swegen 0.1.0__py3-none-any.whl

swegen/tools/harbor_runner.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import shutil
+import subprocess
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+from harbor.models.environment_type import EnvironmentType
+from harbor.models.job.result import JobResult
+from harbor.models.trial.paths import TrialPaths
+from harbor.models.trial.result import TrialResult
+
+
+def harbor_cmd_base() -> list[str]:
+    """Get the base command to invoke Harbor.
+
+    Prefers direct `harbor` binary, falls back to `uv run harbor`.
+    """
+    if shutil.which("harbor"):
+        return ["harbor"]
+    if shutil.which("uv"):
+        return ["uv", "run", "harbor"]
+    return ["python", "-m", "harbor"]
+
+
+def run_harbor_agent(
+    task_id: str,
+    dataset_path: Path,
+    jobs_dir: Path,
+    agent: str,
+    timeout_multiplier: float | None = None,
+    capture_output: bool = False,
+    delete_after: bool = True,
+    environment: EnvironmentType = EnvironmentType.DOCKER,
+) -> tuple[int, Path | None]:
+    """Run a Harbor agent and return (exit_code, job_result_path).
+
+    Args:
+        task_id: The task identifier
+        dataset_path: Path to the Harbor dataset root
+        jobs_dir: Parent directory for job artifacts
+        agent: Agent type ("nop" or "oracle")
+        timeout_multiplier: Optional timeout multiplier for long tasks
+        capture_output: If True, suppress stdout/stderr (for rich console usage)
+        delete_after: If True, delete Docker images after run (default: True)
+                     Set to False to keep images for faster subsequent runs
+        environment: Environment type (docker, daytona, e2b, modal, runloop, gke)
+
+    Returns:
+        Tuple of (exit_code, path_to_result_json or None)
+    """
+    # Create unique job directory to avoid race conditions
+    unique_parent = jobs_dir / f"{task_id}.{agent}.{int(time.time())}"
+    unique_parent.mkdir(parents=True, exist_ok=True)
+    before = set(unique_parent.iterdir())
+
+    cmd = harbor_cmd_base() + [
+        "run",
+        "--agent",
+        agent,
+        "-p",
+        str(dataset_path),
+        "-t",
+        task_id,
+        "--jobs-dir",
+        str(unique_parent),
+        "--env",
+        environment.value,
+    ]
+    if timeout_multiplier is not None:
+        cmd += ["--timeout-multiplier", str(timeout_multiplier)]
+
+    # Control image deletion: --no-delete keeps images for faster subsequent runs
+    if not delete_after:
+        cmd.append("--no-delete")
+
+    proc: subprocess.CompletedProcess[str]
+    if capture_output:
+        proc = subprocess.run(cmd, check=False, capture_output=True, text=True)
+    else:
+        proc_bytes = subprocess.run(cmd, check=False)
+        # Convert to text version for consistent return type
+        proc = subprocess.CompletedProcess(
+            args=proc_bytes.args,
+            returncode=proc_bytes.returncode,
+            stdout="",
+            stderr="",
+        )
+
+    # Check if directory still exists after subprocess
+    if not unique_parent.exists():
+        return proc.returncode, None
+
+    after = set(unique_parent.iterdir())
+    new_dirs = [p for p in (after - before) if p.is_dir()]
+    job_dir = (
+        sorted(new_dirs, key=lambda p: p.stat().st_mtime, reverse=True)[0] if new_dirs else None
+    )
+    job_result = (job_dir / "result.json").resolve() if job_dir else None
+
+    return proc.returncode, job_result
+
+
+@dataclass(frozen=True)
+class HarborOutcome:
+    reward: int | None
+    error: str | None
+
+
+def parse_harbor_outcome(job_result_path: Path | None) -> HarborOutcome:
+    """Parse Harbor job result and return both reward and error (best-effort).
+
+    Uses Harbor's JobResult and TrialResult Pydantic models for type-safe parsing.
+    This automatically handles schema changes and provides better error messages.
+
+    Args:
+        job_result_path: Path to the job-level result.json
+
+    Returns:
+        HarborOutcome with:
+        - reward: 0 or 1 (or None if unavailable)
+        - error: best-effort exception message (or None)
+    """
+    if not job_result_path or not job_result_path.exists():
+        return HarborOutcome(reward=None, error=None)
+
+    try:
+        # Use Harbor's JobResult model for type-safe parsing
+        job_result = JobResult.model_validate_json(job_result_path.read_text())
+
+        # Prefer structured exception info from typed trial results.
+        error: str | None = None
+        for trial_result in job_result.trial_results:
+            if getattr(trial_result, "exception_info", None):
+                exc = trial_result.exception_info
+                msg = getattr(exc, "exception_message", None) or getattr(exc, "exception_type", None)
+                if msg:
+                    error = str(msg)
+                    break
+
+        # Method 1: Check reward_stats in job stats (fastest)
+        if job_result.stats.evals:
+            # Get first eval (typically only one for single-task runs)
+            first_eval = next(iter(job_result.stats.evals.values()))
+
+            # Check reward_stats for "reward" key
+            if first_eval.reward_stats and "reward" in first_eval.reward_stats:
+                reward_map = first_eval.reward_stats["reward"]
+
+                # Check for reward=1 first (oracle success)
+                if 1 in reward_map or 1.0 in reward_map:
+                    return HarborOutcome(reward=1, error=error)
+                # Then check for reward=0 (nop success)
+                if 0 in reward_map or 0.0 in reward_map:
+                    return HarborOutcome(reward=0, error=error)
+
+        # Method 2: Check trial results directly
+        for trial_result in job_result.trial_results:
+            if trial_result.verifier_result and trial_result.verifier_result.rewards:
+                reward_value = trial_result.verifier_result.rewards.get("reward")
+                if reward_value is not None:
+                    return HarborOutcome(reward=int(float(reward_value)), error=error)
+
+        # Method 3: Fallback - scan trial directories using TrialPaths
+        job_root = job_result_path.parent
+        for trial_dir in (p for p in job_root.iterdir() if p.is_dir()):
+            try:
+                trial_paths = TrialPaths(trial_dir)
+                if not trial_paths.result_path.exists():
+                    continue
+                trial_result = TrialResult.model_validate_json(trial_paths.result_path.read_text())
+
+                if error is None and getattr(trial_result, "exception_info", None):
+                    exc = trial_result.exception_info
+                    msg = getattr(exc, "exception_message", None) or getattr(exc, "exception_type", None)
+                    if msg:
+                        error = str(msg)
+
+                if trial_result.verifier_result and trial_result.verifier_result.rewards:
+                    reward_value = trial_result.verifier_result.rewards.get("reward")
+                    if reward_value is not None:
+                        return HarborOutcome(reward=int(float(reward_value)), error=error)
+            except Exception:
+                # Not a valid trial directory, continue searching
+                continue
+
+    except Exception:
+        return HarborOutcome(reward=None, error=None)
+
+    return HarborOutcome(reward=None, error=error)