openadapt-ml 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

openadapt_ml/benchmarks/azure.py

@@ -30,16 +30,13 @@ from __future__ import annotations
 
 import json
 import logging
-import os
-import tempfile
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Callable
+from typing import Callable
 
-from openadapt_ml.benchmarks.agent import BenchmarkAgent
-from openadapt_ml.benchmarks.base import BenchmarkResult, BenchmarkTask
+from openadapt_evals import BenchmarkAgent, BenchmarkResult, BenchmarkTask
 
 logger = logging.getLogger(__name__)
 
@@ -233,7 +230,9 @@ class AzureMLClient:
                 resource_group_name=self.config.resource_group,
                 workspace_name=self.config.workspace_name,
             )
-            logger.info(f"Connected to Azure ML workspace: {self.config.workspace_name}")
+            logger.info(
+                f"Connected to Azure ML workspace: {self.config.workspace_name}"
+            )
         return self._client
 
     def _get_credential(self):
@@ -241,11 +240,13 @@ class AzureMLClient:
         from openadapt_ml.config import settings
 
         # Use service principal if credentials are configured
-        if all([
-            settings.azure_client_id,
-            settings.azure_client_secret,
-            settings.azure_tenant_id,
-        ]):
+        if all(
+            [
+                settings.azure_client_id,
+                settings.azure_client_secret,
+                settings.azure_tenant_id,
+            ]
+        ):
             logger.info("Using service principal authentication")
             return self._ClientSecretCredential(
                 tenant_id=settings.azure_tenant_id,
@@ -301,7 +302,10 @@ class AzureMLClient:
                 f"/providers/Microsoft.ManagedIdentity"
                 f"/userAssignedIdentities/{self.config.managed_identity_name}"
             )
-            compute.identity = {"type": "UserAssigned", "user_assigned_identities": [identity_id]}
+            compute.identity = {
+                "type": "UserAssigned",
+                "user_assigned_identities": [identity_id],
+            }
 
         print(f"      Creating VM: {name}...", end="", flush=True)
         self.client.compute.begin_create_or_update(compute).result()
@@ -355,6 +359,7 @@ class AzureMLClient:
         command: str,
         environment_variables: dict[str, str] | None = None,
         display_name: str | None = None,
+        timeout_hours: float = 4.0,
     ) -> str:
         """Submit a job to a compute instance.
 
@@ -363,6 +368,8 @@ class AzureMLClient:
             command: Command to run.
             environment_variables: Environment variables.
             display_name: Job display name.
+            timeout_hours: Maximum job duration in hours (default: 4). The job
+                will be automatically canceled after this duration.
 
         Returns:
             Job name/ID.
@@ -376,16 +383,28 @@ class AzureMLClient:
             name="waa-agent-env",
         )
 
+        import time
+        import uuid
+
+        timestamp = int(time.time())
+        unique_id = str(uuid.uuid4())[:8]
+        job_name = f"waa-{compute_name}-{timestamp}-{unique_id}"
+
+        # Convert hours to seconds for Azure ML timeout
+        timeout_seconds = int(timeout_hours * 3600)
+
         job = ml_command(
             command=command,
             environment=env,
             compute=compute_name,
+            name=job_name,  # Unique job name for Azure ML
             display_name=display_name or f"waa-job-{compute_name}",
             environment_variables=environment_variables or {},
+            limits={"timeout": timeout_seconds},
         )
 
         submitted = self.client.jobs.create_or_update(job)
-        logger.info(f"Job submitted: {submitted.name}")
+        logger.info(f"Job submitted: {submitted.name} (timeout: {timeout_hours}h)")
         return submitted.name
 
     def wait_for_job(self, job_name: str, timeout_seconds: int = 3600) -> dict:
@@ -458,6 +477,7 @@ class AzureWAAOrchestrator:
         max_steps_per_task: int = 15,
         on_worker_complete: Callable[[WorkerState], None] | None = None,
         cleanup_on_complete: bool = True,
+        timeout_hours: float = 4.0,
     ) -> list[BenchmarkResult]:
         """Run evaluation across multiple Azure VMs.
 
@@ -468,12 +488,14 @@ class AzureWAAOrchestrator:
             max_steps_per_task: Maximum steps per task.
             on_worker_complete: Callback when a worker finishes.
             cleanup_on_complete: Whether to delete VMs after completion.
+            timeout_hours: Maximum job duration in hours (default: 4). Jobs are
+                auto-canceled after this duration to prevent runaway costs.
 
         Returns:
             List of BenchmarkResult for all tasks.
         """
         # Load tasks
-        from openadapt_ml.benchmarks.waa import WAAAdapter
+        from openadapt_evals import WAAMockAdapter as WAAAdapter
 
         adapter = WAAAdapter(waa_repo_path=self.waa_repo_path)
         if task_ids:
@@ -513,17 +535,21 @@ class AzureWAAOrchestrator:
 
         try:
             # Provision VMs in parallel
-            print(f"[2/4] Provisioning {num_workers} Azure VM(s)... (this takes 3-5 minutes)")
+            print(
+                f"[2/4] Provisioning {num_workers} Azure VM(s)... (this takes 3-5 minutes)"
+            )
             self._provision_workers(workers)
-            print(f"      VM(s) ready")
+            print("      VM(s) ready")
 
             # Submit jobs to workers
-            print(f"[3/4] Submitting evaluation jobs...")
-            self._submit_worker_jobs(workers, task_batches, agent, max_steps_per_task)
-            print(f"      Jobs submitted")
+            print("[3/4] Submitting evaluation jobs...")
+            self._submit_worker_jobs(
+                workers, task_batches, agent, max_steps_per_task, timeout_hours
+            )
+            print("      Jobs submitted")
 
             # Wait for completion and collect results
-            print(f"[4/4] Waiting for workers to complete...")
+            print("[4/4] Waiting for workers to complete...")
             results = self._wait_and_collect_results(workers, on_worker_complete)
 
             self._current_run.status = "completed"
@@ -577,8 +603,17 @@ class AzureWAAOrchestrator:
         task_batches: list[list[BenchmarkTask]],
         agent: BenchmarkAgent,
         max_steps: int,
+        timeout_hours: float = 4.0,
     ) -> None:
-        """Submit evaluation jobs to workers."""
+        """Submit evaluation jobs to workers.
+
+        Args:
+            workers: List of worker states.
+            task_batches: Task batches for each worker.
+            agent: Agent to run.
+            max_steps: Maximum steps per task.
+            timeout_hours: Maximum job duration in hours.
+        """
         for worker, tasks in zip(workers, task_batches):
             if worker.status == "failed":
                 continue
@@ -591,7 +626,7 @@ class AzureWAAOrchestrator:
                 # Build command
                 command = self._build_worker_command(task_ids_json, max_steps, agent)
 
-                # Submit job
+                # Submit job with timeout
                 self.ml_client.submit_job(
                     compute_name=worker.compute_name,
                     command=command,
@@ -600,6 +635,7 @@ class AzureWAAOrchestrator:
                         "WAA_MAX_STEPS": str(max_steps),
                     },
                     display_name=f"waa-worker-{worker.worker_id}",
+                    timeout_hours=timeout_hours,
                 )
                 worker.status = "running"
                 worker.start_time = time.time()
@@ -625,9 +661,11 @@ class AzureWAAOrchestrator:
         # TODO: Serialize agent config and pass to remote worker
         # For now, workers use a default agent configuration
         _ = agent  # Reserved for agent serialization
+        # WAA Docker image has client at /client (see Dockerfile-WinArena)
+        # The run.py script is at /client/run.py (not a module, so use python run.py)
         return f"""
-        cd /workspace/WindowsAgentArena && \
-        python -m client.run \
+        cd /client && \
+        python run.py \
             --task_ids '{task_ids_json}' \
             --max_steps {max_steps} \
             --output_dir /outputs