91 lines
4.8 KiB
Python
91 lines
4.8 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import subprocess
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Protocol
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class BackendResult:
|
||
|
|
status: str
|
||
|
|
checkpoint_reference: str = ""
|
||
|
|
checkpoint_hash: str = ""
|
||
|
|
metrics: dict[str, float] | None = None
|
||
|
|
stdout: str = ""
|
||
|
|
stderr: str = ""
|
||
|
|
failure_category: str = ""
|
||
|
|
failure_details: str = ""
|
||
|
|
|
||
|
|
|
||
|
|
class TrainingBackend(Protocol):
|
||
|
|
def estimate_runtime(self, recipe: dict[str, Any]) -> int: ...
|
||
|
|
def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult: ...
|
||
|
|
def validate_checkpoint(self, reference: str) -> bool: ...
|
||
|
|
|
||
|
|
|
||
|
|
class FakeTrainingBackend:
|
||
|
|
"""Deterministic backend for workflow tests; never launches training."""
|
||
|
|
|
||
|
|
def __init__(self, outcomes: list[dict[str, Any]] | None = None) -> None:
|
||
|
|
self.outcomes = list(outcomes or [{"status": "SUCCEEDED", "metrics": {"primary": 0.75}}])
|
||
|
|
|
||
|
|
def estimate_runtime(self, recipe: dict[str, Any]) -> int:
|
||
|
|
return int(recipe.get("estimated_runtime_seconds", 60))
|
||
|
|
|
||
|
|
def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult:
|
||
|
|
outcome = self.outcomes.pop(0) if self.outcomes else {"status": "SUCCEEDED", "metrics": {"primary": 0.75}}
|
||
|
|
status = str(outcome.get("status", "SUCCEEDED"))
|
||
|
|
reference = str(outcome.get("checkpoint_reference", f"fake://{hashlib.sha256(json.dumps(outcome, sort_keys=True).encode()).hexdigest()[:16]}"))
|
||
|
|
return BackendResult(status=status, checkpoint_reference=reference, checkpoint_hash=hashlib.sha256(reference.encode()).hexdigest(), metrics=outcome.get("metrics", {}), failure_category=str(outcome.get("failure_category", "")), failure_details=str(outcome.get("failure_details", "")))
|
||
|
|
|
||
|
|
def validate_checkpoint(self, reference: str) -> bool:
|
||
|
|
return reference.startswith("fake://")
|
||
|
|
|
||
|
|
|
||
|
|
class GuardSubprocessBackend:
|
||
|
|
"""Scoped wrapper for the discovered Guard trainer; commands come from profiles, never an LLM."""
|
||
|
|
|
||
|
|
def estimate_runtime(self, recipe: dict[str, Any]) -> int:
|
||
|
|
return int(recipe.get("estimated_runtime_seconds", 90 * 60))
|
||
|
|
|
||
|
|
def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult:
|
||
|
|
try:
|
||
|
|
completed = subprocess.run(command, cwd=working_directory, capture_output=True, text=True, timeout=timeout_seconds, check=False)
|
||
|
|
except subprocess.TimeoutExpired as exc:
|
||
|
|
return BackendResult(status="TIMEOUT", stdout=exc.stdout or "", stderr=exc.stderr or "", failure_category="TIMEOUT", failure_details=f"Exceeded {timeout_seconds}s")
|
||
|
|
stdout, stderr = completed.stdout or "", completed.stderr or ""
|
||
|
|
if completed.returncode:
|
||
|
|
category = "OOM" if "out of memory" in (stdout + stderr).lower() else "PROCESS_FAILURE"
|
||
|
|
return BackendResult(status=category if category == "OOM" else "FAILED", stdout=stdout, stderr=stderr, failure_category=category, failure_details=f"Exit code {completed.returncode}")
|
||
|
|
return BackendResult(status="SUCCEEDED", stdout=stdout, stderr=stderr)
|
||
|
|
|
||
|
|
def validate_checkpoint(self, reference: str) -> bool:
|
||
|
|
path = Path(reference)
|
||
|
|
return path.is_dir() and any(path.glob("adapter_model.*"))
|
||
|
|
|
||
|
|
|
||
|
|
class SparkGuardBackend(GuardSubprocessBackend):
|
||
|
|
"""Runs only profile-generated Guard commands over the configured Spark SSH alias."""
|
||
|
|
|
||
|
|
def __init__(self, ssh_alias: str = "spark") -> None:
|
||
|
|
self.ssh_alias = ssh_alias
|
||
|
|
|
||
|
|
def launch(self, *, command: list[str], working_directory: str, timeout_seconds: int) -> BackendResult:
|
||
|
|
if not command:
|
||
|
|
return BackendResult(status="FAILED", failure_category="COMMAND_SCOPE", failure_details="Missing profile-generated command.")
|
||
|
|
import shlex
|
||
|
|
|
||
|
|
remote = "cd " + shlex.quote(working_directory) + " && " + " ".join(shlex.quote(part) for part in command)
|
||
|
|
try:
|
||
|
|
completed = subprocess.run(["ssh", self.ssh_alias, remote], capture_output=True, text=True, timeout=timeout_seconds, check=False)
|
||
|
|
except subprocess.TimeoutExpired as exc:
|
||
|
|
return BackendResult(status="TIMEOUT", stdout=exc.stdout or "", stderr=exc.stderr or "", failure_category="TIMEOUT", failure_details=f"Exceeded {timeout_seconds}s")
|
||
|
|
if completed.returncode:
|
||
|
|
combined = (completed.stdout or "") + (completed.stderr or "")
|
||
|
|
category = "OOM" if "out of memory" in combined.lower() else "REMOTE_PROCESS_FAILURE"
|
||
|
|
return BackendResult(status="OOM" if category == "OOM" else "FAILED", stdout=completed.stdout or "", stderr=completed.stderr or "", failure_category=category, failure_details=f"Exit code {completed.returncode}")
|
||
|
|
return BackendResult(status="SUCCEEDED", stdout=completed.stdout or "", stderr=completed.stderr or "")
|