`torchscan.process`¶

The process subpackage exposes explicit workload memory measurement. Legacy get_process_gpu_ram was removed in v0.2 because process snapshots and allocator deltas could not provide a truthful model measurement.

CPU tracked tensors, accelerator allocator peaks, process RSS, and model storage in summary() have distinct scopes. They can overlap, especially on unified memory; do not sum them or equate them to total device use.

Measure one workload's peak memory¶

measure_peak_memory runs an owner-provided callable exactly once. The callable owns the model, tensors, optimizer, gradient mode, device placement, warmup, and cache state.

measure_peak_memory ¶

measure_peak_memory(workload: Callable[[], object], *, device: str | device) -> dict[str, str | int]

Measure peak PyTorch memory used by one workload invocation.

Parameters:

Name	Type	Description	Default
`workload`	`Callable[[], object]`	Zero-argument callable to invoke exactly once.	required
`device`	`str \| device`	Device on which the workload already operates.	required

Returns:

Type	Description
`dict[str, str \| int]`	JSON-serializable backend-specific memory statistics in bytes.

Raises:

Type	Description
`NotImplementedError`	If the backend or its measurement APIs are unsupported.
`RuntimeError`	If the requested device is unavailable.

Measure peak resident process RAM¶

measure_peak_rss runs a command in a fresh child on Linux/macOS. It records the child's whole lifetime, including imports and loading. See workload diagnostics for an example and memory scope boundaries.

measure_peak_rss ¶

measure_peak_rss(command: Sequence[str], *, cwd: str | Path | None = None) -> MetricResult

Measure the OS peak resident RAM of one new child command.

Parameters:

Name	Type	Description	Default
`command`	`Sequence[str]`	Executable and arguments, passed without a shell. The child owns model loading, inputs, device placement, and the workload.	required
`cwd`	`str \| Path \| None`	Optional child working directory. Not retained in the result.	`None`

Returns:

Type	Description
`MetricResult`	Peak resident bytes for the child's whole lifetime, including imports and
`MetricResult`	model loading. This is not an inference-only delta or accelerator memory.

Raises:

Type	Description
`ValueError`	If command is empty or contains non-string arguments.
`NotImplementedError`	If per-child RSS accounting is unavailable.
`CalledProcessError`	If the command exits unsuccessfully.

Notes

Linux/macOS use per-child wait4 accounting, not the parent's cumulative child high-water mark. Standard streams are inherited. Child process trees are not summed, and container/device memory are outside this metric's scope. The command must exit; no deadline or process-tree cancellation is added.

Peak-memory workload examples¶

Inference¶

import os

import torch
from torchscan.process import measure_peak_memory

model = torch.nn.Linear(64, 16)
inputs = torch.randn(32, 64)


def run_inference() -> None:
    with torch.inference_mode():
        model(inputs)


stats = measure_peak_memory(run_inference, device=inputs.device)
owner_budget = int(os.environ["MODEL_MEMORY_BUDGET_BYTES"])
if stats["peak_bytes"] > owner_budget:
    raise RuntimeError("Owner-approved PyTorch memory budget exceeded")

Training step¶

import torch
from torchscan.process import measure_peak_memory

model = torch.nn.Linear(64, 16)
inputs = torch.randn(32, 64)
targets = torch.randn(32, 16)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)


def run_training_step() -> None:
    optimizer.zero_grad(set_to_none=True)
    loss = torch.nn.functional.mse_loss(model(inputs), targets)
    loss.backward()
    optimizer.step()


stats = measure_peak_memory(run_training_step, device=inputs.device)

Limitations¶

CPU pytorch_tensor_bytes covers PyTorch-tracked tensor and operator allocations, not Python heap or process RSS.
Accelerator pytorch_reserved_bytes uses PyTorch caching-allocator state and is not total device or driver memory.
TorchScan does not warm up, repeat, move, reset, or roll back the workload.
Unrelated concurrent allocations can affect process-global metrics.
A mock or skipped accelerator test is not hardware evidence.

torchscan.process¶