distributed-llm 0.4.1__py3-none-any.whl

distllm/__init__.py

@@ -0,0 +1,47 @@
+"""Distributed LLM Inference System — lazy imports to avoid circular dependencies.
+
+Heavy modules are imported on first attribute access via ``__getattr__``.
+This breaks circular import chains and speeds up ``import distllm``.
+"""
+
+__version__ = "0.4.0"
+
+# Map public names to their module paths for lazy loading.
+_LAZY_IMPORTS: dict[str, tuple[str, str]] = {
+    # core
+    "Coordinator": ("distllm.core.coordinator", "Coordinator"),
+    "KVCache": ("distllm.core.kv_cache", "KVCache"),
+    "KVCacheManager": ("distllm.core.kv_cache", "KVCacheManager"),
+    "BatchScheduler": ("distllm.core.batch_scheduler", "BatchScheduler"),
+    "Sequence": ("distllm.core.batch_scheduler", "Sequence"),
+    "SequenceStatus": ("distllm.core.batch_scheduler", "SequenceStatus"),
+    "ScheduledBatch": ("distllm.core.batch_scheduler", "ScheduledBatch"),
+    "JSONSchemaConstraint": ("distllm.core.structured_output", "JSONSchemaConstraint"),
+    "SystemMonitor": ("distllm.core.monitor", "SystemMonitor"),
+    "NodeRegistration": ("distllm.core.resource_manager", "NodeRegistration"),
+    # models
+    "ModelPartitioner": ("distllm.models.partitioner", "ModelPartitioner"),
+    "partition_model_across_nodes": ("distllm.models.partitioner", "partition_model_across_nodes"),
+    "get_model_info": ("distllm.models.partitioner", "get_model_info"),
+    "AdapterManager": ("distllm.models.adapter", "AdapterManager"),
+    # optional
+    "ModelHotSwapManager": ("distllm.core.multi_model_serving", "ModelHotSwapManager"),
+}
+
+__all__ = ["__version__", *sorted(_LAZY_IMPORTS.keys())]
+
+
+def __getattr__(name: str):
+    if name in _LAZY_IMPORTS:
+        module_path, attr_name = _LAZY_IMPORTS[name]
+        import importlib
+        module = importlib.import_module(module_path)
+        attr = getattr(module, attr_name)
+        # Cache on the module so subsequent access is fast
+        globals()[name] = attr
+        return attr
+    raise AttributeError(f"module 'distllm' has no attribute {name!r}")
+
+
+def __dir__():
+    return list(__all__)

distllm/api/__init__.py

@@ -0,0 +1,11 @@
+"""API layer for distributed LLM inference."""
+
+from distllm.api.server import app, create_coordinator, main
+from distllm.api.middleware import AuthMiddleware
+
+__all__ = [
+    "app",
+    "create_coordinator",
+    "main",
+    "AuthMiddleware",
+]

distllm/api/api_settings.py

@@ -0,0 +1,126 @@
+"""Consolidated API settings via Pydantic BaseSettings.
+
+Replaces 20+ scattered ``os.environ.get()`` calls across 15+ files
+with a single typed settings class loaded once.
+
+Usage::
+
+    from distllm.api.api_settings import api_settings
+
+    # Access any setting as a typed attribute:
+    cors_origins = api_settings.cors_origins
+    rate_limit = api_settings.rate_limit_requests
+    enabled = api_settings.prompt_injection_enabled
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+from pydantic import BaseSettings, Field
+
+
+class APISettings(BaseSettings):
+    """Pydantic-based settings for the DistLLM API layer.
+
+    All settings are loaded from environment variables with typed
+    defaults.  To override, set the corresponding env var before
+    starting the server.
+    """
+
+    # ── CORS ──────────────────────────────────────────────────────────
+
+    cors_origins: str = Field(
+        "http://localhost:3000,http://localhost:8080",
+        description="Comma-separated CORS allowed origins",
+        env="DISTLLM_CORS_ORIGINS",
+    )
+    cors_allow_all: bool = Field(
+        False,
+        description="Allow all CORS origins (DANGEROUS — dev only)",
+        env="DISTLLM_CORS_ALLOW_ALL",
+    )
+
+    # ── TLS / Security ────────────────────────────────────────────────
+
+    tls_enabled: bool = Field(False, env="DISTLLM_TLS_ENABLED")
+    trust_proxy_tls: bool = Field(False, env="DISTLLM_TRUST_PROXY_TLS")
+    trust_proxy_headers: bool = Field(False, env="DISTLLM_TRUST_PROXY_HEADERS")
+
+    # ── Rate Limiting ─────────────────────────────────────────────────
+
+    rate_limit_requests: int = Field(
+        1000, ge=0,
+        description="Max requests per IP per 60s window (0=disabled)",
+        env="DISTLLM_RATE_LIMIT_REQUESTS",
+    )
+    rate_limit_auth_attempts: int = Field(
+        30, ge=1,
+        description="Max failed auth attempts per IP per 60s",
+        env="DISTLLM_AUTH_RATE_LIMIT_ATTEMPTS",
+    )
+
+    # ── Prompt Injection ──────────────────────────────────────────────
+
+    prompt_injection_enabled: bool = Field(True, env="DISTLLM_INJECTION_ENABLED")
+    prompt_injection_block_threshold: float = Field(0.9, ge=0.0, le=1.0, env="DISTLLM_INJECTION_BLOCK_THRESHOLD")
+    prompt_injection_sanitize_threshold: float = Field(0.7, ge=0.0, le=1.0, env="DISTLLM_INJECTION_SANITIZE_THRESHOLD")
+    prompt_injection_flag_threshold: float = Field(0.4, ge=0.0, le=1.0, env="DISTLLM_INJECTION_FLAG_THRESHOLD")
+    prompt_injection_model: str = Field("", env="DISTLLM_INJECTION_MODEL")
+    prompt_injection_audit_log: str = Field("", env="DISTLLM_INJECTION_AUDIT_LOG")
+
+    # ── Quota / Usage ──────────────────────────────────────────────────
+
+    quota_enabled: bool = Field(False, env="DISTLLM_QUOTA_ENABLED")
+    usage_db_path: str = Field(".usage.db", env="DISTLLM_USAGE_DB")
+
+    # ── Observability ─────────────────────────────────────────────────
+
+    trace_sample_rate: float = Field(0.1, ge=0.0, le=1.0, env="DISTLLM_TRACE_SAMPLE_RATE")
+    enable_docs: bool = Field(False, env="DISTLLM_ENABLE_DOCS")
+    otlp_endpoint: str = Field("http://localhost:4318/v1/traces", env="DISTLLM_OTLP_ENDPOINT")
+
+    # ── Config / Plugin ────────────────────────────────────────────────
+
+    config_path: str = Field("config.yaml", env="DISTLLM_CONFIG")
+    verify_plugins: bool = Field(False, env="DISTLLM_VERIFY_PLUGINS")
+
+    # ── SSO ────────────────────────────────────────────────────────────
+
+    sso_provider: str = Field("", env="DISTLLM_SSO_PROVIDER")
+    sso_client_id: str = Field("", env="DISTLLM_SSO_CLIENT_ID")
+    sso_client_secret: str = Field("", env="DISTLLM_SSO_CLIENT_SECRET")
+    sso_authority: str = Field("", env="DISTLLM_SSO_AUTHORITY")
+    sso_metadata_url: str = Field("", env="DISTLLM_SSO_METADATA_URL")
+    sso_callback_url: str = Field("", env="DISTLLM_SSO_CALLBACK_URL")
+    sso_jwks_url: str = Field("", env="DISTLLM_SSO_JWKS_URL")
+
+    # ── Feature Flags ─────────────────────────────────────────────────
+
+    enable_debug_routes: bool = Field(False, env="DISTLLM_ENABLE_DEBUG_ROUTES")
+    dev_mode: bool = Field(False, env="DISTLLM_DEV_MODE")
+    api_key: str = Field("", env="API_KEY")
+    cluster_key: str = Field("", env="DISTLLM_CLUSTER_KEY")
+    ha_shared_secret: str = Field("", env="DISTLLM_HA_SECRET")
+
+    class Config:
+        env_file = ".env"
+        env_file_encoding = "utf-8"
+        case_sensitive = False
+
+
+# Global singleton
+_settings: Optional[APISettings] = None
+
+
+def get_api_settings() -> APISettings:
+    global _settings
+    if _settings is None:
+        _settings = APISettings()
+    return _settings
+
+
+# Convenience accessor
+api_settings = get_api_settings()

distllm/api/api_state.py

@@ -0,0 +1,85 @@
+"""Shared application state proxy for cross-module access.
+
+Routes use::
+
+    from distllm.api.api_state import g
+    if g.coordinator is None: ...
+
+All state is held in ``AppState`` (defined in ``app_state.py``).
+``_ServerGlobals`` is a thin proxy that delegates to ``AppState``,
+allowing route modules to access ``g.coordinator`` without
+importing ``server`` directly.
+"""
+
+from distllm.api.app_state import AppState
+
+_state = AppState()
+
+
+class _ServerGlobals:
+    """Proxy delegating to ``AppState`` — the single source of truth."""
+
+    def __init__(self) -> None:
+        pass
+
+    @property
+    def coordinator(self):
+        return _state.coordinator
+
+    @coordinator.setter
+    def coordinator(self, value):
+        _state.coordinator = value
+
+    @property
+    def monitor(self):
+        return _state.monitor
+
+    @monitor.setter
+    def monitor(self, value):
+        _state.monitor = value
+
+    @property
+    def startup_time(self) -> float:
+        return _state.uptime_seconds
+
+    @startup_time.setter
+    def startup_time(self, value: float) -> None:
+        _state.startup_time = value
+
+    @property
+    def metrics_exporter(self):
+        return _state.metrics_exporter
+
+    @metrics_exporter.setter
+    def metrics_exporter(self, value):
+        _state.metrics_exporter = value
+
+    # Generic dict-style access: g.get("name") / g["name"] resolve attributes on
+    # the shared AppState.  Several route modules (routes/exchange.py,
+    # authz/opa.py) rely on this pattern.
+    def get(self, name, default=None):
+        return getattr(_state, name, default)
+
+    def __getitem__(self, name):
+        return getattr(_state, name)
+
+    def __setitem__(self, name, value):
+        setattr(_state, name, value)
+
+    def set(self, name, value) -> None:
+        setattr(_state, name, value)
+
+
+def reset_app_state_for_testing() -> None:
+    """Reset the shared app state to a pristine instance (for tests).
+
+    Mutates the existing AppState in place (rather than rebinding the module
+    global) so imports that captured the old ``_state`` object — and the ``g``
+    proxy — keep pointing at the same, now-reset, instance.  Rebinding would
+    strand old bindings on a stale object (e.g. a test that sets
+    ``_state.prompt_exchange`` after ``_state`` was reassigned elsewhere).
+    """
+    _state.__init__()
+
+
+g = _ServerGlobals()