File size: 5,147 Bytes
7b5e901 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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):
# Process exit releases the model and all CUDA allocations, even after an error.
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:
# Avoid simultaneous checkpoint loads exhausting host RAM.
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()
|