Spaces:
Running on Zero
Running on Zero
| """Maintained Phase60 CUDA pipeline for bulk Arrow screening. | |
| Imports are CUDA-lazy; rich/list operations use cpu_engine. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import importlib | |
| import importlib.util | |
| import json | |
| from pathlib import Path | |
| import sys | |
| import threading | |
| import weakref | |
| _ENGINES = weakref.WeakSet() | |
| _REGISTRY_LOCK = threading.Lock() | |
| _ALLOCATION_LOCK = threading.RLock() | |
| GPU_LIMIT = 2 * 1024**3 | |
| PINNED_LIMIT = 256 * 1024**2 | |
| NATIVE_SHA256 = "8b1dddd69ccc13380cced1f2df4977f2ecb750805c9d127b0c8046fef812563f" | |
| CUBIN_SHA256 = "dd943d9adac702e41c2bb0a2b58f78698498edd8a3dfab16cf9d224c8c6a83fa" | |
| def _sha(path): | |
| digest = hashlib.sha256() | |
| with path.open("rb") as handle: | |
| for chunk in iter(lambda: handle.read(1024**2), b""): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| def _member(root, name): | |
| if not isinstance(name, str) or "\\" in name or ":" in name: | |
| raise ValueError("invalid installed resource path") | |
| parts = name.split("/") | |
| if not parts or any(part in ("", ".", "..") for part in parts): | |
| raise ValueError("invalid installed resource path") | |
| path = root.joinpath(*parts) | |
| if path.is_symlink() or not path.resolve().is_relative_to(root.resolve()): | |
| raise ValueError("installed resource escapes its package") | |
| return path | |
| def _verify_files(root, files): | |
| for name, value in files.items(): | |
| path = _member(root, name) | |
| if not path.is_file() or path.stat().st_size != value["bytes"] or _sha(path) != value["sha256"]: | |
| raise RuntimeError("installed package resource changed: " + name) | |
| def _manifest(root): | |
| path = root / "runtime-manifest.json" | |
| if path.stat().st_size > 1024**2: | |
| raise RuntimeError("installed runtime manifest exceeds its bound") | |
| value = json.loads(path.read_text(encoding="utf-8")) | |
| if (value.get("schema") != "range.cuda.installed-runtime.v1" | |
| or value.get("selected_phase") not in (59, 60) | |
| or value.get("native", {}).get("sha256") != NATIVE_SHA256 | |
| or value.get("cubin", {}).get("sha256") != CUBIN_SHA256 | |
| or value.get("python_tag") != "cp310-cp310-win_amd64" | |
| or value.get("cuda_target") != "sm_75"): | |
| raise RuntimeError("installed companion manifest is incompatible") | |
| _verify_files(root, value["files"]) | |
| return value | |
| def _cpu_api(root, manifest): | |
| expected = root.parent / "range_curate" | |
| spec = importlib.util.find_spec("range_curate") | |
| if spec is None or not spec.origin or Path(spec.origin).resolve() != (expected / "__init__.py").resolve(): | |
| raise RuntimeError("install range-curate beside range-curate-cuda") | |
| _verify_files(expected, manifest["cpu_api_files"]) | |
| module = importlib.import_module("range_curate.examples.general_economy") | |
| if module.__version__ != manifest["cpu_version"]: | |
| raise RuntimeError("CUDA companion requires its matching RANGE release") | |
| return module | |
| def _workers(value): | |
| if isinstance(value, bool) or not isinstance(value, int) or value < 1: | |
| raise ValueError("workers must be a positive integer") | |
| return value | |
| def _validate_arrow(data, row_count, text_column, timestamp_column): | |
| import pyarrow as pa | |
| if not isinstance(data, (pa.Table, pa.RecordBatch)): | |
| raise TypeError("CUDA screening accepts a materialized Arrow Table or RecordBatch") | |
| if row_count is not None and (isinstance(row_count, bool) or not isinstance(row_count, int) or row_count < 0): | |
| raise ValueError("row_count must be a nonnegative integer or None") | |
| if row_count is not None and row_count != data.num_rows: | |
| raise ValueError("row_count differs from Arrow input") | |
| for name in (text_column, timestamp_column): | |
| if data.schema.get_field_index(name) < 0: | |
| raise ValueError("missing or ambiguous Arrow column: " + name) | |
| if not (pa.types.is_string(data.schema.field(text_column).type) | |
| or pa.types.is_large_string(data.schema.field(text_column).type)): | |
| raise TypeError("text must be Arrow string or large_string") | |
| timestamp = data.schema.field(timestamp_column).type | |
| if not pa.types.is_timestamp(timestamp) or timestamp.tz is None: | |
| raise ValueError("timestamp must be a timezone-aware Arrow timestamp") | |
| if data.column(text_column).null_count or data.column(timestamp_column).null_count: | |
| raise ValueError("text and timestamp must not contain nulls") | |
| class CudaRelevanceEngine: | |
| """Opt-in installed research engine for exact Arrow compact screening. | |
| backend='cpu' explicitly uses the private block31 CPU engine. Unsupported | |
| GPU taxonomies also use that exact engine. Missing CUDA and unexpected GPU | |
| failures raise; no partial output or hidden runtime-error fallback is used. | |
| Rich predictions and explanations remain CPU operations via cpu_engine. | |
| """ | |
| def __init__(self, taxonomy=None, workers=10, *, backend="cuda", resource_guard=None): | |
| self.workers = _workers(workers) | |
| if backend not in ("cuda", "cpu"): | |
| raise ValueError("backend must be cuda or cpu") | |
| if resource_guard is not None and not callable(resource_guard): | |
| raise TypeError("resource_guard must be callable or None") | |
| if sys.platform != "win32" or sys.version_info[:2] != (3, 10) or sys.maxsize <= 2**32: | |
| raise RuntimeError("this research companion supports Windows CPython3.10 AMD64 only") | |
| self._root = Path(__file__).resolve().parent | |
| self._manifest = _manifest(self._root) | |
| api = _cpu_api(self._root, self._manifest) | |
| self._cpu_engine = api.RelevanceEngine(taxonomy=taxonomy, workers=self.workers) | |
| # Keep the measured binary private. Never rebind range_curate._native or | |
| # mutate the CPU backend registry; this engine owns its own scanners. | |
| native = importlib.import_module("range_curate_cuda._private._native") | |
| native_path = getattr(native, "__file__", None) | |
| if not native_path or Path(native_path).resolve() != (self._root / self._manifest["native"]["member"]).resolve(): | |
| raise RuntimeError("the private native module was imported outside the installed companion") | |
| from range_curate._core.accelerated.compact import _build_native_specs | |
| indicator = self._cpu_engine._screening_indicator | |
| specs = _build_native_specs(indicator) | |
| indicator._list_scanner = native.NativeCompactFactScanner(*specs) | |
| indicator._arrow_scanner = native.NativeArrowCompactFactScanner(*specs) | |
| self._native = native | |
| self._backend = backend | |
| self._runtime = None | |
| self._machine_ready = False | |
| self._machine_lock = threading.Lock() | |
| self._last_cpu_profile = {} | |
| self._external_guard = resource_guard | |
| with _REGISTRY_LOCK: | |
| _ENGINES.add(self) | |
| if backend == "cuda": | |
| self._construct_runtime() | |
| def _construct_runtime(self): | |
| phase = self._manifest["selected_phase"] | |
| name = "phase60_followup_gpu" if phase == 60 else "phase59_followup_gpu" | |
| module = importlib.import_module("range_curate_cuda._runtime." + name) | |
| for imported_name, imported in tuple(sys.modules.items()): | |
| if imported_name == "range_curate_cuda._runtime" or imported_name.startswith("range_curate_cuda._runtime."): | |
| filename = getattr(imported, "__file__", None) | |
| if not filename or not Path(filename).resolve().is_relative_to((self._root / "_runtime").resolve()): | |
| raise RuntimeError("a runtime module was imported outside the installed companion: " + imported_name) | |
| scanner = self._cpu_engine._screening_indicator._arrow_scanner | |
| def masked(data, mask, *, workers, text_column, timestamp_column): | |
| return scanner.screen_arrow_masked(data, mask, text_column, timestamp_column, workers) | |
| kwargs = {"masked_cpu": masked, "unicode_tables": self._root / "resources" / "unicode13", | |
| "overlap_cpu": True, "diagnostics": False} | |
| if phase == 60: | |
| kind = module.ExactTwoSlotCallWideCudaClassifier | |
| kwargs["resource_guard"] = self._resource_guard | |
| else: | |
| guard = self._resource_guard | |
| class GuardedClassifier(module.ExactWarpPackedMatchingCudaClassifier): | |
| def _guard_allocations(self, sizes): | |
| sizes = tuple(int(size) for size in sizes) | |
| super()._guard_allocations(sizes) | |
| guard(self.cp, additional_device_bytes=sum((size + 511) // 512 * 512 for size in sizes), | |
| additional_pinned_bytes=0) | |
| kind = GuardedClassifier | |
| self._runtime = kind(self._cpu_engine, self._root / self._manifest["cubin"]["member"], **kwargs) | |
| def _resource_guard(self, cp, *, additional_device_bytes=0, additional_pinned_bytes=0): | |
| if additional_device_bytes < 0 or additional_pinned_bytes < 0: | |
| raise ValueError("allocation projection must be nonnegative") | |
| with _REGISTRY_LOCK: | |
| pinned = sum(engine._runtime._pinned_bytes() for engine in tuple(_ENGINES) | |
| if engine._runtime is not None and hasattr(engine._runtime, "_pinned_bytes")) | |
| if int(cp.get_default_memory_pool().total_bytes()) + additional_device_bytes > GPU_LIMIT: | |
| raise MemoryError("companion device pool plus proposed allocations exceeds2GiB") | |
| if pinned + additional_pinned_bytes > PINNED_LIMIT: | |
| raise MemoryError("companion pinned owners plus proposed allocations exceed256MiB") | |
| if self._external_guard is not None: | |
| self._external_guard(cp, additional_device_bytes=additional_device_bytes, | |
| additional_pinned_bytes=additional_pinned_bytes) | |
| def _ensure_machine(self): | |
| if self._runtime is None or self._machine_ready: | |
| return | |
| # Across companion instances admission and the complete fixed owner | |
| # allocation are one serialized operation. Ordinary calls reuse those | |
| # maximum arenas and never rely on an unlocked check-then-allocate gap. | |
| with self._machine_lock, _ALLOCATION_LOCK: | |
| if self._machine_ready: | |
| return | |
| runtime = self._runtime | |
| if runtime._failed_drain or runtime._initialization_failed: | |
| raise RuntimeError("CUDA initialization/drain previously failed; end this worker process") | |
| try: | |
| cp = importlib.import_module("cupy") | |
| except ImportError as error: | |
| raise RuntimeError("CUDA screening requires the separately installed compatible CuPy/CUDA runtime") from error | |
| try: | |
| capability = str(cp.cuda.Device().compute_capability) | |
| except Exception as error: | |
| raise RuntimeError("a usable NVIDIA CUDA device/runtime is required before CUDA screening") from error | |
| if capability != "75": | |
| raise RuntimeError("this bundled research cubin supports the qualified sm75 device only") | |
| if runtime.unsupported_reason: | |
| self._machine_ready = True | |
| return | |
| # Admit the complete maximum owner before any lazy config/workspace | |
| # allocation. Existing live/cached pools and every public pinned | |
| # owner remain charged. No credit is taken for cached reuse. | |
| workspace = int(runtime.backend_report["workspace_requested_max_bytes"]) | |
| configs = sum((max(value.nbytes, value.dtype.itemsize) + 511) // 512 * 512 | |
| for value in runtime._arrays.values()) | |
| # At most34 workspace arrays across two slots, conservatively each | |
| # charged one additional allocator quantum beyond logical bytes. | |
| pinned = 70328320 if self._manifest["selected_phase"] == 60 else 0 | |
| self._resource_guard(cp, additional_device_bytes=workspace + configs + 34 * 512, | |
| additional_pinned_bytes=pinned) | |
| if not runtime._call_lock.acquire(blocking=False): | |
| raise RuntimeError("the installed engine already has an active call") | |
| try: | |
| runtime._initialize() | |
| runtime._reserve(runtime.batch_bytes, runtime.batch_rows) | |
| self._machine_ready = True | |
| finally: | |
| runtime._call_lock.release() | |
| def prepare(self): | |
| """Explicitly validate and allocate the fixed owner outside a call timer.""" | |
| self._ensure_machine() | |
| return self.backend_report | |
| def screen_arrow(self, data, *, text_column="text", timestamp_column="timestamp", workers=None, row_count=None): | |
| workers = self.workers if workers is None else _workers(workers) | |
| if self._runtime is None: | |
| result = self._cpu_engine.screen_arrow(data, text_column=text_column, timestamp_column=timestamp_column, | |
| workers=workers, row_count=row_count) | |
| self._last_cpu_profile = {"backend": "cpu", "execution_path": "explicit_cpu", "rows": result.row_count} | |
| return result | |
| if self._runtime.unsupported_reason: | |
| self._ensure_machine() | |
| result = self._cpu_engine.screen_arrow(data, text_column=text_column, timestamp_column=timestamp_column, | |
| workers=workers, row_count=row_count) | |
| self._runtime.last_profile = {**self._runtime.backend_report, "backend": "cpu", | |
| "execution_path": "unsupported_taxonomy_cpu", "callwide_applied": False, | |
| "callwide_enabled": False, "rows": result.row_count, | |
| "unsupported_taxonomy_reason": self._runtime.unsupported_reason} | |
| return result | |
| _validate_arrow(data, row_count, text_column, timestamp_column) | |
| self._ensure_machine() | |
| result = self._runtime.screen_arrow(data, text_column=text_column, timestamp_column=timestamp_column, | |
| workers=workers, row_count=row_count) | |
| from range_curate._core.accelerated.compact import ArrowCompactFacts | |
| return ArrowCompactFacts(row_count=result.row_count, facts=result.facts) | |
| def cpu_engine(self): | |
| """CPU-only rich/list predictions and explanations for this taxonomy.""" | |
| return self._cpu_engine | |
| def runtime(self): | |
| """Read-only reference for maintainer ownership/parity qualification.""" | |
| return self._runtime | |
| def last_profile(self): | |
| return self._last_cpu_profile if self._runtime is None else self._runtime.last_profile | |
| def backend_report(self): | |
| report = {"backend": "cpu", "execution_path": "explicit_cpu"} if self._runtime is None else dict(self._runtime.backend_report) | |
| if self._runtime is not None and self._runtime.unsupported_reason: | |
| report.update(backend="cpu", execution_path="unsupported_taxonomy_cpu", callwide_applied=False, | |
| unsupported_taxonomy_reason=self._runtime.unsupported_reason) | |
| return {**report, "installed_runtime_phase": self._manifest["selected_phase"], | |
| "distribution": self._manifest["distribution"], "version": self._manifest["version"], | |
| "acceleration_scope": "bulk screen_arrow only; cpu_engine rich/list calls use CPU", | |
| "native_sha256": NATIVE_SHA256, "bundled_cubin_sha256": CUBIN_SHA256, | |
| "cuda_target": "sm_75", "local_dependency_environment_qualification_only": True} | |
| __all__ = ["CudaRelevanceEngine"] | |