`nexuml.data.auto_batch`¶

nexuml.data.auto_batch ¶

Automatic CUDA batch-size resolution helpers.

Result of a single batch-size probe attempt.

Aggregated result of an automatic batch-size probe run.

generate_candidates(config: AutoBatchSizeSpec) -> list[int]

Generate bounded power-of-two candidates in ascending order.

Returns:

Type	Description
`list[int]`	Sorted list of candidate batch sizes.

Raises:

Type	Description
`ValueError`	If the candidate strategy is not `power_of_two`.

select_batch_size(
    successes: list[int], config: AutoBatchSizeSpec
) -> int

Select the effective batch size from successful candidates.

Returns:

Type	Description
`int`	The selected batch size.

Raises:

Type	Description
`ValueError`	If `successes` is empty or the safety policy is unsupported.

is_cuda_oom(exc: BaseException) -> bool

Return True if the exception is a CUDA out-of-memory error.

cuda_device_info() -> dict[str, Any]

Return CUDA device metadata, or {"available": False} if no GPU is present.

resolve_with_probe(
    config: AutoBatchSizeSpec, probe: Callable[[int], None]
) -> BatchProbeResult

Try candidate batch sizes and return selected size plus observable metadata.

Returns:

Type	Description
`BatchProbeResult`	Aggregated probe result with selected batch size.

Raises:

Type	Description
`RuntimeError`	If no candidate succeeds because candidates hit CUDA OOM or exceed the configured margin policy.