cuda/src/range_curate_cuda/api.py

"""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)

    @property
    def cpu_engine(self):
        """CPU-only rich/list predictions and explanations for this taxonomy."""
        return self._cpu_engine

    @property
    def runtime(self):
        """Read-only reference for maintainer ownership/parity qualification."""
        return self._runtime

    @property
    def last_profile(self):
        return self._last_cpu_profile if self._runtime is None else self._runtime.last_profile

    @property
    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"]