syvain-training-utils 0.0.170__tar.gz

syvain_training_utils-0.0.170/PKG-INFO

@@ -0,0 +1,91 @@
+Metadata-Version: 2.3
+Name: syvain-training-utils
+Version: 0.0.170
+Summary: Shared runtime, diagnostics, and checkpoint utilities for Syvain training runs
+Requires-Dist: obstore>=0.11.0,<0.12.0
+Requires-Dist: pydantic>=2.13.4
+Requires-Dist: torch>=2.13.0
+Requires-Dist: pytest>=8.0.0 ; extra == 'dev'
+Requires-Dist: ruff>=0.15.12 ; extra == 'dev'
+Requires-Dist: ty>=0.0.34 ; extra == 'dev'
+Requires-Python: >=3.14, <3.15
+Provides-Extra: dev
+Description-Content-Type: text/markdown
+
+# syvain-training-utils
+
+Internal [Syvain](https://syvain.com/) helpers for small, explicit ML training
+runs. No secret sauce here, just shared runtime, device-diagnostic, and
+checkpoint patterns.
+
+## Install
+
+```bash
+uv add syvain-training-utils
+```
+
+## Runtime setup
+
+```python
+from syvain_training_utils import (
+    generate_run_id,
+    require_torch_compile_toolchain,
+    select_device,
+)
+
+run_id = generate_run_id()
+device = select_device()
+require_torch_compile_toolchain()
+```
+
+## Device smoke test
+
+```python
+import json
+
+from syvain_training_utils import run_device_smoke_test
+
+report = run_device_smoke_test(require_cuda=True)
+print(json.dumps({"smoke": report}, indent=2, sort_keys=True))
+```
+
+## Checkpoint a training run
+
+```python
+from syvain_training_utils import (
+    TrainingLoopState,
+    load_training_checkpoint_if_available,
+    save_model_checkpoint,
+)
+
+loop_state = TrainingLoopState(
+    global_step=global_step,
+    curriculum_stage=curriculum_stage,
+    curriculum_step=curriculum_step,
+)
+
+save_model_checkpoint(
+    object_store=object_store,
+    base_path=base_model_path,
+    experiment_slug=experiment_slug,
+    run_id=run_id,
+    model=model,
+    optimizer=optimizer,
+    scheduler=scheduler,
+    loop_state=loop_state,
+    checkpoint_label=f"step-{global_step:012d}",
+)
+
+resume = load_training_checkpoint_if_available(
+    object_store=object_store,
+    base_path=base_model_path,
+    model=model,
+    optimizer=optimizer,
+    scheduler=scheduler,
+    device=device,
+)
+```
+
+The manifest remains a small pointer to the current checkpoint. Each checkpoint
+contains the model, optimizer, optional scheduler, and PyTorch RNG state needed
+by the common single-file training-run contract.

syvain_training_utils-0.0.170/README.md

@@ -0,0 +1,77 @@
+# syvain-training-utils
+
+Internal [Syvain](https://syvain.com/) helpers for small, explicit ML training
+runs. No secret sauce here, just shared runtime, device-diagnostic, and
+checkpoint patterns.
+
+## Install
+
+```bash
+uv add syvain-training-utils
+```
+
+## Runtime setup
+
+```python
+from syvain_training_utils import (
+    generate_run_id,
+    require_torch_compile_toolchain,
+    select_device,
+)
+
+run_id = generate_run_id()
+device = select_device()
+require_torch_compile_toolchain()
+```
+
+## Device smoke test
+
+```python
+import json
+
+from syvain_training_utils import run_device_smoke_test
+
+report = run_device_smoke_test(require_cuda=True)
+print(json.dumps({"smoke": report}, indent=2, sort_keys=True))
+```
+
+## Checkpoint a training run
+
+```python
+from syvain_training_utils import (
+    TrainingLoopState,
+    load_training_checkpoint_if_available,
+    save_model_checkpoint,
+)
+
+loop_state = TrainingLoopState(
+    global_step=global_step,
+    curriculum_stage=curriculum_stage,
+    curriculum_step=curriculum_step,
+)
+
+save_model_checkpoint(
+    object_store=object_store,
+    base_path=base_model_path,
+    experiment_slug=experiment_slug,
+    run_id=run_id,
+    model=model,
+    optimizer=optimizer,
+    scheduler=scheduler,
+    loop_state=loop_state,
+    checkpoint_label=f"step-{global_step:012d}",
+)
+
+resume = load_training_checkpoint_if_available(
+    object_store=object_store,
+    base_path=base_model_path,
+    model=model,
+    optimizer=optimizer,
+    scheduler=scheduler,
+    device=device,
+)
+```
+
+The manifest remains a small pointer to the current checkpoint. Each checkpoint
+contains the model, optimizer, optional scheduler, and PyTorch RNG state needed
+by the common single-file training-run contract.

syvain_training_utils-0.0.170/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "syvain-training-utils"
+version = "0.0.170"
+description = "Shared runtime, diagnostics, and checkpoint utilities for Syvain training runs"
+readme = "README.md"
+requires-python = ">=3.14,<3.15"
+dependencies = [
+    "obstore>=0.11.0,<0.12.0",
+    "pydantic>=2.13.4",
+    "torch>=2.13.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "ruff>=0.15.12",
+    "ty>=0.0.34",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[build-system]
+requires = ["uv_build>=0.11.9,<0.12"]
+build-backend = "uv_build"

syvain_training_utils-0.0.170/pyproject.toml.orig

@@ -0,0 +1,21 @@
+[project]
+name = "syvain-training-utils"
+version = "0.0.170"
+description = "Shared runtime, diagnostics, and checkpoint utilities for Syvain training runs"
+readme = "README.md"
+requires-python = ">=3.14,<3.15"
+dependencies = [
+    "obstore>=0.11.0,<0.12.0",
+    "pydantic>=2.13.4",
+    "torch>=2.13.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8.0.0", "ruff>=0.15.12", "ty>=0.0.34"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[build-system]
+requires = ["uv_build>=0.11.9,<0.12"]
+build-backend = "uv_build"

syvain_training_utils-0.0.170/src/syvain_training_utils/__init__.py

@@ -0,0 +1,35 @@
+from syvain_training_utils.checkpoint import (
+    CheckpointManifest,
+    ResumeCheckpoint,
+    TrainingLoopState,
+    load_training_checkpoint_if_available,
+    save_model_checkpoint,
+)
+from syvain_training_utils.diagnostics import (
+    DeviceSmokeError,
+    collect_device_diagnostics,
+    nvidia_smi_diagnostics,
+    run_device_smoke_test,
+)
+from syvain_training_utils.runtime import (
+    generate_run_id,
+    require_torch_compile_toolchain,
+    runtime_c_compiler,
+    select_device,
+)
+
+__all__ = [
+    "CheckpointManifest",
+    "DeviceSmokeError",
+    "ResumeCheckpoint",
+    "TrainingLoopState",
+    "collect_device_diagnostics",
+    "generate_run_id",
+    "load_training_checkpoint_if_available",
+    "nvidia_smi_diagnostics",
+    "require_torch_compile_toolchain",
+    "run_device_smoke_test",
+    "runtime_c_compiler",
+    "save_model_checkpoint",
+    "select_device",
+]

syvain_training_utils-0.0.170/src/syvain_training_utils/checkpoint.py

@@ -0,0 +1,301 @@
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Any, Literal
+
+import torch
+from obstore.store import S3Store
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from torch.torch_version import TorchVersion
+
+
+class ResumeCheckpoint(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+
+    experiment_slug: str = Field(min_length=1)
+    run_id: str = Field(min_length=1)
+    checkpoint_label: str = Field(min_length=1)
+    global_step: int = Field(ge=0)
+    curriculum_stage: int = Field(ge=0)
+    curriculum_step: int = Field(ge=0)
+    torch_version: str = Field(min_length=1)
+
+
+class CheckpointManifest(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+
+    version: Literal[1]
+    checkpoint_key: str = Field(min_length=1)
+
+
+class _CheckpointPayload(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+
+    version: Literal[1]
+    checkpoint: ResumeCheckpoint
+    checkpoint_data: dict[str, Any]
+
+
+@dataclass
+class TrainingLoopState:
+    global_step: int
+    curriculum_stage: int
+    curriculum_step: int
+
+
+def _cpu_checkpoint_value(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, dict):
+        return {key: _cpu_checkpoint_value(inner) for key, inner in value.items()}
+    if isinstance(value, list):
+        return [_cpu_checkpoint_value(inner) for inner in value]
+    if isinstance(value, tuple):
+        return tuple(_cpu_checkpoint_value(inner) for inner in value)
+    return value
+
+
+def _checkpoint_module(model: torch.nn.Module) -> torch.nn.Module:
+    original_module = getattr(model, "_orig_mod", None)
+    if isinstance(original_module, torch.nn.Module):
+        return original_module
+    return model
+
+
+def _capture_torch_rng_state() -> dict[str, Any]:
+    state: dict[str, Any] = {
+        "version": 1,
+        "cpu": _cpu_checkpoint_value(torch.random.get_rng_state()),
+    }
+
+    if torch.cuda.is_available():
+        state["cuda"] = _cpu_checkpoint_value(torch.cuda.get_rng_state_all())
+
+    xpu_backend = getattr(torch, "xpu", None)
+    if xpu_backend is not None and bool(xpu_backend.is_available()):
+        state["xpu"] = _cpu_checkpoint_value(xpu_backend.get_rng_state_all())
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    torch_mps = getattr(torch, "mps", None)
+    if (
+        mps_backend is not None
+        and torch_mps is not None
+        and bool(mps_backend.is_available())
+    ):
+        state["mps"] = _cpu_checkpoint_value(torch_mps.get_rng_state())
+
+    return state
+
+
+def _validate_accelerator_rng_states(value: Any, *, backend: str) -> list[torch.Tensor]:
+    if not isinstance(value, list):
+        raise TypeError(f"Checkpoint {backend} RNG state must be a list")
+
+    states: list[torch.Tensor] = []
+    for rng_state in value:
+        if not isinstance(rng_state, torch.Tensor):
+            raise TypeError(f"Checkpoint {backend} RNG states must be tensors")
+        states.append(rng_state.detach().cpu())
+    return states
+
+
+def _restore_torch_rng_state(state: Mapping[str, Any], *, device: torch.device) -> None:
+    if state.get("version") != 1:
+        raise RuntimeError("Checkpoint torch RNG state has unsupported version")
+
+    cpu_state = state.get("cpu")
+    if not isinstance(cpu_state, torch.Tensor):
+        raise TypeError("Checkpoint torch RNG state has no CPU state")
+    torch.random.set_rng_state(cpu_state.detach().cpu())
+
+    if device.type == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA resume requested but CUDA is not available")
+        if "cuda" not in state:
+            raise RuntimeError(
+                "CUDA resume requested but checkpoint has no CUDA RNG state"
+            )
+        cuda_states = _validate_accelerator_rng_states(state["cuda"], backend="CUDA")
+        if len(cuda_states) != torch.cuda.device_count():
+            raise RuntimeError(
+                "Checkpoint CUDA RNG state device count does not match runtime"
+            )
+        torch.cuda.set_rng_state_all(cuda_states)
+
+    xpu_backend = getattr(torch, "xpu", None)
+    xpu_available = xpu_backend is not None and bool(xpu_backend.is_available())
+    if device.type == "xpu":
+        if not xpu_available:
+            raise RuntimeError("XPU resume requested but XPU is not available")
+        if "xpu" not in state:
+            raise RuntimeError(
+                "XPU resume requested but checkpoint has no XPU RNG state"
+            )
+        xpu_states = _validate_accelerator_rng_states(state["xpu"], backend="XPU")
+        if len(xpu_states) != xpu_backend.device_count():
+            raise RuntimeError(
+                "Checkpoint XPU RNG state device count does not match runtime"
+            )
+        xpu_backend.set_rng_state_all(xpu_states)
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    torch_mps = getattr(torch, "mps", None)
+    mps_available = (
+        mps_backend is not None
+        and torch_mps is not None
+        and bool(mps_backend.is_available())
+    )
+    if device.type == "mps":
+        if not mps_available:
+            raise RuntimeError("MPS resume requested but MPS is not available")
+        if "mps" not in state:
+            raise RuntimeError(
+                "MPS resume requested but checkpoint has no MPS RNG state"
+            )
+        mps_state = state["mps"]
+        if not isinstance(mps_state, torch.Tensor):
+            raise RuntimeError("Checkpoint MPS RNG state must be a tensor")
+        torch_mps.set_rng_state(mps_state.detach().cpu())
+
+
+_CHECKPOINT_MANIFEST_FILENAME = "checkpoint-manifest-v1.json"
+
+
+def save_model_checkpoint(
+    *,
+    object_store: S3Store,
+    base_path: str,
+    experiment_slug: str,
+    run_id: str,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    scheduler: torch.optim.lr_scheduler.LRScheduler | None,
+    loop_state: TrainingLoopState,
+    checkpoint_label: str,
+) -> str:
+    """Save model/training state and update the base-path manifest pointer."""
+    checkpoint_filename = f"{checkpoint_label}.pt"
+    checkpoint_key = (
+        f"{base_path.lstrip('/').rstrip('/')}/runs/{run_id}/checkpoints/"
+        f"{checkpoint_filename}"
+    )
+
+    resume_checkpoint = ResumeCheckpoint(
+        experiment_slug=experiment_slug,
+        run_id=run_id,
+        checkpoint_label=checkpoint_label,
+        global_step=loop_state.global_step,
+        curriculum_stage=loop_state.curriculum_stage,
+        curriculum_step=loop_state.curriculum_step,
+        torch_version=str(torch.__version__),
+    )
+    checkpoint = {
+        "version": 1,
+        "checkpoint": resume_checkpoint.model_dump(mode="json"),
+        "checkpoint_data": {
+            "model_state_dict": _cpu_checkpoint_value(
+                _checkpoint_module(model).state_dict()
+            ),
+            "optimizer_state_dict": _cpu_checkpoint_value(optimizer.state_dict()),
+            "scheduler_state_dict": (
+                None
+                if scheduler is None
+                else _cpu_checkpoint_value(scheduler.state_dict())
+            ),
+            "torch_rng_state": _capture_torch_rng_state(),
+        },
+    }
+
+    checkpoint_body = BytesIO()
+    torch.save(checkpoint, checkpoint_body)
+    object_store.put(
+        checkpoint_key,
+        checkpoint_body.getvalue(),
+        attributes={"Content-Type": "application/octet-stream"},
+    )
+
+    manifest_key = (
+        f"{base_path.lstrip('/').rstrip('/')}/{_CHECKPOINT_MANIFEST_FILENAME}"
+    )
+    manifest = CheckpointManifest(version=1, checkpoint_key=checkpoint_key)
+    manifest_body = (manifest.model_dump_json(indent=2) + "\n").encode("utf-8")
+    object_store.put(
+        manifest_key,
+        manifest_body,
+        attributes={"Content-Type": "application/json"},
+    )
+    return checkpoint_key
+
+
+def load_training_checkpoint_if_available(
+    *,
+    object_store: S3Store,
+    base_path: str,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    scheduler: torch.optim.lr_scheduler.LRScheduler | None,
+    device: torch.device,
+) -> ResumeCheckpoint | None:
+    """Restore the manifest-selected checkpoint, or return None when absent."""
+    manifest_key = (
+        f"{base_path.lstrip('/').rstrip('/')}/{_CHECKPOINT_MANIFEST_FILENAME}"
+    )
+    try:
+        manifest_body = bytes(object_store.get(manifest_key).bytes())
+        raw_manifest = json.loads(manifest_body.decode("utf-8"))
+        manifest = CheckpointManifest.model_validate(raw_manifest)
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(
+            f"Checkpoint manifest is not valid JSON: {manifest_key}"
+        ) from exc
+    except ValidationError as exc:
+        raise RuntimeError(
+            f"Checkpoint manifest has invalid shape: {manifest_key}"
+        ) from exc
+    except FileNotFoundError:
+        return None
+
+    checkpoint_key = manifest.checkpoint_key
+    checkpoint_body = BytesIO(bytes(object_store.get(checkpoint_key).bytes()))
+    with torch.serialization.safe_globals([TorchVersion]):
+        checkpoint_object = torch.load(
+            checkpoint_body,
+            map_location=device,
+            weights_only=True,
+        )
+    if not isinstance(checkpoint_object, dict):
+        raise TypeError(f"Checkpoint {checkpoint_key} must contain a mapping")
+
+    try:
+        checkpoint_payload = _CheckpointPayload.model_validate(checkpoint_object)
+    except ValidationError as exc:
+        raise RuntimeError(f"Checkpoint {checkpoint_key} has invalid payload") from exc
+
+    checkpoint_data = checkpoint_payload.checkpoint_data
+    model_state = checkpoint_data.get("model_state_dict")
+    optimizer_state = checkpoint_data.get("optimizer_state_dict")
+    scheduler_state = checkpoint_data.get("scheduler_state_dict")
+    torch_rng_state = checkpoint_data.get("torch_rng_state")
+    if not isinstance(model_state, Mapping):
+        raise TypeError(f"Checkpoint {checkpoint_key} has no model state")
+    if not isinstance(optimizer_state, dict):
+        raise TypeError(f"Checkpoint {checkpoint_key} has no optimizer state")
+    if not isinstance(torch_rng_state, Mapping):
+        raise TypeError(f"Checkpoint {checkpoint_key} has no torch RNG state")
+
+    _checkpoint_module(model).load_state_dict(model_state)
+    optimizer.load_state_dict(optimizer_state)
+    if scheduler is None:
+        if scheduler_state is not None:
+            raise RuntimeError(
+                f"Checkpoint {checkpoint_key} has scheduler state "
+                "but current config disables the scheduler"
+            )
+    else:
+        if not isinstance(scheduler_state, dict):
+            raise RuntimeError(f"Checkpoint {checkpoint_key} has no scheduler state")
+        scheduler.load_state_dict(scheduler_state)
+
+    _restore_torch_rng_state(dict(torch_rng_state), device=device)
+    return checkpoint_payload.checkpoint

syvain_training_utils-0.0.170/src/syvain_training_utils/diagnostics.py

@@ -0,0 +1,197 @@
+import os
+import subprocess
+import time
+from typing import Any
+
+import torch
+
+from syvain_training_utils.runtime import select_device
+
+
+class DeviceSmokeError(RuntimeError):
+    """A device smoke failure with the diagnostics collected before the failure."""
+
+    def __init__(self, message: str, diagnostics: dict[str, Any]) -> None:
+        super().__init__(message)
+        self.diagnostics = diagnostics
+
+
+def nvidia_smi_diagnostics() -> dict[str, object]:
+    """Collect a small, non-failing nvidia-smi report."""
+    try:
+        result = subprocess.run(
+            (
+                "nvidia-smi",
+                "--query-gpu=index,name,driver_version,memory.total",
+                "--format=csv,noheader",
+            ),
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            check=False,
+        )
+    except FileNotFoundError:
+        return {
+            "available": False,
+            "output": "nvidia-smi executable not found",
+        }
+    return {
+        "available": result.returncode == 0,
+        "returncode": result.returncode,
+        "output": result.stdout.strip(),
+    }
+
+
+def collect_device_diagnostics() -> dict[str, Any]:
+    """Describe visible CUDA, XPU, and MPS devices without selecting one."""
+    cuda_available = torch.cuda.is_available()
+    diagnostics: dict[str, Any] = {
+        "torch_version": torch.__version__,
+        "cuda_version": torch.version.cuda,
+        "environment": {
+            "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES"),
+            "NVIDIA_VISIBLE_DEVICES": os.environ.get("NVIDIA_VISIBLE_DEVICES"),
+            "SKYPILOT_NUM_GPUS_PER_NODE": os.environ.get("SKYPILOT_NUM_GPUS_PER_NODE"),
+        },
+        "nvidia_smi": nvidia_smi_diagnostics(),
+        "cuda": {
+            "available": cuda_available,
+            "device_count": torch.cuda.device_count(),
+            "devices": [],
+        },
+        "mps": {
+            "built": False,
+            "available": False,
+        },
+        "xpu": {
+            "available": False,
+            "device_count": 0,
+            "devices": [],
+        },
+    }
+
+    if cuda_available:
+        for index in range(torch.cuda.device_count()):
+            properties = torch.cuda.get_device_properties(index)
+            cuda_device: dict[str, Any] = {
+                "index": index,
+                "name": torch.cuda.get_device_name(index),
+                "compute_capability": f"{properties.major}.{properties.minor}",
+                "total_memory_gb": round(properties.total_memory / (1024**3), 3),
+            }
+            try:
+                free_bytes, total_bytes = torch.cuda.mem_get_info(index)
+                cuda_device["free_memory_gb"] = round(free_bytes / (1024**3), 3)
+                cuda_device["visible_memory_gb"] = round(total_bytes / (1024**3), 3)
+            except RuntimeError as exc:
+                cuda_device["memory_info_error"] = str(exc)
+            diagnostics["cuda"]["devices"].append(cuda_device)
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    if mps_backend is not None:
+        diagnostics["mps"]["built"] = bool(mps_backend.is_built())
+        diagnostics["mps"]["available"] = bool(mps_backend.is_available())
+
+    xpu_backend = getattr(torch, "xpu", None)
+    if xpu_backend is not None:
+        diagnostics["xpu"]["available"] = bool(xpu_backend.is_available())
+        if diagnostics["xpu"]["available"]:
+            device_count = int(xpu_backend.device_count())
+            diagnostics["xpu"]["device_count"] = device_count
+            for index in range(device_count):
+                xpu_device: dict[str, Any] = {"index": index}
+                try:
+                    xpu_device["name"] = xpu_backend.get_device_name(index)
+                except RuntimeError as exc:
+                    xpu_device["name_error"] = str(exc)
+                diagnostics["xpu"]["devices"].append(xpu_device)
+
+    return diagnostics
+
+
+def _synchronize_device(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+        return
+    if device.type == "xpu":
+        xpu_backend = getattr(torch, "xpu", None)
+        if xpu_backend is not None:
+            xpu_backend.synchronize(device)
+
+
+def run_device_smoke_test(
+    *,
+    require_cuda: bool = False,
+    matrix_size: int = 1024,
+    repeats: int = 3,
+) -> dict[str, Any]:
+    """Run a deterministic matrix multiply on the selected device."""
+    if matrix_size < 1:
+        raise ValueError("matrix_size must be at least 1")
+    if repeats < 1:
+        raise ValueError("repeats must be at least 1")
+
+    diagnostics = collect_device_diagnostics()
+    diagnostics["require_cuda"] = require_cuda
+    if require_cuda and not diagnostics["cuda"]["available"]:
+        raise DeviceSmokeError(
+            "CUDA is required for this smoke run but PyTorch could not initialize it",
+            diagnostics,
+        )
+
+    device = select_device()
+    dtype = torch.float32
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(0)
+    left = torch.randn(
+        matrix_size,
+        matrix_size,
+        generator=generator,
+        dtype=dtype,
+        device="cpu",
+    ).to(device)
+    right = torch.randn(
+        matrix_size,
+        matrix_size,
+        generator=generator,
+        dtype=dtype,
+        device="cpu",
+    ).to(device)
+
+    _synchronize_device(device)
+    start = time.perf_counter()
+    result = left @ right
+    for _ in range(repeats - 1):
+        result = left @ right
+    _synchronize_device(device)
+    elapsed_seconds = time.perf_counter() - start
+
+    is_finite = bool(torch.isfinite(result).all().cpu().item())
+    if not is_finite:
+        raise DeviceSmokeError(
+            f"Smoke matmul produced non-finite values on {device}", diagnostics
+        )
+
+    selected_device: dict[str, Any] = {
+        "accelerator": device.type,
+        "device": str(device),
+        "dtype": str(dtype),
+    }
+    diagnostics["selected_device"] = selected_device
+    diagnostics["matmul_smoke"] = {
+        "matrix_size": matrix_size,
+        "repeats": repeats,
+        "elapsed_seconds": round(elapsed_seconds, 6),
+        "mean_checksum": float(result.detach().float().mean().cpu().item()),
+        "finite": True,
+    }
+
+    if device.type == "cuda":
+        selected_device["memory_allocated_gb"] = round(
+            torch.cuda.memory_allocated(device) / (1024**3), 3
+        )
+        selected_device["max_memory_allocated_gb"] = round(
+            torch.cuda.max_memory_allocated(device) / (1024**3), 3
+        )
+
+    return diagnostics

syvain_training_utils-0.0.170/src/syvain_training_utils/runtime.py

@@ -0,0 +1,54 @@
+import datetime as dt
+import logging
+import os
+import shutil
+
+import torch
+
+
+def generate_run_id() -> str:
+    """Return a compact UTC timestamp suitable for a training-run identifier."""
+    return dt.datetime.now(dt.UTC).strftime("%Y%m%d-%H%M%S")
+
+
+def runtime_c_compiler() -> str | None:
+    """Find the configured or first conventional C compiler on PATH."""
+    configured_cc = os.getenv("CC")
+    if configured_cc:
+        return shutil.which(configured_cc)
+
+    for compiler in ("cc", "gcc", "clang"):
+        resolved = shutil.which(compiler)
+        if resolved is not None:
+            return resolved
+    return None
+
+
+def require_torch_compile_toolchain(*, logger: logging.Logger | None = None) -> str:
+    """Require the runtime compiler used by TorchInductor and return its path."""
+    compiler = runtime_c_compiler()
+    if compiler is None:
+        raise RuntimeError(
+            "torch.compile uses TorchInductor/Triton and requires a C compiler "
+            "at runtime. Use an image with cc/gcc/clang installed or "
+            "set CC to a valid compiler path."
+        )
+    if logger is not None:
+        logger.info("Using C compiler for torch.compile: %s", compiler)
+    return compiler
+
+
+def select_device() -> torch.device:
+    """Select the best available PyTorch device in CUDA, XPU, MPS, CPU order."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+
+    xpu_backend = getattr(torch, "xpu", None)
+    if xpu_backend is not None and bool(xpu_backend.is_available()):
+        return torch.device("xpu")
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    if mps_backend is not None and bool(mps_backend.is_available()):
+        return torch.device("mps")
+
+    return torch.device("cpu")