sie-sdk 0.1.7__py3-none-any.whl

sie_sdk/__init__.py

@@ -0,0 +1,93 @@
+"""SIE SDK - Python client for Search Inference Engine.
+
+Example:
+    >>> from sie_sdk import SIEClient
+    >>> client = SIEClient("http://localhost:8080")
+    >>> result = client.encode("bge-m3", {"text": "Hello world"})
+    >>> result["dense"]  # np.ndarray, shape [1024]
+
+For ColBERT/late interaction models, use the scoring module:
+    >>> from sie_sdk.scoring import maxsim
+    >>> scores = maxsim(query_multivector, doc_multivectors)
+"""
+
+from sie_sdk.client import (
+    LoraLoadingError,
+    ModelLoadingError,
+    PoolError,
+    ProvisioningError,
+    RequestError,
+    ServerError,
+    SIEAsyncClient,
+    SIEClient,
+    SIEConnectionError,
+    SIEError,
+)
+from sie_sdk.types import (
+    # Response types
+    AssignedWorkerInfo,
+    CapacityInfo,
+    ClusterStatusMessage,
+    ClusterSummary,
+    EncodeResult,
+    Entity,
+    ExtractResult,
+    HealthResponse,
+    Item,
+    ModelInfo,
+    ModelSummary,
+    PoolInfo,
+    PoolListItem,
+    PoolResponse,
+    PoolSpec,
+    PoolSpecResponse,
+    PoolStatusInfo,
+    ScoreResult,
+    SparseResult,
+    StatusMessage,
+    TimingInfo,
+    WorkerInfo,
+    WorkerStatusMessage,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Response types
+    "AssignedWorkerInfo",
+    "CapacityInfo",
+    "ClusterStatusMessage",
+    "ClusterSummary",
+    "EncodeResult",
+    "Entity",
+    "ExtractResult",
+    "HealthResponse",
+    # Input types
+    "Item",
+    # Exceptions
+    "LoraLoadingError",
+    "ModelInfo",
+    "ModelLoadingError",
+    "ModelSummary",
+    "PoolError",
+    "PoolInfo",
+    "PoolListItem",
+    "PoolResponse",
+    "PoolSpec",
+    "PoolSpecResponse",
+    "PoolStatusInfo",
+    "ProvisioningError",
+    "RequestError",
+    # Clients
+    "SIEAsyncClient",
+    "SIEClient",
+    "SIEConnectionError",
+    "SIEError",
+    "ScoreResult",
+    "ServerError",
+    "SparseResult",
+    "StatusMessage",
+    "TimingInfo",
+    "WorkerInfo",
+    "WorkerStatusMessage",
+]

sie_sdk/bundle_utils.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+def _scan_model_adapters(models_dir: Path) -> dict[str, set[str]]:
+    """Scan model config YAMLs and return adapter modules per model.
+
+    Args:
+        models_dir: Path to the models directory containing *.yaml configs.
+
+    Returns:
+        Dict mapping model name to set of adapter module paths.
+    """
+    result: dict[str, set[str]] = {}
+    if not models_dir.exists():
+        return result
+
+    for model_path in sorted(models_dir.glob("*.yaml")):
+        try:
+            model_data = yaml.safe_load(model_path.read_text()) or {}
+        except Exception:
+            logger.exception("Failed to parse model config %s", model_path.name)
+            continue
+        model_name = model_data.get("sie_id", model_path.stem.replace("__", "/"))
+        modules: set[str] = set()
+        for profile in model_data.get("profiles", {}).values():
+            adapter_path = profile.get("adapter_path", "")
+            module_path = adapter_path.split(":", maxsplit=1)[0]
+            if module_path:
+                modules.add(module_path)
+        if modules:
+            result[model_name] = modules
+
+    return result
+
+
+def match_bundle_models(bundle_path: Path, models_dir: Path) -> list[str]:
+    """Match models to a bundle by adapter module paths.
+
+    Loads the bundle YAML to get its adapter module list, then scans
+    model config YAMLs to find models whose adapter_path module matches.
+
+    Args:
+        bundle_path: Path to the bundle YAML file.
+        models_dir: Path to the models directory containing *.yaml configs.
+
+    Returns:
+        List of model names (sie_id or derived from filename) whose adapters
+        match the bundle's adapter list.
+    """
+    with bundle_path.open() as f:
+        data = yaml.safe_load(f) or {}
+
+    adapter_modules = set(data.get("adapters", []))
+    if not adapter_modules:
+        return []
+
+    model_adapters = _scan_model_adapters(models_dir)
+    return [name for name, modules in model_adapters.items() if modules & adapter_modules]
+
+
+def find_bundle_for_models(
+    model_names: list[str],
+    bundles_dir: Path,
+    models_dir: Path,
+) -> str | None:
+    """Find the best bundle whose adapters cover the given models.
+
+    Scans all bundle YAMLs in bundles_dir and returns the one whose adapter
+    set covers all requested models with the fewest extra adapters (most
+    specific match). Ties are broken by bundle priority (lower = higher
+    priority).
+
+    Args:
+        model_names: List of model names to match.
+        bundles_dir: Path to the bundles directory.
+        models_dir: Path to the models directory containing *.yaml configs.
+
+    Returns:
+        Bundle name (without .yaml) of the best match, or None if no bundle
+        covers all requested models.
+    """
+    if not model_names or not bundles_dir.exists() or not models_dir.exists():
+        return None
+
+    # Collect adapter modules needed by the requested models
+    model_adapters = _scan_model_adapters(models_dir)
+    needed_adapters: set[str] = set()
+    for name in model_names:
+        needed_adapters |= model_adapters.get(name, set())
+
+    if not needed_adapters:
+        return None
+
+    # Score each bundle: must cover all needed adapters
+    best_name: str | None = None
+    best_extra = float("inf")
+    best_priority = float("inf")
+
+    for bundle_path in sorted(bundles_dir.glob("*.yaml")):
+        try:
+            data = yaml.safe_load(bundle_path.read_text()) or {}
+        except Exception:
+            logger.exception("Failed to parse bundle %s", bundle_path.name)
+            continue
+        bundle_adapters = set(data.get("adapters", []))
+        if not needed_adapters <= bundle_adapters:
+            continue  # doesn't cover all needed adapters
+        extra = len(bundle_adapters - needed_adapters)
+        priority = data.get("priority", 50)
+        if extra < best_extra or (extra == best_extra and priority < best_priority):
+            best_name = bundle_path.stem
+            best_extra = extra
+            best_priority = priority
+
+    return best_name

sie_sdk/cache.py

@@ -0,0 +1,374 @@
+"""Model weight caching with hierarchical lookup.
+
+Implements the caching hierarchy for model weights:
+1. Local cache (HF_HOME/hub by default)
+2. Cluster cache (S3/GCS bucket)
+3. HuggingFace Hub fallback (if enabled)
+
+The caching is transparent to adapters - they always see files in local cache.
+This module pre-populates local cache from cluster cache if needed.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from sie_sdk.exceptions import GatedModelError
+from sie_sdk.storage import get_storage_backend, join_path
+
+if TYPE_CHECKING:
+    from sie_sdk.storage import StorageBackend
+
+logger = logging.getLogger(__name__)
+
+HF_IGNORE_PATTERNS = [
+    "*.md",
+    "README*",
+    "LICENSE*",
+    "docs/*",
+    ".github/*",
+    "*.onnx",
+    "onnx/*",
+    # training-only
+    "optimizer.*",
+    "scheduler.*",
+    "trainer_state.json",
+    "training_args.bin",
+    "checkpoint-*/*",
+]
+HF_LEGACY_WEIGHT_PATTERNS = ["*.bin", "*.ckpt", "*.msgpack", "*.h5", "*.ot"]
+
+
+@dataclass
+class CacheConfig:
+    """Configuration for the weight cache."""
+
+    local_cache: Path
+    """Local cache directory (usually HF_HOME/hub)."""
+
+    cluster_cache: str | None = None
+    """Cluster cache URL (s3:// or gs://), or None if not configured."""
+
+    hf_fallback: bool = True
+    """Whether to fallback to HuggingFace Hub for downloads."""
+
+
+def get_cache_config() -> CacheConfig:
+    """Get cache configuration from environment variables.
+
+    Reads:
+        SIE_LOCAL_CACHE: Local cache directory (default: HF_HOME/hub)
+        SIE_CLUSTER_CACHE: Cluster cache URL (s3:// or gs://)
+        SIE_HF_FALLBACK: Whether to enable HF Hub fallback (default: true)
+
+    Returns:
+        CacheConfig with resolved paths.
+    """
+    # Local cache: explicit SIE_LOCAL_CACHE, or HF_HOME, or default
+    local_cache_env = os.environ.get("SIE_LOCAL_CACHE")
+    if local_cache_env:
+        local_cache = Path(local_cache_env)
+    else:
+        hf_home = os.environ.get("HF_HOME")
+        if hf_home:
+            local_cache = Path(hf_home) / "hub"
+        else:
+            local_cache = Path.home() / ".cache" / "huggingface" / "hub"
+
+    # Cluster cache: S3/GCS URL
+    cluster_cache = os.environ.get("SIE_CLUSTER_CACHE")
+
+    # HF fallback: default true
+    hf_fallback_str = os.environ.get("SIE_HF_FALLBACK", "true").lower()
+    hf_fallback = hf_fallback_str in ("true", "1", "yes")
+
+    return CacheConfig(
+        local_cache=local_cache,
+        cluster_cache=cluster_cache,
+        hf_fallback=hf_fallback,
+    )
+
+
+def is_model_cached(model_id: str, config: CacheConfig | None = None) -> bool:
+    """Check if a model is already in local cache.
+
+    Uses HuggingFace Hub's cache structure to check for the model.
+
+    Args:
+        model_id: HuggingFace model ID (e.g., "BAAI/bge-m3").
+        config: Cache configuration. If None, reads from environment.
+
+    Returns:
+        True if model appears to be cached locally.
+    """
+    if config is None:
+        config = get_cache_config()
+
+    # HF Hub cache structure: models--{org}--{model}/snapshots/{revision}/
+    cache_dir = config.local_cache / f"models--{model_id.replace('/', '--')}"
+    if not cache_dir.exists():
+        return False
+
+    # Check for any snapshot directory with files
+    snapshots_dir = cache_dir / "snapshots"
+    if not snapshots_dir.exists():
+        return False
+
+    return any(snapshot.is_dir() and any(snapshot.iterdir()) for snapshot in snapshots_dir.iterdir())
+
+
+def ensure_model_cached(
+    model_id: str,
+    config: CacheConfig | None = None,
+    revision: str | None = None,
+) -> Path:
+    """Ensure model weights are in local cache.
+
+    Implements the caching hierarchy:
+    1. Check local cache - return path if found
+    2. Check cluster cache - download to local if found
+    3. Download from HuggingFace Hub (if fallback enabled)
+
+    Args:
+        model_id: HuggingFace model ID (e.g., "BAAI/bge-m3").
+        config: Cache configuration. If None, reads from environment.
+        revision: Specific model revision (commit hash, branch, or tag).
+
+    Returns:
+        Path to cached model in local cache.
+
+    Raises:
+        GatedModelError: If model is gated and authentication fails.
+        RuntimeError: If model not found in any cache tier and HF fallback disabled.
+    """
+    if config is None:
+        config = get_cache_config()
+
+    # Check local cache first
+    if is_model_cached(model_id, config):
+        logger.debug("Model %s found in local cache", model_id)
+        return _get_model_cache_path(model_id, config)
+
+    # Try cluster cache if configured
+    if config.cluster_cache:
+        if _download_from_cluster_cache(model_id, config):
+            logger.info("Downloaded %s from cluster cache", model_id)
+            return _get_model_cache_path(model_id, config)
+
+    # Download from HuggingFace Hub if fallback enabled
+    if config.hf_fallback:
+        logger.info("Downloading %s from HuggingFace Hub", model_id)
+        return _download_from_huggingface(model_id, config, revision)
+
+    # Not in any cache and HF fallback disabled
+    raise RuntimeError(
+        f"Model '{model_id}' not found in local or cluster cache, "
+        f"and HuggingFace fallback is disabled (SIE_HF_FALLBACK=false)"
+    )
+
+
+def _get_model_cache_path(model_id: str, config: CacheConfig) -> Path:
+    """Get the local cache path for a model.
+
+    Args:
+        model_id: HuggingFace model ID.
+        config: Cache configuration.
+
+    Returns:
+        Path to the model's cache directory.
+    """
+    return config.local_cache / f"models--{model_id.replace('/', '--')}"
+
+
+def _download_from_cluster_cache(model_id: str, config: CacheConfig) -> bool:
+    """Download model from cluster cache to local cache.
+
+    Mirrors the HuggingFace Hub cache structure from cluster to local.
+
+    Args:
+        model_id: HuggingFace model ID.
+        config: Cache configuration (must have cluster_cache set).
+
+    Returns:
+        True if successfully downloaded, False if not found in cluster cache.
+    """
+    if not config.cluster_cache:
+        return False
+
+    backend = get_storage_backend(config.cluster_cache)
+
+    # Construct cluster cache path
+    model_folder = f"models--{model_id.replace('/', '--')}"
+    cluster_model_path = join_path(config.cluster_cache, model_folder)
+
+    # Check if model exists in cluster cache
+    if not backend.exists(join_path(cluster_model_path, "snapshots")):
+        logger.debug("Model %s not found in cluster cache", model_id)
+        return False
+
+    # Create local cache directory
+    local_model_path = config.local_cache / model_folder
+    local_model_path.mkdir(parents=True, exist_ok=True)
+
+    # Download the entire model directory structure
+    _download_directory(backend, cluster_model_path, local_model_path)
+
+    return True
+
+
+def _download_directory(backend: StorageBackend, src_path: str, dst_path: Path) -> None:
+    """Recursively download a directory from cloud storage.
+
+    Args:
+        backend: Storage backend instance.
+        src_path: Source path in cloud storage.
+        dst_path: Destination local path.
+    """
+    # Download files in current directory
+    for filename in backend.list_files(src_path):
+        src_file = join_path(src_path, filename)
+        dst_file = dst_path / filename
+        backend.download_file(src_file, dst_file)
+
+    # Recursively download subdirectories
+    for dirname in backend.list_dirs(src_path):
+        src_subdir = join_path(src_path, dirname)
+        dst_subdir = dst_path / dirname
+        dst_subdir.mkdir(parents=True, exist_ok=True)
+        _download_directory(backend, src_subdir, dst_subdir)
+
+
+def _get_hf_ignore_patterns(
+    model_id: str,
+    revision: str | None,
+    token: str | None,
+) -> list[str]:
+    """Check if a HF repo has safetensors files; This avoids downloading duplicate weight files (e.g., both .safetensors and .bin)."""
+    from huggingface_hub import file_exists
+    from huggingface_hub.errors import (
+        GatedRepoError,
+        HfHubHTTPError,
+    )
+
+    try:
+        has_safetensors = any(
+            file_exists(model_id, filename, revision=revision, token=token)
+            for filename in ("model.safetensors", "model.safetensors.index.json")
+        )
+        if has_safetensors:
+            return HF_IGNORE_PATTERNS + HF_LEGACY_WEIGHT_PATTERNS
+    except (HfHubHTTPError, GatedRepoError) as e:
+        logger.warning(
+            "Failed to list HF repo files; falling back to default ignore patterns",
+            extra={
+                "model_id": model_id,
+                "revision": revision,
+                "status_code": getattr(e.response, "status_code", None),
+            },
+        )
+    return HF_IGNORE_PATTERNS
+
+
+def _download_from_huggingface(
+    model_id: str,
+    config: CacheConfig,
+    revision: str | None = None,
+) -> Path:
+    """Download model from HuggingFace Hub to local cache.
+
+    Uses huggingface_hub's snapshot_download which automatically uses
+    the standard HF cache structure (~/.cache/huggingface/hub/).
+
+    Args:
+        model_id: HuggingFace model ID.
+        config: Cache configuration.
+        revision: Specific model revision.
+
+    Returns:
+        Path to the model in local cache.
+
+    Raises:
+        GatedModelError: If model is gated and authentication fails.
+    """
+    from huggingface_hub import snapshot_download
+    from huggingface_hub.errors import (
+        GatedRepoError,
+        HfHubHTTPError,
+        RepositoryNotFoundError,
+    )
+
+    # Get token from environment
+    token = os.environ.get("HF_TOKEN")
+
+    try:
+        # snapshot_download automatically uses HF cache structure
+        # cache_dir should point to HF_HOME/hub where model dirs are created
+        snapshot_download(
+            model_id,
+            revision=revision,
+            token=token,
+            cache_dir=str(config.local_cache),  # Points to HF_HOME/hub
+            ignore_patterns=_get_hf_ignore_patterns(model_id, revision, token),
+        )
+        return _get_model_cache_path(model_id, config)
+
+    except GatedRepoError as e:
+        # User-friendly error for gated models
+        raise GatedModelError(model_id, e) from e
+
+    except RepositoryNotFoundError as e:
+        # Check if this might be a gated model accessed without auth
+        if token is None:
+            msg = (
+                f"Model '{model_id}' not found. This could mean:\n"
+                f"  1. The model ID is incorrect\n"
+                f"  2. The model is private/gated and requires authentication\n\n"
+                f"If this is a gated model, set HF_TOKEN environment variable:\n"
+                f"  export HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxx\n\n"
+                f"Original error: {e}"
+            )
+            raise RuntimeError(msg) from e
+        raise
+
+    except HfHubHTTPError as e:
+        # Handle 401/403 errors that might indicate auth issues
+        status_code = getattr(e.response, "status_code", None) if hasattr(e, "response") else None
+        if status_code in (401, 403):
+            raise GatedModelError(model_id, e) from e
+        raise
+
+
+def populate_cluster_cache(
+    model_id: str,
+    config: CacheConfig | None = None,
+) -> bool:
+    """Upload model from local cache to cluster cache.
+
+    Used by sie-admin to pre-populate cluster cache.
+
+    Args:
+        model_id: HuggingFace model ID.
+        config: Cache configuration.
+
+    Returns:
+        True if successfully uploaded, False if model not in local cache.
+    """
+    if config is None:
+        config = get_cache_config()
+
+    if not config.cluster_cache:
+        logger.warning("No cluster cache configured")
+        return False
+
+    if not is_model_cached(model_id, config):
+        logger.warning("Model %s not in local cache, cannot populate cluster", model_id)
+        return False
+
+    # TODO: Implement upload to cluster cache
+    # This requires adding upload methods to storage backends
+    logger.info("populate_cluster_cache not yet implemented")
+    return False

sie_sdk/client/__init__.py

@@ -0,0 +1,30 @@
+"""SIE SDK Client package.
+
+Re-exports all client classes and errors for backwards compatibility.
+"""
+
+from sie_sdk.client.async_ import SIEAsyncClient
+from sie_sdk.client.errors import (
+    LoraLoadingError,
+    ModelLoadingError,
+    PoolError,
+    ProvisioningError,
+    RequestError,
+    ServerError,
+    SIEConnectionError,
+    SIEError,
+)
+from sie_sdk.client.sync import SIEClient
+
+__all__ = [
+    "LoraLoadingError",
+    "ModelLoadingError",
+    "PoolError",
+    "ProvisioningError",
+    "RequestError",
+    "SIEAsyncClient",
+    "SIEClient",
+    "SIEConnectionError",
+    "SIEError",
+    "ServerError",
+]