Functions
	_import_torch ()
bool	cuda_available ()
bool	mps_available ()
bool	is_rocm_build ()
str	cuda_backend_label ()
str\|None	_env_override ()
str	resolve_compute_device (*, str\|None prefer=None)
Any	default_dtype_for_device (str\|None device=None)
tuple[str, Any]	device_and_dtype (*, str\|None prefer=None)
bool	accelerator_available ()
bool	allow_cpu_load ()
None	empty_device_cache (str\|None device=None)
None	synchronize_device (str\|None device=None)
bool	is_oom_error (BaseException exc)
list[dict[str, Any]]	_cuda_gpu_entries ()
bool	_is_windows ()
str\|None	_rocm_install_hint ()
str\|None	torch_version_label ()
bool	torch_is_cpu_wheel ()
list[str]	nvidia_smi_gpu_names ()
str\|None	_nvidia_cpu_torch_hint ()
str\|None	_best_fix_hint ()
str\|None	_windows_amd_hint ()
str\|None	_darwin_cpu_hint ()
str	load_blocked_message (str model_id)
str	require_load_device (str model_id, dict[str, Any]\|None cfg=None)
dict[str, Any]	probe_backend ()
str	format_backend_summary ()
str	probe_backend_json ()

Variables
	_VALID_OVERRIDES = frozenset({"cuda", "mps", "cpu", "auto"})

Function Documentation

◆ _best_fix_hint()

str | None _best_fix_hint ( )

protected

Definition at line 242 of file device.py.

References _darwin_cpu_hint(), _nvidia_cpu_torch_hint(), _rocm_install_hint(), _windows_amd_hint(), and nvidia_smi_gpu_names().

Referenced by load_blocked_message(), and probe_backend().

◆ _cuda_gpu_entries()

list[dict[str, Any]] _cuda_gpu_entries ( )

protected

Definition at line 163 of file device.py.

References _import_torch().

Referenced by probe_backend().

◆ _darwin_cpu_hint()

str | None _darwin_cpu_hint ( )

protected

Definition at line 262 of file device.py.

References mps_available().

Referenced by _best_fix_hint().

◆ _env_override()

str | None _env_override ( )

protected

Definition at line 55 of file device.py.

Referenced by resolve_compute_device().

◆ _import_torch()

_import_torch ( )

protected

Definition at line 20 of file device.py.

Referenced by _cuda_gpu_entries(), cuda_available(), default_dtype_for_device(), empty_device_cache(), is_rocm_build(), mps_available(), synchronize_device(), and torch_version_label().

◆ _is_windows()

bool _is_windows ( )

protected

Definition at line 181 of file device.py.

Referenced by _windows_amd_hint().

◆ _nvidia_cpu_torch_hint()

str | None _nvidia_cpu_torch_hint ( )

protected

Definition at line 226 of file device.py.

References cuda_available(), nvidia_smi_gpu_names(), torch_is_cpu_wheel(), and torch_version_label().

Referenced by _best_fix_hint().

◆ _rocm_install_hint()

str | None _rocm_install_hint ( )

protected

Definition at line 185 of file device.py.

References cuda_available().

Referenced by _best_fix_hint().

◆ _windows_amd_hint()

str | None _windows_amd_hint ( )

protected

Definition at line 252 of file device.py.

References _is_windows(), cuda_available(), and mps_available().

Referenced by _best_fix_hint().

◆ accelerator_available()

bool accelerator_available ( )

Definition at line 109 of file device.py.

References resolve_compute_device().

Referenced by probe_backend().

◆ allow_cpu_load()

bool allow_cpu_load ( )

Definition at line 113 of file device.py.

Referenced by probe_backend(), and require_load_device().

◆ cuda_available()

bool cuda_available ( )

Definition at line 26 of file device.py.

References _import_torch().

Referenced by _nvidia_cpu_torch_hint(), _rocm_install_hint(), _windows_amd_hint(), probe_backend(), and resolve_compute_device().

◆ cuda_backend_label()

str cuda_backend_label ( )

PyTorch still uses device type 'cuda' for ROCm builds.

Definition at line 50 of file device.py.

References is_rocm_build().

Referenced by probe_backend().

◆ default_dtype_for_device()

Any default_dtype_for_device ( str | None device = None )

Definition at line 93 of file device.py.

References _import_torch(), and resolve_compute_device().

Referenced by device_and_dtype().

◆ device_and_dtype()

tuple[str, Any] device_and_dtype	(	*	,
		str \| None	prefer = None )

Convenience for train/eval modules that need both together.

Definition at line 103 of file device.py.

References default_dtype_for_device(), and resolve_compute_device().

◆ empty_device_cache()

None empty_device_cache ( str | None device = None )

Definition at line 122 of file device.py.

References _import_torch(), mps_available(), and resolve_compute_device().

◆ format_backend_summary()

str format_backend_summary ( )

Definition at line 352 of file device.py.

References probe_backend().

◆ is_oom_error()

bool is_oom_error ( BaseException exc )

Definition at line 150 of file device.py.

◆ is_rocm_build()

bool is_rocm_build ( )

Definition at line 42 of file device.py.

References _import_torch().

Referenced by cuda_backend_label(), and probe_backend().

◆ load_blocked_message()

str load_blocked_message ( str model_id )

Definition at line 270 of file device.py.

References _best_fix_hint(), and resolve_compute_device().

Referenced by require_load_device().

◆ mps_available()

bool mps_available ( )

Definition at line 33 of file device.py.

References _import_torch().

Referenced by _darwin_cpu_hint(), _windows_amd_hint(), empty_device_cache(), probe_backend(), resolve_compute_device(), and synchronize_device().

◆ nvidia_smi_gpu_names()

list[str] nvidia_smi_gpu_names ( )

Definition at line 208 of file device.py.

Referenced by _best_fix_hint(), _nvidia_cpu_torch_hint(), and probe_backend().

◆ probe_backend()

dict[str, Any] probe_backend ( )

JSON-friendly backend snapshot for status / diagnostics.

Definition at line 300 of file device.py.

References _best_fix_hint(), _cuda_gpu_entries(), accelerator_available(), allow_cpu_load(), cuda_available(), cuda_backend_label(), is_rocm_build(), mps_available(), nvidia_smi_gpu_names(), resolve_compute_device(), torch_is_cpu_wheel(), and torch_version_label().

Referenced by format_backend_summary(), and probe_backend_json().

◆ probe_backend_json()

str probe_backend_json ( )

Definition at line 356 of file device.py.

References probe_backend().

◆ require_load_device()

str require_load_device	(	str	model_id,
		dict[str, Any] \| None	cfg = None )

Device for aquin load model. Raises RuntimeError when only CPU and not allowed.

Definition at line 288 of file device.py.

References allow_cpu_load(), load_blocked_message(), and resolve_compute_device().

◆ resolve_compute_device()

str resolve_compute_device	(	*	,
		str \| None	prefer = None )

Pick torch device string: cuda → mps → cpu.
Override with AQUIN_DEVICE=cuda|mps|cpu|auto or `prefer`.

Definition at line 64 of file device.py.

References _env_override(), cuda_available(), and mps_available().

Referenced by accelerator_available(), default_dtype_for_device(), device_and_dtype(), empty_device_cache(), load_blocked_message(), probe_backend(), require_load_device(), and synchronize_device().

◆ synchronize_device()

None synchronize_device ( str | None device = None )

Definition at line 136 of file device.py.

References _import_torch(), mps_available(), and resolve_compute_device().

◆ torch_is_cpu_wheel()

bool torch_is_cpu_wheel ( )

Definition at line 203 of file device.py.

References torch_version_label().

Referenced by _nvidia_cpu_torch_hint(), and probe_backend().

◆ torch_version_label()

str | None torch_version_label ( )

Definition at line 196 of file device.py.

References _import_torch().

Referenced by _nvidia_cpu_torch_hint(), probe_backend(), and torch_is_cpu_wheel().

Variable Documentation

◆ _VALID_OVERRIDES

aquin.compute.device._VALID_OVERRIDES = frozenset({"cuda", "mps", "cpu", "auto"})

protected

Definition at line 17 of file device.py.

Functions

Variables

Function Documentation

◆ _best_fix_hint()

◆ _cuda_gpu_entries()

◆ _darwin_cpu_hint()

◆ _env_override()

◆ _import_torch()

◆ _is_windows()

◆ _nvidia_cpu_torch_hint()

◆ _rocm_install_hint()

◆ _windows_amd_hint()

◆ accelerator_available()

◆ allow_cpu_load()

◆ cuda_available()

◆ cuda_backend_label()

◆ default_dtype_for_device()

◆ device_and_dtype()

◆ empty_device_cache()

◆ format_backend_summary()

◆ is_oom_error()

◆ is_rocm_build()

◆ load_blocked_message()

◆ mps_available()

◆ nvidia_smi_gpu_names()

◆ probe_backend()

◆ probe_backend_json()

◆ require_load_device()

◆ resolve_compute_device()

◆ synchronize_device()

◆ torch_is_cpu_wheel()

◆ torch_version_label()

Variable Documentation

◆ _VALID_OVERRIDES