ltx/runtime_utils.py

from __future__ import annotations

import copy
import hashlib
import importlib.metadata as importlib_metadata
import json
import logging
import platform
import resource
import shutil
import subprocess
import sys
import time
import uuid
from dataclasses import dataclass
from pathlib import Path

import torch


def open_request_logger(log_path: Path, request_id: str, formatter: logging.Formatter):
    """Open one append-only FileHandler for one UUID request; caller must close it."""
    logger = logging.getLogger(f"ltx25.request.{request_id}")
    logger.setLevel(logging.DEBUG)
    logger.propagate = True
    handler = logging.FileHandler(log_path, mode="a", encoding="utf-8", delay=False)
    handler.setLevel(logging.DEBUG)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    return logger, handler


def close_request_logger(logger, handler) -> None:
    if logger is None or handler is None:
        return
    try:
        handler.flush()
    finally:
        logger.removeHandler(handler)
        handler.close()


def is_uuid_hex(value: str) -> bool:
    value = str(value or "").strip().lower()
    return len(value) == 32 and all(ch in "0123456789abcdef" for ch in value)


def hf_repo_url(repo_id: str, repo_type: str = "model") -> str:
    """Return the canonical Hugging Face Hub page for one repository."""
    repo_id = str(repo_id or "").strip().strip("/")
    repo_type = str(repo_type or "model").strip().lower()
    prefix = {"model": "", "dataset": "datasets/", "space": "spaces/"}.get(repo_type)
    if not repo_id or prefix is None:
        return ""
    return f"https://huggingface.co/{prefix}{repo_id}"


def hf_repo_markdown_link(repo_id: str, repo_type: str = "model") -> str:
    """Render a repo ID as a beginner-friendly Markdown link when possible."""
    repo_id = str(repo_id or "").strip()
    url = hf_repo_url(repo_id, repo_type)
    return f"[{repo_id}]({url})" if url else f"`{repo_id}`"


def create_session_id() -> str:
    return uuid.uuid4().hex


def normalize_session_id(value) -> str:
    value = str(value or "").strip().lower()
    return value if is_uuid_hex(value) else create_session_id()


def cleanup_session_results(worker_root: Path, session_id) -> None:
    value = str(session_id or "").strip().lower()
    if not is_uuid_hex(value):
        return
    root = worker_root / value
    try:
        shutil.rmtree(root, ignore_errors=True)
    except Exception:
        pass


@dataclass(frozen=True)
class RequestPaths:
    session_id: str
    request_id: str
    root: Path
    video: Path
    diagnostics: Path
    run_info: Path
    log: Path
    probe: Path


def create_request_paths(worker_root: Path, session_id) -> RequestPaths:
    session_id = normalize_session_id(session_id)
    request_id = uuid.uuid4().hex
    root = worker_root / session_id / request_id
    root.mkdir(parents=True, mode=0o700, exist_ok=False)
    return RequestPaths(
        session_id=session_id,
        request_id=request_id,
        root=root,
        video=root / f"ltx25_{request_id}.mp4",
        diagnostics=root / f"ltx25_request_{request_id}.json",
        run_info=root / f"ltx25_run_{request_id}.json",
        log=root / f"ltx25_request_{request_id}.log",
        probe=root / f"ltx25_probe_{request_id}.zip",
    )



def monotonic_elapsed(started: float) -> float:
    """Return monotonic elapsed seconds for Probe/runtime phase timing."""
    return time.monotonic() - float(started)


def process_rss_kib() -> int:
    """Return process max RSS using the platform value already used by runtime diagnostics."""
    return int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)


def reset_gpu_peak_memory() -> None:
    """Reset CUDA peak counters when CUDA is available; safe on CPU-only runtimes."""
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()

def disk_state() -> dict:
    d = shutil.disk_usage(Path.home())
    return {"total": d.total, "used": d.used, "free": d.free}


def gpu_state() -> dict:
    if not torch.cuda.is_available():
        return {"cuda_available": False}
    free, total = torch.cuda.mem_get_info()
    return {
        "cuda_available": True,
        "device": torch.cuda.get_device_name(0),
        "compute_capability": list(torch.cuda.get_device_capability(0)),
        "torch_cuda": torch.version.cuda,
        "free_bytes": int(free),
        "total_bytes": int(total),
        "allocated_bytes": int(torch.cuda.memory_allocated()),
        "max_allocated_bytes": int(torch.cuda.max_memory_allocated()),
        "reserved_bytes": int(torch.cuda.memory_reserved()),
        "max_reserved_bytes": int(torch.cuda.max_memory_reserved()),
    }



def _nvidia_driver_version() -> str | None:
    if not torch.cuda.is_available():
        return None
    try:
        proc = subprocess.run(
            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            timeout=2.0,
            check=False,
        )
        if proc.returncode == 0:
            first = (proc.stdout or "").strip().splitlines()
            if first:
                return first[0].strip() or None
    except Exception:
        pass
    return None


def runtime_environment_identity() -> dict:
    """Return process-level runtime identity.

    On ZeroGPU this may run during startup CUDA emulation, so its GPU fields are
    preload context rather than authoritative execution-GPU evidence.
    """
    identity = {
        "python": {
            "version": platform.python_version(),
            "implementation": platform.python_implementation(),
            "compiler": platform.python_compiler(),
        },
        "platform": {
            "system": platform.system(),
            "release": platform.release(),
            "version": platform.version(),
            "machine": platform.machine(),
            "platform": platform.platform(),
        },
        "process": {
            "python_executable": sys.executable,
        },
    }
    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        identity["cuda"] = {
            "torch_cuda": torch.version.cuda,
            "driver_version": _nvidia_driver_version(),
        }
        identity["gpu"] = {
            "device": torch.cuda.get_device_name(0),
            "compute_capability": list(torch.cuda.get_device_capability(0)),
            "total_bytes": int(total),
            "free_bytes_at_capture": int(free),
        }
    else:
        identity["cuda"] = {"torch_cuda": torch.version.cuda, "driver_version": None}
        identity["gpu"] = {"cuda_available": False}
    return identity


def execution_environment_identity(preload_environment: dict | None = None) -> dict:
    """Return runtime identity with GPU fields captured inside the active GPU callback.

    ZeroGPU exposes CUDA emulation during module startup and a real allocated GPU
    only inside ``@spaces.GPU``.  Preserve stable process/platform fields from the
    preload capture, but replace CUDA/GPU identity with the callback-visible device.
    """
    identity = copy.deepcopy(preload_environment or runtime_environment_identity())
    identity["capture_scope"] = "gpu_callback_execution"
    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        cuda_record = dict(identity.get("cuda") or {})
        cuda_record["torch_cuda"] = torch.version.cuda
        identity["cuda"] = cuda_record
        identity["gpu"] = {
            "device": torch.cuda.get_device_name(0),
            "compute_capability": list(torch.cuda.get_device_capability(0)),
            "total_bytes": int(total),
            "free_bytes_at_capture": int(free),
        }
    else:
        identity["gpu"] = {"cuda_available": False}
    return identity


def resolved_revision_from_hub_path(path: str | Path) -> str | None:
    """Return the snapshot commit encoded in a Hugging Face Hub cache path."""
    parts = list(Path(path).parts)
    if "snapshots" in parts:
        snap_idx = parts.index("snapshots")
        if snap_idx + 1 < len(parts):
            return parts[snap_idx + 1]
    return None


def package_identity(name: str) -> dict:
    out = {"name": name}
    try:
        dist = importlib_metadata.distribution(name)
        out["version"] = dist.version
        direct = dist.read_text("direct_url.json")
        if direct:
            payload = json.loads(direct)
            out["direct_url"] = payload.get("url")
            vcs = payload.get("vcs_info") or {}
            if vcs.get("commit_id"):
                out["vcs_commit"] = vcs["commit_id"]
            if vcs.get("requested_revision"):
                out["requested_revision"] = vcs["requested_revision"]
    except Exception as exc:
        out["identity_error"] = f"{type(exc).__name__}: {exc}"
    return out


def sha256_file(path: str | Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()