File size: 2,658 Bytes
e1acc15 | 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 | """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
|