config/clearml_enrich.py

"""Extra ClearML polish: env tags, config snapshot, output model metadata."""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path


def project_root() -> Path:
    return Path(__file__).resolve().parent.parent


def active_config_path() -> Path:
    env = os.environ.get("FOCUSGUARD_CONFIG")
    if env:
        return Path(env).expanduser()
    return Path(__file__).resolve().parent / "default.yaml"


def enrich_task(task, *, role: str) -> None:
    """Tags for filtering in the UI (Python, OS, torch device, git revision)."""
    tags = [
        role,
        f"py{sys.version_info.major}{sys.version_info.minor}",
        sys.platform.replace(" ", "_"),
    ]
    try:
        import torch

        ver = torch.__version__.split("+")[0].replace(".", "_")
        tags.append(f"torch_{ver}")
        tags.append("cuda" if torch.cuda.is_available() else "cpu")
    except ImportError:
        tags.append("no_torch")
    rev = _git_short_rev()
    if rev:
        tags.append(f"git_{rev}")
    task.add_tags(tags)


def _git_short_rev() -> str | None:
    root = project_root()
    try:
        p = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            cwd=str(root),
            capture_output=True,
            text=True,
            timeout=6,
            check=False,
        )
        if p.returncode == 0 and p.stdout:
            return p.stdout.strip()
    except (OSError, subprocess.TimeoutExpired):
        pass
    return None


def upload_repro_artifacts(task) -> None:
    """Pin the exact YAML + requirements file used for this run."""
    cfg = active_config_path()
    if cfg.is_file():
        task.upload_artifact(name="config_yaml", artifact_object=str(cfg))
    req = project_root() / "requirements.txt"
    if req.is_file():
        task.upload_artifact(name="requirements_txt", artifact_object=str(req))


def attach_output_metrics(output_model, metrics: dict[str, float | str]) -> None:
    """Surface headline metrics on the registered model card."""
    for k, v in metrics.items():
        key = str(k).replace("/", "_")
        try:
            output_model.set_metadata(key, str(v))
        except Exception:
            pass


def task_done_summary(task, summary: str) -> None:
    setter = getattr(task, "set_comment", None)
    if callable(setter):
        try:
            setter(summary)
        except Exception:
            pass