API Reference

Autogenerated docs for the public modules most users interact with. The CLI module exposes the Typer entry point, while the controllers and utilities power both the command-line and Python recipes.

For global sessions, gpu_ids=None means all visible GPUs. Explicit values are visible device ordinals after CUDA or ROCm visibility filtering. Empty, duplicate, or out-of-range lists are invalid; lists with more than 64 entries are also invalid. Startup raises ValueError if discovery resolves to zero devices. Telemetry may expose metadata such as physical_id, but those fields are not accepted as selection IDs. On CUDA, NVML telemetry records are returned only for visible ordinals that Torch CUDA can select, and surviving visible IDs are not compacted after filtering. NVML-only devices are not exposed as public gpu_ids. On ROCm, telemetry records are returned only for visible ordinals that Torch can select; nullable memory fields mean memory telemetry is unavailable after successful selection.

GlobalGPUController validates local constructor inputs (gpu_ids, interval, busy_threshold, and vram_to_keep) before platform or hardware probing. Visible-count checks for explicit IDs still run after backend discovery because they depend on the current visible device count. Its omitted vram_to_keep default is the shared low-power public default, 1GiB, matching the CLI and service APIs. Direct CudaGPUController, RocmGPUController, and MacMGPUController constructors use the same omitted vram_to_keep="1GiB" default. Public interval values must be finite positive seconds, including fractional seconds, capped by the Python runtime wait limit. Public VRAM byte-equivalent values below 4 bytes or above 1 PiB are rejected, and internal tensor element counts round up to cover the requested byte-equivalent amount.

Direct CUDA, ROCm, and Mac M controller rank values are public visible device ordinals. CUDA and ROCm ranks must be plain integers within 0..torch.cuda.device_count()-1 in the current process environment; Mac M ranks must be the plain integer 0. Non-integer ranks raise TypeError, and negative or out-of-range ranks raise ValueError during construction, before KeepGPU creates a torch.device, checks backend availability, calls backend device selection, starts a worker, or queries telemetry. CUDA and ROCm call torch.cuda.device_count() only after the rank is known to be a plain integer.

CUDA telemetry resolves visible ordinals through CUDA_VISIBLE_DEVICES before querying NVML. Numeric tokens, full UUID tokens, and unique UUID prefixes are supported, and parsing stops at -1 after any valid preceding tokens. Malformed, duplicate/equivalent, ambiguous, out-of-range, or unresolved mappings report unavailable utilization instead of guessing a physical device. ROCm telemetry resolves ROCR_VISIBLE_DEVICES as the base mask and one matching HIP_VISIBLE_DEVICES/CUDA_VISIBLE_DEVICES overlay before querying ROCm SMI. Unresolved mappings report unavailable utilization instead of falling back to a possibly wrong physical device.

Public Python controller, CLI, REST, JSON-RPC, and MCP defaults use vram_to_keep/vram 1GiB and busy_threshold=25, so omitted settings target a modest VRAM signal only when utilization backoff permits; busy or unavailable telemetry sleeps before allocating keep tensors or running compute. Pass busy_threshold=-1 only when you intentionally want unconditional keepalive compute without utilization backoff.

CUDA and ROCm keep() calls wait for fatal backend startup setup to succeed before reporting success. Startup failures such as device-selection errors are raised synchronously, while normal low-power allocation can still be deferred by utilization backoff after startup succeeds.

Single-GPU workload iteration controls must be positive integers. CUDA exposes relu_iterations; ROCm and Mac M expose iterations. Non-integer values raise TypeError, and non-positive values raise ValueError before a worker can start, rather than creating a silent no-op keep loop or a background thread crash.

For service session IDs, job_id=None is the only omitted/all-sessions sentinel. Custom IDs must be non-empty strings containing only letters, digits, ., _, -, or ~, and may not be standalone . or ..; invalid values raise ValueError before session state changes.

CLI entrypoint for KeepGPU.

ServiceRPCError

Bases: RuntimeError

Raised when the local service returns a JSON-RPC error.

Source code in src/keep_gpu/cli.py

class ServiceRPCError(RuntimeError):
    """Raised when the local service returns a JSON-RPC error."""

    def __init__(self, message: str, code: Optional[int] = None) -> None:
        super().__init__(message)
        self.code = code

ServiceResponseError

Bases: RuntimeError

Raised when the local service responds with an invalid HTTP payload.

Source code in src/keep_gpu/cli.py

class ServiceResponseError(RuntimeError):
    """Raised when the local service responds with an invalid HTTP payload."""

ServiceUnreachableError

Bases: RuntimeError

Raised when the local service cannot be reached.

Source code in src/keep_gpu/cli.py

class ServiceUnreachableError(RuntimeError):
    """Raised when the local service cannot be reached."""

list_gpus(host=typer.Option(DEFAULT_SERVICE_HOST, '--host', help='Service host.'), port=typer.Option(str(DEFAULT_SERVICE_PORT), '--port', help='Service port.'))

List GPU telemetry from local service.

Source code in src/keep_gpu/cli.py

@app.command("list-gpus")
def list_gpus(
    host: str = typer.Option(DEFAULT_SERVICE_HOST, "--host", help="Service host."),
    port: str = typer.Option(
        str(DEFAULT_SERVICE_PORT),
        "--port",
        help="Service port.",
    ),
):
    """List GPU telemetry from local service."""
    try:
        host, port = _validate_cli_service_endpoint(host, port)
        result = _rpc_call("list_gpus", {}, host, port)
        result = _validate_list_gpus_result(result)
        _print_machine_json(result)
    except (RuntimeError, typer.BadParameter) as exc:
        _print_machine_json({"error": str(exc)})
        raise typer.Exit(code=1) from exc

main(ctx, interval=typer.Option('300', metavar='NUMBER', help='Interval in seconds between GPU usage checks (blocking mode only).'), gpu_ids=typer.Option(None, help='Comma-separated visible GPU ordinals for blocking mode (default: all).'), vram=typer.Option('1GiB', '--vram', help='Amount of VRAM to keep occupied (blocking mode).'), legacy_threshold=typer.Option(None, '--threshold', hidden=True, help='Deprecated alias: numeric maps to busy-threshold, string maps to vram.'), busy_threshold=typer.Option(str(DEFAULT_BUSY_THRESHOLD), '--busy-threshold', '--util-threshold', metavar='INTEGER', help='Back off when utilization is above this 0..100 percent threshold or telemetry is unavailable; -1 disables utilization backoff (blocking mode).'))

Run blocking keep-alive mode when no subcommand is provided.

Source code in src/keep_gpu/cli.py

@app.callback(invoke_without_command=True)
def main(
    ctx: typer.Context,
    interval: str = typer.Option(
        "300",
        metavar="NUMBER",
        help="Interval in seconds between GPU usage checks (blocking mode only).",
    ),
    gpu_ids: Optional[str] = typer.Option(
        None,
        help="Comma-separated visible GPU ordinals for blocking mode (default: all).",
    ),
    vram: str = typer.Option(
        "1GiB",
        "--vram",
        help="Amount of VRAM to keep occupied (blocking mode).",
    ),
    legacy_threshold: Optional[str] = typer.Option(
        None,
        "--threshold",
        hidden=True,
        help="Deprecated alias: numeric maps to busy-threshold, string maps to vram.",
    ),
    busy_threshold: str = typer.Option(
        str(DEFAULT_BUSY_THRESHOLD),
        "--busy-threshold",
        "--util-threshold",
        metavar="INTEGER",
        help=(
            "Back off when utilization is above this 0..100 percent threshold "
            "or telemetry is unavailable; -1 disables utilization backoff "
            "(blocking mode)."
        ),
    ),
):
    """Run blocking keep-alive mode when no subcommand is provided."""
    try:
        if ctx.invoked_subcommand is not None:
            _reject_root_blocking_options_before_subcommand(ctx)
            return
        interval = _validate_cli_interval(interval)
        busy_threshold = _validate_cli_busy_threshold(busy_threshold)
        _parse_gpu_ids(gpu_ids)
        legacy_vram, legacy_busy_threshold, _ = _apply_legacy_threshold(
            vram, legacy_threshold, busy_threshold
        )
        _validate_cli_vram(legacy_vram)
        _validate_cli_busy_threshold(legacy_busy_threshold)
        _run_blocking(interval, gpu_ids, vram, legacy_threshold, busy_threshold)
    except typer.BadParameter as exc:
        if _subcommand_outputs_machine_json(ctx):
            _print_machine_json({"error": str(exc)})
        else:
            console.print(f"[bold red]Error: {exc}[/bold red]")
        raise typer.Exit(code=1) from exc

serve(host=typer.Option(DEFAULT_SERVICE_HOST, help='Host interface for KeepGPU local service.'), port=typer.Option(str(DEFAULT_SERVICE_PORT), help='Port for KeepGPU local service.'))

Run KeepGPU local service (HTTP + JSON-RPC + dashboard).

Source code in src/keep_gpu/cli.py

@app.command("serve")
def serve(
    host: str = typer.Option(
        DEFAULT_SERVICE_HOST,
        help="Host interface for KeepGPU local service.",
    ),
    port: str = typer.Option(
        str(DEFAULT_SERVICE_PORT),
        help="Port for KeepGPU local service.",
    ),
):
    """Run KeepGPU local service (HTTP + JSON-RPC + dashboard)."""
    try:
        host, port = _validate_cli_service_endpoint(host, port)
    except typer.BadParameter as exc:
        console.print(f"[bold red]Error: {exc}[/bold red]")
        raise typer.Exit(code=1) from exc

    from keep_gpu.mcp.server import KeepGPUServer, run_http

    console.print(f"[bold cyan]Service URL:[/bold cyan] http://{host}:{port}/")
    console.print(
        "[dim]Press Ctrl+C to stop the foreground service, or use `keep-gpu service-stop` for auto-started daemons.[/dim]"
    )
    run_http(KeepGPUServer(), host=host, port=port)

service_stop(host=typer.Option(DEFAULT_SERVICE_HOST, '--host', help='Service host.'), port=typer.Option(str(DEFAULT_SERVICE_PORT), '--port', help='Service port.'), force=typer.Option(False, '--force', help='Stop service even if tracked sessions exist.'))

Stop local KeepGPU service daemon started by auto-start logic.

Source code in src/keep_gpu/cli.py

@app.command("service-stop")
def service_stop(
    host: str = typer.Option(DEFAULT_SERVICE_HOST, "--host", help="Service host."),
    port: str = typer.Option(str(DEFAULT_SERVICE_PORT), "--port", help="Service port."),
    force: bool = typer.Option(
        False,
        "--force",
        help="Stop service even if tracked sessions exist.",
    ),
):
    """Stop local KeepGPU service daemon started by auto-start logic."""
    try:
        host, port = _validate_cli_service_endpoint(host, port)
        if force:
            stopped = _stop_service_process(host, port)
            if not stopped:
                raise RuntimeError(
                    "No managed daemon PID found. If service was started in foreground, stop it with Ctrl+C in that terminal."
                )
            console.print(
                f"[bold green]Force-stopped KeepGPU service daemon[/bold green] at http://{host}:{port}/"
            )
            return

        try:
            status = _rpc_call("status", {}, host, port)
            status = _validate_status_result(status, single_job=False)
        except ServiceUnreachableError as exc:
            service_stop_force_command = _keep_gpu_hint_command(
                "service-stop", host, port, "--force"
            )
            raise RuntimeError(
                f"KeepGPU service is unavailable at {host}:{port}. Non-force service-stop must verify no tracked keep sessions before stopping the daemon. "
                "For an unresponsive auto-started daemon, run "
                f"`{service_stop_force_command}`."
            ) from exc

        stop_all_command = _keep_gpu_hint_command("stop", host, port, "--all")
        service_stop_force_command = _keep_gpu_hint_command(
            "service-stop", host, port, "--force"
        )
        _require_no_active_jobs_for_service_stop(
            status,
            stop_all_command=stop_all_command,
            force_command=service_stop_force_command,
        )
        stop_result = _rpc_call("stop_keep", {}, host, port, timeout=45.0)
        stop_result = _validate_stop_keep_result(stop_result)
        _require_clean_stop_keep_for_service_stop(
            stop_result,
            force_command=_keep_gpu_hint_command("service-stop", host, port, "--force"),
        )
        status = _rpc_call("status", {}, host, port)
        status = _validate_status_result(status, single_job=False)
        _require_no_active_jobs_for_service_stop(
            status,
            stop_all_command=stop_all_command,
            force_command=service_stop_force_command,
        )

        stopped = _stop_service_process(host, port)
        if not stopped:
            raise RuntimeError(
                "No managed daemon PID found. If service was started in foreground, stop it with Ctrl+C in that terminal."
            )

        console.print(
            f"[bold green]Stopped KeepGPU service daemon[/bold green] at http://{host}:{port}/"
        )
    except (RuntimeError, typer.BadParameter) as exc:
        console.print(f"[bold red]Error: {exc}[/bold red]")
        raise typer.Exit(code=1) from exc

start(gpu_ids=typer.Option(None, help='Comma-separated visible GPU ordinals.'), vram=typer.Option('1GiB', '--vram', help='VRAM to keep per GPU.'), interval=typer.Option('300', metavar='NUMBER', help='Interval in seconds between checks.'), busy_threshold=typer.Option(str(DEFAULT_BUSY_THRESHOLD), '--busy-threshold', '--util-threshold', metavar='INTEGER', help='Back off when utilization is above this 0..100 percent threshold or telemetry is unavailable; -1 disables utilization backoff.'), job_id=typer.Option(None, help='Optional custom job id. Auto-generated when omitted.'), host=typer.Option(DEFAULT_SERVICE_HOST, '--host', help='KeepGPU service host.'), port=typer.Option(str(DEFAULT_SERVICE_PORT), '--port', help='KeepGPU service port.'), auto_start=typer.Option(True, '--auto-start/--no-auto-start', help='Auto-start local service when unavailable.'))

Start a non-blocking keep session and return a job id.

Use keep-gpu stop --job-id <id> to release this session and keep-gpu service-stop to stop the local service daemon.

Source code in src/keep_gpu/cli.py

@app.command("start")
def start(
    gpu_ids: Optional[str] = typer.Option(
        None,
        help="Comma-separated visible GPU ordinals.",
    ),
    vram: str = typer.Option("1GiB", "--vram", help="VRAM to keep per GPU."),
    interval: str = typer.Option(
        "300", metavar="NUMBER", help="Interval in seconds between checks."
    ),
    busy_threshold: str = typer.Option(
        str(DEFAULT_BUSY_THRESHOLD),
        "--busy-threshold",
        "--util-threshold",
        metavar="INTEGER",
        help=(
            "Back off when utilization is above this 0..100 percent threshold "
            "or telemetry is unavailable; -1 disables utilization backoff."
        ),
    ),
    job_id: Optional[str] = typer.Option(
        None,
        help="Optional custom job id. Auto-generated when omitted.",
    ),
    host: str = typer.Option(
        DEFAULT_SERVICE_HOST,
        "--host",
        help="KeepGPU service host.",
    ),
    port: str = typer.Option(
        str(DEFAULT_SERVICE_PORT),
        "--port",
        help="KeepGPU service port.",
    ),
    auto_start: bool = typer.Option(
        True,
        "--auto-start/--no-auto-start",
        help="Auto-start local service when unavailable.",
    ),
):
    """Start a non-blocking keep session and return a job id.

    Use `keep-gpu stop --job-id <id>` to release this session and
    `keep-gpu service-stop` to stop the local service daemon.
    """
    auto_started = False
    try:
        host, port = _validate_cli_service_endpoint(host, port)
        interval = _validate_cli_interval(interval)
        busy_threshold = _validate_cli_busy_threshold(busy_threshold)
        parsed_gpu_ids = _parse_gpu_ids(gpu_ids)
        _validate_cli_vram(vram)
        job_id = _validate_cli_job_id(job_id)
        auto_started = _ensure_service_running(host, port, auto_start=auto_start)
        result = _rpc_call(
            "start_keep",
            {
                "gpu_ids": parsed_gpu_ids,
                "vram": vram,
                "interval": interval,
                "busy_threshold": busy_threshold,
                "job_id": job_id,
            },
            host,
            port,
        )
        result_job_id = _validate_start_keep_result(result, expected_job_id=job_id)
        if auto_started:
            console.print(
                f"[bold cyan]Auto-started KeepGPU service[/bold cyan] at http://{host}:{port}/"
            )
        console.print(
            f"[bold green]Started keep session[/bold green] job_id={result_job_id}"
        )
        console.print(f"[cyan]Dashboard:[/cyan] http://{host}:{port}/")
        status_command = _keep_gpu_hint_command(
            "status", host, port, "--job-id", result_job_id
        )
        stop_command = _keep_gpu_hint_command(
            "stop", host, port, "--job-id", result_job_id
        )
        service_stop_command = _keep_gpu_hint_command("service-stop", host, port)
        console.print(f"[dim]Next: {status_command} | {stop_command}[/dim]")
        console.print(
            f"[dim]When all sessions are done, stop daemon with: {service_stop_command}[/dim]"
        )
    except ServiceRPCError as exc:
        _rollback_auto_started_service_on_startup_unavailable(
            auto_started, exc, host, port
        )
        console.print(f"[bold red]Error: {exc}[/bold red]")
        raise typer.Exit(code=1) from exc
    except (RuntimeError, typer.BadParameter) as exc:
        console.print(f"[bold red]Error: {exc}[/bold red]")
        raise typer.Exit(code=1) from exc

status(job_id=typer.Option(None, help='Session id to inspect.'), host=typer.Option(DEFAULT_SERVICE_HOST, '--host', help='Service host.'), port=typer.Option(str(DEFAULT_SERVICE_PORT), '--port', help='Service port.'))

Show session status from KeepGPU local service.

Source code in src/keep_gpu/cli.py

@app.command("status")
def status(
    job_id: Optional[str] = typer.Option(None, help="Session id to inspect."),
    host: str = typer.Option(DEFAULT_SERVICE_HOST, "--host", help="Service host."),
    port: str = typer.Option(
        str(DEFAULT_SERVICE_PORT),
        "--port",
        help="Service port.",
    ),
):
    """Show session status from KeepGPU local service."""
    try:
        host, port = _validate_cli_service_endpoint(host, port)
        job_id = _validate_cli_job_id(job_id)
        result = _rpc_call(
            "status",
            {} if job_id is None else {"job_id": job_id},
            host,
            port,
        )
        result = _validate_status_result(
            result,
            single_job=job_id is not None,
            expected_job_id=job_id,
        )
        _print_machine_json(result)
    except (RuntimeError, typer.BadParameter) as exc:
        _print_machine_json({"error": str(exc)})
        raise typer.Exit(code=1) from exc

stop(job_id=typer.Option(None, help='Session id to stop. Omit with --all to stop every session.'), all_sessions=typer.Option(False, '--all', help='Stop all sessions.'), host=typer.Option(DEFAULT_SERVICE_HOST, '--host', help='Service host.'), port=typer.Option(str(DEFAULT_SERVICE_PORT), '--port', help='Service port.'))

Stop one session or all sessions.

Source code in src/keep_gpu/cli.py

@app.command("stop")
def stop(
    job_id: Optional[str] = typer.Option(
        None,
        help="Session id to stop. Omit with --all to stop every session.",
    ),
    all_sessions: bool = typer.Option(
        False,
        "--all",
        help="Stop all sessions.",
    ),
    host: str = typer.Option(DEFAULT_SERVICE_HOST, "--host", help="Service host."),
    port: str = typer.Option(
        str(DEFAULT_SERVICE_PORT),
        "--port",
        help="Service port.",
    ),
):
    """Stop one session or all sessions."""
    try:
        if job_id is not None and all_sessions:
            raise RuntimeError("Use either --job-id or --all, not both.")
        if job_id is None and not all_sessions:
            raise RuntimeError("Provide --job-id or use --all.")
        host, port = _validate_cli_service_endpoint(host, port)
        job_id = _validate_cli_job_id(job_id)
        if all_sessions:
            result = _stop_all_sessions_with_fallback(host, port)
        else:
            result = _rpc_call(
                "stop_keep",
                {"job_id": job_id},
                host,
                port,
                timeout=45.0,
            )
            result = _validate_stop_keep_result(result, expected_job_id=job_id)
        _print_machine_json(result)
    except (RuntimeError, typer.BadParameter) as exc:
        _print_machine_json({"error": str(exc)})
        raise typer.Exit(code=1) from exc

CudaGPUController

Bases: BaseGPUController

CudaGPUController Keep a single CUDA GPU busy by repeatedly running lightweight elementwise workloads in a background thread.

Typical usage:

ctrl = CudaGPUController(rank=0, interval=0.5)
ctrl.keep()          # occupy GPU while you do CPU-only work
dataset.process()
ctrl.release()        # give GPU memory back
model.train_start()   # now run real GPU training

Or as a context manager:

with CudaGPUController(rank=0, interval=0.5):
    dataset.process()  # GPU occupied inside this block
model.train_start()    # GPU free after exiting block

Source code in src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

class CudaGPUController(BaseGPUController):
    """CudaGPUController
    Keep a single CUDA GPU busy by repeatedly running lightweight
    elementwise workloads in a background thread.

    Typical usage:

    ```python
    ctrl = CudaGPUController(rank=0, interval=0.5)
    ctrl.keep()          # occupy GPU while you do CPU-only work
    dataset.process()
    ctrl.release()        # give GPU memory back
    model.train_start()   # now run real GPU training
    ```

    Or as a context manager:

    ```python
    with CudaGPUController(rank=0, interval=0.5):
        dataset.process()  # GPU occupied inside this block
    model.train_start()    # GPU free after exiting block
    ```
    """

    def __init__(
        self,
        *,
        rank: int,
        interval: float = 1.0,
        relu_iterations: int = 5000,
        vram_to_keep: str | int = "1GiB",
        busy_threshold: int = DEFAULT_BUSY_THRESHOLD,
    ):
        """
        Args:
            rank (int): Local CUDA device index to occupy.
            interval (float, optional): Sleep time (seconds) between workload
                batches. Defaults to 1.0.
            relu_iterations (int, optional): Number of in-place ReLU ops per
                batch.
            vram_to_keep (int or str, optional): Amount of VRAM to keep busy,
                e.g. `"1GiB"`, `"20 GB"`, or an integer like `1000 * 1000`.
                This represents the total size of the matrix allocated to
                occupy the GPU.
            busy_threshold (int, optional): Defaults to 25. If current
                utilization (%) exceeds this threshold, or utilization is
                unavailable, the worker inserts extra sleeps to avoid hogging
                the GPU. Use -1 only to disable utilization backoff.

        """
        super().__init__(vram_to_keep=vram_to_keep, interval=interval)
        self.relu_iterations = validate_positive_integer(
            relu_iterations, "relu_iterations"
        )
        self.busy_threshold = validate_busy_threshold(busy_threshold)
        rank = validate_rank_type(rank)
        rank = validate_visible_rank(rank, visible_torch_device_count())
        self.rank = rank
        self.device = torch.device(f"cuda:{rank}")
        self.interval = interval
        self.platform = ComputingPlatform.CUDA

        self._stop_evt: Optional[threading.Event] = None
        self._thread: Optional[threading.Thread] = None
        self._failure_exc: Optional[Exception] = None
        self._num_elements: Optional[int] = None

    @staticmethod
    def parse_size(text: str) -> int:
        return parse_size(text)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def keep(self) -> None:
        """Launch the background thread that keeps the GPU busy."""
        if self._thread and self._thread.is_alive():
            if self._stop_evt is not None and self._stop_evt.is_set():
                raise RuntimeError(
                    f"rank {self.rank}: previous keep thread startup did not complete"
                )
            logger.warning("rank %s: keep thread already running", self.rank)
            return

        self._failure_exc = None
        self._num_elements = int(self.vram_to_keep)
        if self._num_elements <= 0:
            raise ValueError("vram_to_keep must be positive")

        self._stop_evt = threading.Event()
        startup_evt = threading.Event()
        startup_errors: list[Exception] = []
        self._thread = threading.Thread(
            target=self._keep_loop,
            args=(startup_evt, startup_errors),
            name=f"gpu-keeper-{self.rank}",
            daemon=True,  # daemon so program can exit cleanly
        )
        try:
            self._thread.start()
        except Exception:  # noqa: BLE001 - preserve original thread-start failure
            self._thread = None
            self._stop_evt = None
            raise
        startup_timeout = 5.0
        if not startup_evt.wait(startup_timeout):
            stop_evt = self._stop_evt
            if stop_evt is not None:
                stop_evt.set()
            self._thread.join(timeout=1.0)
            if self._thread.is_alive():
                raise RuntimeError(
                    f"rank {self.rank}: keep thread did not complete startup within "
                    f"{startup_timeout:.1f}s"
                )
            self._thread = None
            self._stop_evt = None
            raise RuntimeError(
                f"rank {self.rank}: keep thread exited before startup completed"
            )
        if startup_errors:
            self._thread.join(timeout=1.0)
            self._thread = None
            self._stop_evt = None
            raise startup_errors[0]
        logger.info("rank %s: keep thread started", self.rank)

    def release(self) -> None:
        """
        Stop the background thread and clear CUDA cache so the memory
        becomes immediately available to other code.
        """
        thread = self._thread
        if thread is not None and not thread.is_alive():
            stop_evt = self._stop_evt
            if (stop_evt is not None and stop_evt.is_set()) or getattr(
                self, "_failure_exc", None
            ) is not None:
                torch.cuda.empty_cache()
                self._thread = None
                self._stop_evt = None
                logger.info(
                    "rank %s: previously stopping keep thread exited; cache cleared",
                    self.rank,
                )
                return

        if not (thread and thread.is_alive()):
            logger.warning("rank %s: keep thread not running", self.rank)
            return

        stop_evt = self._stop_evt
        if stop_evt is None:
            raise RuntimeError(f"rank {self.rank}: stop event missing")
        assert stop_evt is not None

        stop_evt.set()
        join_timeout = max(2.0, min(float(self.interval) + 2.0, 30.0))
        thread.join(timeout=join_timeout)
        if thread.is_alive():
            logger.warning(
                "rank %s: keep thread did not stop within %.1fs",
                self.rank,
                join_timeout,
            )
            raise TimeoutError(
                f"rank {self.rank}: keep thread did not stop within {join_timeout:.1f}s"
            )
        torch.cuda.empty_cache()
        self._thread = None
        self._stop_evt = None
        logger.info("rank %s: keep thread stopped & cache cleared", self.rank)

    # Context-manager helpers -------------------------------------------------
    def __enter__(self):
        self.keep()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.release()

    # ------------------------------------------------------------------
    # Background loop
    # ------------------------------------------------------------------
    def _keep_loop(
        self,
        startup_evt: Optional[threading.Event] = None,
        startup_errors: Optional[list[Exception]] = None,
    ) -> None:
        """Internal: run workloads until stop event is set."""
        startup_confirmed = startup_evt is None

        def confirm_startup() -> None:
            nonlocal startup_confirmed
            if startup_confirmed:
                return
            startup_confirmed = True
            assert startup_evt is not None
            startup_evt.set()

        def record_startup_failure(exc: Exception) -> None:
            if not startup_confirmed and startup_errors is not None:
                startup_errors.append(exc)
            else:
                self._failure_exc = exc
            confirm_startup()

        def record_worker_failure(exc: Exception) -> None:
            failure = RuntimeError(
                f"rank {self.rank}: unexpected CUDA keep worker failure: {exc}"
            )
            record_startup_failure(failure)
            logger.exception("%s", failure)

        stop_evt = self._stop_evt
        if stop_evt is None:
            exc = RuntimeError(f"rank {self.rank}: stop event not initialized")
            logger.error("%s", exc)
            record_startup_failure(exc)
            return
        assert stop_evt is not None

        try:
            torch.cuda.set_device(self.rank)
        except Exception as exc:  # noqa: BLE001 - surface backend startup failure
            logger.error("rank %s: CUDA startup failed: %s", self.rank, exc)
            record_startup_failure(exc)
            return
        num_elements = self._num_elements if self._num_elements is not None else 0
        if num_elements <= 0:
            exc = RuntimeError(
                f"rank {self.rank}: invalid vram_to_keep={self.vram_to_keep}"
            )
            logger.error("%s", exc)
            record_startup_failure(exc)
            return
        matrix = None
        while not stop_evt.is_set():
            try:
                gpu_utilization = self._monitor_utilization(self.rank)
                if not self._should_run_batch(gpu_utilization, self.busy_threshold):
                    confirm_startup()
                    logger.debug(
                        "rank %s: GPU utilization unavailable or busy (%s), deferring allocation",
                        self.rank,
                        "n/a" if gpu_utilization is None else f"{gpu_utilization}%",
                    )
                    if stop_evt.wait(self.interval):
                        return
                    continue
                matrix = torch.rand(
                    num_elements,
                    device=self.device,
                    dtype=torch.float32,
                    requires_grad=False,
                )
                confirm_startup()
                break
            except RuntimeError as e:
                logger.error("rank %s: failed to allocate matrix: %s", self.rank, e)
                if "out of memory" in str(e).lower():
                    torch.cuda.empty_cache()
                    confirm_startup()
                    if stop_evt.wait(self.interval):
                        return
                    continue
                record_worker_failure(e)
                return
            except Exception as exc:  # noqa: BLE001 - surface backend startup failure
                record_worker_failure(exc)
                return
        if matrix is None:
            if not startup_confirmed:
                exc = RuntimeError(
                    f"rank {self.rank}: stopped before CUDA startup allocation"
                )
                logger.error("%s", exc)
                record_startup_failure(exc)
            else:
                logger.debug(
                    "rank %s: exiting before CUDA allocation after startup confirmation",
                    self.rank,
                )
            return
        while not stop_evt.is_set():
            try:
                gpu_utilization = self._monitor_utilization(self.rank)
                if not self._should_run_batch(gpu_utilization, self.busy_threshold):
                    logger.debug(
                        "rank %s: GPU utilization unavailable or busy (%s), sleeping longer",
                        self.rank,
                        "n/a" if gpu_utilization is None else f"{gpu_utilization}%",
                    )
                else:
                    self._run_relu_batch(matrix)
                if stop_evt.wait(self.interval):
                    break
            except RuntimeError as e:
                # Handle OOM by clearing cache; then sleep and continue
                if "out of memory" in str(e).lower():
                    torch.cuda.empty_cache()
                    if stop_evt.wait(self.interval):
                        break
                    continue
                self._failure_exc = RuntimeError(
                    f"rank {self.rank}: unexpected CUDA keep worker failure: {e}"
                )
                logger.exception("%s", self._failure_exc)
                return
            except Exception as exc:
                self._failure_exc = RuntimeError(
                    f"rank {self.rank}: unexpected CUDA keep worker failure: {exc}"
                )
                logger.exception("%s", self._failure_exc)
                return

    def allocation_status(self) -> Optional[Exception]:
        """
        Return fatal worker failure captured after startup, if any.

        The reference assignment/read is thread-safe for CPython's GIL model.
        """
        return self._failure_exc

    # ------------------------------------------------------------------
    # Workload implementation
    # ------------------------------------------------------------------
    @torch.no_grad()
    def _run_relu_batch(self, matrix: torch.Tensor) -> None:
        """Run a batch of in-place ReLU ops to keep GPU busy."""
        stop_evt = self._stop_evt

        tic = time.monotonic()
        for _ in range(self.relu_iterations):
            torch.relu_(matrix)
            if stop_evt is not None and stop_evt.is_set():
                break
        torch.cuda.synchronize()
        toc = time.monotonic()

        logger.debug(
            "rank %s: relu ops batch done - avg %.2f ms",
            self.rank,
            (toc - tic) * 1000 / max(1, self.relu_iterations),
        )

    # ------------------------------------------------------------------
    # Utilization monitor
    # ------------------------------------------------------------------
    @staticmethod
    def _monitor_utilization(rank: int) -> Optional[int]:
        """
        Return current GPU utilization (%) for `rank`.
        Returns None when telemetry is unavailable.
        """
        return get_gpu_utilization(rank)

__init__(*, rank, interval=1.0, relu_iterations=5000, vram_to_keep='1GiB', busy_threshold=DEFAULT_BUSY_THRESHOLD)

Parameters:

Name	Type	Description	Default
rank	int	Local CUDA device index to occupy.	required
interval	float	Sleep time (seconds) between workload batches. Defaults to 1.0.	1.0
relu_iterations	int	Number of in-place ReLU ops per batch.	5000
vram_to_keep	int or str	Amount of VRAM to keep busy, e.g. "1GiB", "20 GB", or an integer like 1000 * 1000. This represents the total size of the matrix allocated to occupy the GPU.	'1GiB'
busy_threshold	int	Defaults to 25. If current utilization (%) exceeds this threshold, or utilization is unavailable, the worker inserts extra sleeps to avoid hogging the GPU. Use -1 only to disable utilization backoff.	DEFAULT_BUSY_THRESHOLD

Source code in src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

def __init__(
    self,
    *,
    rank: int,
    interval: float = 1.0,
    relu_iterations: int = 5000,
    vram_to_keep: str | int = "1GiB",
    busy_threshold: int = DEFAULT_BUSY_THRESHOLD,
):
    """
    Args:
        rank (int): Local CUDA device index to occupy.
        interval (float, optional): Sleep time (seconds) between workload
            batches. Defaults to 1.0.
        relu_iterations (int, optional): Number of in-place ReLU ops per
            batch.
        vram_to_keep (int or str, optional): Amount of VRAM to keep busy,
            e.g. `"1GiB"`, `"20 GB"`, or an integer like `1000 * 1000`.
            This represents the total size of the matrix allocated to
            occupy the GPU.
        busy_threshold (int, optional): Defaults to 25. If current
            utilization (%) exceeds this threshold, or utilization is
            unavailable, the worker inserts extra sleeps to avoid hogging
            the GPU. Use -1 only to disable utilization backoff.

    """
    super().__init__(vram_to_keep=vram_to_keep, interval=interval)
    self.relu_iterations = validate_positive_integer(
        relu_iterations, "relu_iterations"
    )
    self.busy_threshold = validate_busy_threshold(busy_threshold)
    rank = validate_rank_type(rank)
    rank = validate_visible_rank(rank, visible_torch_device_count())
    self.rank = rank
    self.device = torch.device(f"cuda:{rank}")
    self.interval = interval
    self.platform = ComputingPlatform.CUDA

    self._stop_evt: Optional[threading.Event] = None
    self._thread: Optional[threading.Thread] = None
    self._failure_exc: Optional[Exception] = None
    self._num_elements: Optional[int] = None

allocation_status()

Return fatal worker failure captured after startup, if any.

The reference assignment/read is thread-safe for CPython's GIL model.

Source code in src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

def allocation_status(self) -> Optional[Exception]:
    """
    Return fatal worker failure captured after startup, if any.

    The reference assignment/read is thread-safe for CPython's GIL model.
    """
    return self._failure_exc

keep()

Launch the background thread that keeps the GPU busy.

Source code in src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

def keep(self) -> None:
    """Launch the background thread that keeps the GPU busy."""
    if self._thread and self._thread.is_alive():
        if self._stop_evt is not None and self._stop_evt.is_set():
            raise RuntimeError(
                f"rank {self.rank}: previous keep thread startup did not complete"
            )
        logger.warning("rank %s: keep thread already running", self.rank)
        return

    self._failure_exc = None
    self._num_elements = int(self.vram_to_keep)
    if self._num_elements <= 0:
        raise ValueError("vram_to_keep must be positive")

    self._stop_evt = threading.Event()
    startup_evt = threading.Event()
    startup_errors: list[Exception] = []
    self._thread = threading.Thread(
        target=self._keep_loop,
        args=(startup_evt, startup_errors),
        name=f"gpu-keeper-{self.rank}",
        daemon=True,  # daemon so program can exit cleanly
    )
    try:
        self._thread.start()
    except Exception:  # noqa: BLE001 - preserve original thread-start failure
        self._thread = None
        self._stop_evt = None
        raise
    startup_timeout = 5.0
    if not startup_evt.wait(startup_timeout):
        stop_evt = self._stop_evt
        if stop_evt is not None:
            stop_evt.set()
        self._thread.join(timeout=1.0)
        if self._thread.is_alive():
            raise RuntimeError(
                f"rank {self.rank}: keep thread did not complete startup within "
                f"{startup_timeout:.1f}s"
            )
        self._thread = None
        self._stop_evt = None
        raise RuntimeError(
            f"rank {self.rank}: keep thread exited before startup completed"
        )
    if startup_errors:
        self._thread.join(timeout=1.0)
        self._thread = None
        self._stop_evt = None
        raise startup_errors[0]
    logger.info("rank %s: keep thread started", self.rank)

release()

Stop the background thread and clear CUDA cache so the memory becomes immediately available to other code.

Source code in src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

def release(self) -> None:
    """
    Stop the background thread and clear CUDA cache so the memory
    becomes immediately available to other code.
    """
    thread = self._thread
    if thread is not None and not thread.is_alive():
        stop_evt = self._stop_evt
        if (stop_evt is not None and stop_evt.is_set()) or getattr(
            self, "_failure_exc", None
        ) is not None:
            torch.cuda.empty_cache()
            self._thread = None
            self._stop_evt = None
            logger.info(
                "rank %s: previously stopping keep thread exited; cache cleared",
                self.rank,
            )
            return

    if not (thread and thread.is_alive()):
        logger.warning("rank %s: keep thread not running", self.rank)
        return

    stop_evt = self._stop_evt
    if stop_evt is None:
        raise RuntimeError(f"rank {self.rank}: stop event missing")
    assert stop_evt is not None

    stop_evt.set()
    join_timeout = max(2.0, min(float(self.interval) + 2.0, 30.0))
    thread.join(timeout=join_timeout)
    if thread.is_alive():
        logger.warning(
            "rank %s: keep thread did not stop within %.1fs",
            self.rank,
            join_timeout,
        )
        raise TimeoutError(
            f"rank {self.rank}: keep thread did not stop within {join_timeout:.1f}s"
        )
    torch.cuda.empty_cache()
    self._thread = None
    self._stop_evt = None
    logger.info("rank %s: keep thread stopped & cache cleared", self.rank)

RocmGPUController

Bases: BaseGPUController

Keep a single ROCm GPU busy by running lightweight elementwise ops in a background thread. Requires a ROCm-enabled torch build.

Source code in src/keep_gpu/single_gpu_controller/rocm_gpu_controller.py

class RocmGPUController(BaseGPUController):
    """
    Keep a single ROCm GPU busy by running lightweight elementwise ops
    in a background thread. Requires a ROCm-enabled torch build.
    """

    def __init__(
        self,
        *,
        rank: int,
        interval: float = 1.0,
        vram_to_keep: str | int = "1GiB",
        busy_threshold: int = DEFAULT_BUSY_THRESHOLD,
        iterations: int = 5000,
        max_allocation_retries: Optional[int] = None,
    ):
        super().__init__(vram_to_keep=vram_to_keep, interval=interval)
        self.busy_threshold = validate_busy_threshold(busy_threshold)
        self.iterations = validate_positive_integer(iterations, "iterations")
        self.max_allocation_retries = (
            None
            if max_allocation_retries is None
            else validate_positive_integer(
                max_allocation_retries, "max_allocation_retries"
            )
        )
        rank = validate_rank_type(rank)
        rank = validate_visible_rank(rank, visible_torch_device_count())
        self.rank = rank
        self.device = torch.device(f"cuda:{rank}")
        self._stop_evt: Optional[threading.Event] = None
        self._thread: Optional[threading.Thread] = None
        self._failure_exc: Optional[Exception] = None
        self._num_elements: Optional[int] = None
        self._rocm_smi_initialized = False

        # Lazy rocm_smi import; keep handle for reuse
        try:
            import rocm_smi  # type: ignore

            self._rocm_smi = rocm_smi
        except Exception as exc:  # pragma: no cover - env-specific
            logger.debug("rocm_smi not available: %s", exc)
            self._rocm_smi = None

    def keep(self) -> None:
        if self._thread and self._thread.is_alive():
            if self._stop_evt is not None and self._stop_evt.is_set():
                raise RuntimeError(
                    f"rank {self.rank}: previous keep thread startup did not complete"
                )
            logger.warning("rank %s: keep thread already running", self.rank)
            return
        self._failure_exc = None
        self._num_elements = int(self.vram_to_keep)
        if self._num_elements <= 0:
            raise ValueError("vram_to_keep must be positive")
        self._ensure_rocm_smi_initialized()

        self._stop_evt = threading.Event()
        startup_evt = threading.Event()
        startup_errors: list[Exception] = []
        self._thread = threading.Thread(
            target=self._keep_loop,
            args=(startup_evt, startup_errors),
            name=f"gpu-keeper-rocm-{self.rank}",
            daemon=True,
        )
        try:
            self._thread.start()
        except Exception:  # noqa: BLE001 - preserve original thread-start failure
            self._thread = None
            self._stop_evt = None
            self._shutdown_rocm_smi()
            raise
        startup_timeout = 5.0
        if not startup_evt.wait(startup_timeout):
            stop_evt = self._stop_evt
            if stop_evt is not None:
                stop_evt.set()
            self._thread.join(timeout=1.0)
            if self._thread.is_alive():
                self._shutdown_rocm_smi()
                raise RuntimeError(
                    f"rank {self.rank}: ROCm keep thread did not complete startup "
                    f"within {startup_timeout:.1f}s"
                )
            self._thread = None
            self._stop_evt = None
            self._shutdown_rocm_smi()
            raise RuntimeError(
                f"rank {self.rank}: ROCm keep thread exited before startup completed"
            )
        if startup_errors:
            self._thread.join(timeout=1.0)
            self._thread = None
            self._stop_evt = None
            self._shutdown_rocm_smi()
            raise startup_errors[0]
        logger.info("rank %s: ROCm keep thread started", self.rank)

    def _shutdown_rocm_smi(self) -> None:
        if not self._rocm_smi:
            return
        try:
            self._rocm_smi.rsmi_shut_down()
        except Exception as exc:  # noqa: BLE001  # pragma: no cover - best effort
            logger.debug("rsmi_shut_down failed: %s", exc)
        finally:
            self._rocm_smi_initialized = False

    def _ensure_rocm_smi_initialized(self) -> bool:
        if not self._rocm_smi:
            return False
        if self._rocm_smi_initialized:
            return True
        try:
            self._rocm_smi.rsmi_init()
        except Exception as exc:  # noqa: BLE001  # pragma: no cover - env-specific
            logger.debug("rsmi_init failed: %s", exc)
            return False
        self._rocm_smi_initialized = True
        return True

    def release(self) -> None:
        try:
            thread = self._thread
            if thread is not None and not thread.is_alive():
                stop_evt = self._stop_evt
                if (stop_evt is not None and stop_evt.is_set()) or getattr(
                    self, "_failure_exc", None
                ) is not None:
                    torch.cuda.empty_cache()
                    self._thread = None
                    self._stop_evt = None
                    logger.info(
                        "rank %s: previously stopping keep thread exited; cache cleared",
                        self.rank,
                    )
                    return

            if thread and thread.is_alive():
                stop_evt = self._stop_evt
                if stop_evt is None:
                    raise RuntimeError(f"rank {self.rank}: stop event missing")
                assert stop_evt is not None
                stop_evt.set()
                join_timeout = max(2.0, min(float(self.interval) + 2.0, 30.0))
                thread.join(timeout=join_timeout)
                if thread.is_alive():
                    logger.warning(
                        "rank %s: ROCm keep thread did not stop within %.1fs",
                        self.rank,
                        join_timeout,
                    )
                    raise TimeoutError(
                        f"rank {self.rank}: ROCm keep thread did not stop within {join_timeout:.1f}s"
                    )
                torch.cuda.empty_cache()
                self._thread = None
                self._stop_evt = None
            else:
                logger.warning("rank %s: keep thread not running", self.rank)
                return
        finally:
            self._shutdown_rocm_smi()
        logger.info("rank %s: keep thread stopped & cache cleared", self.rank)

    def __enter__(self):
        self.keep()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.release()

    def _query_utilization(self) -> Optional[int]:
        for attempt in range(2):
            if not self._ensure_rocm_smi_initialized():
                return None
            try:
                smi_index = resolve_rocm_visible_rank_to_smi_index(
                    self.rank,
                    self._rocm_smi,
                )
                if smi_index is None:
                    return None
                util = self._rocm_smi.rsmi_dev_busy_percent_get(smi_index)
                utilization = normalize_utilization_percent(util)
                return int(utilization) if utilization is not None else None
            except Exception as exc:  # noqa: BLE001  # pragma: no cover - env-specific
                logger.debug("ROCm utilization query failed: %s", exc)
                if attempt == 0:
                    self._rocm_smi_initialized = False
                    continue
                return None
        return None

    def _keep_loop(
        self,
        startup_evt: Optional[threading.Event] = None,
        startup_errors: Optional[list[Exception]] = None,
    ) -> None:
        startup_confirmed = startup_evt is None

        def confirm_startup() -> None:
            nonlocal startup_confirmed
            if startup_confirmed:
                return
            startup_confirmed = True
            assert startup_evt is not None
            startup_evt.set()

        def record_startup_failure(exc: Exception) -> None:
            if not startup_confirmed and startup_errors is not None:
                startup_errors.append(exc)
            else:
                self._failure_exc = exc
            confirm_startup()

        def record_worker_failure(exc: Exception) -> None:
            failure = RuntimeError(
                f"rank {self.rank}: unexpected ROCm keep worker failure: {exc}"
            )
            record_startup_failure(failure)
            logger.exception("%s", failure)

        stop_evt = self._stop_evt
        if stop_evt is None:
            exc = RuntimeError(f"rank {self.rank}: stop event not initialized")
            logger.error("%s", exc)
            record_startup_failure(exc)
            return
        assert stop_evt is not None

        try:
            torch.cuda.set_device(self.rank)
        except Exception as exc:  # noqa: BLE001 - surface backend startup failure
            logger.error("rank %s: ROCm startup failed: %s", self.rank, exc)
            record_startup_failure(exc)
            return
        tensor = None
        attempts = 0
        num_elements = self._num_elements if self._num_elements is not None else 0
        if num_elements <= 0:
            exc = RuntimeError(
                f"rank {self.rank}: invalid vram_to_keep={self.vram_to_keep}"
            )
            logger.error("%s", exc)
            record_startup_failure(exc)
            return
        while not stop_evt.is_set():
            try:
                util = self._query_utilization()
                if not self._should_run_batch(util, self.busy_threshold):
                    confirm_startup()
                    logger.debug(
                        "rank %s: GPU utilization unavailable or busy (%s), deferring allocation",
                        self.rank,
                        "n/a" if util is None else f"{util}%",
                    )
                    if stop_evt.wait(self.interval):
                        return
                    continue
                tensor = torch.rand(
                    num_elements,
                    device=self.device,
                    dtype=torch.float32,
                    requires_grad=False,
                )
                confirm_startup()
                break
            except RuntimeError as exc:
                attempts += 1
                logger.error(
                    "rank %s: failed to allocate tensor (attempt %d%s): %s",
                    self.rank,
                    attempts,
                    (
                        f"/{self.max_allocation_retries}"
                        if self.max_allocation_retries is not None
                        else ""
                    ),
                    exc,
                )
                if "out of memory" not in str(exc).lower():
                    record_worker_failure(exc)
                    return
                if (
                    self.max_allocation_retries is not None
                    and attempts >= self.max_allocation_retries
                ):
                    failure = RuntimeError(
                        f"rank {self.rank}: failed to allocate tensor after {attempts} attempts"
                    )
                    record_startup_failure(failure)
                    logger.error("%s", failure)
                    return
                torch.cuda.empty_cache()
                confirm_startup()
                if stop_evt.wait(self.interval):
                    return
            except Exception as exc:
                record_worker_failure(exc)
                return

        if tensor is None:
            if not startup_confirmed:
                exc = RuntimeError(
                    f"rank {self.rank}: stopped before ROCm startup allocation"
                )
                logger.error("%s", exc)
                record_startup_failure(exc)
            else:
                logger.debug(
                    "rank %s: exiting before ROCm allocation after startup confirmation",
                    self.rank,
                )
            return
        assert tensor is not None

        while not stop_evt.is_set():
            try:
                util = self._query_utilization()
                if self._should_run_batch(util, self.busy_threshold):
                    self._run_batch(tensor)
                else:
                    logger.debug(
                        "rank %s: GPU utilization unavailable or busy (%s), sleeping",
                        self.rank,
                        "n/a" if util is None else f"{util}%",
                    )
                if stop_evt.wait(self.interval):
                    break
            except RuntimeError as exc:
                if "out of memory" in str(exc).lower():
                    torch.cuda.empty_cache()
                    if stop_evt.wait(self.interval):
                        break
                    continue
                self._failure_exc = RuntimeError(
                    f"rank {self.rank}: unexpected ROCm keep worker failure: {exc}"
                )
                logger.exception("%s", self._failure_exc)
                return
            except Exception as exc:
                self._failure_exc = RuntimeError(
                    f"rank {self.rank}: unexpected ROCm keep worker failure: {exc}"
                )
                logger.exception("%s", self._failure_exc)
                return

    def allocation_status(self) -> Optional[Exception]:
        """
        Return allocation failure captured in the worker thread, if any.

        The reference assignment/read is thread-safe for CPython's GIL model.
        """
        return self._failure_exc

    @torch.no_grad()
    def _run_batch(self, tensor: torch.Tensor) -> None:
        stop_evt = self._stop_evt
        tic = time.monotonic()
        for _ in range(self.iterations):
            torch.relu_(tensor)
            if stop_evt is not None and stop_evt.is_set():
                break
        torch.cuda.synchronize()
        toc = time.monotonic()
        logger.debug(
            "rank %s: elementwise batch done - avg %.2f ms",
            self.rank,
            (toc - tic) * 1000 / max(1, self.iterations),
        )

allocation_status()

Return allocation failure captured in the worker thread, if any.

The reference assignment/read is thread-safe for CPython's GIL model.

Source code in src/keep_gpu/single_gpu_controller/rocm_gpu_controller.py

def allocation_status(self) -> Optional[Exception]:
    """
    Return allocation failure captured in the worker thread, if any.

    The reference assignment/read is thread-safe for CPython's GIL model.
    """
    return self._failure_exc

MPSBackendUnavailableError

Bases: RuntimeError

PyTorch MPS cannot be used for a Mac M keep session.

Source code in src/keep_gpu/single_gpu_controller/macm_gpu_controller.py

class MPSBackendUnavailableError(RuntimeError):
    """PyTorch MPS cannot be used for a Mac M keep session."""

MacMGPUController

Bases: BaseGPUController

Source code in src/keep_gpu/single_gpu_controller/macm_gpu_controller.py

class MacMGPUController(BaseGPUController):
    def __init__(
        self,
        *,
        rank: int = 0,
        interval: float = 1.0,
        vram_to_keep: str | int = "1GiB",
        busy_threshold: int = DEFAULT_BUSY_THRESHOLD,
        iterations: int = 5000,
    ):
        super().__init__(vram_to_keep=vram_to_keep, interval=interval)
        busy_threshold = validate_busy_threshold(busy_threshold)
        iterations = validate_positive_integer(iterations, "iterations")
        rank = validate_visible_rank(rank, 1)
        try:
            mps_available = torch.backends.mps.is_available()
        except Exception as exc:
            raise MPSBackendUnavailableError(
                f"PyTorch MPS backend availability check failed: {exc}"
            ) from exc
        if not mps_available:
            raise MPSBackendUnavailableError("PyTorch MPS backend is not available")

        self.rank = rank
        self.device = torch.device("mps")
        self.busy_threshold = busy_threshold
        self.iterations = iterations
        self.platform = ComputingPlatform.MACM

        self._stop_evt: Optional[threading.Event] = None
        self._thread: Optional[threading.Thread] = None
        self._failure_exc: Optional[Exception] = None
        self._num_elements: Optional[int] = None

    def keep(self) -> None:
        if self._thread and self._thread.is_alive():
            if self._stop_evt is not None and self._stop_evt.is_set():
                raise RuntimeError(
                    f"rank {self.rank}: previous keep thread is still stopping"
                )
            logger.warning("rank %s: keep thread already running", self.rank)
            return

        self._failure_exc = None
        self._num_elements = int(self.vram_to_keep)
        if self._num_elements <= 0:
            raise ValueError("vram_to_keep must be positive")

        self._stop_evt = threading.Event()
        startup_evt = threading.Event()
        startup_errors: list[Exception] = []
        self._thread = threading.Thread(
            target=self._keep_loop,
            args=(startup_evt, startup_errors),
            name=f"gpu-keeper-macm-{self.rank}",
            daemon=True,
        )
        try:
            self._thread.start()
        except Exception:  # noqa: BLE001 - preserve original thread-start failure
            self._thread = None
            self._stop_evt = None
            raise
        startup_timeout = 5.0
        if not startup_evt.wait(startup_timeout):
            stop_evt = self._stop_evt
            if stop_evt is not None:
                stop_evt.set()
            self._thread.join(timeout=1.0)
            if self._thread.is_alive():
                raise RuntimeError(
                    f"rank {self.rank}: MPS keep thread did not complete startup "
                    f"within {startup_timeout:.1f}s"
                )
            self._thread = None
            self._stop_evt = None
            raise RuntimeError(
                f"rank {self.rank}: MPS keep thread exited before startup completed"
            )
        if startup_errors:
            self._thread.join(timeout=1.0)
            self._thread = None
            self._stop_evt = None
            raise startup_errors[0]
        logger.info("rank %s: MPS keep thread started", self.rank)

    def release(self) -> None:
        thread = self._thread
        if thread is not None and not thread.is_alive():
            stop_evt = self._stop_evt
            if (stop_evt is not None and stop_evt.is_set()) or getattr(
                self, "_failure_exc", None
            ) is not None:
                torch.mps.empty_cache()
                gc.collect()
                self._thread = None
                self._stop_evt = None
                logger.info(
                    "rank %s: previously stopping keep thread exited; cache cleared",
                    self.rank,
                )
                return

        if not (thread and thread.is_alive()):
            logger.warning("rank %s: keep thread not running", self.rank)
            return

        stop_evt = self._stop_evt
        if stop_evt is None:
            raise RuntimeError(f"rank {self.rank}: stop event missing")

        stop_evt.set()
        join_timeout = max(2.0, min(float(self.interval) + 2.0, 30.0))
        thread.join(timeout=join_timeout)
        if thread.is_alive():
            logger.warning(
                "rank %s: MPS keep thread did not stop within %.1fs",
                self.rank,
                join_timeout,
            )
            raise TimeoutError(
                f"rank {self.rank}: MPS keep thread did not stop within {join_timeout:.1f}s"
            )

        torch.mps.empty_cache()
        gc.collect()
        self._thread = None
        self._stop_evt = None
        logger.info("rank %s: keep thread stopped & cache cleared", self.rank)

    def __enter__(self):
        self.keep()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.release()

    def _keep_loop(
        self,
        startup_evt: Optional[threading.Event] = None,
        startup_errors: Optional[list[Exception]] = None,
    ) -> None:
        stop_evt = self._stop_evt
        if stop_evt is None:
            exc = RuntimeError(f"rank {self.rank}: stop event not initialized")
            logger.error("%s", exc)
            if startup_errors is not None:
                startup_errors.append(exc)
            else:
                self._failure_exc = exc
            if startup_evt is not None:
                startup_evt.set()
            return

        num_elements = self._num_elements if self._num_elements is not None else 0
        if num_elements <= 0:
            exc = RuntimeError(
                f"rank {self.rank}: invalid vram_to_keep={self.vram_to_keep}"
            )
            logger.error("%s", exc)
            if startup_errors is not None:
                startup_errors.append(exc)
            else:
                self._failure_exc = exc
            if startup_evt is not None:
                startup_evt.set()
            return

        tensor = None
        while not stop_evt.is_set():
            try:
                if not self._should_run_batch(None, self.busy_threshold):
                    logger.debug(
                        "rank %s: MPS utilization unavailable; deferring allocation because busy_threshold=%s",
                        self.rank,
                        self.busy_threshold,
                    )
                    if startup_evt is not None:
                        startup_evt.set()
                    if stop_evt.wait(self.interval):
                        return
                    continue
                tensor = torch.rand(
                    num_elements,
                    device=self.device,
                    dtype=torch.float32,
                    requires_grad=False,
                )
                if startup_evt is not None:
                    startup_evt.set()
                break
            except RuntimeError as exc:
                logger.error("rank %s: failed to allocate tensor: %s", self.rank, exc)
                if "out of memory" in str(exc).lower():
                    torch.mps.empty_cache()
                    gc.collect()
                    if startup_evt is not None:
                        startup_evt.set()
                    if stop_evt.wait(self.interval):
                        return
                    continue
                failure = RuntimeError(
                    f"rank {self.rank}: unexpected MPS keep worker failure: {exc}"
                )
                if startup_evt is not None and not startup_evt.is_set():
                    if startup_errors is not None:
                        startup_errors.append(failure)
                    else:
                        self._failure_exc = failure
                    startup_evt.set()
                else:
                    self._failure_exc = failure
                logger.exception("%s", failure)
                return
            except Exception as exc:
                failure = RuntimeError(
                    f"rank {self.rank}: unexpected MPS keep worker failure: {exc}"
                )
                if startup_evt is not None and not startup_evt.is_set():
                    if startup_errors is not None:
                        startup_errors.append(failure)
                    else:
                        self._failure_exc = failure
                    startup_evt.set()
                else:
                    self._failure_exc = failure
                logger.exception("%s", failure)
                return

        if tensor is None:
            if startup_evt is not None and not startup_evt.is_set():
                exc = RuntimeError(
                    f"rank {self.rank}: stopped before MPS startup allocation"
                )
                logger.error("%s", exc)
                if startup_errors is not None:
                    startup_errors.append(exc)
                else:
                    self._failure_exc = exc
                startup_evt.set()
            else:
                logger.debug(
                    "rank %s: exiting before MPS allocation after startup confirmation",
                    self.rank,
                )
            return

        while not stop_evt.is_set():
            try:
                if self._should_run_batch(None, self.busy_threshold):
                    self._run_batch(tensor)
                else:
                    logger.debug(
                        "rank %s: MPS utilization unavailable; sleeping because busy_threshold=%s",
                        self.rank,
                        self.busy_threshold,
                    )
                if stop_evt.wait(self.interval):
                    break
            except RuntimeError as exc:
                if "out of memory" in str(exc).lower():
                    torch.mps.empty_cache()
                    gc.collect()
                    if stop_evt.wait(self.interval):
                        break
                    continue
                self._failure_exc = RuntimeError(
                    f"rank {self.rank}: unexpected MPS keep worker failure: {exc}"
                )
                logger.exception("%s", self._failure_exc)
                return
            except Exception as exc:
                self._failure_exc = RuntimeError(
                    f"rank {self.rank}: unexpected MPS keep worker failure: {exc}"
                )
                logger.exception("%s", self._failure_exc)
                return

    def allocation_status(self) -> Optional[Exception]:
        """
        Return fatal worker failure captured after startup, if any.

        The reference assignment/read is thread-safe for CPython's GIL model.
        """
        return self._failure_exc

    @torch.no_grad()
    def _run_batch(self, tensor: torch.Tensor) -> None:
        stop_evt = self._stop_evt

        tic = time.monotonic()
        for _ in range(self.iterations):
            torch.relu_(tensor)
            if stop_evt is not None and stop_evt.is_set():
                break
        torch.mps.synchronize()
        toc = time.monotonic()

        logger.debug(
            "rank %s: elementwise batch done - avg %.2f ms",
            self.rank,
            (toc - tic) * 1000 / max(1, self.iterations),
        )

allocation_status()

Return fatal worker failure captured after startup, if any.

The reference assignment/read is thread-safe for CPython's GIL model.

Source code in src/keep_gpu/single_gpu_controller/macm_gpu_controller.py

def allocation_status(self) -> Optional[Exception]:
    """
    Return fatal worker failure captured after startup, if any.

    The reference assignment/read is thread-safe for CPython's GIL model.
    """
    return self._failure_exc

ControllerStartupUnavailable

Bases: Exception

Expected hardware/platform unavailability during controller startup.

Source code in src/keep_gpu/global_gpu_controller/global_gpu_controller.py

class ControllerStartupUnavailable(Exception):
    """Expected hardware/platform unavailability during controller startup."""

GlobalGPUController

Source code in src/keep_gpu/global_gpu_controller/global_gpu_controller.py

class GlobalGPUController:
    def __init__(
        self,
        gpu_ids: Optional[List[int]] = None,
        interval: Union[int, float] = 300,
        vram_to_keep: Union[int, str] = "1GiB",
        busy_threshold: int = DEFAULT_BUSY_THRESHOLD,
    ):
        self.interval = validate_interval(interval)
        self.busy_threshold = validate_busy_threshold(busy_threshold)
        parse_vram_to_elements(vram_to_keep)
        self.vram_to_keep = vram_to_keep
        gpu_ids = validate_gpu_ids(gpu_ids)
        self.computing_platform = get_platform()
        controller_unavailable_errors = (
            VisibleRankValidationError,
            DeviceEnumerationUnavailableError,
        )
        if self.computing_platform == ComputingPlatform.CUDA:
            from keep_gpu.single_gpu_controller.cuda_gpu_controller import (
                CudaGPUController,
            )

            controller_cls = CudaGPUController
        elif self.computing_platform == ComputingPlatform.ROCM:
            from keep_gpu.single_gpu_controller.rocm_gpu_controller import (
                RocmGPUController,
            )

            controller_cls = RocmGPUController
        elif self.computing_platform == ComputingPlatform.MACM:
            from keep_gpu.single_gpu_controller.macm_gpu_controller import (
                MacMGPUController,
                MPSBackendUnavailableError,
            )

            controller_cls = MacMGPUController
            controller_unavailable_errors = (
                *controller_unavailable_errors,
                MPSBackendUnavailableError,
            )
        else:
            raise UnsupportedControllerPlatformError(
                f"GlobalGPUController not implemented for platform {self.computing_platform}"
            )

        if self.computing_platform == ComputingPlatform.MACM:
            if gpu_ids is None:
                self.gpu_ids = [0]
            elif gpu_ids == [0]:
                self.gpu_ids = gpu_ids
            else:
                raise InvalidVisibleGPUSelectionError(
                    f"MACM platform only supports gpu_ids=[0] or None, got {gpu_ids}"
                )
        elif self.computing_platform in (
            ComputingPlatform.CUDA,
            ComputingPlatform.ROCM,
        ):
            self.gpu_ids = _resolve_visible_gpu_ids(gpu_ids)
        else:
            self.gpu_ids = gpu_ids

        if not self.gpu_ids:
            raise NoGPUAvailableError("No GPUs available for GlobalGPUController")

        try:
            self.controllers = [
                controller_cls(
                    rank=i,
                    interval=self.interval,
                    vram_to_keep=self.vram_to_keep,
                    busy_threshold=self.busy_threshold,
                )
                for i in self.gpu_ids
            ]
        except controller_unavailable_errors as exc:
            raise NoGPUAvailableError(str(exc)) from exc

    def keep(self) -> None:
        started = []
        for ctrl in self.controllers:
            already_running = _controller_worker_alive(ctrl)
            try:
                ctrl.keep()
            except Exception:
                if _controller_has_worker_state(ctrl):
                    try:
                        ctrl.release()
                    except Exception as cleanup_exc:
                        logger.warning(
                            "Failed to clean up failed controller rank %s after start "
                            "failure: %s",
                            getattr(ctrl, "rank", "unknown"),
                            cleanup_exc,
                        )
                for started_ctrl in reversed(started):
                    try:
                        started_ctrl.release()
                    except Exception as cleanup_exc:
                        logger.warning(
                            "Failed to roll back controller rank %s after start failure: %s",
                            getattr(started_ctrl, "rank", "unknown"),
                            cleanup_exc,
                        )
                raise
            if not already_running:
                started.append(ctrl)

    @staticmethod
    def parse_size(text: str) -> int:
        return parse_size(text)

    def release(self) -> None:
        threads = []
        errors = []
        errors_lock = threading.Lock()

        def _release_controller(ctrl) -> None:
            try:
                ctrl.release()
            except Exception as exc:
                with errors_lock:
                    errors.append((getattr(ctrl, "rank", "unknown"), exc))

        for ctrl in self.controllers:
            t = threading.Thread(target=_release_controller, args=(ctrl,))
            t.start()
            threads.append(t)
        for t in threads:
            t.join()
        if errors:
            details = "; ".join(
                (
                    str(exc)
                    if str(exc).startswith(f"rank {rank}:")
                    else f"rank {rank}: {exc}"
                )
                for rank, exc in errors
            )
            raise RuntimeError(f"Failed to release GPU controllers: {details}")

    def runtime_error(self) -> Optional[Exception]:
        """Return the first terminal child-controller runtime error, if any."""
        for ctrl in self.controllers:
            allocation_status = getattr(ctrl, "allocation_status", None)
            if not callable(allocation_status):
                continue
            try:
                error = allocation_status()
            except Exception as exc:  # noqa: BLE001 - health hook failure is health
                error = exc
            if error is None:
                continue
            rank = getattr(ctrl, "rank", "unknown")
            message = str(error)
            if message.startswith(f"rank {rank}:"):
                return error
            return RuntimeError(f"rank {rank}: {message}")
        return None

    def __enter__(self) -> "GlobalGPUController":
        self.keep()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.release()

runtime_error()

Return the first terminal child-controller runtime error, if any.

Source code in src/keep_gpu/global_gpu_controller/global_gpu_controller.py

def runtime_error(self) -> Optional[Exception]:
    """Return the first terminal child-controller runtime error, if any."""
    for ctrl in self.controllers:
        allocation_status = getattr(ctrl, "allocation_status", None)
        if not callable(allocation_status):
            continue
        try:
            error = allocation_status()
        except Exception as exc:  # noqa: BLE001 - health hook failure is health
            error = exc
        if error is None:
            continue
        rank = getattr(ctrl, "rank", "unknown")
        message = str(error)
        if message.startswith(f"rank {rank}:"):
            return error
        return RuntimeError(f"rank {rank}: {message}")
    return None

InvalidVisibleGPUSelectionError

Bases: ValueError

Public gpu_ids selection is not valid for the visible device set.

Source code in src/keep_gpu/global_gpu_controller/global_gpu_controller.py

class InvalidVisibleGPUSelectionError(ValueError):
    """Public gpu_ids selection is not valid for the visible device set."""

NoGPUAvailableError

Bases: ControllerStartupUnavailable, ValueError

No visible GPU devices are available for a global keep session.

Source code in src/keep_gpu/global_gpu_controller/global_gpu_controller.py

class NoGPUAvailableError(ControllerStartupUnavailable, ValueError):
    """No visible GPU devices are available for a global keep session."""

UnsupportedControllerPlatformError

Bases: ControllerStartupUnavailable, NotImplementedError

The current platform cannot run the global GPU controller.

Source code in src/keep_gpu/global_gpu_controller/global_gpu_controller.py

class UnsupportedControllerPlatformError(
    ControllerStartupUnavailable, NotImplementedError
):
    """The current platform cannot run the global GPU controller."""

Utilities for KeepGPU project.