Skip to content

KeepGPU

API Reference

Wangmerlyn/KeepGPU

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.

CLI entrypoint for KeepGPU.

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

List GPU telemetry from local service.

Source code in keep_gpu/cli.py

@app.command("list-gpus")
def list_gpus(
    host: str = typer.Option(DEFAULT_SERVICE_HOST, "--host", help="Service host."),
    port: int = typer.Option(DEFAULT_SERVICE_PORT, "--port", help="Service port."),
):
    """List GPU telemetry from local service."""
    try:
        result = _rpc_call("list_gpus", {}, host, port)
        console.print_json(data=json.dumps(result))
    except RuntimeError as exc:
        console.print(f"[bold red]Error: {exc}[/bold red]")
        raise typer.Exit(code=1) from exc

main(ctx, interval=typer.Option(300, help='Interval in seconds between GPU usage checks (blocking mode only).'), gpu_ids=typer.Option(None, help='Comma-separated GPU IDs 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(-1, '--busy-threshold', '--util-threshold', help='Max utilization threshold before backing off (blocking mode).')) ¶

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

Source code in keep_gpu/cli.py

@app.callback(invoke_without_command=True)
def main(
    ctx: typer.Context,
    interval: int = typer.Option(
        300,
        help="Interval in seconds between GPU usage checks (blocking mode only).",
    ),
    gpu_ids: Optional[str] = typer.Option(
        None,
        help="Comma-separated GPU IDs 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: int = typer.Option(
        -1,
        "--busy-threshold",
        "--util-threshold",
        help="Max utilization threshold before backing off (blocking mode).",
    ),
):
    """Run blocking keep-alive mode when no subcommand is provided."""
    if ctx.invoked_subcommand is not None:
        return
    try:
        _run_blocking(interval, gpu_ids, vram, legacy_threshold, busy_threshold)
    except typer.BadParameter as exc:
        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(DEFAULT_SERVICE_PORT, help='Port for KeepGPU local service.'))` ¶

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

Source code in keep_gpu/cli.py

@app.command("serve")
def serve(
    host: str = typer.Option(
        DEFAULT_SERVICE_HOST,
        help="Host interface for KeepGPU local service.",
    ),
    port: int = typer.Option(
        DEFAULT_SERVICE_PORT,
        help="Port for KeepGPU local service.",
    ),
):
    """Run KeepGPU local service (HTTP + JSON-RPC + dashboard)."""
    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(DEFAULT_SERVICE_PORT, '--port', help='Service port.'), force=typer.Option(False, '--force', help='Stop service even if active sessions exist.'))` ¶

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

Source code in keep_gpu/cli.py

@app.command("service-stop")
def service_stop(
    host: str = typer.Option(DEFAULT_SERVICE_HOST, "--host", help="Service host."),
    port: int = typer.Option(DEFAULT_SERVICE_PORT, "--port", help="Service port."),
    force: bool = typer.Option(
        False,
        "--force",
        help="Stop service even if active sessions exist.",
    ),
):
    """Stop local KeepGPU service daemon started by auto-start logic."""
    try:
        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

        if _service_available(host, port):
            status = _rpc_call("status", {}, host, port)
            active_jobs = status.get("active_jobs", [])
            if active_jobs:
                raise RuntimeError(
                    "Active keep sessions detected. Stop sessions first (`keep-gpu stop --all`) or re-run with --force."
                )
            _rpc_call("stop_keep", {}, host, port, timeout=45.0)

        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 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 GPU IDs.'), vram=typer.Option('1GiB', '--vram', help='VRAM to keep per GPU.'), interval=typer.Option(300, help='Interval in seconds between checks.'), busy_threshold=typer.Option(-1, '--busy-threshold', '--util-threshold', help='Back off when utilization exceeds this percent.'), 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(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 keep_gpu/cli.py

@app.command("start")
def start(
    gpu_ids: Optional[str] = typer.Option(None, help="Comma-separated GPU IDs."),
    vram: str = typer.Option("1GiB", "--vram", help="VRAM to keep per GPU."),
    interval: int = typer.Option(300, help="Interval in seconds between checks."),
    busy_threshold: int = typer.Option(
        -1,
        "--busy-threshold",
        "--util-threshold",
        help="Back off when utilization exceeds this percent.",
    ),
    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: int = typer.Option(
        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.
    """
    try:
        auto_started = _ensure_service_running(host, port, auto_start=auto_start)
        result = _rpc_call(
            "start_keep",
            {
                "gpu_ids": _parse_gpu_ids(gpu_ids),
                "vram": vram,
                "interval": interval,
                "busy_threshold": busy_threshold,
                "job_id": job_id,
            },
            host,
            port,
        )
        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}/")
        console.print(
            f"[dim]Next: keep-gpu status --job-id {result['job_id']} | keep-gpu stop --job-id {result['job_id']}[/dim]"
        )
        console.print(
            "[dim]When all sessions are done, stop daemon with: keep-gpu service-stop[/dim]"
        )
    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(DEFAULT_SERVICE_PORT, '--port', help='Service port.'))` ¶

Show session status from KeepGPU local service.

Source code in 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: int = typer.Option(DEFAULT_SERVICE_PORT, "--port", help="Service port."),
):
    """Show session status from KeepGPU local service."""
    try:
        result = _rpc_call(
            "status",
            {"job_id": job_id} if job_id else {},
            host,
            port,
        )
        console.print_json(data=json.dumps(result))
    except RuntimeError as exc:
        console.print(f"[bold red]Error: {exc}[/bold red]")
        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(DEFAULT_SERVICE_PORT, '--port', help='Service port.'))` ¶

Stop one session or all sessions.

Source code in 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: int = typer.Option(DEFAULT_SERVICE_PORT, "--port", help="Service port."),
):
    """Stop one session or all sessions."""
    try:
        if not job_id and not all_sessions:
            raise RuntimeError("Provide --job-id or use --all.")
        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,
            )
        console.print_json(data=json.dumps(result))
    except RuntimeError as exc:
        console.print(f"[bold red]Error: {exc}[/bold red]")
        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 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,
        matmul_iterations: Optional[int] = None,
        vram_to_keep: str | int = "1000 MB",
        busy_threshold: int = 10,
    ):
        """
        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.
            matmul_iterations (int, optional): Legacy alias for
                `relu_iterations`. When set, it overrides `relu_iterations`.
            vram_to_keep (int or str, optional): Amount of VRAM to keep busy,
                e.g. `"1000 MB"`, `"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): If current utilisation (%) exceeds
                this threshold, the worker will insert extra sleeps to avoid
                hogging the GPU.

        """
        super().__init__(vram_to_keep=vram_to_keep, interval=interval)
        self.rank = rank
        self.device = torch.device(f"cuda:{rank}")
        self.interval = interval
        if matmul_iterations is not None:
            relu_iterations = matmul_iterations
        if relu_iterations <= 0:
            raise ValueError("relu_iterations must be positive")
        self.relu_iterations = relu_iterations
        self.busy_threshold = busy_threshold
        self.platform = ComputingPlatform.CUDA

        self._stop_evt: Optional[threading.Event] = None
        self._thread: Optional[threading.Thread] = 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():
            logger.warning("rank %s: keep thread already running", self.rank)
            return

        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()
        self._thread = threading.Thread(
            target=self._keep_loop,
            name=f"gpu-keeper-{self.rank}",
            daemon=True,  # daemon so program can exit cleanly
        )
        self._thread.start()
        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.
        """
        if not (self._thread and self._thread.is_alive()):
            logger.warning("rank %s: keep thread not running", self.rank)
            return

        stop_evt = self._stop_evt
        if stop_evt is None:
            logger.warning("rank %s: stop event missing; skipping release", self.rank)
            return
        assert stop_evt is not None

        stop_evt.set()
        join_timeout = max(2.0, min(float(self.interval) + 2.0, 30.0))
        self._thread.join(timeout=join_timeout)
        if self._thread.is_alive():
            logger.warning(
                "rank %s: keep thread did not stop within %.1fs",
                self.rank,
                join_timeout,
            )
            return
        torch.cuda.empty_cache()
        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) -> None:
        """Internal: run workloads until stop event is set."""
        stop_evt = self._stop_evt
        if stop_evt is None:
            logger.error("rank %s: stop event not initialized", self.rank)
            return
        assert stop_evt is not None

        torch.cuda.set_device(self.rank)
        num_elements = self._num_elements if self._num_elements is not None else 0
        if num_elements <= 0:
            logger.error(
                "rank %s: invalid vram_to_keep=%s", self.rank, self.vram_to_keep
            )
            return
        matrix = None
        while not stop_evt.is_set():
            try:
                matrix = torch.rand(
                    num_elements,
                    device=self.device,
                    dtype=torch.float32,
                    requires_grad=False,
                )
                break
            except RuntimeError as e:
                logger.error("rank %s: failed to allocate matrix: %s", self.rank, e)
                if stop_evt.wait(self.interval):
                    return
        if matrix is None:
            logger.error("rank %s: failed to allocate matrix, exiting loop", self.rank)
            return
        while not stop_evt.is_set():
            try:
                gpu_utilization = self._monitor_utilization(self.rank)
                if gpu_utilization > self.busy_threshold:
                    logger.debug(
                        "rank %s: GPU busy (%d%%), sleeping longer",
                        self.rank,
                        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
            except Exception:
                # Log unexpected exceptions but keep running
                logger.exception("rank %s: unexpected error", self.rank)
                if stop_evt.wait(self.interval):
                    break

    # ------------------------------------------------------------------
    # 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.time()
        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.time()

        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) -> int:
        """
        Return current GPU utilization (%) for `rank`.
        Falls back to 0 when NVML is unavailable.
        """
        utilization = get_gpu_utilization(rank)
        return utilization if utilization is not None else 0

`init(*, rank, interval=1.0, relu_iterations=5000, matmul_iterations=None, vram_to_keep='1000 MB', busy_threshold=10)` ¶

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`
`matmul_iterations`	`int`	Legacy alias for `relu_iterations`. When set, it overrides `relu_iterations`.	`None`
`vram_to_keep`	`int or str`	Amount of VRAM to keep busy, e.g. `"1000 MB"`, `"20 GB"`, or an integer like `1000 * 1000`. This represents the total size of the matrix allocated to occupy the GPU.	`'1000 MB'`
`busy_threshold`	`int`	If current utilisation (%) exceeds this threshold, the worker will insert extra sleeps to avoid hogging the GPU.	`10`

Source code in keep_gpu/single_gpu_controller/cuda_gpu_controller.py

def __init__(
    self,
    *,
    rank: int,
    interval: float = 1.0,
    relu_iterations: int = 5000,
    matmul_iterations: Optional[int] = None,
    vram_to_keep: str | int = "1000 MB",
    busy_threshold: int = 10,
):
    """
    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.
        matmul_iterations (int, optional): Legacy alias for
            `relu_iterations`. When set, it overrides `relu_iterations`.
        vram_to_keep (int or str, optional): Amount of VRAM to keep busy,
            e.g. `"1000 MB"`, `"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): If current utilisation (%) exceeds
            this threshold, the worker will insert extra sleeps to avoid
            hogging the GPU.

    """
    super().__init__(vram_to_keep=vram_to_keep, interval=interval)
    self.rank = rank
    self.device = torch.device(f"cuda:{rank}")
    self.interval = interval
    if matmul_iterations is not None:
        relu_iterations = matmul_iterations
    if relu_iterations <= 0:
        raise ValueError("relu_iterations must be positive")
    self.relu_iterations = relu_iterations
    self.busy_threshold = busy_threshold
    self.platform = ComputingPlatform.CUDA

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

`keep()` ¶

Launch the background thread that keeps the GPU busy.

Source code in 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():
        logger.warning("rank %s: keep thread already running", self.rank)
        return

    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()
    self._thread = threading.Thread(
        target=self._keep_loop,
        name=f"gpu-keeper-{self.rank}",
        daemon=True,  # daemon so program can exit cleanly
    )
    self._thread.start()
    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 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.
    """
    if not (self._thread and self._thread.is_alive()):
        logger.warning("rank %s: keep thread not running", self.rank)
        return

    stop_evt = self._stop_evt
    if stop_evt is None:
        logger.warning("rank %s: stop event missing; skipping release", self.rank)
        return
    assert stop_evt is not None

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

Utilities for KeepGPU project.