| from __future__ import annotations |
|
|
| from dataclasses import fields |
| import os |
|
|
| import torch |
|
|
| from config import V2Config |
|
|
|
|
| def resolve_device(requested: str | torch.device) -> torch.device: |
| name = str(requested).lower() |
| if name == "auto": |
| if torch.cuda.is_available(): |
| return torch.device("cuda") |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): |
| return torch.device("mps") |
| return torch.device("cpu") |
| if name.startswith("cuda"): |
| return torch.device(requested if torch.cuda.is_available() else "cpu") |
| if name == "mps": |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): |
| return torch.device("mps") |
| return torch.device("cpu") |
| return torch.device(requested) |
|
|
|
|
| def sync_device(device: torch.device) -> None: |
| if device.type == "cuda": |
| torch.cuda.synchronize(device) |
| elif device.type == "mps": |
| torch.mps.synchronize() |
|
|
|
|
| def configure_threads(cpu_threads: int | None = None) -> int: |
| count = int(cpu_threads or os.cpu_count() or 1) |
| count = max(1, count) |
| for key in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_MAX_THREADS"): |
| os.environ.setdefault(key, str(count)) |
| torch.set_num_threads(count) |
| |
| torch.set_num_interop_threads(1) |
| return count |
|
|
|
|
| def make_config(**kwargs: object) -> V2Config: |
| allowed = {field.name for field in fields(V2Config)} |
| return V2Config(**{key: value for key, value in kwargs.items() if key in allowed}) |
|
|
|
|
| def normalize_config(config: V2Config) -> V2Config: |
| defaults = V2Config() |
| values = { |
| field.name: getattr(config, field.name, getattr(defaults, field.name)) |
| for field in fields(V2Config) |
| } |
| return V2Config(**values) |
|
|
