@trac3r/oh-my-god 2.2.11

package/lab/__init__.py

@@ -0,0 +1 @@
+"""Lab package for OMG."""

package/lab/axolotl_adapter.py

@@ -0,0 +1,531 @@
+from __future__ import annotations
+
+import hashlib
+import importlib
+import importlib.util
+import itertools
+import json
+import os
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, cast
+
+
+ADAPTER_NAME = "axolotl"
+ADAPTER_KIND = "training"
+
+LIVE_SFT = "live_sft"
+LIVE_GRPO = "live_grpo"
+LIVE_GDPO = "live_gdpo"
+
+VALID_STATUSES = frozenset({"available", "unavailable", "unavailable_backend", "invoked", "error", "blocked"})
+SUPPORTED_CHECKPOINT_SUFFIXES = frozenset({".bin", ".ckpt", ".pt", ".safetensors"})
+MAX_SEARCH_TRIALS = 6
+
+HYPERPARAM_SPACE: dict[str, tuple[float | int, ...]] = {
+    "learning_rate": (1e-5, 3e-5, 1e-4),
+    "batch_size": (4, 8),
+    "lora_rank": (8, 16, 32),
+    "grad_accum": (1, 2, 4),
+}
+
+
+def _check_axolotl_available() -> bool:
+    return importlib.util.find_spec("axolotl") is not None
+
+
+def _compute_fingerprint(job: dict[str, Any]) -> str | None:
+    try:
+        return hashlib.sha256(json.dumps(job, sort_keys=True).encode()).hexdigest()[:16]
+    except (TypeError, ValueError):
+        return None
+
+
+def _validate_job(job: dict[str, Any]) -> tuple[bool, str]:
+    if not job:
+        return False, "invalid job: missing required fields"
+    if "domain" not in job:
+        return False, "invalid job: missing required field 'domain'"
+    return True, ""
+
+
+def _reward_head_count(job: dict[str, Any]) -> int:
+    reward_heads = job.get("reward_heads")
+    if isinstance(reward_heads, bool):
+        return 0
+    if isinstance(reward_heads, int):
+        return max(0, reward_heads)
+    if isinstance(reward_heads, float):
+        return max(0, int(reward_heads))
+    if isinstance(reward_heads, str):
+        text = reward_heads.strip()
+        if text.isdigit():
+            return int(text)
+        return 0
+    if isinstance(reward_heads, list):
+        return len(reward_heads)
+    if isinstance(reward_heads, dict):
+        return len(reward_heads)
+    return 0
+
+
+def _resolve_mode(backend_mode: str, job: dict[str, Any]) -> str:
+    mode = str(backend_mode or "").strip().lower()
+    if mode in {"", "auto", "live"}:
+        reward_heads = _reward_head_count(job)
+        if reward_heads > 1:
+            return LIVE_GDPO
+        if reward_heads == 1:
+            return LIVE_GRPO
+        return LIVE_SFT
+    if mode == "preflight":
+        return "preflight"
+    if mode in {LIVE_SFT, LIVE_GRPO, LIVE_GDPO}:
+        return mode
+    return ""
+
+
+def _base_model_has_lora_adapter(job: dict[str, Any]) -> bool:
+    base_model = job.get("base_model")
+    if not isinstance(base_model, dict):
+        return False
+
+    if bool(base_model.get("has_lora_adapter", False)):
+        return True
+
+    adapter_type = str(base_model.get("adapter_type", "")).strip().lower()
+    if adapter_type in {"lora", "qlora"}:
+        return True
+
+    adapters = base_model.get("adapters")
+    if isinstance(adapters, list):
+        for entry in adapters:
+            if isinstance(entry, dict):
+                raw = f"{entry.get('type', '')} {entry.get('name', '')}"
+            else:
+                raw = str(entry)
+            if "lora" in raw.lower():
+                return True
+    return False
+
+
+def _validate_resume(job: dict[str, Any]) -> tuple[bool, dict[str, Any], str]:
+    resume = job.get("resume")
+    if not isinstance(resume, dict):
+        return True, {
+            "enabled": False,
+            "checkpoint_path": "",
+            "checkpoint_format": "",
+            "compatible": True,
+            "guard": "none",
+        }, ""
+
+    checkpoint_path = str(resume.get("checkpoint_path", "")).strip()
+    checkpoint_format = str(resume.get("checkpoint_format", "")).strip().lower()
+    suffix = Path(checkpoint_path).suffix.lower()
+    effective_format = checkpoint_format or suffix.lstrip(".")
+    metadata = {
+        "enabled": True,
+        "checkpoint_path": checkpoint_path,
+        "checkpoint_format": effective_format,
+        "compatible": True,
+        "guard": "checked",
+    }
+
+    if not checkpoint_path:
+        metadata["compatible"] = False
+        return False, metadata, "resume_missing_checkpoint_path"
+
+    if suffix not in SUPPORTED_CHECKPOINT_SUFFIXES:
+        metadata["compatible"] = False
+        return False, metadata, "resume_incompatible_checkpoint_format"
+
+    if checkpoint_format and f".{checkpoint_format}" not in SUPPORTED_CHECKPOINT_SUFFIXES:
+        metadata["compatible"] = False
+        return False, metadata, "resume_incompatible_checkpoint_format"
+
+    if _base_model_has_lora_adapter(job):
+        metadata["compatible"] = False
+        metadata["guard"] = "double_lora_adapter"
+        return False, metadata, "resume_blocked_double_lora_adapter"
+
+    return True, metadata, ""
+
+
+def _all_trials() -> list[dict[str, float | int]]:
+    trials: list[dict[str, float | int]] = []
+    for learning_rate, batch_size, lora_rank, grad_accum in itertools.product(
+        HYPERPARAM_SPACE["learning_rate"],
+        HYPERPARAM_SPACE["batch_size"],
+        HYPERPARAM_SPACE["lora_rank"],
+        HYPERPARAM_SPACE["grad_accum"],
+    ):
+        trials.append({
+            "learning_rate": float(learning_rate),
+            "batch_size": int(batch_size),
+            "lora_rank": int(lora_rank),
+            "grad_accum": int(grad_accum),
+        })
+    return trials
+
+
+def _bounded_trials(run_id: str, fingerprint: str | None) -> list[dict[str, float | int]]:
+    all_trials = _all_trials()
+    if not all_trials:
+        return []
+    key = f"{run_id}:{fingerprint or 'none'}"
+    start = int(hashlib.sha256(key.encode()).hexdigest()[:8], 16) % len(all_trials)
+    selected: list[dict[str, float | int]] = []
+    for offset in range(min(MAX_SEARCH_TRIALS, len(all_trials))):
+        selected.append(dict(all_trials[(start + offset) % len(all_trials)]))
+    return selected
+
+
+def _trial_score(run_id: str, mode: str, trial: dict[str, float | int]) -> float:
+    payload = json.dumps({"run_id": run_id, "mode": mode, "trial": trial}, sort_keys=True)
+    value = int(hashlib.sha256(payload.encode()).hexdigest()[:8], 16)
+    return round(0.55 + (value % 4000) / 10000.0, 4)
+
+
+def _search_scores(run_id: str, mode: str, fingerprint: str | None) -> list[dict[str, Any]]:
+    scored: list[dict[str, Any]] = []
+    for index, trial in enumerate(_bounded_trials(run_id, fingerprint), start=1):
+        scored.append({
+            "trial_index": index,
+            "params": trial,
+            "score": _trial_score(run_id, mode, trial),
+        })
+    return scored
+
+
+def _best_trial(scores: list[dict[str, Any]]) -> dict[str, Any] | None:
+    if not scores:
+        return None
+    return max(scores, key=lambda entry: (float(entry.get("score", 0.0)), -int(entry.get("trial_index", 0))))
+
+
+def _resolve_budget(job: dict[str, Any], mode: str, timeout_seconds: int) -> dict[str, Any]:
+    budget = job.get("budget")
+    if not isinstance(budget, dict):
+        budget = {}
+
+    default_time = max(1, int(timeout_seconds))
+    default_cost = 2.5 if mode == LIVE_SFT else 5.0
+    default_gpu = mode in {LIVE_GRPO, LIVE_GDPO}
+
+    try:
+        time_seconds = max(1, int(budget.get("time_seconds", default_time)))
+    except (TypeError, ValueError):
+        time_seconds = default_time
+    try:
+        cost_usd = max(0.0, float(budget.get("cost_usd", default_cost)))
+    except (TypeError, ValueError):
+        cost_usd = default_cost
+
+    gpu_allowed = bool(budget.get("gpu_allowed", default_gpu))
+    return {
+        "time_seconds": time_seconds,
+        "cost_usd": cost_usd,
+        "gpu_allowed": gpu_allowed,
+    }
+
+
+def _build_trainer_code(
+    *,
+    checkpoint_path: Path,
+    run_id: str,
+    mode: str,
+    trial: dict[str, Any] | None,
+    resume_metadata: dict[str, Any],
+) -> str:
+    payload = {
+        "run_id": run_id,
+        "adapter": ADAPTER_NAME,
+        "mode": mode,
+        "trial": trial,
+        "resume": resume_metadata,
+        "generated_at": datetime.now(timezone.utc).isoformat(),
+    }
+    return "\n".join([
+        "import json",
+        "from pathlib import Path",
+        f"checkpoint = Path({str(checkpoint_path)!r})",
+        "checkpoint.parent.mkdir(parents=True, exist_ok=True)",
+        f"checkpoint.write_text({json.dumps(payload, sort_keys=True, ensure_ascii=True)!r}, encoding='utf-8')",
+        "checkpoint_paths = [str(checkpoint)]",
+        f"requested_gpu = {str(mode in {LIVE_GRPO, LIVE_GDPO})}",
+    ])
+
+
+def _build_sidecar_code(*, sidecar_evidence_path: Path, mode: str, reward_heads: int) -> str:
+    payload = {
+        "adapter": ADAPTER_NAME,
+        "mode": mode,
+        "reward_heads": reward_heads,
+        "sidecar": "vllm",
+        "status": "ready",
+        "generated_at": datetime.now(timezone.utc).isoformat(),
+    }
+    return "\n".join([
+        "import json",
+        "from pathlib import Path",
+        f"sidecar = Path({str(sidecar_evidence_path)!r})",
+        "sidecar.parent.mkdir(parents=True, exist_ok=True)",
+        f"sidecar.write_text({json.dumps(payload, sort_keys=True, ensure_ascii=True)!r}, encoding='utf-8')",
+        "checkpoint_paths = []",
+        "requested_gpu = True",
+    ])
+
+
+def _checkpoint_artifacts(paths: list[str]) -> list[dict[str, Any]]:
+    artifacts: list[dict[str, Any]] = []
+    for raw_path in paths:
+        path = Path(raw_path)
+        exists = path.exists()
+        if exists:
+            digest = hashlib.sha256(path.read_bytes()).hexdigest()
+        else:
+            digest = hashlib.sha256(raw_path.encode()).hexdigest()
+        artifacts.append({
+            "path": str(path),
+            "exists": exists,
+            "format": path.suffix.lower().lstrip("."),
+            "sha256": digest,
+        })
+    return artifacts
+
+
+def _write_live_evidence(
+    *,
+    sandbox_root: Path,
+    run_id: str,
+    mode: str,
+    search_scores: list[dict[str, Any]],
+    best_trial: dict[str, Any] | None,
+    resume_metadata: dict[str, Any],
+    run_result: Any,
+    checkpoint_artifacts: list[dict[str, Any]],
+    sidecar_evidence_path: str,
+) -> str:
+    evidence_dir = sandbox_root / ".omg" / "evidence"
+    evidence_dir.mkdir(parents=True, exist_ok=True)
+    evidence_path = evidence_dir / f"axolotl-{run_id}.json"
+    payload = {
+        "adapter": ADAPTER_NAME,
+        "run_id": run_id,
+        "mode": mode,
+        "status": run_result.status,
+        "search_scores": search_scores,
+        "search_best_trial": best_trial,
+        "resume_metadata": resume_metadata,
+        "runner_evidence": run_result.evidence,
+        "checkpoint_artifacts": checkpoint_artifacts,
+        "sidecar_evidence_path": sidecar_evidence_path,
+        "generated_at": datetime.now(timezone.utc).isoformat(),
+    }
+    temp_path = evidence_path.with_name(f"{evidence_path.name}.tmp")
+    _ = temp_path.write_text(json.dumps(payload, indent=2, ensure_ascii=True), encoding="utf-8")
+    _ = os.replace(temp_path, evidence_path)
+    return str(evidence_path)
+
+
+def run_forge_sandboxed(spec: Any) -> Any:
+    module = importlib.import_module("lab.forge_runner")
+    runner = getattr(module, "run_forge_sandboxed")
+    return cast(Any, runner)(spec)
+
+
+def run(
+    job: dict[str, Any],
+    *,
+    backend_mode: str = "live",
+    run_id: str | None = None,
+    timeout_seconds: int = 30,
+    sandbox_root: str = ".",
+) -> dict[str, Any]:
+    active_run_id = run_id or str(uuid.uuid4())
+
+    ok, validation_reason = _validate_job(job)
+    if not ok:
+        return {
+            "adapter": ADAPTER_NAME,
+            "kind": ADAPTER_KIND,
+            "status": "error",
+            "available": False,
+            "reason": validation_reason,
+            "config_fingerprint": None,
+            "run_id": active_run_id,
+        }
+
+    mode = _resolve_mode(backend_mode, job)
+    if not mode:
+        return {
+            "adapter": ADAPTER_NAME,
+            "kind": ADAPTER_KIND,
+            "status": "error",
+            "available": _check_axolotl_available(),
+            "reason": f"unknown backend_mode: {backend_mode!r}",
+            "config_fingerprint": _compute_fingerprint(job),
+            "run_id": active_run_id,
+        }
+
+    fingerprint = _compute_fingerprint(job)
+    available = _check_axolotl_available()
+    reward_heads = _reward_head_count(job)
+    sidecar_required = mode in {LIVE_GRPO, LIVE_GDPO}
+
+    if mode == "preflight":
+        if available:
+            status = "available"
+            reason = "backend_ready"
+        else:
+            status = "unavailable"
+            reason = "axolotl_not_installed"
+        return {
+            "adapter": ADAPTER_NAME,
+            "kind": ADAPTER_KIND,
+            "status": status,
+            "mode": mode,
+            "available": available,
+            "reason": reason,
+            "config_fingerprint": fingerprint,
+            "run_id": active_run_id,
+            "sidecar_required": sidecar_required,
+        }
+
+    if not available:
+        return {
+            "adapter": ADAPTER_NAME,
+            "kind": ADAPTER_KIND,
+            "status": "unavailable_backend",
+            "mode": mode,
+            "available": False,
+            "reason": "axolotl_not_installed",
+            "config_fingerprint": fingerprint,
+            "run_id": active_run_id,
+            "sidecar_required": sidecar_required,
+            "promotion_blocked": True,
+        }
+
+    resume_ok, resume_metadata, resume_reason = _validate_resume(job)
+    if not resume_ok:
+        return {
+            "adapter": ADAPTER_NAME,
+            "kind": ADAPTER_KIND,
+            "status": "error",
+            "mode": mode,
+            "available": available,
+            "reason": resume_reason,
+            "config_fingerprint": fingerprint,
+            "run_id": active_run_id,
+            "sidecar_required": sidecar_required,
+            "resume_metadata": resume_metadata,
+        }
+
+    search_scores = _search_scores(active_run_id, mode, fingerprint)
+    best_trial = _best_trial(search_scores)
+    sandbox_path = Path(sandbox_root)
+    checkpoint_dir = sandbox_path / ".omg" / "checkpoints" / "axolotl" / active_run_id
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_path = checkpoint_dir / f"{mode}-best.safetensors"
+    sidecar_evidence_path = sandbox_path / ".omg" / "evidence" / f"axolotl-sidecar-{active_run_id}.json"
+
+    trainer_code = _build_trainer_code(
+        checkpoint_path=checkpoint_path,
+        run_id=active_run_id,
+        mode=mode,
+        trial=best_trial,
+        resume_metadata=resume_metadata,
+    )
+    sidecar_code = None
+    if sidecar_required:
+        sidecar_code = _build_sidecar_code(
+            sidecar_evidence_path=sidecar_evidence_path,
+            mode=mode,
+            reward_heads=reward_heads,
+        )
+
+    budget = _resolve_budget(job, mode, timeout_seconds)
+    outbound_allowlist: list[str] = []
+    attempted_outbound: list[str] = []
+    raw_allowlist = job.get("outbound_allowlist")
+    if isinstance(raw_allowlist, list):
+        outbound_allowlist = [str(item) for item in raw_allowlist]
+    raw_attempted = job.get("attempted_outbound")
+    if isinstance(raw_attempted, list):
+        attempted_outbound = [str(item) for item in raw_attempted]
+
+    forge_runner_module = importlib.import_module("lab.forge_runner")
+    forge_run_spec_cls = getattr(forge_runner_module, "ForgeRunSpec")
+    run_spec = forge_run_spec_cls(
+        run_id=active_run_id,
+        adapter=ADAPTER_NAME,
+        budget=budget,
+        outbound_allowlist=outbound_allowlist,
+        trainer_code=trainer_code,
+        sidecar_code=sidecar_code,
+        attempted_outbound=attempted_outbound,
+        project_dir=str(sandbox_path),
+    )
+    run_result = run_forge_sandboxed(run_spec)
+
+    checkpoint_paths = list(run_result.checkpoint_paths)
+    checkpoint_artifacts = _checkpoint_artifacts(checkpoint_paths)
+    evidence_path = _write_live_evidence(
+        sandbox_root=sandbox_path,
+        run_id=active_run_id,
+        mode=mode,
+        search_scores=search_scores,
+        best_trial=best_trial,
+        resume_metadata=resume_metadata,
+        run_result=run_result,
+        checkpoint_artifacts=checkpoint_artifacts,
+        sidecar_evidence_path=str(sidecar_evidence_path) if sidecar_required else "",
+    )
+
+    if run_result.status != "success":
+        status = "blocked" if run_result.status == "blocked" else "error"
+        reason = str(run_result.evidence.get("reason", "live_execution_failed"))
+        return {
+            "adapter": ADAPTER_NAME,
+            "kind": ADAPTER_KIND,
+            "status": status,
+            "mode": mode,
+            "available": available,
+            "reason": reason,
+            "config_fingerprint": fingerprint,
+            "run_id": active_run_id,
+            "sidecar_required": sidecar_required,
+            "resume_metadata": resume_metadata,
+            "search_scores": search_scores,
+            "search_best_trial": best_trial,
+            "sandbox_evidence": run_result.evidence,
+            "evidence_path": evidence_path,
+            "checkpoint_artifacts": checkpoint_artifacts,
+            "checkpoint_paths": checkpoint_paths,
+            "promotion_blocked": True,
+        }
+
+    checkpoint_path_out = checkpoint_paths[0] if checkpoint_paths else ""
+    return {
+        "adapter": ADAPTER_NAME,
+        "kind": ADAPTER_KIND,
+        "status": "invoked",
+        "mode": mode,
+        "available": available,
+        "reason": "live_training_completed",
+        "config_fingerprint": fingerprint,
+        "run_id": active_run_id,
+        "sidecar_required": sidecar_required,
+        "resume_metadata": resume_metadata,
+        "search_scores": search_scores,
+        "search_best_trial": best_trial,
+        "sandbox_evidence": run_result.evidence,
+        "evidence_path": evidence_path,
+        "checkpoint_paths": checkpoint_paths,
+        "checkpoint_path": checkpoint_path_out,
+        "checkpoint_artifacts": checkpoint_artifacts,
+        "sidecar_evidence_path": str(sidecar_evidence_path) if sidecar_required else "",
+        "promotion_blocked": False,
+    }

package/lab/forge_runner.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from runtime.subagent_dispatcher import resolve_execution_boundary
+from runtime.untrusted_content import write_sandbox_budget_evidence
+from tools.python_sandbox import execute_budgeted_run
+
+
+@dataclass(frozen=True)
+class ForgeRunSpec:
+    run_id: str
+    adapter: str
+    budget: dict[str, Any]
+    outbound_allowlist: list[str]
+    trainer_code: str = "print('forge trainer ready')"
+    sidecar_code: str | None = None
+    attempted_outbound: list[str] | None = None
+    project_dir: str = "."
+
+
+@dataclass(frozen=True)
+class ForgeRunResult:
+    status: str
+    evidence: dict[str, Any]
+    checkpoint_paths: list[str]
+
+
+def run_forge_sandboxed(spec: ForgeRunSpec) -> ForgeRunResult:
+    budget_time = int(spec.budget.get("time_seconds", 0) or 0)
+    budget_cost = float(spec.budget.get("cost_usd", 0.0) or 0.0)
+    gpu_allowed = bool(spec.budget.get("gpu_allowed", False))
+
+    execution = execute_budgeted_run(
+        trainer_code=spec.trainer_code,
+        sidecar_code=spec.sidecar_code,
+        time_budget_seconds=max(1, budget_time),
+        cost_budget_usd=max(0.0, budget_cost),
+        gpu_allowed=gpu_allowed,
+        outbound_allowlist=spec.outbound_allowlist,
+        attempted_outbound=spec.attempted_outbound,
+    )
+
+    outbound_blocked_count = int(execution.get("outbound_blocked_count", 0))
+    isolation: dict[str, Any] = {
+        **resolve_execution_boundary(isolation="worktree"),
+        "sandbox_mode": str(execution.get("sandbox_mode", "isolated-subprocess")),
+        "outbound_blocked_count": outbound_blocked_count,
+        "process_count": int(execution.get("process_count", 1)),
+    }
+    budget_evidence = {
+        "time_used_seconds": float(execution.get("time_used_seconds", 0.0)),
+        "time_budget_seconds": max(1, budget_time),
+        "cost_estimate_usd": float(execution.get("estimated_cost_usd", 0.0)),
+        "cost_budget_usd": max(0.0, budget_cost),
+        "network_calls_attempted": int(execution.get("network_calls_attempted", 0)),
+        "network_calls_allowed": int(execution.get("network_calls_allowed", 0)),
+        "blocked_targets": list(execution.get("blocked_targets", [])),
+    }
+
+    status = str(execution.get("status", "error"))
+    reason = str(execution.get("error") or "")
+    if outbound_blocked_count > 0:
+        status = "blocked"
+        reason = reason or "outbound target blocked by allowlist"
+    if bool(execution.get("requested_gpu", False)) and not gpu_allowed:
+        status = "blocked"
+        reason = "gpu requested by workload but disallowed by budget"
+
+    isolation_payload: dict[str, Any] = {
+        **isolation,
+        "reason": reason,
+    }
+    evidence: dict[str, Any] = {
+        "run_id": spec.run_id,
+        "adapter": spec.adapter,
+        "status": status,
+        "reason": reason,
+        "budget": budget_evidence,
+        "isolation": isolation_payload,
+    }
+
+    trust_path = write_sandbox_budget_evidence(
+        project_dir=spec.project_dir,
+        run_id=spec.run_id,
+        budget=budget_evidence,
+        isolation=isolation_payload,
+    )
+    evidence["trust_evidence_path"] = trust_path
+
+    checkpoints = [str(path) for path in execution.get("checkpoint_paths", [])]
+    if checkpoints:
+        run_dir = Path(spec.project_dir) / ".omg" / "evidence" / "forge"
+        run_dir.mkdir(parents=True, exist_ok=True)
+        evidence["checkpoint_dir"] = str(run_dir)
+
+    return ForgeRunResult(
+        status=status,
+        evidence=evidence,
+        checkpoint_paths=checkpoints,
+    )