inference update (e4500b6e) · Commits · 张泽凯 / VAGEN

scripts/exps/prompting/frozenlake/inference_config.yaml

+15 −12

Original line number	Diff line number	Diff line
		batch_size: 32
		max_steps: 10
		num_workers: 4

		dataset_dir: data
		use_split: test

		server_url: http://localhost:5000
		# Server configuration
		server_url: "http://localhost:5000"
		server_timeout: 600
		server_max_workers: 48

		output_dir: inference_results
		show_progress: true
		# Inference parameters
		batch_size: 32
		max_steps: 10
		split: "test"
		debug: false

		# Output configuration
		output_dir: "inference_outputs"

		# WandB configuration
		use_wandb: true
		wandb_project: vagen-prompting-frozenlake
		wandb_project: "vagen-inference"

		debug: false
		No newline at end of file
		# Display settings
		show_progress: true
		val_generations_to_log_to_wandb: 10
		No newline at end of file

scripts/exps/prompting/frozenlake/model_config.yaml

+2 −9

Original line number	Diff line number	Diff line
		models:
		qwen_0.5b:
		provider: vllm
		model_name: Qwen/Qwen2.5-0.5B-Instruct
		max_tokens: 1024
		temperature: 0.7
		top_p: 0.9
		top_k: 50
		tensor_parallel_size: 1
		gpu_memory_utilization: 0.9
		dtype: bfloat16
		trust_remote_code: true

		# qwen_vl_3b:
		# provider: vllm
		# model_name: Qwen/Qwen2.5-VL-3B-Instruct
		# max_tokens: 1024
		# temperature: 0.7
		# top_p: 0.9
		# top_k: 50
		# tensor_parallel_size: 1
		# tensor_parallel_size: 2
		# gpu_memory_utilization: 0.9
		No newline at end of file
		# dtype: bfloat16
		# trust_remote_code: true
		No newline at end of file

test_frozenlake/frozenlake_0.png

0 → 100644

+6.07 KiB

Loading image diff...

vagen/inference/run_inference.py

+191 −157

Original line number	Diff line number	Diff line
		# vagen/inference/run_inference.py

		import os
		import sys
		import argparse
		import logging
		import yaml
		import json
		import wandb
		from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
		import pandas as pd
		import numpy as np
		from datetime import datetime
		from typing import Dict, List, Any
		from pathlib import Path

		from ..mllm_agent.model_interface.factory_model import ModelFactory
		from ..mllm_agent.inference_rollout.inference_rollout_service import InferenceRolloutService
		from .utils.environment import load_environment_configs_from_parquet
		from .utils.logging import setup_wandb_for_model, log_metrics_to_wandb
		from .utils.config import load_yaml_config
		from vagen.mllm_agent.model_interface.factory_model import ModelFactory
		from vagen.mllm_agent.inference_rollout.inference_rollout_service import InferenceRolloutService
		from vagen.inference.utils.metrics import calculate_aggregate_metrics
		from vagen.inference.utils.logging import maybe_log_val_generations_to_wandb, log_metrics_by_config_id

		logger = logging.getLogger(__name__)

		def parse_args():
		"""Parse command line arguments."""
		parser = argparse.ArgumentParser(description="Run inference with multiple models")
		parser = argparse.ArgumentParser(description="Run inference with models")

		parser.add_argument("--inference_config_path", type=str, required=True,
		help="Path to inference configuration YAML")
		@@ -28,8 +33,132 @@ def parse_args():

		return parser.parse_args()

		def load_yaml_config(config_path: str) -> Dict[str, Any]:
		"""Load configuration from YAML file."""
		with open(config_path, 'r') as f:
		config = yaml.safe_load(f)
		return config

		def load_environment_configs_from_parquet(val_files_path: str) -> List[Dict]:
		"""Load environment configurations from parquet file."""
		df = pd.read_parquet(val_files_path)
		env_configs = []

		for idx, row in df.iterrows():
		extra_info = row.get('extra_info', {})
		config = {
		"env_name": extra_info.get("env_name"),
		"env_config": extra_info.get("env_config", {}),
		"seed": extra_info.get("seed", 42)
		}
		env_configs.append(config)

		return env_configs

		def setup_wandb(model_name: str, model_config: Dict, inference_config: Dict) -> None:
		"""Initialize wandb run."""
		timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
		run_name = f"{model_name}_inference_{timestamp}"

		wandb.init(
		project=inference_config.get('wandb_project', 'vagen-inference'),
		name=run_name,
		config={
		"model_name": model_name,
		"model_config": model_config,
		"inference_config": inference_config
		}
		)

		def run_inference_batch(
		service: InferenceRolloutService,
		env_configs: List[Dict],
		batch_size: int,
		max_steps: int
		) -> List[Dict]:
		"""Run inference on a batch of environments."""
		all_results = []

		# Process environments in batches
		for i in range(0, len(env_configs), batch_size):
		batch_configs = env_configs[i:i + batch_size]
		logger.info(f"Processing batch {i//batch_size + 1}/{(len(env_configs) + batch_size - 1)//batch_size}")

		# Reset environments for this batch
		service.reset(batch_configs)

		# Run inference
		service.run(max_steps=max_steps)

		# Get results
		batch_results = service.recording_to_log()
		all_results.extend(batch_results)

		# Log batch progress
		if wandb.run:
		wandb.log({
		"batch_progress": (i + batch_size) / len(env_configs) * 100,
		"num_environments_processed": i + len(batch_configs)
		})

		return all_results

		def log_results_to_wandb(results: List[Dict], global_step: int = 0) -> None:
		"""Log results to wandb using the same format as training."""
		# Log metrics by config_id (same as training)
		metrics = log_metrics_by_config_id(results, mode='val') # Use 'val' to match training prefix
		wandb.log(metrics)

		# Log generation table (same as training)
		generations_to_log = 10 # You can make this configurable
		maybe_log_val_generations_to_wandb(results, generations_to_log, global_step)

		# Log overall statistics
		success_count = sum(1 for r in results if r['metrics'].get('success', 0) > 0)
		done_count = sum(1 for r in results if r['metrics'].get('done', 0) > 0)
		total_score = sum(r['metrics'].get('score', 0) for r in results)
		avg_steps = sum(r['metrics'].get('step', 0) for r in results) / len(results) if results else 0

		wandb.log({
		"val/success_rate": success_count / len(results) if results else 0,
		"val/completion_rate": done_count / len(results) if results else 0,
		"val/average_score": total_score / len(results) if results else 0,
		"val/average_steps": avg_steps
		})

		def save_results(results: List[Dict], output_dir: str) -> None:
		"""Save results to disk."""
		os.makedirs(output_dir, exist_ok=True)

		# Custom JSON encoder to handle NumPy types
		class NumpyEncoder(json.JSONEncoder):
		def default(self, obj):
		import numpy as np
		if isinstance(obj, np.integer):
		return int(obj)
		elif isinstance(obj, np.floating):
		return float(obj)
		elif isinstance(obj, np.bool_):
		return bool(obj)
		elif isinstance(obj, np.ndarray):
		return obj.tolist()
		return super(NumpyEncoder, self).default(obj)

		# Save raw results
		results_file = os.path.join(output_dir, "results.json")
		with open(results_file, "w") as f:
		json.dump(results, f, indent=2, cls=NumpyEncoder)

		# Save summary
		summary = calculate_aggregate_metrics(results)
		summary_file = os.path.join(output_dir, "summary.json")
		with open(summary_file, "w") as f:
		json.dump(summary, f, indent=2, cls=NumpyEncoder)

		logger.info(f"Results saved to {output_dir}")

		def main():
		"""Main entry point for parallel multi-model inference."""
		"""Main entry point for inference."""
		args = parse_args()

		# Load configurations
		@@ -42,183 +171,88 @@ def main():
		format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
		)

		logger.info("Starting multi-model inference pipeline")

		try:
		run_parallel_model_inference(
		inference_config=inference_config,
		model_config=model_config,
		val_files_path=args.val_files_path
		)
		except Exception as e:
		logger.error(f"Inference pipeline failed: {str(e)}")
		raise
		finally:
		logger.info("Inference pipeline completed")
		logger.info("Starting inference pipeline")

		def run_parallel_model_inference(
		inference_config: Dict,
		model_config: Dict,
		val_files_path: str
		) -> None:
		"""
		Runs inference for all models in parallel.
		Each model processes all environments independently.
		Uses thread pool for API models, process pool for local models.
		"""
		# Load environment configs once for all models
		env_configs = load_environment_configs_from_parquet(val_files_path)
		# Load environment configurations
		env_configs = load_environment_configs_from_parquet(args.val_files_path)
		logger.info(f"Loaded {len(env_configs)} environment configurations")

		# Extract models from config
		# Process each model
		models = model_config.get('models', {})

		# Determine which executor to use based on model types
		use_processes = any(
		model_cfg.get('provider') == "vllm"
		for model_cfg in models.values()
		)

		executor_class = ProcessPoolExecutor if use_processes else ThreadPoolExecutor
		max_workers = inference_config.get('num_workers', 4)

		logger.info(f"Using {executor_class.__name__} with {max_workers} workers")

		# Run inference for each model in parallel
		with executor_class(max_workers=max_workers) as executor:
		futures = []

		for model_name, model_cfg in models.items():
		future = executor.submit(
		run_single_model_inference,
		model_name,
		model_cfg,
		env_configs,
		inference_config
		)
		futures.append(future)

		# Wait for all models to complete
		for future in futures:
		try:
		future.result()
		except Exception as e:
		logger.error(f"Model inference failed: {str(e)}")

		def run_single_model_inference(
		model_name: str,
		model_config: Dict,
		env_configs: List[Dict],
		inference_config: Dict
		) -> None:
		"""
		Runs inference for a single model on all environments.
		Logs results directly to wandb, no local saving.
		"""
		logger.info(f"Starting inference for model: {model_name}")
		logger.info(f"Running inference for model: {model_name}")

		# Setup wandb for this model
		if inference_config.get('use_wandb', True):
		setup_wandb_for_model(model_name, model_config, inference_config)
		setup_wandb(model_name, model_cfg, inference_config)

		try:
		# Create service for this model
		service = create_model_service(model_name, model_config, inference_config)

		# Process environments in batches
		batch_size = inference_config.get('batch_size', 32)
		total_batches = (len(env_configs) + batch_size - 1) // batch_size
		# Create model interface
		model_interface = ModelFactory.create(model_cfg)

		all_results = []
		# Create inference service
		service = InferenceRolloutService(
		config=inference_config,
		model_interface=model_interface,
		base_url=inference_config.get('server_url', 'http://localhost:5000'),
		timeout=inference_config.get('server_timeout', 600),
		max_workers=inference_config.get('server_max_workers', 48),
		split=inference_config.get('split', 'test'),
		debug=inference_config.get('debug', False)
		)

		for batch_idx in range(total_batches):
		start_idx = batch_idx * batch_size
		end_idx = min(start_idx + batch_size, len(env_configs))
		batch_configs = env_configs[start_idx:end_idx]
		# Run inference
		results = run_inference_batch(
		service=service,
		env_configs=env_configs,
		batch_size=inference_config.get('batch_size', 32),
		max_steps=inference_config.get('max_steps', 10)
		)

		logger.info(f"Model {model_name}: Processing batch {batch_idx+1}/{total_batches}")
		# Log results to wandb (aligned with training format)
		if inference_config.get('use_wandb', True):
		log_results_to_wandb(results, global_step=0)

		# Reset environments for this batch
		service.reset(batch_configs)
		# Save results
		output_dir = os.path.join(
		inference_config.get('output_dir', 'inference_outputs'),
		model_name,
		datetime.now().strftime("%Y%m%d_%H%M%S")
		)
		save_results(results, output_dir)

		# Run inference
		service.run(max_steps=inference_config.get('max_steps', 10))
		# Print summary
		print(f"\n===== Results for {model_name} =====")
		print(f"Total environments: {len(results)}")

		# Get results
		batch_results = service.recording_to_log()
		all_results.extend(batch_results)
		success_count = sum(1 for r in results if r['metrics'].get('success', 0) > 0)
		done_count = sum(1 for r in results if r['metrics'].get('done', 0) > 0)
		total_score = sum(r['metrics'].get('score', 0) for r in results)

		# Log metrics to wandb
		if inference_config.get('use_wandb', True):
		log_metrics_to_wandb(batch_results, model_name)
		print(f"Success rate: {success_count / len(results) * 100:.1f}%")
		print(f"Completion rate: {done_count / len(results) * 100:.1f}%")
		print(f"Average score: {total_score / len(results):.4f}")

		# Print summary
		print_model_summary(model_name, all_results)
		# Print by config_id results (like training)
		metrics_by_config = log_metrics_by_config_id(results, mode='val')
		print("\nMetrics by config:")
		for metric_name, value in metrics_by_config.items():
		print(f" {metric_name}: {value:.4f}")

		except Exception as e:
		logger.error(f"Error during inference for model {model_name}: {str(e)}")
		raise

		finally:
		# Cleanup
		if 'service' in locals():
		cleanup_model_resources(model_name, service)

		def print_model_summary(
		model_name: str,
		results: List[Dict]
		) -> None:
		"""
		Prints a simple summary for the model at the end of inference.
		Just basic stats like completion count, not per-environment details.
		"""
		total_envs = len(results)
		completed_envs = sum(1 for r in results if r['metrics'].get('done', False))

		print(f"\n===== Model: {model_name} =====")
		print(f"Total environments: {total_envs}")
		print(f"Completed environments: {completed_envs}")
		print(f"Completion rate: {completed_envs/total_envs*100:.1f}%")

		def create_model_service(
		model_name: str,
		model_config: Dict,
		inference_config: Dict
		) -> InferenceRolloutService:
		"""
		Creates InferenceRolloutService for a specific model.
		"""
		# Create model interface using factory
		model_interface = ModelFactory.create(model_config)

		# Create and return service
		service = InferenceRolloutService(
		config=inference_config,
		model_interface=model_interface,
		base_url=inference_config.get('server_url', 'http://localhost:5000'),
		timeout=inference_config.get('server_timeout', 600),
		max_workers=inference_config.get('server_max_workers', 48),
		split=inference_config.get('use_split', 'test'),
		debug=inference_config.get('debug', False)
		)

		return service

		def cleanup_model_resources(
		model_name: str,
		service: InferenceRolloutService
		) -> None:
		"""
		Cleans up resources for a specific model.
		Closes environments and finishes wandb run.
		"""
		logger.info(f"Cleaning up resources for model: {model_name}")

		# Close environments
		service.close()

		# Finish wandb run if active
		if wandb.run is not None:
		# Finish wandb run
		if wandb.run:
		wandb.finish()

		logger.info("Inference pipeline completed")

		if __name__ == "__main__":
		main()
		No newline at end of file

vagen/inference/utils/logging.py

+143 −75

File changed.

Preview size limit exceeded, changes collapsed.

Admin message