| """Production loading of the p2p extension from a prebuilt, pinned build (S10 P-1; R2-RUNTIME). |
| |
| The development path (p2p.ext() without DSPARK_P2P_PREBUILT) compiles p2p.cu with torch.utils.cpp_extension at the first |
| use. A production unit must not run the compiler at start, nor pick up a build it was not qualified with. With |
| DSPARK_P2P_PREBUILT=<dir> the worker instead loads <dir>/<src16>/dspark_p2p.so, where <src16> is the first 16 hex digits |
| of the sha256 of the p2p.cu it runs, and only when the file's sha256 equals <dir>/<src16>/dspark_p2p.so.sha256 (written |
| by run/prod_tp2.sh build-p2p when it built that source). Anything else is a refusal naming the fix, never a fallback |
| to compiling. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import importlib.machinery |
| import importlib.util |
| import os |
| import sys |
|
|
| SO_NAME = "dspark_p2p.so" |
| MODULE = "dspark_p2p" |
|
|
|
|
| class PrebuiltError(RuntimeError): |
| pass |
|
|
|
|
| def _sha256(path: str) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 20), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def source_key(src: str) -> str: |
| return _sha256(src)[:16] |
|
|
|
|
| def prebuilt_path(root: str, src: str) -> str: |
| return os.path.join(os.path.expanduser(root), source_key(src), SO_NAME) |
|
|
|
|
| def load_prebuilt(root: str, src: str): |
| """The extension module from the pinned build of this p2p.cu under root; PrebuiltError otherwise.""" |
| so = prebuilt_path(root, src) |
| pin = so + ".sha256" |
| if not os.path.isfile(so) or not os.path.isfile(pin): |
| raise PrebuiltError(f"no prebuilt p2p extension for {src} (source key {source_key(src)}) under {root}: " |
| f"build it with offload/dspark/worker/run/prod_tp2.sh build-p2p (expected {so} + .sha256)") |
| with open(pin) as f: |
| words = f.read().split() |
| if not words: |
| raise PrebuiltError(f"{pin} is empty") |
| want = words[0].lower() |
| got = _sha256(so) |
| if got != want: |
| raise PrebuiltError(f"{so}: sha256 {got} is not the pinned {want} ({pin})") |
| if MODULE in sys.modules: |
| mod = sys.modules[MODULE] |
| if getattr(mod, "__file__", None) != so: |
| raise PrebuiltError(f"{MODULE} already loaded from {getattr(mod, '__file__', None)}, not {so}") |
| return mod |
| loader = importlib.machinery.ExtensionFileLoader(MODULE, so) |
| spec = importlib.util.spec_from_file_location(MODULE, so, loader=loader) |
| if spec is None: |
| raise PrebuiltError(f"cannot load {so}") |
| mod = importlib.util.module_from_spec(spec) |
| loader.exec_module(mod) |
| sys.modules[MODULE] = mod |
| return mod |
|
|
