| """Spawn one D3 process per GPU so FLA/Triton mutable state is not shared.""" |
|
|
| import multiprocessing |
| import threading |
| import traceback |
| from concurrent.futures import ProcessPoolExecutor |
|
|
| from ._compat import Engine, NativeAbstention, Unsupported |
|
|
| TARGET = f"{__package__}.engine:D3Engine" |
| RUN_SEED = 20260919 |
|
|
| _ENGINE = None |
| _FAILURE = None |
| _START_LOCK = threading.Lock() |
|
|
|
|
| class WorkerError(RuntimeError): |
| def __init__(self, record): |
| super().__init__(record["error"]) |
| self.exception_name = record["exception"] |
| self.remote_traceback = record["traceback"] |
|
|
|
|
| def _invoke(target, options, method, args, kwargs): |
| global _ENGINE, _FAILURE |
| if _FAILURE is not None: |
| return False, _FAILURE |
| try: |
| if _ENGINE is None: |
| from importlib import import_module |
|
|
| if target == TARGET: |
| import torch |
|
|
| torch.manual_seed(RUN_SEED) |
| module, name = target.split(":") |
| _ENGINE = getattr(import_module(module), name)(**options) |
| if method == "info": |
| value = (_ENGINE.provenance, _ENGINE.runtime(), _ENGINE.latency) |
| else: |
| value = getattr(_ENGINE, method)(*args, **kwargs) |
| if method == "__call__": |
| _ENGINE.synchronize() |
| return True, value |
| except Exception as exc: |
| record = {"error": str(exc), "exception": type(exc).__name__, |
| "traceback": traceback.format_exc()} |
| if isinstance(exc, NativeAbstention): |
| record.update(kind="abstained", raw=exc.raw) |
| elif isinstance(exc, Unsupported): |
| record["kind"] = "unsupported" |
| if _ENGINE is None or record["exception"] in ("OutOfMemoryError", "AcceleratorError") or "device-side assert" in str(exc): |
| _FAILURE = record |
| return False, record |
|
|
|
|
| class _Child: |
| def __init__(self, options, target=TARGET): |
| self.options = options |
| self.target = target |
| self.lock = threading.Lock() |
| self.pool = ProcessPoolExecutor(max_workers=1, mp_context=multiprocessing.get_context("spawn")) |
| try: |
| self.info = self.call("info") |
| except BaseException: |
| self.pool.shutdown(wait=True, cancel_futures=True) |
| raise |
|
|
| def call(self, method, *args, **kwargs): |
| with self.lock: |
| ok, value = self.pool.submit(_invoke, self.target, self.options, method, args, kwargs).result() |
| if ok: |
| return value |
| if value.get("kind") == "unsupported": |
| raise Unsupported(value["error"]) |
| if value.get("kind") == "abstained": |
| raise NativeAbstention(value["raw"]) |
| raise WorkerError(value) |
|
|
| def close(self): |
| |
| self.pool.shutdown(wait=True, cancel_futures=True) |
|
|
|
|
| class D3ProcessEngine(Engine): |
| name = "d3" |
|
|
| def __init__(self, device=None, replicas=None, **options): |
| import torch |
|
|
| gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0 |
| if replicas is not None and int(replicas) > 1: |
| if int(replicas) > gpus: |
| raise ValueError(f"replicas={replicas} needs more than {gpus} CUDA devices") |
| self.devices = [f"cuda:{i}" for i in range(int(replicas))] |
| elif device in (None, "cuda"): |
| self.devices = [f"cuda:{i}" for i in range(gpus)] or ["cpu"] |
| else: |
| self.devices = [str(device)] |
| super().__init__(device=device, **options) |
| self._options = options |
| self._children = {} |
| self._local = threading.local() |
| self._lock = threading.Lock() |
| self._next = 0 |
| child = self._child(0) |
| provenance, runtime, self.latency = child.info |
| self.provenance = {**provenance, "gpu_execution": "one spawned process per GPU"} |
| self._runtime = runtime |
|
|
| def _child(self, slot): |
| child = self._children.get(slot) |
| if child is not None: |
| return child |
| with self._lock: |
| if slot not in self._children: |
| |
| with _START_LOCK: |
| self._children[slot] = _Child({**self._options, "device": self.devices[slot]}) |
| return self._children[slot] |
|
|
| def __call__(self, *args, **kwargs): |
| slot = getattr(self._local, "slot", None) |
| if slot is None: |
| with self._lock: |
| slot = self._next |
| self._next = (self._next + 1) % len(self.devices) |
| self._local.slot = slot |
| return self._child(slot).call("__call__", *args, **kwargs) |
|
|
| def warmup(self): |
| self._child(0).call("warmup") |
|
|
| def synchronize(self): |
| for child in list(self._children.values()): |
| child.call("synchronize") |
|
|
| def runtime(self): |
| return {**self._runtime, "gpu_execution": "spawn", "devices": self.devices} |
|
|
| def close(self): |
| for child in self._children.values(): |
| child.close() |
| self._children.clear() |
|
|
