152 lines
4.7 KiB
Python
152 lines
4.7 KiB
Python
"""Credential-free hardware metadata for remote execution targets."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
if __package__:
|
|
from .execution_environments import EnvironmentCapabilities
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from execution_environments import EnvironmentCapabilities
|
|
|
|
_MODAL_GPU_VRAM_GB: dict[str, float] = {
|
|
"T4": 16,
|
|
"L4": 24,
|
|
"A10": 24,
|
|
"L40S": 48,
|
|
"A100": 40,
|
|
"A100-40GB": 40,
|
|
"A100-80GB": 80,
|
|
"RTX-PRO-6000": 96,
|
|
"H100": 80,
|
|
"H100!": 80,
|
|
"H200": 141,
|
|
"B200": 180,
|
|
"B200+": 180,
|
|
"B300": 288,
|
|
}
|
|
# Modal public pricing verified 2026-08-22; CPU and memory charges are excluded.
|
|
_MODAL_GPU_COST_USD_PER_SECOND: dict[str, float] = {
|
|
"T4": 0.000164,
|
|
"L4": 0.000222,
|
|
"A10": 0.000306,
|
|
"L40S": 0.000542,
|
|
"A100": 0.000583,
|
|
"A100-40GB": 0.000583,
|
|
"A100-80GB": 0.000694,
|
|
"RTX-PRO-6000": 0.000842,
|
|
"H100": 0.001097,
|
|
"H100!": 0.001097,
|
|
"H200": 0.001261,
|
|
"B200": 0.001736,
|
|
"B200+": 0.001736,
|
|
"B300": 0.001972,
|
|
}
|
|
_HBM_GPU_NAME_MARKERS = ("A100", "H100", "H200", "B200", "B300")
|
|
|
|
|
|
def _gpu_memory_kind(machine_type: str) -> str:
|
|
"""Return the best-known GPU memory technology for a machine label."""
|
|
normalized = machine_type.upper().replace(" ", "-")
|
|
return (
|
|
"HBM"
|
|
if any(marker in normalized for marker in _HBM_GPU_NAME_MARKERS)
|
|
else "VRAM"
|
|
)
|
|
|
|
|
|
def _hardware_payload(
|
|
*,
|
|
machine_type: str,
|
|
gpu_count: int,
|
|
gpu_memory_bytes_per_device: int,
|
|
gpu_memory_bytes_total: int,
|
|
ram_bytes: int = 0,
|
|
ram_available_bytes: int | None = None,
|
|
ram_capacity_label: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Return compact credential-free target hardware metadata."""
|
|
payload: dict[str, Any] = {
|
|
"machine_type": machine_type or "Unknown GPU",
|
|
"gpu_count": max(0, gpu_count),
|
|
"gpu_memory_kind": _gpu_memory_kind(machine_type),
|
|
"gpu_memory_bytes_per_device": max(0, gpu_memory_bytes_per_device),
|
|
"gpu_memory_bytes_total": max(0, gpu_memory_bytes_total),
|
|
}
|
|
if ram_bytes > 0:
|
|
payload["ram_bytes"] = ram_bytes
|
|
if ram_available_bytes is not None and ram_available_bytes >= 0:
|
|
payload["ram_available_bytes"] = ram_available_bytes
|
|
if ram_capacity_label:
|
|
payload["ram_capacity_label"] = ram_capacity_label
|
|
return payload
|
|
|
|
|
|
def _capabilities_hardware_payload(
|
|
capabilities: EnvironmentCapabilities | None,
|
|
) -> dict[str, Any] | None:
|
|
"""Summarize a probed host's GPUs and system memory for the UI."""
|
|
if capabilities is None:
|
|
return None
|
|
gpu_names = tuple(dict.fromkeys(gpu.name for gpu in capabilities.gpus))
|
|
machine_type = " + ".join(gpu_names) if gpu_names else "CPU-only"
|
|
total_gpu_memory = sum(gpu.total_vram_bytes for gpu in capabilities.gpus)
|
|
per_device_memory = max(
|
|
(gpu.total_vram_bytes for gpu in capabilities.gpus),
|
|
default=0,
|
|
)
|
|
return _hardware_payload(
|
|
machine_type=machine_type,
|
|
gpu_count=len(capabilities.gpus),
|
|
gpu_memory_bytes_per_device=per_device_memory,
|
|
gpu_memory_bytes_total=total_gpu_memory,
|
|
ram_bytes=capabilities.total_ram_bytes,
|
|
ram_available_bytes=capabilities.available_ram_bytes,
|
|
)
|
|
|
|
|
|
def _modal_hardware_payload(gpu_type: str) -> dict[str, Any]:
|
|
"""Summarize known Modal GPU capacity without guessing dynamic host RAM."""
|
|
gpu_memory_bytes = int(_MODAL_GPU_VRAM_GB.get(gpu_type, 0.0) * 1024**3)
|
|
return _hardware_payload(
|
|
machine_type=gpu_type,
|
|
gpu_count=1,
|
|
gpu_memory_bytes_per_device=gpu_memory_bytes,
|
|
gpu_memory_bytes_total=gpu_memory_bytes,
|
|
ram_capacity_label="Provider managed",
|
|
)
|
|
|
|
|
|
def _vast_hardware_payload(resource: Any) -> dict[str, Any]:
|
|
"""Summarize a Vast offer or lease using its advertised capacities."""
|
|
gpu_count = max(
|
|
0,
|
|
int(
|
|
getattr(resource, "gpu_count", None)
|
|
or getattr(resource, "num_gpus", 0)
|
|
or 0
|
|
),
|
|
)
|
|
per_device_memory = max(
|
|
0,
|
|
int(getattr(resource, "gpu_ram_mb", 0) or 0) * 1024**2,
|
|
)
|
|
advertised_total_mb = int(getattr(resource, "gpu_total_ram_mb", 0) or 0)
|
|
total_gpu_memory = (
|
|
advertised_total_mb * 1024**2
|
|
if advertised_total_mb > 0
|
|
else per_device_memory * max(gpu_count, 1)
|
|
)
|
|
return _hardware_payload(
|
|
machine_type=str(
|
|
getattr(resource, "gpu_name", "Unknown GPU") or "Unknown GPU"
|
|
),
|
|
gpu_count=gpu_count,
|
|
gpu_memory_bytes_per_device=per_device_memory,
|
|
gpu_memory_bytes_total=total_gpu_memory,
|
|
ram_bytes=max(
|
|
0,
|
|
int(getattr(resource, "cpu_ram_mb", 0) or 0) * 1024**2,
|
|
),
|
|
)
|