Snakepit.Hardware (Snakepit v0.8.7)

View Source

Hardware abstraction layer for Snakepit.

Provides unified hardware detection and device selection for ML workloads. Supports CPU, NVIDIA CUDA, Apple MPS, and AMD ROCm accelerators.

Features

  • Automatic Detection: Detects available hardware at startup
  • Device Selection: Intelligent device selection with fallback strategies
  • Caching: Results are cached for performance
  • Lock File Support: Identity map for lock file generation

Usage

# Detect all hardware
info = Snakepit.Hardware.detect()
# => %{accelerator: :cuda, cpu: %{...}, cuda: %{...}, ...}

# Check capabilities
caps = Snakepit.Hardware.capabilities()
# => %{cuda: true, mps: false, avx2: true, ...}

# Select device
{:ok, device} = Snakepit.Hardware.select(:auto)
# => {:ok, {:cuda, 0}}

# Select with fallback
{:ok, device} = Snakepit.Hardware.select_with_fallback([:cuda, :mps, :cpu])
# => {:ok, :cpu}

Identity Map

The identity/0 function returns a map suitable for lock file generation:

identity = Snakepit.Hardware.identity()
# => %{"platform" => "linux-x86_64", "accelerator" => "cuda", ...}

This can be serialized to JSON/YAML for lock files that need to track the hardware environment.

Summary

Types

capabilities()
device()
device_preference()
hardware_info()

Functions

capabilities()

Returns hardware capability flags.

clear_cache()

Clears the hardware detection cache.

detect()

Detects all hardware information.

device_info(device)

Returns device information for a selected device.

identity()

Returns a hardware identity map for lock files.

info()

Alias for detect/0.

select(preference)

Selects a device based on preference.

select_with_fallback(preferences)

Selects the first available device from a preference list.

Types

capabilities()

@type capabilities() :: Snakepit.Hardware.Detector.capabilities()

device()

@type device() :: Snakepit.Hardware.Selector.device()

device_preference()

@type device_preference() :: Snakepit.Hardware.Selector.device_preference()

hardware_info()

@type hardware_info() :: Snakepit.Hardware.Detector.hardware_info()

Functions

capabilities()

@spec capabilities() :: capabilities()

Returns hardware capability flags.

Returns a map of boolean flags for quick feature checks:

  • :cuda - CUDA available
  • :mps - Apple MPS available
  • :rocm - AMD ROCm available
  • :avx - AVX instruction set available
  • :avx2 - AVX2 instruction set available
  • :avx512 - AVX-512 instruction set available
  • :cuda_version - CUDA version string or nil
  • :cudnn_version - cuDNN version string or nil
  • :cudnn - cuDNN available

Examples

caps = Snakepit.Hardware.capabilities()
if caps.cuda do
  cuda_version = caps.cuda_version
end

clear_cache()

@spec clear_cache() :: :ok

Clears the hardware detection cache.

Forces re-detection on next call. Useful after hardware changes or for testing.

Examples

Snakepit.Hardware.clear_cache()
:ok

detect()

@spec detect() :: hardware_info()

Detects all hardware information.

Returns a map with:

  • :accelerator - Primary accelerator type (:cpu, :cuda, :mps, :rocm)
  • :cpu - CPU information (cores, threads, model, features, memory)
  • :cuda - NVIDIA CUDA info or nil
  • :mps - Apple MPS info or nil
  • :rocm - AMD ROCm info or nil
  • :platform - Platform string (e.g., "linux-x86_64")

Examples

info = Snakepit.Hardware.detect()
info.accelerator
#=> :cuda

info.cpu.cores
#=> 8

device_info(device)

@spec device_info(device()) :: map()

Returns device information for a selected device.

Returns a map with device-specific details useful for logging, telemetry, and diagnostics.

Examples

info = Snakepit.Hardware.device_info({:cuda, 0})
# => %{type: :cuda, device_id: 0, name: "NVIDIA GeForce RTX 3080", ...}

identity()

@spec identity() :: map()

Returns a hardware identity map for lock files.

The identity map contains string keys and is suitable for serialization to JSON/YAML lock files that need to track the hardware environment.

Keys

  • "platform" - Platform string (e.g., "linux-x86_64")
  • "accelerator" - Primary accelerator type as string
  • "cpu_features" - List of CPU feature strings
  • "gpu_count" - Number of GPUs detected

Examples

identity = Snakepit.Hardware.identity()
Jason.encode!(identity)
# => "{\"platform\":\"linux-x86_64\",\"accelerator\":\"cuda\",...}"

info()

@spec info() :: hardware_info()

Alias for detect/0.

Returns the same hardware info map as detect/0.

select(preference)

@spec select(device_preference()) :: {:ok, device()} | {:error, :device_not_available}

Selects a device based on preference.

Options

  • :auto - Automatically select best available accelerator
  • :cpu - Select CPU (always available)
  • :cuda - Select CUDA (fails if not available)
  • :mps - Select MPS (fails if not macOS with Apple Silicon)
  • :rocm - Select ROCm (fails if not available)
  • {:cuda, device_id} - Select specific CUDA device by ID

Returns

  • {:ok, device} on success
  • {:error, :device_not_available} if requested device is unavailable

Examples

# Auto-select best device
{:ok, device} = Snakepit.Hardware.select(:auto)

# Request specific device
case Snakepit.Hardware.select(:cuda) do
  {:ok, {:cuda, 0}} -> :ok
  {:error, :device_not_available} -> :error
end

select_with_fallback(preferences)

@spec select_with_fallback([device_preference()]) ::
  {:ok, device()} | {:error, :no_device}

Selects the first available device from a preference list.

Tries each device in order until one is available. This is useful for graceful degradation strategies.

Examples

# Prefer CUDA, fall back to MPS, then CPU
{:ok, device} = Snakepit.Hardware.select_with_fallback([:cuda, :mps, :cpu])

# Returns :cpu if CUDA and MPS are unavailable