GPU Temperature Guard

GPU temperature supervisor for CUDA training safety.

Provides a CPU-side supervisor that owns the lifecycle of a GPU training subprocess. The supervisor itself holds no GPU memory and runs purely on the CPU, polling GPU temperature via NVML (primary) or nvidia-smi (fallback).

State machine

The supervisor re-launches the training child (resuming from the latest rolling checkpoint) and watches the GPU temperature every poll_interval_seconds:

  • temp >= critical_threshold_c:

    • switch_on_thermal == False (default): permanent abort — kill the child and return a non-zero exit code (hardware unsafe).

    • switch_on_thermal == True: CPU continuation — checkpoint the child via SIGUSR1, kill it, then re-launch with --device cpu so training continues on CPU. The supervisor then stops polling temperature (no GPU in use) and waits for the CPU child to finish.

  • pause_threshold_c <= temp < critical_threshold_c: thermal pause — send SIGUSR1 (request a rolling checkpoint flush), wait up to checkpoint_grace_seconds for the child to acknowledge via a sentinel file, then SIGKILL the process group. Wait for the GPU to cool below resume_threshold_c and re-launch on GPU with --resume-from-checkpoint auto.

  • child exits cleanly (rc==0): training done, return 0.

  • child crashes (rc!=0, not killed by us): return the child’s exit code.

Backward compatibility

GPUTemperatureGuard is kept as a thin shim that exposes the temperature-reading helpers (read_temperature_c, threshold properties) so that existing tests and lightweight in-process callers continue to work. The supervisor (GPUTempSupervisor) is the new primary class.

exception src.utils.gpu_temp_guard.GPUTelemetryError[source]

Bases: RuntimeError

Fatal telemetry error indicating training must stop to avoid unsafe execution.

class src.utils.gpu_temp_guard.GPUTempCheckResult(temp_c: float, paused: bool, pause_duration_s: float, checks_during_pause: int, critical_seen: bool)[source]

Bases: object

Result of a single GPU temperature check cycle.

temp_c

Current GPU temperature in degrees Celsius.

Type:

float

paused

Whether training was paused during this check.

Type:

bool

pause_duration_s

Total seconds spent paused (0 if not paused).

Type:

float

checks_during_pause

Number of temperature polls performed while paused.

Type:

int

critical_seen

Whether the critical threshold was reached at any point.

Type:

bool

checks_during_pause: int
critical_seen: bool
pause_duration_s: float
paused: bool
property repair_action: str

Human-readable label describing the thermal action taken.

temp_c: float
class src.utils.gpu_temp_guard.GPUTempSupervisor(*, enabled: bool, device: str, nvml_device_index: int = 0, pause_threshold_c: float = 90.0, resume_threshold_c: float = 80.0, critical_threshold_c: float | None = None, poll_interval_seconds: float = 30.0, checkpoint_dir: str = 'checkpoints', checkpoint_grace_seconds: float = 30.0, switch_on_thermal: bool = False, child_argv: ~typing.List[str] | None = None, child_env: dict | None = None, log_fn: ~typing.Callable[[str], None] = <function warning>)[source]

Bases: object

CPU-side supervisor that kills/restarts a GPU training subprocess on thermal events.

The supervisor spawns the training command as a direct child process (start_new_session=True so it leads its own process group). It polls GPU temperature every poll_interval_seconds and, on a thermal event:

  • Sends SIGUSR1 to the child to request an immediate rolling checkpoint flush (the child acknowledges by writing a sentinel file in checkpoint_dir).

  • Waits up to checkpoint_grace_seconds for the sentinel, then SIGKILL the child’s process group.

  • Waits for the GPU to cool below resume_threshold_c and re-launches the child with --resume-from-checkpoint auto (the child loads the latest rolling checkpoint itself).

At the critical threshold the behaviour depends on switch_on_thermal:

  • switch_on_thermal == False: kill and abort permanently (return 2).

  • switch_on_thermal == True: checkpoint, kill, and re-launch the child with --device cpu (resume on CPU). The supervisor then stops polling temperature and just waits for the CPU child to finish.

is_active

Whether the supervisor is active (enabled + CUDA available).

pause_threshold_c

Temperature that triggers a kill+cooldown+restart.

resume_threshold_c

Temperature below which the GPU child is re-launched.

critical_threshold_c

Permanent-abort (or CPU-switch) threshold.

poll_interval_seconds

Seconds between temperature polls while the GPU child is running.

checkpoint_grace_seconds

Max seconds to wait for the child’s checkpoint acknowledgement before SIGKILL.

switch_on_thermal

Whether critical temp triggers CPU continuation.

total_pause_events

Cumulative kill+restart events.

total_paused_seconds

Cumulative seconds spent cooling down.

__init__(*, enabled: bool, device: str, nvml_device_index: int = 0, pause_threshold_c: float = 90.0, resume_threshold_c: float = 80.0, critical_threshold_c: float | None = None, poll_interval_seconds: float = 30.0, checkpoint_dir: str = 'checkpoints', checkpoint_grace_seconds: float = 30.0, switch_on_thermal: bool = False, child_argv: ~typing.List[str] | None = None, child_env: dict | None = None, log_fn: ~typing.Callable[[str], None] = <function warning>)[source]

Initialize the supervisor.

Parameters:

  • enabled – Whether the supervisor is enabled.

  • device – Target device string (supervisor only activates for CUDA).

  • nvml_device_index – GPU index to monitor.

  • pause_threshold_c – Kill threshold in Celsius. Must be > 0.

  • resume_threshold_c – Re-launch threshold in Celsius. Must be > 0 and < pause_threshold_c.

  • critical_threshold_c – Optional permanent-abort / CPU-switch threshold. Must be > 0 when provided.

  • poll_interval_seconds – Seconds between temperature polls. > 0.

  • checkpoint_dir – Directory where the child writes rolling checkpoints and the checkpoint-done sentinel file.

  • checkpoint_grace_seconds – Grace window after SIGUSR1 before SIGKILL. Must be > 0.

  • switch_on_thermal – If True, critical temp triggers CPU continuation instead of permanent abort.

  • child_argv – Argument vector used to (re)launch the training child. The supervisor prepends sys.executable -u if the first element is not already an interpreter path.

  • child_env – Optional environment override for the child. If None, os.environ is used (with the supervisor-child marker added).

  • log_fn – Callable used for log messages (defaults to logging.warning).

Raises:

ValueError – If threshold / interval / grace constraints are violated.

property checkpoint_grace_seconds: float
property critical_threshold_c: float | None
property is_active: bool
property pause_threshold_c: float
property poll_interval_seconds: float
read_temperature_c() → float[source]

Read the current GPU temperature via the probe.

Returns:

Temperature in Celsius, or 0.0 if the supervisor is inactive.

property resume_threshold_c: float
run() → int[source]

Run the supervisor loop until the child finishes or a permanent abort.

Returns:

0 on successful completion, 2 on permanent thermal abort, or the child’s non-zero exit code on crash.

property switch_on_thermal: bool
class src.utils.gpu_temp_guard.GPUTemperatureGuard(*, enabled: bool, device: str, nvml_device_index: int = 0, pause_threshold_c: float = 90.0, resume_threshold_c: float = 80.0, critical_threshold_c: float | None = None, poll_interval_seconds: float = 30.0)[source]

Bases: object

Backward-compatible temperature reader / threshold holder.

This class no longer implements the in-process pause/resume loop (that responsibility moved to GPUTempSupervisor). It still exposes read_temperature_c and the threshold properties so legacy callers and tests that only need a temperature probe can keep working unchanged.

is_active

Whether temperature probing is available (CUDA available and guard enabled).

pause_threshold_c

Temperature above which the supervisor pauses.

resume_threshold_c

Temperature below which the supervisor resumes.

critical_threshold_c

Optional critical / permanent-abort threshold.

poll_interval_seconds

Seconds between temperature polls.

last_temperature_c

Most recently read temperature, or None.

total_pause_events

Cumulative pause events (kept for API parity; the supervisor owns the real counter now).

total_paused_seconds

Cumulative paused seconds (API parity only).

__init__(*, enabled: bool, device: str, nvml_device_index: int = 0, pause_threshold_c: float = 90.0, resume_threshold_c: float = 80.0, critical_threshold_c: float | None = None, poll_interval_seconds: float = 30.0)[source]

Initialize the temperature guard shim.

Parameters:

  • enabled – Whether the guard is enabled.

  • device – Target device string (probing only activates for CUDA).

  • nvml_device_index – NVML device index for multi-GPU systems.

  • pause_threshold_c – Pause threshold in Celsius. Must be > 0.

  • resume_threshold_c – Resume threshold in Celsius. Must be > 0 and < pause_threshold_c.

  • critical_threshold_c – Optional critical threshold. Must be > 0 when provided.

  • poll_interval_seconds – Seconds between temperature polls. > 0.

Raises:

ValueError – If threshold or poll interval constraints are violated.

property critical_threshold_c: float | None
property is_active: bool
last_temperature_c: float | None
property pause_threshold_c: float
property poll_interval_seconds: float
read_temperature_c() → float[source]

Read the current GPU temperature in Celsius.

Attempts NVML first, then nvidia-smi. Returns 0.0 when the guard is not active (CPU device or disabled).

Raises:

GPUTelemetryError – If both backends fail to provide a valid temperature reading.

property resume_threshold_c: float
wait_until_safe(*, context: str = '') → GPUTempCheckResult[source]

Block until GPU temperature drops below the resume threshold.

Kept for backward compatibility with callers/tests that drive an in-process pause loop. The supervisor does not use this method; it kills and restarts the child instead.

Parameters:

context – Optional human-readable label for log messages.

Returns:

A GPUTempCheckResult describing the check outcome.

src.utils.gpu_temp_guard.read_temperature_from_nvidia_smi(nvml_device_index: int = 0) → float[source]

Read GPU temperature via nvidia-smi subprocess.

Parameters:

nvml_device_index – GPU index to query.

Returns:

Temperature in Celsius.

Raises:

GPUTelemetryError – If nvidia-smi is unavailable or returns an invalid payload.