code/self_healing_finetuning/providers/hf_managed.py

from __future__ import annotations

import os
import subprocess
from typing import Any

from budget.limits import check_budget
from n21.config import write_json
from n21.settings import SHFT_WORKSPACE_ROOT
from providers.base import BaseProvider
from n21.types import ValidationResult

HF_CLI_TIMEOUT_SECONDS = int(os.environ.get("SHFT_HF_CLI_TIMEOUT_SECONDS", "120"))


class HFManagedProvider(BaseProvider):
    name = "hf_managed"
    required_env = ["HF_TOKEN"]

    def validate_config(self, config: dict[str, Any], *, live: bool = False) -> ValidationResult:
        result = super().validate_config(config, live=False)
        warnings = list(result["warnings"])
        errors = list(result["errors"])
        if live:
            warnings = []
            if not os.environ.get("HF_TOKEN") and not self._hf_cli_authenticated():
                errors.append("missing Hugging Face credential: set HF_TOKEN or run `hf auth login`")
            if not self._hf_cli_command_available("jobs"):
                errors.append("Hugging Face CLI does not expose `hf jobs`; upgrade huggingface_hub CLI in the active environment")
            if not self._hf_cli_command_available("buckets"):
                errors.append("Hugging Face CLI does not expose `hf buckets`; upgrade huggingface_hub CLI in the active environment")
            base_model_id = str(config.get("jobs", {}).get("base_model_id", "meta-llama/Meta-Llama-3-8B"))
            if not self._hf_model_file_accessible(base_model_id, "config.json"):
                errors.append(
                    f"missing Hugging Face access to gated base model `{base_model_id}`; "
                    f"request access on the model page and ensure HF_TOKEN is authorized"
                )
        return {"ok": not errors, "errors": errors, "warnings": warnings}

    @staticmethod
    def _hf_cli_authenticated() -> bool:
        proc = _run_hf(["hf", "auth", "whoami"])
        return proc.returncode == 0

    @staticmethod
    def _hf_model_file_accessible(repo_id: str, filename: str) -> bool:
        proc = _run_hf(["hf", "download", repo_id, filename, "--dry-run"])
        return proc.returncode == 0

    @staticmethod
    def _hf_cli_command_available(command: str) -> bool:
        proc = _run_hf(["hf", command, "--help"])
        return proc.returncode == 0

    def start_train(self, run_manifest: dict[str, Any]) -> dict[str, Any]:
        if run_manifest["execution"].get("dry_run", True):
            return super().start_train(run_manifest)

        provider_config = run_manifest.get("provider_config", {})
        command = self._build_jobs_command(run_manifest, provider_config)
        estimate = self.estimate_cost({"provider_config": provider_config})
        budget_errors = check_budget(estimate)
        submitted = os.environ.get("SHFT_SUBMIT_HF_JOB", "").lower() == "true"
        status = "blocked_budget" if budget_errors else ("submitted" if submitted else "live_plan_ready")
        proc_result: dict[str, object] | None = None

        if submitted and not budget_errors:
            proc = _run_hf(command)
            proc_result = {
                "returncode": proc.returncode,
                "stdout": proc.stdout,
                "stderr": proc.stderr,
            }
            if proc.returncode != 0:
                status = "submit_failed"

        handle = {
            "run_id": run_manifest["run_id"],
            "provider_job_id": f"hf-live-plan-{run_manifest['run_id']}",
            "status": status,
            "logs_uri": None,
            "live": True,
            "submitted": submitted and status == "submitted",
            "submit_guard": "set SHFT_SUBMIT_HF_JOB=true to submit the planned HF Job",
            "command": command,
            "cost_estimate": estimate,
            "budget_errors": budget_errors,
            "provider_result": proc_result,
        }
        out_dir = SHFT_WORKSPACE_ROOT / "runs" / run_manifest["run_id"] / "provider_plans"
        write_json(out_dir / "hf_train_job_plan.json", handle)
        return handle

    def estimate_cost(self, plan: dict[str, Any]) -> dict[str, Any]:
        config = plan.get("provider_config", {})
        jobs = config.get("jobs", {})
        timeout = str(jobs.get("timeout", "8h"))
        gpu_hours = _parse_hours(timeout)
        flavor = str(jobs.get("flavor", "a100-large"))
        hourly = _rough_hourly_rate(flavor)
        return {
            "provider": self.name,
            "flavor": flavor,
            "gpu_hours": gpu_hours,
            "cost_usd": round(gpu_hours * hourly, 2),
            "hourly_rate_usd_assumption": hourly,
            "estimation_method": "static_preflight_estimate_not_provider_billing",
        }

    @staticmethod
    def _build_jobs_command(run_manifest: dict[str, Any], config: dict[str, Any]) -> list[str]:
        jobs = config.get("jobs", {})
        storage = config.get("storage", {})
        namespace = str(config.get("namespace", "linvest21"))
        image = str(jobs.get("image", "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel"))
        base_model_id = str(jobs.get("base_model_id", "meta-llama/Meta-Llama-3-8B"))
        training = jobs.get("training", {})
        max_steps = str(os.environ.get("SHFT_TRAIN_MAX_STEPS") or training.get("max_steps", 20))
        batch_size = str(os.environ.get("SHFT_TRAIN_BATCH_SIZE") or training.get("per_device_train_batch_size", 1))
        grad_accum = str(os.environ.get("SHFT_TRAIN_GRAD_ACCUM") or training.get("gradient_accumulation_steps", 8))
        learning_rate = str(os.environ.get("SHFT_TRAIN_LEARNING_RATE") or training.get("learning_rate", 0.00008))
        lora_r = str(os.environ.get("SHFT_TRAIN_LORA_R") or training.get("lora_r", 16))
        lora_alpha = str(os.environ.get("SHFT_TRAIN_LORA_ALPHA") or training.get("lora_alpha", 32))
        lora_dropout = str(os.environ.get("SHFT_TRAIN_LORA_DROPOUT") or training.get("lora_dropout", 0.05))
        max_seq_length = str(os.environ.get("SHFT_TRAIN_MAX_SEQ_LENGTH") or training.get("max_seq_length", 2048))
        logging_steps = str(os.environ.get("SHFT_LOGGING_STEPS") or training.get("logging_steps", 5))
        checkpoint_steps = str(os.environ.get("SHFT_CHECKPOINT_STEPS") or training.get("checkpoint_steps", 50))
        eval_steps = str(os.environ.get("SHFT_EVAL_STEPS") or training.get("eval_steps", checkpoint_steps))
        save_total_limit = str(os.environ.get("SHFT_SAVE_TOTAL_LIMIT") or training.get("save_total_limit", 4))
        metric_for_best = str(os.environ.get("SHFT_METRIC_FOR_BEST_MODEL") or training.get("metric_for_best_model", "eval_loss"))
        overfit_tolerance = str(os.environ.get("SHFT_OVERFIT_TOLERANCE") or training.get("overfit_tolerance", 0.10))
        min_steps = str(os.environ.get("SHFT_MIN_PRODUCTION_STEPS") or training.get("min_production_steps", 100))
        min_train_records = str(
            os.environ.get("SHFT_MIN_PRODUCTION_TRAIN_RECORDS") or training.get("min_production_train_records", 100)
        )
        flavor = str(jobs.get("flavor", "a100-large"))
        timeout = str(jobs.get("timeout", "8h"))
        artifact_bucket = str(storage.get("bucket", "linvest21/shft-artifacts"))
        dataset_repo = str(storage.get("dataset_repo", "linvest21/shft-datasets"))
        model_candidate = str(run_manifest["model_candidate"])
        training_start = run_manifest.get("training_start", {})
        finetune_start_policy = str(training_start.get("policy", "bootstrap"))
        start_adapter = str(training_start.get("start_adapter") or "")
        run_id = str(run_manifest["run_id"])
        dataset_stage = _load_dataset_stage(run_id)
        dataset_dir = str(dataset_stage.get("job_dataset_dir") or f"/data/runs/{run_id}")
        split_sha256 = dataset_stage.get("split_sha256") if isinstance(dataset_stage.get("split_sha256"), dict) else {}
        dataset_manifest_sha256 = str(dataset_stage.get("dataset_manifest_sha256") or "")
        remote_code_dir = "/artifacts/code/self_healing_finetuning"
        remote_output_dir = f"/artifacts/runs/{run_id}"
        entrypoint_script = f"{remote_code_dir}/training/hf_job_entrypoint.py"
        command = [
            "hf",
            "jobs",
            "run",
            "--detach",
            "--namespace",
            namespace,
            "--flavor",
            flavor,
            "--timeout",
            timeout,
            "--secrets",
            "HF_TOKEN",
            "--env",
            f"SHFT_RUN_ID={run_id}",
            "--env",
            f"SHFT_MODEL_CANDIDATE={model_candidate}",
            "--env",
            f"SHFT_FINETUNE_START_POLICY={finetune_start_policy}",
            "--env",
            f"PYTHONPATH={remote_code_dir}",
            "--env",
            f"PYTORCH_CUDA_ALLOC_CONF={os.environ.get('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')}",
            "--env",
            f"SHFT_MIN_PRODUCTION_STEPS={min_steps}",
            "--env",
            f"SHFT_MIN_PRODUCTION_TRAIN_RECORDS={min_train_records}",
            "--env",
            f"SHFT_EXPECTED_DATASET_DIR={dataset_dir}",
            "--volume",
            f"hf://{base_model_id}:/models/base:ro",
            "--volume",
            f"hf://datasets/{dataset_repo}:/data:ro",
            "--volume",
            f"hf://buckets/{artifact_bucket}:/artifacts",
            image,
            "python",
            entrypoint_script,
            "--run-id",
            run_id,
            "--model-candidate",
            model_candidate,
            "--finetune-start-policy",
            finetune_start_policy,
            "--dataset-dir",
            dataset_dir,
            "--output-dir",
            remote_output_dir,
            "--base-model-id",
            base_model_id,
            "--max-steps",
            max_steps,
            "--per-device-train-batch-size",
            batch_size,
            "--gradient-accumulation-steps",
            grad_accum,
            "--learning-rate",
            learning_rate,
            "--lora-r",
            lora_r,
            "--lora-alpha",
            lora_alpha,
            "--lora-dropout",
            lora_dropout,
            "--max-seq-length",
            max_seq_length,
            "--logging-steps",
            logging_steps,
            "--checkpoint-steps",
            checkpoint_steps,
            "--eval-steps",
            eval_steps,
            "--save-total-limit",
            save_total_limit,
            "--metric-for-best-model",
            metric_for_best,
            "--overfit-tolerance",
            overfit_tolerance,
        ]
        optional_hash_args = [
            ("--expected-dataset-manifest-sha256", dataset_manifest_sha256),
            ("--expected-train-sha256", str(split_sha256.get("train") or "")),
            ("--expected-valid-sha256", str(split_sha256.get("valid") or "")),
            ("--expected-test-sha256", str(split_sha256.get("test") or "")),
        ]
        for flag, value in optional_hash_args:
            if value:
                command.extend([flag, value])
        if start_adapter:
            marker = command.index("--finetune-start-policy")
            command[marker:marker] = ["--start-adapter", start_adapter]
        return command


def _parse_hours(timeout: str) -> float:
    raw = timeout.strip().lower()
    if raw.endswith("h"):
        return float(raw[:-1])
    if raw.endswith("m"):
        return float(raw[:-1]) / 60.0
    if raw.endswith("d"):
        return float(raw[:-1]) * 24.0
    if raw.endswith("s"):
        return float(raw[:-1]) / 3600.0
    return float(raw) / 3600.0


def _rough_hourly_rate(flavor: str) -> float:
    if "h200" in flavor:
        return 8.0
    if "a100" in flavor:
        return 4.0
    if "l40" in flavor:
        return 2.5
    if "a10g" in flavor:
        return 1.5
    if "l4" in flavor:
        return 1.0
    if "t4" in flavor:
        return 0.6
    return 0.0


def _load_dataset_stage(run_id: str) -> dict[str, Any]:
    path = SHFT_WORKSPACE_ROOT / "runs" / run_id / "provider_plans" / "hf_dataset_stage_result.json"
    if not path.exists():
        return {}
    try:
        import json

        return json.loads(path.read_text(encoding="utf-8-sig"))
    except (OSError, ValueError):
        return {}


def _run_hf(command: list[str]) -> subprocess.CompletedProcess[str]:
    env = os.environ.copy()
    env.setdefault("PYTHONUTF8", "1")
    env.setdefault("PYTHONIOENCODING", "utf-8")
    env.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
    try:
        return subprocess.run(
            command,
            text=True,
            capture_output=True,
            check=False,
            env=env,
            encoding="utf-8",
            errors="replace",
            timeout=HF_CLI_TIMEOUT_SECONDS,
        )
    except subprocess.TimeoutExpired as exc:
        return subprocess.CompletedProcess(
            command,
            124,
            stdout=exc.stdout or "",
            stderr=f"HF CLI command timed out after {HF_CLI_TIMEOUT_SECONDS}s: {' '.join(command)}",
        )