Spaces:
Running on Zero
Running on Zero
File size: 5,592 Bytes
f8768df c4aee46 f8768df 8947dc2 f8768df 8947dc2 f8768df 8947dc2 f8768df 8947dc2 f8768df 8947dc2 f8768df | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | """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"]
|