Spaces:
Running on Zero
Running on Zero
| """Backend selection for the supported CPU implementation.""" | |
| from __future__ import annotations | |
| from importlib.util import find_spec | |
| from typing import TYPE_CHECKING, Protocol | |
| from ._core.accelerated.compact import ( | |
| ArrowCompactFacts, | |
| ColumnArticleBatch, | |
| CompactColumnScreenBatch, | |
| CompactRelevanceIndicator, | |
| ) | |
| from .models import BackendInfo | |
| if TYPE_CHECKING: | |
| from .customization import Taxonomy | |
| def _native_available() -> bool: | |
| try: | |
| return find_spec("range_curate._native") is not None | |
| except (ImportError, ModuleNotFoundError, ValueError): | |
| return False | |
| class _Backend(Protocol): | |
| name: str | |
| def describe( | |
| self, | |
| indicator: CompactRelevanceIndicator, | |
| ) -> BackendInfo: ... | |
| def screen_batch( | |
| self, | |
| indicator: CompactRelevanceIndicator, | |
| batch: ColumnArticleBatch, | |
| *, | |
| workers: int, | |
| ) -> CompactColumnScreenBatch: ... | |
| def screen_arrow( | |
| self, | |
| indicator: CompactRelevanceIndicator, | |
| taxonomy: "Taxonomy", | |
| data, | |
| *, | |
| text_column: str, | |
| timestamp_column: str, | |
| workers: int, | |
| row_count: int | None, | |
| ) -> ArrowCompactFacts: ... | |
| class _CpuBackend: | |
| name = "cpu" | |
| def describe( | |
| self, | |
| indicator: CompactRelevanceIndicator, | |
| ) -> BackendInfo: | |
| native_list = indicator.list_available | |
| native_arrow = indicator.arrow_available | |
| if native_list and native_arrow: | |
| detail = "Rust list and Arrow screening are available" | |
| elif native_list: | |
| detail = "Rust list screening is available; Arrow screening is unavailable" | |
| else: | |
| detail = ( | |
| "Python prediction is available; rebuild the native extension " | |
| "for compact screening" | |
| ) | |
| return BackendInfo( | |
| name=self.name, | |
| available=True, | |
| accelerated=native_list, | |
| detail=detail, | |
| ) | |
| def screen_batch( | |
| self, | |
| indicator: CompactRelevanceIndicator, | |
| batch: ColumnArticleBatch, | |
| *, | |
| workers: int, | |
| ) -> CompactColumnScreenBatch: | |
| if not indicator.list_available: | |
| raise RuntimeError( | |
| "Compact screening requires the packaged Rust extension; " | |
| "install a wheel for this platform or rebuild the package" | |
| ) | |
| return indicator.screen_column_batch(batch, workers=workers) | |
| def screen_arrow( | |
| self, | |
| indicator: CompactRelevanceIndicator, | |
| taxonomy: "Taxonomy", | |
| data, | |
| *, | |
| text_column: str, | |
| timestamp_column: str, | |
| workers: int, | |
| row_count: int | None, | |
| ) -> ArrowCompactFacts: | |
| del taxonomy | |
| if not indicator.arrow_available: | |
| raise RuntimeError( | |
| "Arrow screening requires the packaged Rust extension and the " | |
| "'arrow' extra" | |
| ) | |
| return indicator.screen_arrow( | |
| data, | |
| text_column=text_column, | |
| timestamp_column=timestamp_column, | |
| workers=workers, | |
| row_count=row_count, | |
| ) | |
| class _CudaBackend(_CpuBackend): | |
| """CUDA bulk Arrow screening; rich/list operations remain CPU work.""" | |
| name = "cuda" | |
| def __init__(self): | |
| import threading | |
| if find_spec("range_curate_cuda") is None: | |
| raise RuntimeError("Install the matching range-curate-cuda wheel for CUDA Arrow screening") | |
| self._engine = None | |
| self._lock = threading.Lock() | |
| def describe(self, indicator): | |
| return BackendInfo(name="cuda", available=True, accelerated=True, | |
| detail="Phase60 CUDA for bulk Arrow screening; rich/list operations use CPU. Device support is validated on first Arrow call.") | |
| def screen_arrow(self, indicator, taxonomy, data, *, text_column, timestamp_column, workers, row_count): | |
| if self._engine is None: | |
| with self._lock: | |
| if self._engine is None: | |
| from range_curate_cuda import CudaRelevanceEngine | |
| self._engine = CudaRelevanceEngine(taxonomy=taxonomy, workers=workers) | |
| return self._engine.screen_arrow(data, text_column=text_column, | |
| timestamp_column=timestamp_column, workers=workers, row_count=row_count) | |
| def create_backend(name: str) -> _Backend: | |
| normalized = name.strip().lower() if isinstance(name, str) else "" | |
| if normalized == "cpu": | |
| return _CpuBackend() | |
| if normalized == "cuda": | |
| return _CudaBackend() | |
| raise ValueError("backend must be 'cpu' or 'cuda'") | |
| def available_backends() -> tuple[BackendInfo, ...]: | |
| """Discover packages without importing CuPy or initializing CUDA.""" | |
| native = _native_available() | |
| cuda = find_spec("range_curate_cuda") is not None | |
| return ( | |
| BackendInfo(name="cpu", available=True, accelerated=native, | |
| detail="Native compact screening" if native else "Python rich predictions; native screening unavailable"), | |
| BackendInfo(name="cuda", available=cuda, accelerated=cuda, | |
| detail="Optional Phase60 Arrow runtime installed; device support checked on use" if cuda | |
| else "Install the matching range-curate-cuda wheel; Windows CPython 3.10 and sm75 only"), | |
| ) | |
| __all__ = ["available_backends", "create_backend"] | |