Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import csv | |
| import json | |
| import multiprocessing | |
| from collections.abc import Callable | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from dataclasses import dataclass | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from time import perf_counter | |
| from typing import Any | |
| from data_agent_baseline.agents.model import AzureOpenAIModelAdapter, OpenAIModelAdapter | |
| from data_agent_baseline.agents.react import ReActAgent, ReActAgentConfig | |
| from data_agent_baseline.benchmark.dataset import DABenchPublicDataset | |
| from data_agent_baseline.config import AppConfig | |
| from data_agent_baseline.tools.registry import ToolRegistry, create_default_tool_registry | |
| from data_agent_baseline import logger | |
| class TaskRunArtifacts: | |
| task_id: str | |
| task_output_dir: Path | |
| prediction_csv_path: Path | None | |
| trace_path: Path | |
| succeeded: bool | |
| failure_reason: str | None | |
| elapsed_seconds: float | None = None | |
| def to_dict(self) -> dict[str, Any]: | |
| return { | |
| "task_id": self.task_id, | |
| "task_output_dir": str(self.task_output_dir), | |
| "prediction_csv_path": str(self.prediction_csv_path) if self.prediction_csv_path else None, | |
| "trace_path": str(self.trace_path), | |
| "succeeded": self.succeeded, | |
| "failure_reason": self.failure_reason, | |
| "elapsed_seconds": self.elapsed_seconds, | |
| } | |
| def create_run_id() -> str: | |
| return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") | |
| def resolve_run_id(run_id: str | None = None) -> str: | |
| if run_id is None: | |
| return create_run_id() | |
| normalized = run_id.strip() | |
| if not normalized: | |
| raise ValueError("run_id must not be empty.") | |
| if normalized in {".", ".."} or "/" in normalized or "\\" in normalized: | |
| raise ValueError("run_id must be a single directory name, not a path.") | |
| return normalized | |
| def create_run_output_dir(output_root: Path, *, run_id: str | None = None) -> tuple[str, Path]: | |
| effective_run_id = resolve_run_id(run_id) | |
| run_output_dir = output_root / effective_run_id | |
| run_output_dir.mkdir(parents=True, exist_ok=False) | |
| return effective_run_id, run_output_dir | |
| # TODO: Add support for other model providers by making this function more generic and configurable. | |
| def build_model_adapter(config: AppConfig): | |
| """ | |
| Build a model adapter based on the configuration. | |
| If Azure OpenAI configuration is present (azure_endpoint AND azure_deployment are set), | |
| use AzureOpenAIModelAdapter. Otherwise, use OpenAIModelAdapter. | |
| """ | |
| # Check if Azure configuration is provided (need both endpoint and deployment) | |
| has_azure_config = ( | |
| config.agent.azure_endpoint | |
| and config.agent.azure_deployment | |
| and config.agent.azure_api_key | |
| ) | |
| if has_azure_config: | |
| logger.debug( | |
| "Initializing Azure OpenAI model adapter", | |
| endpoint=config.agent.azure_endpoint, | |
| deployment=config.agent.azure_deployment, | |
| api_version=config.agent.azure_api_version, | |
| temperature=config.agent.temperature, | |
| custom_headers=bool(config.agent.azure_custom_headers) | |
| ) | |
| return AzureOpenAIModelAdapter( | |
| azure_endpoint=config.agent.azure_endpoint, | |
| azure_api_key=config.agent.azure_api_key, | |
| azure_api_version=config.agent.azure_api_version, | |
| azure_deployment=config.agent.azure_deployment, | |
| temperature=config.agent.temperature, | |
| azure_custom_headers=config.agent.azure_custom_headers, | |
| ) | |
| # Default to OpenAI configuration | |
| logger.debug( | |
| "Initializing OpenAI model adapter", | |
| model=config.agent.model, | |
| api_base=config.agent.api_base, | |
| temperature=config.agent.temperature | |
| ) | |
| return OpenAIModelAdapter( | |
| model=config.agent.model, | |
| api_base=config.agent.api_base, | |
| api_key=config.agent.api_key, | |
| temperature=config.agent.temperature, | |
| ) | |
| def _write_json(path: Path, payload: dict[str, Any]) -> None: | |
| path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n") | |
| def _write_csv(path: Path, columns: list[str], rows: list[list[Any]]) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| with path.open("w", newline="") as handle: | |
| writer = csv.writer(handle) | |
| writer.writerow(columns) | |
| for row in rows: | |
| writer.writerow(row) | |
| def _failure_run_result_payload(task_id: str, failure_reason: str) -> dict[str, Any]: | |
| return { | |
| "task_id": task_id, | |
| "answer": None, | |
| "steps": [], | |
| "failure_reason": failure_reason, | |
| "succeeded": False, | |
| } | |
| def _run_single_task_core( | |
| *, | |
| task_id: str, | |
| config: AppConfig, | |
| model=None, | |
| tools: ToolRegistry | None = None, | |
| ) -> dict[str, Any]: | |
| logger.debug(f"Starting task execution", task_id=task_id) | |
| public_dataset = DABenchPublicDataset(config.dataset.root_path) | |
| task = public_dataset.get_task(task_id) | |
| logger.debug( | |
| "Task loaded", | |
| task_id=task_id, | |
| difficulty=task.difficulty, | |
| context_dir=str(task.context_dir) | |
| ) | |
| agent = ReActAgent( | |
| model=model or build_model_adapter(config), | |
| tools=tools or create_default_tool_registry(), | |
| config=ReActAgentConfig(max_steps=config.agent.max_steps), | |
| ) | |
| logger.debug(f"Running agent", task_id=task_id, max_steps=config.agent.max_steps) | |
| run_result = agent.run(task) | |
| logger.debug( | |
| "Task completed", | |
| task_id=task_id, | |
| succeeded=run_result.succeeded, | |
| steps=len(run_result.steps) | |
| ) | |
| return run_result.to_dict() | |
| def _run_single_task_in_subprocess(task_id: str, config: AppConfig, queue: multiprocessing.Queue[Any]) -> None: | |
| try: | |
| queue.put( | |
| { | |
| "ok": True, | |
| "run_result": _run_single_task_core(task_id=task_id, config=config), | |
| } | |
| ) | |
| except BaseException as exc: # noqa: BLE001 | |
| queue.put( | |
| { | |
| "ok": False, | |
| "error": str(exc), | |
| } | |
| ) | |
| def _run_single_task_with_timeout(*, task_id: str, config: AppConfig) -> dict[str, Any]: | |
| timeout_seconds = config.run.task_timeout_seconds | |
| if timeout_seconds <= 0: | |
| return _run_single_task_core(task_id=task_id, config=config) | |
| queue: multiprocessing.Queue[Any] = multiprocessing.Queue() | |
| process = multiprocessing.Process( | |
| target=_run_single_task_in_subprocess, | |
| args=(task_id, config, queue), | |
| ) | |
| process.start() | |
| process.join(timeout_seconds) | |
| if process.is_alive(): | |
| process.terminate() | |
| process.join(timeout=1.0) | |
| if process.is_alive(): | |
| process.kill() | |
| process.join() | |
| return _failure_run_result_payload(task_id, f"Task timed out after {timeout_seconds} seconds.") | |
| if queue.empty(): | |
| exit_code = process.exitcode | |
| if exit_code not in (None, 0): | |
| return _failure_run_result_payload( | |
| task_id, | |
| f"Task exited unexpectedly with exit code {exit_code}.", | |
| ) | |
| return _failure_run_result_payload(task_id, "Task exited without returning a result.") | |
| result = queue.get() | |
| if result.get("ok"): | |
| return dict(result["run_result"]) | |
| return _failure_run_result_payload(task_id, f"Task failed with uncaught error: {result['error']}") | |
| def _write_task_outputs(task_id: str, run_output_dir: Path, run_result: dict[str, Any]) -> TaskRunArtifacts: | |
| task_output_dir = run_output_dir / task_id | |
| task_output_dir.mkdir(parents=True, exist_ok=True) | |
| trace_path = task_output_dir / "trace.json" | |
| _write_json(trace_path, run_result) | |
| prediction_csv_path: Path | None = None | |
| answer = run_result.get("answer") | |
| if isinstance(answer, dict): | |
| prediction_csv_path = task_output_dir / "prediction.csv" | |
| _write_csv( | |
| prediction_csv_path, | |
| list(answer.get("columns", [])), | |
| [list(row) for row in answer.get("rows", [])], | |
| ) | |
| return TaskRunArtifacts( | |
| task_id=task_id, | |
| task_output_dir=task_output_dir, | |
| prediction_csv_path=prediction_csv_path, | |
| trace_path=trace_path, | |
| succeeded=bool(run_result.get("succeeded")), | |
| failure_reason=run_result.get("failure_reason"), | |
| elapsed_seconds=run_result.get("e2e_elapsed_seconds"), | |
| ) | |
| def run_single_task( | |
| *, | |
| task_id: str, | |
| config: AppConfig, | |
| run_output_dir: Path, | |
| model=None, | |
| tools: ToolRegistry | None = None, | |
| ) -> TaskRunArtifacts: | |
| started_at = perf_counter() | |
| if model is None and tools is None: | |
| run_result = _run_single_task_with_timeout(task_id=task_id, config=config) | |
| else: | |
| run_result = _run_single_task_core(task_id=task_id, config=config, model=model, tools=tools) | |
| run_result["e2e_elapsed_seconds"] = round(perf_counter() - started_at, 3) | |
| return _write_task_outputs(task_id, run_output_dir, run_result) | |
| def run_benchmark( | |
| *, | |
| config: AppConfig, | |
| model=None, | |
| tools: ToolRegistry | None = None, | |
| limit: int | None = None, | |
| progress_callback: Callable[[TaskRunArtifacts], None] | None = None, | |
| ) -> tuple[Path, list[TaskRunArtifacts]]: | |
| benchmark_start = perf_counter() | |
| effective_run_id, run_output_dir = create_run_output_dir(config.run.output_dir, run_id=config.run.run_id) | |
| dataset = DABenchPublicDataset(config.dataset.root_path) | |
| tasks = dataset.iter_tasks() | |
| if limit is not None: | |
| tasks = tasks[:limit] | |
| effective_workers = config.run.max_workers | |
| if effective_workers < 1: | |
| raise ValueError("max_workers must be at least 1.") | |
| if model is not None or tools is not None: | |
| effective_workers = 1 | |
| task_ids = [task.task_id for task in tasks] | |
| task_artifacts: list[TaskRunArtifacts] | |
| if effective_workers == 1: | |
| shared_model = model or build_model_adapter(config) | |
| shared_tools = tools or create_default_tool_registry() | |
| task_artifacts = [] | |
| for task_id in task_ids: | |
| artifact = run_single_task( | |
| task_id=task_id, | |
| config=config, | |
| run_output_dir=run_output_dir, | |
| model=shared_model, | |
| tools=shared_tools, | |
| ) | |
| task_artifacts.append(artifact) | |
| if progress_callback is not None: | |
| progress_callback(artifact) | |
| else: | |
| with ThreadPoolExecutor(max_workers=effective_workers) as executor: | |
| future_to_index = { | |
| executor.submit( | |
| run_single_task, | |
| task_id=task_id, | |
| config=config, | |
| run_output_dir=run_output_dir, | |
| ): index | |
| for index, task_id in enumerate(task_ids) | |
| } | |
| indexed_artifacts: list[TaskRunArtifacts | None] = [None] * len(task_ids) | |
| for future in as_completed(future_to_index): | |
| artifact = future.result() | |
| indexed_artifacts[future_to_index[future]] = artifact | |
| if progress_callback is not None: | |
| progress_callback(artifact) | |
| task_artifacts = [artifact for artifact in indexed_artifacts if artifact is not None] | |
| summary_path = run_output_dir / "summary.json" | |
| total_elapsed = round(perf_counter() - benchmark_start, 3) | |
| _write_json( | |
| summary_path, | |
| { | |
| "run_id": effective_run_id, | |
| "task_count": len(task_artifacts), | |
| "succeeded_task_count": sum(1 for artifact in task_artifacts if artifact.succeeded), | |
| "max_workers": effective_workers, | |
| "total_elapsed_seconds": total_elapsed, | |
| "tasks": [artifact.to_dict() for artifact in task_artifacts], | |
| }, | |
| ) | |
| return run_output_dir, task_artifacts | |