diff --git a/CLAUDE.md b/CLAUDE.md index eed22f5db..8d712a074 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -145,6 +145,14 @@ Skulk now treats model capability handling as two layers: This capability spine is the source of truth for model-aware reasoning defaults, prompt rendering, output parsing, tool-call handling, and additive `/v1/models` metadata consumed by the dashboard. +### Planned LARQL Slice Mode +Phase 1 LARQL ADRs live in `docs/adr/`. Phase 2 adds internal runner +supervision/readiness, but not slice placement: +- `LarqlRunner` is a Worker-managed runner subtype that supervises an upstream `larql serve` child process. +- The MLX runner remains the head runtime; LARQL peers serve cold FFN / expert slices. +- The MLX head never loads a vindex. Vindexes are consumed from HuggingFace and produced by the separate `skulk-vindex-publisher` repo. +- Phase 4 slice-placement code is blocked until the Phase 3 MLX FFN delegation spike confirms the design. + ### Logging & Observability Centralized logging uses a three-layer stack: - **Structured JSON stdout**: When `logging.enabled` is `true` and `logging.ingest_url` is set in `skulk.yaml` (or dashboard Settings), skulk emits one JSON object per line on stdout (alongside human-readable stderr). Settings sync to all nodes via gossipsub. Configured in `src/exo/shared/logging.py`. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d6118ce5a..2568c7f39 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -76,6 +76,7 @@ This starts a Vite dev server on port 3000 with hot reload. The dev server proxi - `src/exo/worker/` — Worker node (inference, runner management, download coordination) - `src/exo/store/` — Model store (registry, downloads, config, model optimizer) - `src/exo/shared/` — Shared types, constants, topology +- `docs/adr/` — Architecture decision records for changes that affect runner taxonomy, placement, event-sourced state, APIs, or operator workflow - `website/docs/` — Docusaurus documentation source, including API guide and model-capability docs ## Development Guidelines diff --git a/docs/adr/0001-larql-runner-type.md b/docs/adr/0001-larql-runner-type.md new file mode 100644 index 000000000..30d4c1983 --- /dev/null +++ b/docs/adr/0001-larql-runner-type.md @@ -0,0 +1,62 @@ +# ADR 0001: LARQL Runner Type + +Status: Accepted for planning +Date: 2026-05-12 +Parent roadmap: https://github.com/Foxlight-Foundation/Skulk/issues/173 +Tracking issue: https://github.com/Foxlight-Foundation/Skulk/issues/153 +Source plan: https://github.com/Foxlight-Foundation/Skulk/blob/claude/understand-larql-repo-sJqA1/docs/slice-placement-and-vindex-publisher-plan.md + +## Context + +Skulk's current execution model is centered on worker-managed runner +subprocesses. A worker observes event-sourced placement state, downloads or +stages the assigned model artifacts, starts a runner subprocess, and reports +lifecycle transitions through the same event stream every node applies. + +LARQL introduces a different execution role: a process that serves vindex-backed +FFN or expert slices over HTTP. That process still needs the same operational +properties Skulk expects from model runners: deterministic startup, supervised +shutdown, crash visibility, logging, readiness state, and eventual placement +metadata. + +## Decision + +Skulk will model LARQL as a first-class `LarqlRunner` runner type managed by +the worker. The worker will supervise a child `larql-server` process alongside +the existing MLX runner subprocesses. + +The `LarqlRunner` will use Skulk's runner lifecycle conventions: + +- worker-owned process supervision +- stdout/stderr forwarding into Skulk logging +- readiness and failure reporting through event-sourced state +- shutdown driven by instance and runner lifecycle events + +The initial implementation will treat `larql-server` as an upstream binary +dependency, not as in-tree Rust or Python code. + +## Consequences + +`LarqlRunner` becomes part of Skulk's runner taxonomy. Future implementation +work must add explicit runner metadata rather than overloading MLX shard +metadata or treating LARQL as an external sidecar. + +Operators should be able to reason about LARQL-backed slices through the same +dashboard, state, diagnostics, and logging surfaces used for MLX runners. + +Skulk remains insulated from LARQL internals. The integration boundary is the +LARQL server process and its HTTP contract. + +## Rejected Alternatives + +### Sidecar + +A sidecar would be quick to prototype, but it moves process lifecycle, +readiness, logs, and crash recovery outside Skulk. That creates a second +operator workflow and makes slice placement harder to explain and diagnose. + +### In-tree Port + +Reimplementing LARQL's slice protocol inside Skulk would be a large, slow fork +of upstream LARQL. It would also make it harder to pick up future LARQL +improvements in vindex format, server behavior, and FFN/expert endpoints. diff --git a/docs/adr/0002-head-mlx-cold-larql.md b/docs/adr/0002-head-mlx-cold-larql.md new file mode 100644 index 000000000..3720c1743 --- /dev/null +++ b/docs/adr/0002-head-mlx-cold-larql.md @@ -0,0 +1,73 @@ +# ADR 0002: MLX Head With LARQL Cold Tier + +Status: Accepted for planning; gated by Phase 3 feasibility +Date: 2026-05-12 +Parent roadmap: https://github.com/Foxlight-Foundation/Skulk/issues/173 +Tracking issue: https://github.com/Foxlight-Foundation/Skulk/issues/154 +Gate issues: https://github.com/Foxlight-Foundation/Skulk/issues/161 and https://github.com/Foxlight-Foundation/Skulk/issues/162 +Source plan: https://github.com/Foxlight-Foundation/Skulk/blob/claude/understand-larql-repo-sJqA1/docs/slice-placement-and-vindex-publisher-plan.md + +## Context + +Skulk's strongest runtime path is MLX on Apple Silicon. Slice placement should +extend that path instead of replacing it. The desired architecture lets a Mac +head node keep the hot attention/router path local while RAM-rich commodity +peers serve cold FFN or expert weights from LARQL vindexes. + +This decision is expensive to reverse once placement state, runner assignment, +and API surfaces begin to encode slice responsibilities. + +## Decision + +The MLX runner remains the head runtime for slice mode. It owns the standard +MLX weights needed for the head role: embeddings, attention, norms, router, and +any locally assigned layers. It delegates selected per-layer FFN or expert work +to `LarqlRunner` peers over HTTP. + +The MLX head never loads a vindex. Vindexes are cold-tier artifacts consumed by +LARQL peers. + +The default wire format for delegated tensors is f16. i8 remains an explicit +future opt-in only where the LARQL contract supports it and Skulk can preserve +correctness. + +## Feasibility Gate + +This ADR is accepted for planning, not yet accepted for irreversible runtime +implementation. Phase 3 must prove that Skulk's MLX path can delegate a +per-layer FFN or expert step and continue generation with acceptable overhead. + +If MLX does not expose usable hooks and a manual forward-pass split is too +fragile, too invasive, or too slow, this ADR must be superseded before Phase 4 +slice-placement work starts. + +## Consequences + +Existing MLX single-node and MLX pipeline placement remain the default path for +models that fit on the selected head node. + +Slice placement is additive. It is only considered when the normal MLX path +cannot fit the model or when the operator explicitly chooses a slice-mode flow +in future UI/API work. + +The slice plan must identify which LARQL peer serves which preset, layer range, +expert range, and vindex URI so the MLX head can dispatch remote FFN/expert +calls deterministically. + +## Rejected Alternatives + +### Replace the Head Runtime With LARQL + +Replacing MLX would discard Skulk's current strongest execution path and make +Apple Silicon performance dependent on a new serving stack. + +### Load Vindexes on the Head + +Loading vindexes on the head duplicates cold-tier storage and undermines the +purpose of using commodity RAM-rich peers for dormant weights. + +### General Remote-Compute Abstraction + +The v1 design targets LARQL's concrete FFN/expert server contract. A generic +remote execution abstraction would add surface area before Skulk has proven the +basic slice-mode value proposition. diff --git a/docs/adr/0003-vindex-provenance.md b/docs/adr/0003-vindex-provenance.md new file mode 100644 index 000000000..a1d3d57aa --- /dev/null +++ b/docs/adr/0003-vindex-provenance.md @@ -0,0 +1,55 @@ +# ADR 0003: Vindex Provenance + +Status: Accepted for planning +Date: 2026-05-12 +Parent roadmap: https://github.com/Foxlight-Foundation/Skulk/issues/173 +Tracking issue: https://github.com/Foxlight-Foundation/Skulk/issues/155 +Publisher issue: https://github.com/Foxlight-Foundation/Skulk/issues/156 +Source plan: https://github.com/Foxlight-Foundation/Skulk/blob/claude/understand-larql-repo-sJqA1/docs/slice-placement-and-vindex-publisher-plan.md + +## Context + +LARQL vindexes are directory-shaped artifacts derived from source model +weights. Extracting them can require substantial scratch disk, time, and +toolchain setup. Running that extraction on every Skulk user's machine would +make first use slow and operationally fragile. + +Skulk already has model-store and download concepts for consuming artifacts. +The clean boundary is to make Skulk a vindex consumer and move extraction into +a dedicated publisher workflow. + +## Decision + +Skulk will consume vindexes from HuggingFace URIs such as `hf://...`. Skulk will +not extract vindexes in-tree. + +Extraction, publication, manifest curation, and catalogue governance live in a +separate sibling repository: `skulk-vindex-publisher`. + +The publisher repository owns scheduled LARQL extraction and publication jobs. +Skulk owns runtime consumption, local caching/staging, and placement metadata. + +## Consequences + +Skulk does not add a Rust toolchain or LARQL extraction dependency to its normal +runtime setup. + +Future model-store work must support directory-shaped vindex artifacts, but it +does not need to know how to produce them. + +The published vindex URI convention becomes part of the contract between the +publisher repo and Skulk's placement/runtime code. + +## Rejected Alternatives + +### Extract Inside Skulk + +This would push heavyweight extraction work onto every operator machine, +including laptops that only need to run inference. It also expands Skulk's +runtime dependency surface for a build-time artifact-production task. + +### No Curated Catalogue + +Without a curated catalogue, users would need to find community vindexes or +produce their own before slice mode is useful. That weakens the operator +experience and makes supported-model behavior harder to reproduce. diff --git a/docs/adr/README.md b/docs/adr/README.md new file mode 100644 index 000000000..c3c61596e --- /dev/null +++ b/docs/adr/README.md @@ -0,0 +1,27 @@ +# Architecture Decision Records + +Architecture Decision Records capture choices that are expensive to reverse +once they affect Skulk's event-sourced state, runner taxonomy, placement +contracts, public APIs, or operator workflow. + +## Format + +Each ADR uses this structure: + +- **Status:** Proposed, Accepted for planning, Accepted, Superseded, or Rejected. +- **Context:** The forces that made the decision necessary. +- **Decision:** The choice Skulk will implement. +- **Consequences:** Operational and implementation effects of the decision. +- **Rejected alternatives:** Options considered and why they were not chosen. + +## Numbering + +Use monotonically increasing four-digit filenames: + +```text +0001-short-title.md +0002-short-title.md +``` + +Do not renumber existing ADRs. If a decision changes, add a new ADR that +supersedes the old one. diff --git a/docs/architecture.md b/docs/architecture.md index bd6ad7cfd..1cdf103c4 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -66,6 +66,13 @@ This is where Skulk selects inference behavior such as: - MLX execution path - KV cache backend choice +Skulk's accepted LARQL planning ADRs add a runner subtype: `LarqlRunner`. +Phase 2 implements its internal worker-managed supervision path. It starts and +readiness-checks an upstream `larql serve` child process for directory-shaped +vindex artifacts, but placement does not create LARQL runners yet. Phase 3 of +the LARQL roadmap must first prove MLX can delegate per-layer FFN work with +acceptable overhead. + ### API The API server exposes: @@ -153,6 +160,29 @@ Instead, it changes how model artifacts are sourced: The rest of the system still uses the same master, worker, API, and placement model. +The LARQL roadmap extends this artifact model to vindexes. Skulk will consume +HuggingFace-hosted vindex directories, while extraction and publication live in +the separate `skulk-vindex-publisher` repository. Skulk does not extract +vindexes in-tree. + +## Planned LARQL Slice Mode + +The accepted planning ADRs in `docs/adr/` define the intended slice-mode shape: + +- `LarqlRunner` is a first-class worker-managed runner type, not an unmanaged + sidecar. +- The MLX runner remains the head runtime. It owns the hot attention/router + path and delegates selected FFN or expert work to LARQL peers. +- The MLX head never loads a vindex; vindexes are cold-tier artifacts served by + LARQL runners. +- Slice mode is additive. Existing MLX placement remains unchanged for models + that fit on the selected head node. + +Phase 2 implements internal `LarqlRunner` supervision and readiness state, but +slice mode is still not an operator-visible placement mode. Phase 4 placement +work is blocked until the Phase 3 MLX delegation spike confirms that the design +is viable. + ## Where the Dashboard Fits The dashboard is not a separate product or service. diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 34df2f74a..bd52e7be8 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -13,8 +13,10 @@ from exo.download.download_utils import ( RepoDownloadProgress, delete_model, + is_vindex_directory_complete, map_repo_download_progress_to_download_progress_data, resolve_model_in_path, + resolve_vindex_location, ) from exo.download.shard_downloader import ShardDownloader from exo.shared.constants import EXO_MODELS_DIR @@ -40,7 +42,11 @@ DownloadPending, DownloadProgress, ) -from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata +from exo.shared.types.worker.shards import ( + LarqlShardMetadata, + PipelineShardMetadata, + ShardMetadata, +) from exo.store.config import resolve_config_path from exo.utils.channels import Receiver, Sender from exo.utils.task_group import TaskGroup @@ -342,7 +348,12 @@ async def _start_download(self, shard: ShardMetadata) -> None: # clear the stale status and re-download from scratch. if isinstance(status, DownloadCompleted) and status.model_directory: model_dir = Path(status.model_directory) - if not model_dir.is_dir() or not (model_dir / "config.json").exists(): + directory_is_complete = ( + is_vindex_directory_complete(model_dir, shard.vindex_uri) + if isinstance(shard, LarqlShardMetadata) + else model_dir.is_dir() and (model_dir / "config.json").exists() + ) + if not directory_is_complete: logger.info( f"DownloadCoordinator: {model_id} was DownloadCompleted but " f"model directory {status.model_directory} no longer exists, re-downloading" @@ -367,7 +378,16 @@ async def _start_download(self, shard: ShardMetadata) -> None: return # Check EXO_MODELS_PATH for pre-downloaded models - found_path = resolve_model_in_path(model_id) + vindex_location = ( + resolve_vindex_location(model_id, shard.vindex_uri) + if isinstance(shard, LarqlShardMetadata) + else None + ) + found_path = ( + vindex_location.path + if vindex_location is not None + else resolve_model_in_path(model_id) + ) if found_path is not None: logger.info( f"DownloadCoordinator: Model {model_id} found in EXO_MODELS_PATH at {found_path}" @@ -377,7 +397,9 @@ async def _start_download(self, shard: ShardMetadata) -> None: node_id=self.node_id, total=shard.model_card.storage_size, model_directory=str(found_path), - read_only=True, + read_only=vindex_location.read_only + if vindex_location is not None + else True, ) self.download_status[model_id] = completed await self.event_sender.send( diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 4897df840..954590a3e 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -1,12 +1,14 @@ import asyncio import hashlib +import json import os import shutil import ssl import time import traceback from collections.abc import Awaitable -from datetime import timedelta +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Callable, Literal from urllib.parse import urljoin @@ -21,6 +23,7 @@ from loguru import logger from pydantic import ( TypeAdapter, + ValidationError, ) from exo.download.huggingface_utils import ( @@ -43,6 +46,17 @@ ) from exo.shared.types.worker.shards import ShardMetadata +VINDEX_COMPLETE_MARKER = ".skulk-vindex-complete.json" +VINDEX_COMPLETE_MARKER_ADAPTER = TypeAdapter(dict[str, object]) + + +@dataclass(frozen=True) +class VindexPathResolution: + """Resolved local LARQL vindex path plus ownership metadata.""" + + path: Path + read_only: bool + class HuggingFaceAuthenticationError(Exception): """Raised when HuggingFace returns 401/403 for a model download.""" @@ -111,11 +125,12 @@ def map_repo_download_progress_to_download_progress_data( def resolve_model_in_path(model_id: ModelId) -> Path | None: - """Search EXO_MODELS_PATH directories for a pre-existing model. + """Search read-only EXO_MODELS_PATH directories for a pre-existing model. Checks each directory for the normalized name (org--model). A candidate is only returned if ``is_model_directory_complete`` confirms all weight - files are present. + files are present. Writable cache directories are intentionally excluded + because callers mark returned paths as externally managed/read-only. Reads the search path dynamically from ``exo.shared.constants`` so that paths added at runtime (e.g. by the model store) are picked up. @@ -131,6 +146,108 @@ def resolve_model_in_path(model_id: ModelId) -> Path | None: return None +def _read_vindex_complete_marker(vindex_dir: Path) -> dict[str, object] | None: + marker = vindex_dir / VINDEX_COMPLETE_MARKER + if not marker.is_file(): + return None + try: + return VINDEX_COMPLETE_MARKER_ADAPTER.validate_json( + marker.read_text(encoding="utf-8") + ) + except (OSError, ValidationError): + return None + + +def is_vindex_directory_complete( + vindex_dir: Path, expected_vindex_uri: str | None = None +) -> bool: + """Return whether a staged LARQL vindex directory has usable contents.""" + + if not vindex_dir.is_dir(): + return False + marker = _read_vindex_complete_marker(vindex_dir) + if marker is None: + return False + if expected_vindex_uri is not None and marker.get("vindex_uri") != expected_vindex_uri: + return False + + has_metadata = False + has_payload = False + for path in vindex_dir.rglob("*"): + if not path.is_file() or path.name == VINDEX_COMPLETE_MARKER: + continue + has_metadata = has_metadata or path.suffix == ".json" + has_payload = has_payload or path.suffix == ".bin" + if has_metadata and has_payload: + return True + return False + + +def mark_vindex_directory_complete(vindex_dir: Path, vindex_uri: str) -> None: + """Persist the success marker for a completed LARQL vindex pull.""" + + marker = vindex_dir / VINDEX_COMPLETE_MARKER + marker.write_text( + json.dumps( + { + "vindex_uri": vindex_uri, + "completed_at": datetime.now(tz=timezone.utc).isoformat(), + }, + sort_keys=True, + ), + encoding="utf-8", + ) + + +def build_vindex_path(vindex_id: ModelId) -> Path: + """Return the configured writable cache path for a LARQL vindex artifact.""" + + return EXO_MODELS_DIR / vindex_id.normalize() + + +def _vindex_search_path() -> tuple[VindexPathResolution, ...]: + import exo.shared.constants as _constants + + configured_paths = _constants.EXO_MODELS_PATH or () + candidates = ( + *((path, True) for path in configured_paths), + (EXO_MODELS_DIR, False), + (Path.home() / ".exo" / "models", False), + (Path.home() / ".exo" / "staging", False), + ) + unique_paths: list[VindexPathResolution] = [] + seen_paths: set[Path] = set() + for candidate, read_only in candidates: + expanded = candidate.expanduser() + if expanded in seen_paths: + continue + seen_paths.add(expanded) + unique_paths.append(VindexPathResolution(path=expanded, read_only=read_only)) + return tuple(unique_paths) + + +def resolve_vindex_location( + vindex_id: ModelId, expected_vindex_uri: str | None = None +) -> VindexPathResolution | None: + """Search model paths for a complete LARQL vindex and return ownership.""" + + normalized = vindex_id.normalize() + for search_root in _vindex_search_path(): + candidate = search_root.path / normalized + if is_vindex_directory_complete(candidate, expected_vindex_uri): + return VindexPathResolution(path=candidate, read_only=search_root.read_only) + return None + + +def resolve_vindex_in_path( + vindex_id: ModelId, expected_vindex_uri: str | None = None +) -> Path | None: + """Search model paths for a complete directory-shaped LARQL vindex.""" + + found = resolve_vindex_location(vindex_id, expected_vindex_uri) + return found.path if found is not None else None + + def build_model_path(model_id: ModelId) -> Path: """Resolve a local filesystem path for *model_id*. diff --git a/src/exo/download/impl_shard_downloader.py b/src/exo/download/impl_shard_downloader.py index 0820380f1..7d298a11e 100644 --- a/src/exo/download/impl_shard_downloader.py +++ b/src/exo/download/impl_shard_downloader.py @@ -1,6 +1,9 @@ import asyncio +import shutil +import uuid from asyncio import create_task from collections.abc import Awaitable +from datetime import timedelta from pathlib import Path from typing import AsyncIterator, Callable @@ -8,7 +11,11 @@ from exo.download.download_utils import ( RepoDownloadProgress, + build_vindex_path, download_shard, + is_vindex_directory_complete, + mark_vindex_directory_complete, + resolve_vindex_in_path, ) from exo.download.shard_downloader import ShardDownloader from exo.shared.models.model_cards import ( @@ -19,6 +26,7 @@ ) from exo.shared.types.memory import Memory from exo.shared.types.worker.shards import ( + LarqlShardMetadata, PipelineShardMetadata, ShardMetadata, ) @@ -115,6 +123,9 @@ def on_progress( async def ensure_shard( self, shard: ShardMetadata, config_only: bool = False ) -> Path: + if isinstance(shard, LarqlShardMetadata): + return await self._ensure_larql_vindex(shard) + allow_patterns = ["config.json"] if config_only else None target_dir, _ = await download_shard( @@ -158,6 +169,78 @@ async def ensure_shard( return target_dir + async def _ensure_larql_vindex(self, shard: LarqlShardMetadata) -> Path: + """Pull or reuse one directory-shaped LARQL vindex artifact.""" + + found = resolve_vindex_in_path(shard.model_card.model_id, shard.vindex_uri) + if found is not None: + await self.on_progress_wrapper( + shard, + RepoDownloadProgress( + repo_id=shard.vindex_uri, + repo_revision="main", + shard=shard, + completed_files=1, + total_files=1, + downloaded=shard.model_card.storage_size, + downloaded_this_session=Memory.from_bytes(0), + total=shard.model_card.storage_size, + overall_speed=0, + overall_eta=timedelta(seconds=0), + status="complete", + ), + ) + return found + + if self.offline: + raise FileNotFoundError( + f"LARQL vindex {shard.vindex_uri} is not available locally" + ) + + target_dir = build_vindex_path(shard.model_card.model_id) + target_dir.parent.mkdir(parents=True, exist_ok=True) + temp_dir = target_dir.with_name(f".{target_dir.name}.partial-{uuid.uuid4().hex}") + process = await asyncio.create_subprocess_exec( + "larql", + "pull", + shard.vindex_uri, + "--output", + str(temp_dir), + ) + return_code = await process.wait() + if return_code != 0: + await asyncio.to_thread(shutil.rmtree, temp_dir, ignore_errors=True) + raise RuntimeError( + f"larql pull failed for {shard.vindex_uri} with exit code {return_code}" + ) + if not temp_dir.is_dir(): + await asyncio.to_thread(shutil.rmtree, temp_dir, ignore_errors=True) + raise RuntimeError(f"larql pull did not create a vindex at {temp_dir}") + mark_vindex_directory_complete(temp_dir, shard.vindex_uri) + if not is_vindex_directory_complete(temp_dir): + await asyncio.to_thread(shutil.rmtree, temp_dir, ignore_errors=True) + raise RuntimeError(f"larql pull produced an incomplete vindex at {temp_dir}") + if target_dir.exists(): + await asyncio.to_thread(shutil.rmtree, target_dir, ignore_errors=True) + temp_dir.rename(target_dir) + await self.on_progress_wrapper( + shard, + RepoDownloadProgress( + repo_id=shard.vindex_uri, + repo_revision="main", + shard=shard, + completed_files=1, + total_files=1, + downloaded=shard.model_card.storage_size, + downloaded_this_session=shard.model_card.storage_size, + total=shard.model_card.storage_size, + overall_speed=0, + overall_eta=timedelta(seconds=0), + status="complete", + ), + ) + return target_dir + async def get_shard_download_status( self, ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]: @@ -195,6 +278,24 @@ async def download_with_semaphore( async def get_shard_download_status_for_shard( self, shard: ShardMetadata ) -> RepoDownloadProgress: + if isinstance(shard, LarqlShardMetadata): + complete = ( + resolve_vindex_in_path(shard.model_card.model_id, shard.vindex_uri) + is not None + ) + return RepoDownloadProgress( + repo_id=shard.vindex_uri, + repo_revision="main", + shard=shard, + completed_files=1 if complete else 0, + total_files=1, + downloaded=shard.model_card.storage_size if complete else Memory(), + downloaded_this_session=Memory(), + total=shard.model_card.storage_size, + overall_speed=0, + overall_eta=timedelta(seconds=0), + status="complete" if complete else "not_started", + ) _, progress = await download_shard( shard, self.on_progress_wrapper, diff --git a/src/exo/download/tests/test_download_verification.py b/src/exo/download/tests/test_download_verification.py index 2d8d076de..bf4a71b07 100644 --- a/src/exo/download/tests/test_download_verification.py +++ b/src/exo/download/tests/test_download_verification.py @@ -12,8 +12,14 @@ from pydantic import TypeAdapter from exo.download.download_utils import ( + build_model_path, + build_vindex_path, delete_model, fetch_file_list_with_cache, + mark_vindex_directory_complete, + resolve_model_in_path, + resolve_vindex_in_path, + resolve_vindex_location, ) from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory @@ -34,6 +40,125 @@ async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]: yield models_dir +def _write_complete_model_directory(model_dir: Path) -> None: + model_dir.mkdir(parents=True) + (model_dir / "config.json").write_text("{}", encoding="utf-8") + (model_dir / "model.safetensors").write_bytes(b"weights") + (model_dir / "model.safetensors.index.json").write_text( + '{"metadata": {}, "weight_map": {"layer.weight": "model.safetensors"}}', + encoding="utf-8", + ) + + +class TestModelPathResolution: + """Tests for read-only search path semantics.""" + + def test_resolve_model_in_path_excludes_writable_models_dir( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """Writable cache models are loadable but not treated as read-only hits.""" + + model_dir = temp_models_dir / model_id.normalize() + _write_complete_model_directory(model_dir) + + with ( + patch("exo.download.download_utils.EXO_MODELS_DIR", temp_models_dir), + patch("exo.shared.constants.EXO_MODELS_PATH", None), + ): + assert resolve_model_in_path(model_id) is None + assert build_model_path(model_id) == model_dir + + def test_resolve_model_in_path_returns_configured_read_only_path( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """Explicit search paths remain externally managed read-only hits.""" + + search_root = tmp_path / "read-only-models" + model_dir = search_root / model_id.normalize() + _write_complete_model_directory(model_dir) + + with patch("exo.shared.constants.EXO_MODELS_PATH", (search_root,)): + assert resolve_model_in_path(model_id) == model_dir + + +class TestVindexPathResolution: + """Tests for directory-shaped LARQL vindex cache discovery.""" + + def test_resolve_vindex_searches_default_models_dir( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """Vindexes in the writable default models dir are reusable offline.""" + + vindex_dir = temp_models_dir / model_id.normalize() + vindex_dir.mkdir(parents=True) + (vindex_dir / "manifest.json").write_text("{}", encoding="utf-8") + (vindex_dir / "weights.bin").write_bytes(b"vindex") + mark_vindex_directory_complete(vindex_dir, "hf://test-org/test-model") + + with ( + patch("exo.download.download_utils.EXO_MODELS_DIR", temp_models_dir), + patch("exo.shared.constants.EXO_MODELS_PATH", None), + ): + assert build_vindex_path(model_id) == vindex_dir + assert resolve_vindex_in_path(model_id, "hf://test-org/test-model") == vindex_dir + resolved = resolve_vindex_location(model_id, "hf://test-org/test-model") + assert resolved is not None + assert resolved.path == vindex_dir + assert not resolved.read_only + + def test_resolve_vindex_rejects_marker_uri_mismatch( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """A cache entry is not reusable for a different vindex URI.""" + + vindex_dir = temp_models_dir / model_id.normalize() + vindex_dir.mkdir(parents=True) + (vindex_dir / "manifest.json").write_text("{}", encoding="utf-8") + (vindex_dir / "weights.bin").write_bytes(b"vindex") + mark_vindex_directory_complete(vindex_dir, "hf://test-org/old-vindex") + + with ( + patch("exo.download.download_utils.EXO_MODELS_DIR", temp_models_dir), + patch("exo.shared.constants.EXO_MODELS_PATH", None), + ): + assert resolve_vindex_in_path(model_id, "hf://test-org/new-vindex") is None + + def test_resolve_vindex_reports_configured_search_path_read_only( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """Explicit vindex search roots remain protected from deletion.""" + + search_root = tmp_path / "read-only-vindexes" + vindex_dir = search_root / model_id.normalize() + vindex_dir.mkdir(parents=True) + (vindex_dir / "manifest.json").write_text("{}", encoding="utf-8") + (vindex_dir / "weights.bin").write_bytes(b"vindex") + mark_vindex_directory_complete(vindex_dir, "hf://test-org/test-model") + + with patch("exo.shared.constants.EXO_MODELS_PATH", (search_root,)): + resolved = resolve_vindex_location(model_id, "hf://test-org/test-model") + + assert resolved is not None + assert resolved.path == vindex_dir + assert resolved.read_only + + def test_resolve_vindex_rejects_unmarked_partial_directory( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """Partial vindex directories are not reused without a success marker.""" + + vindex_dir = temp_models_dir / model_id.normalize() + vindex_dir.mkdir(parents=True) + (vindex_dir / "manifest.json").write_text("{}", encoding="utf-8") + (vindex_dir / "weights.bin").write_bytes(b"partial") + + with ( + patch("exo.download.download_utils.EXO_MODELS_DIR", temp_models_dir), + patch("exo.shared.constants.EXO_MODELS_PATH", None), + ): + assert resolve_vindex_in_path(model_id) is None + + class TestFileVerification: """Tests for file size verification in _download_file.""" diff --git a/src/exo/download/tests/test_larql_vindex_download.py b/src/exo/download/tests/test_larql_vindex_download.py new file mode 100644 index 000000000..2a1243e7a --- /dev/null +++ b/src/exo/download/tests/test_larql_vindex_download.py @@ -0,0 +1,147 @@ +"""Tests for LARQL vindex download path handling.""" + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from exo.download.coordinator import DownloadCoordinator +from exo.download.download_utils import VindexPathResolution +from exo.download.impl_shard_downloader import ResumableShardDownloader +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.types.commands import ForwarderDownloadCommand +from exo.shared.types.common import NodeId +from exo.shared.types.events import Event, NodeDownloadProgress +from exo.shared.types.memory import Memory +from exo.shared.types.worker.downloads import DownloadCompleted +from exo.shared.types.worker.shards import LarqlShardMetadata +from exo.utils.channels import channel + + +class _SuccessfulProcess: + async def wait(self) -> int: + return 0 + + +class _FailedProcess: + async def wait(self) -> int: + return 1 + + +def _larql_shard() -> LarqlShardMetadata: + return LarqlShardMetadata( + model_card=ModelCard( + model_id=ModelId("skulk/test-vindex"), + storage_size=Memory.from_mb(128), + n_layers=12, + hidden_size=2048, + supports_tensor=False, + tasks=[ModelTask.TextGeneration], + ), + device_rank=0, + world_size=1, + start_layer=0, + end_layer=12, + n_layers=12, + vindex_uri="hf://skulk/test-vindex", + preset="expert-server", + ) + + +@pytest.mark.asyncio +async def test_larql_vindex_pull_uses_configured_models_dir(tmp_path: Path) -> None: + """Pulled vindexes are written under the configured writable models dir.""" + + models_dir = tmp_path / "configured-models" + shard = _larql_shard() + recorded_commands: list[tuple[str, ...]] = [] + + async def fake_create_subprocess_exec(*args: str) -> _SuccessfulProcess: + recorded_commands.append(tuple(args)) + output_dir = Path(args[-1]) + output_dir.mkdir(parents=True) + (output_dir / "manifest.json").write_text("{}", encoding="utf-8") + (output_dir / "weights.bin").write_bytes(b"vindex") + return _SuccessfulProcess() + + with ( + patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir), + patch( + "exo.download.impl_shard_downloader.resolve_vindex_in_path", + return_value=None, + ), + patch("asyncio.create_subprocess_exec", new=fake_create_subprocess_exec), + ): + result = await ResumableShardDownloader().ensure_shard(shard) + + target_dir = models_dir / shard.model_card.model_id.normalize() + assert result == target_dir + assert target_dir.is_dir() + assert recorded_commands[0][:4] == ("larql", "pull", shard.vindex_uri, "--output") + assert recorded_commands[0][4].startswith(str(models_dir / f".{target_dir.name}")) + + +@pytest.mark.asyncio +async def test_larql_vindex_pull_cleans_partial_directory_on_failure( + tmp_path: Path, +) -> None: + """Failed pulls leave no reusable partial vindex directory behind.""" + + models_dir = tmp_path / "configured-models" + shard = _larql_shard() + partial_dir: Path | None = None + + async def fake_create_subprocess_exec(*args: str) -> _FailedProcess: + nonlocal partial_dir + partial_dir = Path(args[-1]) + partial_dir.mkdir(parents=True) + (partial_dir / "manifest.json").write_text("{}", encoding="utf-8") + (partial_dir / "weights.bin").write_bytes(b"partial") + return _FailedProcess() + + with ( + patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir), + patch( + "exo.download.impl_shard_downloader.resolve_vindex_in_path", + return_value=None, + ), + patch("asyncio.create_subprocess_exec", new=fake_create_subprocess_exec), + pytest.raises(RuntimeError, match="larql pull failed"), + ): + await ResumableShardDownloader().ensure_shard(shard) + + target_dir = models_dir / shard.model_card.model_id.normalize() + assert partial_dir is not None + assert not partial_dir.exists() + assert not target_dir.exists() + + +@pytest.mark.asyncio +async def test_larql_vindex_writable_cache_hit_remains_deletable( + tmp_path: Path, +) -> None: + """Coordinator preserves writable ownership when reusing local vindexes.""" + + _, command_receiver = channel[ForwarderDownloadCommand]() + event_sender, event_receiver = channel[Event]() + shard = _larql_shard() + vindex_dir = tmp_path / "configured-models" / shard.model_card.model_id.normalize() + coordinator = DownloadCoordinator( + node_id=NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"), + shard_downloader=ResumableShardDownloader(), + download_command_receiver=command_receiver, + event_sender=event_sender, + ) + + with patch( + "exo.download.coordinator.resolve_vindex_location", + return_value=VindexPathResolution(path=vindex_dir, read_only=False), + ): + await coordinator._start_download(shard) # pyright: ignore[reportPrivateUsage] + + event = await event_receiver.receive() + + assert isinstance(event, NodeDownloadProgress) + assert isinstance(event.download_progress, DownloadCompleted) + assert event.download_progress.model_directory == str(vindex_dir) + assert not event.download_progress.read_only diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index efceeb058..68060a7cf 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -14,6 +14,7 @@ InputChunkReceived, InstanceCreated, InstanceDeleted, + LarqlRunnerReadinessUpdated, NodeDownloadProgress, NodeGatheredInfo, NodeTimedOut, @@ -43,6 +44,7 @@ from exo.shared.types.topology import Connection, RDMAConnection from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId +from exo.shared.types.worker.larql import LarqlRunnerReadiness from exo.shared.types.worker.runners import RunnerId, RunnerShutdown, RunnerStatus from exo.utils.info_gatherer.info_gatherer import ( MacmonMetrics, @@ -85,6 +87,8 @@ def event_apply(event: Event, state: State) -> State: return apply_node_gathered_info(event, state) case RunnerStatusUpdated(): return apply_runner_status_updated(event, state) + case LarqlRunnerReadinessUpdated(): + return apply_larql_runner_readiness_updated(event, state) case TaskCreated(): return apply_task_created(event, state) case TaskDeleted(): @@ -201,10 +205,24 @@ def apply_instance_created(event: InstanceCreated, state: State) -> State: def apply_instance_deleted(event: InstanceDeleted, state: State) -> State: + instance = state.instances.get(event.instance_id) new_instances: Mapping[InstanceId, Instance] = { iid: inst for iid, inst in state.instances.items() if iid != event.instance_id } - return state.model_copy(update={"instances": new_instances}) + if instance is None: + return state.model_copy(update={"instances": new_instances}) + removed_runner_ids = set(instance.shard_assignments.runner_to_shard) + new_larql_readiness: Mapping[RunnerId, LarqlRunnerReadiness] = { + rid: readiness + for rid, readiness in state.larql_runner_readiness.items() + if rid not in removed_runner_ids + } + return state.model_copy( + update={ + "instances": new_instances, + "larql_runner_readiness": new_larql_readiness, + } + ) def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State: @@ -212,7 +230,17 @@ def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> Sta new_runners: Mapping[RunnerId, RunnerStatus] = { rid: rs for rid, rs in state.runners.items() if rid != event.runner_id } - return state.model_copy(update={"runners": new_runners}) + new_larql_readiness: Mapping[RunnerId, LarqlRunnerReadiness] = { + rid: readiness + for rid, readiness in state.larql_runner_readiness.items() + if rid != event.runner_id + } + return state.model_copy( + update={ + "runners": new_runners, + "larql_runner_readiness": new_larql_readiness, + } + ) new_runners = { **state.runners, event.runner_id: event.runner_status, @@ -220,6 +248,19 @@ def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> Sta return state.model_copy(update={"runners": new_runners}) +def apply_larql_runner_readiness_updated( + event: LarqlRunnerReadinessUpdated, state: State +) -> State: + """Store the latest replay-safe LARQL readiness record for one runner.""" + + readiness = event.readiness + new_readiness: Mapping[RunnerId, LarqlRunnerReadiness] = { + **state.larql_runner_readiness, + readiness.runner_id: readiness, + } + return state.model_copy(update={"larql_runner_readiness": new_readiness}) + + def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: topology = copy.deepcopy(state.topology) topology.remove_node(event.node_id) diff --git a/src/exo/shared/tests/test_apply/test_larql_readiness.py b/src/exo/shared/tests/test_apply/test_larql_readiness.py new file mode 100644 index 000000000..d0a8d3840 --- /dev/null +++ b/src/exo/shared/tests/test_apply/test_larql_readiness.py @@ -0,0 +1,107 @@ +from exo.shared.apply import apply +from exo.shared.models.model_cards import ModelId +from exo.shared.types.common import NodeId +from exo.shared.types.events import ( + IndexedEvent, + InstanceCreated, + InstanceDeleted, + LarqlRunnerReadinessUpdated, + RunnerStatusUpdated, +) +from exo.shared.types.memory import Memory +from exo.shared.types.state import State +from exo.shared.types.worker.instances import InstanceId +from exo.shared.types.worker.larql import LarqlRunnerReadiness +from exo.shared.types.worker.runners import RunnerShutdown +from exo.worker.tests.constants import RUNNER_1_ID +from exo.worker.tests.unittests.conftest import ( + get_mlx_ring_instance, + get_pipeline_shard_metadata, +) + + +def test_larql_readiness_updates_state() -> None: + readiness = LarqlRunnerReadiness( + runner_id=RUNNER_1_ID, + vindex_uri="hf://skulk/gemma-4-26b-a4b-expert-server-q4-k-vindex", + preset="expert-server", + start_layer=4, + end_layer=12, + expert_range=None, + units_manifest_path=None, + host="127.0.0.1", + port=49152, + status="ready", + ram_footprint=Memory.from_mb(128), + ) + + state = apply( + State(), + IndexedEvent( + idx=0, + event=LarqlRunnerReadinessUpdated(readiness=readiness), + ), + ) + + assert state.larql_runner_readiness[RUNNER_1_ID] == readiness + + +def test_larql_readiness_is_removed_on_runner_shutdown() -> None: + readiness = LarqlRunnerReadiness( + runner_id=RUNNER_1_ID, + vindex_uri="hf://skulk/gemma", + preset="full", + start_layer=0, + end_layer=1, + host="127.0.0.1", + port=49152, + status="ready", + ) + state = State(larql_runner_readiness={RUNNER_1_ID: readiness}) + + updated = apply( + state, + IndexedEvent( + idx=0, + event=RunnerStatusUpdated( + runner_id=RUNNER_1_ID, + runner_status=RunnerShutdown(), + ), + ), + ) + + assert RUNNER_1_ID not in updated.larql_runner_readiness + + +def test_larql_readiness_is_removed_on_instance_delete() -> None: + instance = get_mlx_ring_instance( + instance_id=InstanceId("instance-a"), + model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"), + node_to_runner={NodeId("node-a"): RUNNER_1_ID}, + runner_to_shard={ + RUNNER_1_ID: get_pipeline_shard_metadata( + ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"), 0 + ) + }, + ) + readiness = LarqlRunnerReadiness( + runner_id=RUNNER_1_ID, + vindex_uri="hf://skulk/gemma", + preset="full", + start_layer=0, + end_layer=1, + host="127.0.0.1", + port=49152, + status="ready", + ) + state = apply(State(), IndexedEvent(idx=0, event=InstanceCreated(instance=instance))) + state = state.model_copy( + update={"larql_runner_readiness": {RUNNER_1_ID: readiness}} + ) + + updated = apply( + state, + IndexedEvent(idx=1, event=InstanceDeleted(instance_id=instance.instance_id)), + ) + + assert RUNNER_1_ID not in updated.larql_runner_readiness diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index 3d26a8ea6..5e0f26ee7 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -11,6 +11,7 @@ from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId +from exo.shared.types.worker.larql import LarqlRunnerReadiness from exo.shared.types.worker.runners import RunnerId, RunnerStatus from exo.utils.info_gatherer.info_gatherer import GatheredInfo from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel @@ -75,6 +76,13 @@ class RunnerStatusUpdated(BaseEvent): runner_status: RunnerStatus +@final +class LarqlRunnerReadinessUpdated(BaseEvent): + """Cluster-visible readiness state for a supervised LARQL server.""" + + readiness: LarqlRunnerReadiness + + class NodeTimedOut(BaseEvent): node_id: NodeId @@ -168,6 +176,7 @@ class TracingStateChanged(BaseEvent): | InstanceCreated | InstanceDeleted | RunnerStatusUpdated + | LarqlRunnerReadinessUpdated | NodeTimedOut | NodeGatheredInfo | NodeDownloadProgress diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py index fa0b0f5e2..27c12fbfa 100644 --- a/src/exo/shared/types/state.py +++ b/src/exo/shared/types/state.py @@ -20,6 +20,7 @@ from exo.shared.types.tasks import Task, TaskId from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId +from exo.shared.types.worker.larql import LarqlRunnerReadiness from exo.shared.types.worker.runners import RunnerId, RunnerStatus from exo.utils.pydantic_ext import CamelCaseModel @@ -42,6 +43,7 @@ class State(CamelCaseModel): ) instances: Mapping[InstanceId, Instance] = {} runners: Mapping[RunnerId, RunnerStatus] = {} + larql_runner_readiness: Mapping[RunnerId, LarqlRunnerReadiness] = {} downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {} tasks: Mapping[TaskId, Task] = {} last_seen: Mapping[NodeId, datetime] = {} diff --git a/src/exo/shared/types/worker/larql.py b/src/exo/shared/types/worker/larql.py new file mode 100644 index 000000000..0f6f064c2 --- /dev/null +++ b/src/exo/shared/types/worker/larql.py @@ -0,0 +1,40 @@ +from typing import Literal, final + +from pydantic import Field + +from exo.shared.types.memory import Memory +from exo.shared.types.worker.runners import RunnerId +from exo.shared.types.worker.shards import LarqlExpertRange, LarqlPreset +from exo.utils.pydantic_ext import CamelCaseModel + +LarqlReadinessStatus = Literal["ready", "not_ready", "failed"] + + +@final +class LarqlRunnerReadiness(CamelCaseModel): + """Event-sourced readiness metadata for one supervised LARQL server.""" + + runner_id: RunnerId = Field(description="Runner that owns this LARQL server.") + vindex_uri: str = Field(description="Immutable source URI for the vindex artifact.") + preset: LarqlPreset = Field(description="LARQL serving preset.") + start_layer: int = Field(ge=0, description="Inclusive first served layer.") + end_layer: int = Field(ge=0, description="Exclusive final served layer.") + expert_range: LarqlExpertRange | None = Field( + default=None, + description="Optional half-open expert range for expert-server slices.", + ) + units_manifest_path: str | None = Field( + default=None, + description="Optional local LARQL units manifest path.", + ) + host: str = Field(description="Host where the local LARQL server listens.") + port: int = Field(ge=1, le=65535, description="TCP port for the LARQL server.") + status: LarqlReadinessStatus = Field(description="Current readiness state.") + ram_footprint: Memory | None = Field( + default=None, + description="Measured resident memory for the LARQL process, when known.", + ) + error_message: str | None = Field( + default=None, + description="Failure or readiness error detail, when unavailable.", + ) diff --git a/src/exo/shared/types/worker/shards.py b/src/exo/shared/types/worker/shards.py index 59a6c54eb..a10b702d6 100644 --- a/src/exo/shared/types/worker/shards.py +++ b/src/exo/shared/types/worker/shards.py @@ -1,10 +1,28 @@ from enum import Enum -from typing import TypeAlias, final +from typing import Literal, TypeAlias, final -from pydantic import Field +from pydantic import Field, model_validator from exo.shared.models.model_cards import ModelCard -from exo.utils.pydantic_ext import TaggedModel +from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel + +LarqlPreset = Literal["full", "expert-server"] + + +@final +class LarqlExpertRange(CamelCaseModel): + """Half-open expert range served by one LARQL expert-server runner.""" + + start_expert: int = Field(ge=0, description="Inclusive first expert index.") + end_expert: int = Field(ge=0, description="Exclusive final expert index.") + + @model_validator(mode="after") + def validate_non_empty(self) -> "LarqlExpertRange": + """Require a non-empty expert interval.""" + + if self.end_expert <= self.start_expert: + raise ValueError("end_expert must be greater than start_expert") + return self class Sharding(str, Enum): @@ -79,6 +97,56 @@ class TensorShardMetadata(BaseShardMetadata): pass +@final +class LarqlShardMetadata(BaseShardMetadata): + """Shard metadata for a future worker-managed LARQL cold-tier runner.""" + + vindex_uri: str = Field( + description="Immutable URI for the vindex directory artifact." + ) + preset: LarqlPreset = Field(description="LARQL serving preset for this shard.") + local_vindex_path: str | None = Field( + default=None, + description="Resolved local vindex directory after staging, if known.", + ) + server_host: str = Field( + default="127.0.0.1", + description="Local bind host for the supervised LARQL HTTP server.", + ) + server_port: int | None = Field( + default=None, + ge=1, + le=65535, + description="Requested LARQL port; omitted means allocate a free local port.", + ) + expert_range: LarqlExpertRange | None = Field( + default=None, + description="Optional expert range for expert-server slices.", + ) + units_manifest_path: str | None = Field( + default=None, + description="Optional LARQL units manifest path; mutually exclusive with expert_range.", + ) + max_crash_restarts: int = Field( + default=3, + ge=0, + description="Maximum ordinary crash restarts before terminal failure.", + ) + readiness_timeout_seconds: float = Field( + default=30.0, + gt=0, + description="Maximum time to wait for LARQL readiness after process start.", + ) + + @model_validator(mode="after") + def validate_slice_arguments(self) -> "LarqlShardMetadata": + """Reject ambiguous expert selection for LARQL serve commands.""" + + if self.expert_range is not None and self.units_manifest_path is not None: + raise ValueError("expert_range and units_manifest_path are mutually exclusive") + return self + + ShardMetadata: TypeAlias = ( - PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata + PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata | LarqlShardMetadata ) diff --git a/src/exo/store/model_store.py b/src/exo/store/model_store.py index 7029d7bc9..47272064b 100644 --- a/src/exo/store/model_store.py +++ b/src/exo/store/model_store.py @@ -94,6 +94,7 @@ class StoreModelEntry(BaseModel): model_config = ConfigDict(frozen=True, strict=True, extra="forbid") model_id: str + artifact_kind: Literal["model", "vindex"] = "model" store_path: str files: list[str] downloaded_at: str @@ -226,6 +227,7 @@ def register_model( model_path: Path, files: list[str], total_bytes: int, + artifact_kind: Literal["model", "vindex"] = "model", ) -> None: """Add or update *model_id* in the registry. @@ -238,10 +240,13 @@ def register_model( Must be inside ``store_path``. files: List of file paths relative to *model_path*. total_bytes: Sum of file sizes in bytes. + artifact_kind: Whether this entry contains MLX weights or a LARQL + vindex directory. """ relative_path = str(model_path.relative_to(self._store_path)) entry = StoreModelEntry( model_id=model_id, + artifact_kind=artifact_kind, store_path=relative_path, files=files, downloaded_at=datetime.now(tz=timezone.utc).isoformat(), @@ -253,6 +258,23 @@ def register_model( f"({total_bytes:,} bytes, {len(files)} files)" ) + def register_vindex( + self, + vindex_id: str, + vindex_path: Path, + files: list[str], + total_bytes: int, + ) -> None: + """Add or update a directory-shaped LARQL vindex artifact.""" + + self.register_model( + vindex_id, + vindex_path, + files, + total_bytes, + artifact_kind="vindex", + ) + def list_files_for_model(self, model_id: str) -> list[str] | None: """Return the file list for *model_id* from the registry, or ``None``. diff --git a/src/exo/store/tests/test_model_store_vindex.py b/src/exo/store/tests/test_model_store_vindex.py new file mode 100644 index 000000000..5a3c18ff9 --- /dev/null +++ b/src/exo/store/tests/test_model_store_vindex.py @@ -0,0 +1,27 @@ +from pathlib import Path + +from exo.store.model_store import ModelStore + + +def test_model_store_registers_vindex_directory(tmp_path: Path) -> None: + root = tmp_path / "store" + vindex = root / "skulk--gemma-vindex" + vindex.mkdir(parents=True) + (vindex / "metadata.json").write_text("{}") + (vindex / "weights.bin").write_bytes(b"abc") + + store = ModelStore(root) + store.register_vindex( + "skulk/gemma-vindex", + vindex, + ["metadata.json", "weights.bin"], + 5, + ) + + entries = store.list_models() + + assert len(entries) == 1 + assert entries[0].artifact_kind == "vindex" + assert entries[0].model_id == "skulk/gemma-vindex" + assert store.get_store_path("skulk/gemma-vindex") == vindex + diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 1082f1fce..85613fb1a 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -62,6 +62,7 @@ ) from exo.shared.types.worker.shards import ( CfgShardMetadata, + LarqlShardMetadata, PipelineShardMetadata, ShardMetadata, TensorShardMetadata, @@ -753,6 +754,8 @@ def shard_and_load( "CfgShardMetadata is not supported for text model loading - " "this metadata type is only for image generation models" ) + case LarqlShardMetadata(): + raise ValueError("LarqlShardMetadata is not loaded by the MLX runtime") # TODO: Do we need this? mx.eval(model) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index db9b5337e..56235a9e7 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -11,7 +11,7 @@ from PIL import Image from exo.api.types import ImageEditsTaskParams -from exo.download.download_utils import resolve_model_in_path +from exo.download.download_utils import resolve_model_in_path, resolve_vindex_location from exo.shared.apply import apply from exo.shared.constants import EXO_IMAGE_TRANSPORT_DEBUG from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card @@ -57,7 +57,7 @@ from exo.shared.types.topology import Connection, SocketConnection from exo.shared.types.worker.downloads import DownloadCompleted, DownloadPending from exo.shared.types.worker.runners import RunnerId -from exo.shared.types.worker.shards import ShardMetadata +from exo.shared.types.worker.shards import LarqlShardMetadata, ShardMetadata from exo.store.config import StagingNodeConfig from exo.store.model_store_client import ModelStoreClient from exo.utils.channels import Receiver, Sender, channel @@ -66,8 +66,11 @@ from exo.utils.keyed_backoff import KeyedBackoff from exo.utils.task_group import TaskGroup from exo.worker.plan import plan +from exo.worker.runner.larql_supervisor import LarqlRunnerSupervisor from exo.worker.runner.runner_supervisor import RunnerSupervisor +WorkerRunnerSupervisor = RunnerSupervisor | LarqlRunnerSupervisor + def _summarize_worker_task(task: Task) -> str: """Return a compact task summary for worker lifecycle logs.""" @@ -238,7 +241,7 @@ def __init__( self._staging_config = staging_config self.state: State = State() - self.runners: dict[RunnerId, RunnerSupervisor] = {} + self.runners: dict[RunnerId, WorkerRunnerSupervisor] = {} self._tg: TaskGroup = TaskGroup() self._system_id = SystemId() @@ -366,7 +369,16 @@ async def plan_step(self): model_id = shard.model_card.model_id self._download_backoff.record_attempt(model_id) - found_path = resolve_model_in_path(model_id) + vindex_location = ( + resolve_vindex_location(model_id, shard.vindex_uri) + if isinstance(shard, LarqlShardMetadata) + else None + ) + found_path = ( + vindex_location.path + if vindex_location is not None + else resolve_model_in_path(model_id) + ) if found_path is not None: logger.info( f"Model {model_id} found in EXO_MODELS_PATH at {found_path}" @@ -378,7 +390,9 @@ async def plan_step(self): shard_metadata=shard, model_directory=str(found_path), total=shard.model_card.storage_size, - read_only=True, + read_only=vindex_location.read_only + if vindex_location is not None + else True, ) ) ) @@ -579,7 +593,7 @@ async def _start_runner_task(self, task: Task): ) await self.runners[runner_id].start_task(task) - def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor: + def _create_supervisor(self, task: CreateRunner) -> WorkerRunnerSupervisor: """Creates and stores a new AssignedRunner with initial downloading status.""" shard = task.bound_instance.bound_shard logger.info( @@ -591,10 +605,16 @@ def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor: f"world_size={shard.world_size}, " f"layers={shard.start_layer}:{shard.end_layer})" ) - runner = RunnerSupervisor.create( - bound_instance=task.bound_instance, - event_sender=self.event_sender.clone(), - ) + if isinstance(shard, LarqlShardMetadata): + runner: WorkerRunnerSupervisor = LarqlRunnerSupervisor.create( + bound_instance=task.bound_instance, + event_sender=self.event_sender.clone(), + ) + else: + runner = RunnerSupervisor.create( + bound_instance=task.bound_instance, + event_sender=self.event_sender.clone(), + ) self.runners[task.bound_instance.bound_runner_id] = runner self._tg.start_soon(runner.run) return runner diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 3116cf4c0..bf8478aa8 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -42,13 +42,16 @@ RunnerWarmingUp, ) from exo.shared.types.worker.shards import ShardMetadata +from exo.worker.runner.larql_supervisor import LarqlRunnerSupervisor from exo.worker.runner.runner_supervisor import RunnerSupervisor +WorkerRunnerSupervisor = RunnerSupervisor | LarqlRunnerSupervisor + def plan( node_id: NodeId, # Runners is expected to be FRESH and so should not come from state - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], instances: Mapping[InstanceId, Instance], all_runners: Mapping[RunnerId, RunnerStatus], # all global @@ -73,7 +76,7 @@ def plan( def _kill_runner( - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], instances: Mapping[InstanceId, Instance], ) -> Shutdown | None: @@ -97,7 +100,7 @@ def _kill_runner( def _create_runner( node_id: NodeId, - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], instances: Mapping[InstanceId, Instance], ) -> CreateRunner | None: for instance in instances.values(): @@ -121,7 +124,7 @@ def _create_runner( def _model_needs_download( node_id: NodeId, - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], ) -> DownloadModel | None: local_downloads = global_download_status.get(node_id, []) @@ -146,7 +149,7 @@ def _model_needs_download( def _init_distributed_backend( - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], ): for runner in runners.values(): @@ -196,7 +199,7 @@ def _init_distributed_backend( def _load_model( - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], ) -> LoadModel | None: @@ -237,7 +240,7 @@ def _load_model( def _ready_to_warmup( - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], ) -> StartWarmup | None: for runner in runners.values(): @@ -303,7 +306,7 @@ def _uses_independent_distributed_warmup(shard: ShardMetadata) -> bool: def _pending_tasks( - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], tasks: Mapping[TaskId, Task], all_runners: Mapping[RunnerId, RunnerStatus], input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]] | None, @@ -346,7 +349,7 @@ def _pending_tasks( def _cancel_tasks( - runners: Mapping[RunnerId, RunnerSupervisor], + runners: Mapping[RunnerId, WorkerRunnerSupervisor], tasks: Mapping[TaskId, Task], ) -> Task | None: for task in tasks.values(): diff --git a/src/exo/worker/runner/larql_supervisor.py b/src/exo/worker/runner/larql_supervisor.py new file mode 100644 index 000000000..dda25c704 --- /dev/null +++ b/src/exo/worker/runner/larql_supervisor.py @@ -0,0 +1,498 @@ +from __future__ import annotations + +import contextlib +import socket +import subprocess +import time +from collections import deque +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import IO, Literal, cast + +import anyio +from anyio import BrokenResourceError, ClosedResourceError, to_thread +from loguru import logger + +from exo.download.download_utils import ( + build_vindex_path, + create_http_session, + resolve_vindex_in_path, +) +from exo.shared.types.diagnostics import ( + RunnerFlightRecorderEntry, + RunnerLifecycleMilestone, + RunnerPhaseName, + RunnerSupervisorDiagnostics, + RunnerTaskDiagnostics, +) +from exo.shared.types.events import ( + Event, + LarqlRunnerReadinessUpdated, + RunnerStatusUpdated, + TaskAcknowledged, + TaskStatusUpdated, +) +from exo.shared.types.memory import Memory +from exo.shared.types.tasks import ( + LoadModel, + Shutdown, + StartWarmup, + Task, + TaskId, + TaskStatus, +) +from exo.shared.types.worker.instances import BoundInstance +from exo.shared.types.worker.larql import LarqlRunnerReadiness +from exo.shared.types.worker.runners import ( + RunnerFailed, + RunnerIdle, + RunnerLoading, + RunnerReady, + RunnerShutdown, + RunnerShuttingDown, + RunnerStatus, +) +from exo.shared.types.worker.shards import LarqlShardMetadata +from exo.utils.channels import Sender +from exo.utils.task_group import TaskGroup + +ProcessFactory = Callable[[Sequence[str]], subprocess.Popen[str]] + + +def _now_utc_iso() -> str: + """Return the current UTC time for lifecycle diagnostics.""" + + return datetime.now(tz=timezone.utc).isoformat() + + +def allocate_larql_port(host: str = "127.0.0.1") -> int: + """Reserve and release a free local TCP port for a LARQL server process.""" + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((host, 0)) + _, port = cast(tuple[str, int], sock.getsockname()) + return port + + +def _format_larql_inclusive_range( + range_name: Literal["layers", "experts"], start: int, exclusive_end: int +) -> str: + """Convert Skulk half-open intervals to LARQL CLI inclusive ranges.""" + + if exclusive_end <= start: + raise ValueError(f"LARQL {range_name} range must be non-empty") + return f"{start}-{exclusive_end - 1}" + + +def build_larql_serve_command( + shard: LarqlShardMetadata, + *, + vindex_path: Path, + port: int, +) -> tuple[str, ...]: + """Build the deterministic LARQL serve command for one cold-tier shard.""" + + command = [ + "larql", + "serve", + str(vindex_path), + "--host", + shard.server_host, + "--port", + str(port), + "--ffn-only", + "--layers", + _format_larql_inclusive_range("layers", shard.start_layer, shard.end_layer), + "--preset", + shard.preset, + ] + if shard.expert_range is not None: + command.extend( + [ + "--experts", + _format_larql_inclusive_range( + "experts", + shard.expert_range.start_expert, + shard.expert_range.end_expert, + ), + ] + ) + if shard.units_manifest_path is not None: + command.extend(["--units", shard.units_manifest_path]) + return tuple(command) + + +def _default_process_factory(command: Sequence[str]) -> subprocess.Popen[str]: + """Start a LARQL child process with text stdout/stderr pipes.""" + + return subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +def _resolve_larql_vindex_path(shard: LarqlShardMetadata) -> Path: + """Resolve the staged local vindex path for a LARQL shard.""" + + if shard.local_vindex_path is not None: + return Path(shard.local_vindex_path).expanduser() + found = resolve_vindex_in_path(shard.model_card.model_id, shard.vindex_uri) + if found is not None: + return found + return build_vindex_path(shard.model_card.model_id) + + +@dataclass(eq=False) +class LarqlRunnerSupervisor: + """Worker-managed supervisor for one upstream `larql serve` child process.""" + + bound_instance: BoundInstance + shard_metadata: LarqlShardMetadata + _event_sender: Sender[Event] + _process_factory: ProcessFactory = _default_process_factory + _readiness_poll_interval: float = 0.25 + _tg: TaskGroup = field(default_factory=TaskGroup, init=False) + status: RunnerStatus = field(default_factory=RunnerIdle, init=False) + pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False) + in_progress: dict[TaskId, Task] = field(default_factory=dict, init=False) + completed: set[TaskId] = field(default_factory=set, init=False) + cancelled: set[TaskId] = field(default_factory=set, init=False) + _process: subprocess.Popen[str] | None = field(default=None, init=False) + _port: int | None = field(default=None, init=False) + _shutdown_requested: bool = field(default=False, init=False) + _status_since: str = field(default_factory=_now_utc_iso, init=False) + _status_since_monotonic: float = field(default_factory=time.monotonic, init=False) + _phase: RunnerPhaseName = field(default="created", init=False) + _phase_started_at: str = field(default_factory=_now_utc_iso, init=False) + _phase_started_monotonic: float = field(default_factory=time.monotonic, init=False) + _last_progress_at: str | None = field(default=None, init=False) + _last_task_sent_at: str | None = field(default=None, init=False) + _last_event_received_at: str | None = field(default=None, init=False) + _last_event_type: str | None = field(default=None, init=False) + _phase_detail: str | None = field(default=None, init=False) + _milestones: deque[RunnerLifecycleMilestone] = field( + default_factory=lambda: deque(maxlen=32), init=False + ) + _flight_recorder: deque[RunnerFlightRecorderEntry] = field( + default_factory=lambda: deque(maxlen=128), init=False + ) + + @classmethod + def create( + cls, + *, + bound_instance: BoundInstance, + event_sender: Sender[Event], + process_factory: ProcessFactory = _default_process_factory, + ) -> "LarqlRunnerSupervisor": + """Construct a LARQL supervisor for an assigned LARQL shard.""" + + shard = bound_instance.bound_shard + if not isinstance(shard, LarqlShardMetadata): + raise TypeError("LarqlRunnerSupervisor requires LarqlShardMetadata") + return cls( + bound_instance=bound_instance, + shard_metadata=shard, + _event_sender=event_sender, + _process_factory=process_factory, + ) + + async def run(self) -> None: + """Keep the supervisor task group alive until the worker shuts it down.""" + + self._record_milestone("supervisor_created", self.status.__class__.__name__) + try: + async with self._tg: + await anyio.sleep_forever() + finally: + await self._terminate_child() + with contextlib.suppress(ClosedResourceError): + self._event_sender.close() + + def shutdown(self) -> None: + """Request local LARQL child shutdown and cancel supervisor tasks.""" + + self._shutdown_requested = True + self._record_milestone("shutdown_requested") + self._tg.cancel_tasks() + + async def start_task(self, task: Task) -> None: + """Execute lifecycle tasks understood by the LARQL supervisor.""" + + self._last_task_sent_at = _now_utc_iso() + self.pending[task.task_id] = anyio.Event() + self.in_progress[task.task_id] = task + await self._send_event(TaskAcknowledged(task_id=task.task_id)) + self.pending.pop(task.task_id, None) + try: + if isinstance(task, LoadModel): + await self._start_until_ready() + self.completed.add(task.task_id) + await self._send_event( + TaskStatusUpdated( + task_id=task.task_id, + task_status=TaskStatus.Complete, + ) + ) + elif isinstance(task, StartWarmup): + self.completed.add(task.task_id) + await self._send_event( + TaskStatusUpdated( + task_id=task.task_id, + task_status=TaskStatus.Complete, + ) + ) + elif isinstance(task, Shutdown): + self._shutdown_requested = True + await self._terminate_child() + self.completed.add(task.task_id) + await self._send_event( + TaskStatusUpdated( + task_id=task.task_id, + task_status=TaskStatus.Complete, + ) + ) + await self._set_status(RunnerShutdown()) + else: + raise RuntimeError( + f"LarqlRunner does not execute {task.__class__.__name__} tasks" + ) + except Exception as exc: + await self._mark_failed(str(exc)) + await self._send_event( + TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Failed) + ) + finally: + self.in_progress.pop(task.task_id, None) + + async def cancel_task(self, task_id: TaskId) -> None: + """Mark a task cancellation request for diagnostics.""" + + self.cancelled.add(task_id) + + async def _start_until_ready(self) -> None: + vindex_path = _resolve_larql_vindex_path(self.shard_metadata) + port = self.shard_metadata.server_port or allocate_larql_port( + self.shard_metadata.server_host + ) + self._port = port + command = build_larql_serve_command( + self.shard_metadata, + vindex_path=vindex_path, + port=port, + ) + for attempt in range(self.shard_metadata.max_crash_restarts + 1): + await self._set_status(RunnerLoading()) + self._record_milestone("larql_start_requested", " ".join(command)) + self._process = self._process_factory(command) + self._tg.start_soon(self._forward_stream, self._process.stdout, "stdout") + self._tg.start_soon(self._forward_stream, self._process.stderr, "stderr") + try: + await self._wait_until_ready() + await self._set_status(RunnerReady()) + await self._send_readiness("ready") + assert self._process is not None + self._tg.start_soon(self._watch_process_exit, self._process) + return + except Exception as exc: + await self._terminate_child() + if attempt >= self.shard_metadata.max_crash_restarts: + raise RuntimeError( + "LARQL child failed readiness after " + f"{attempt + 1} attempt(s): {exc}" + ) from exc + self._record_milestone( + "larql_restart", + f"attempt={attempt + 1}: {type(exc).__name__}", + ) + + async def _wait_until_ready(self) -> None: + deadline = time.monotonic() + self.shard_metadata.readiness_timeout_seconds + while time.monotonic() < deadline: + process = self._process + if process is None: + raise RuntimeError("LARQL process was not started") + if process.poll() is not None: + raise RuntimeError(f"LARQL process exited with {process.returncode}") + if await self._health_check(): + return + await anyio.sleep(self._readiness_poll_interval) + raise TimeoutError("Timed out waiting for LARQL readiness") + + async def _watch_process_exit(self, process: subprocess.Popen[str]) -> None: + """Report a ready LARQL child that exits outside intentional shutdown.""" + + return_code = await to_thread.run_sync(process.wait, abandon_on_cancel=True) + if self._shutdown_requested or self._process is not process: + return + message = f"LARQL child exited unexpectedly with exit code {return_code}" + self._record_milestone("larql_exited", message) + await self._mark_failed(message) + + async def _health_check(self) -> bool: + port = self._port + if port is None: + return False + url = f"http://{self.shard_metadata.server_host}:{port}/v1/health" + try: + async with ( + create_http_session(timeout_profile="short") as session, + session.get(url) as response, + ): + return response.status == 200 + except Exception: + return False + + async def _send_readiness( + self, + status: Literal["ready", "not_ready", "failed"], + error_message: str | None = None, + ) -> None: + port = self._port or 0 + readiness = LarqlRunnerReadiness( + runner_id=self.bound_instance.bound_runner_id, + vindex_uri=self.shard_metadata.vindex_uri, + preset=self.shard_metadata.preset, + start_layer=self.shard_metadata.start_layer, + end_layer=self.shard_metadata.end_layer, + expert_range=self.shard_metadata.expert_range, + units_manifest_path=self.shard_metadata.units_manifest_path, + host=self.shard_metadata.server_host, + port=max(port, 1), + status=status, + ram_footprint=await self._process_memory(), + error_message=error_message, + ) + await self._send_event(LarqlRunnerReadinessUpdated(readiness=readiness)) + + async def _process_memory(self) -> Memory | None: + process = self._process + if process is None: + return None + try: + import psutil + + info = await to_thread.run_sync( + lambda: psutil.Process(process.pid).memory_info() + ) + return Memory.from_bytes(int(info.rss)) + except Exception: + return None + + async def _forward_stream( + self, + stream: IO[str] | None, + stream_name: Literal["stdout", "stderr"], + ) -> None: + if stream is None: + return + while True: + line = await to_thread.run_sync(stream.readline, abandon_on_cancel=True) + if not line: + return + logger.info(f"larql-server[{stream_name}]: {line.rstrip()}") + + async def _terminate_child(self) -> None: + process = self._process + if process is None or process.poll() is not None: + return + await self._set_status(RunnerShuttingDown()) + process.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + await to_thread.run_sync( + lambda: process.wait(timeout=5), + abandon_on_cancel=True, + ) + if process.poll() is None: + process.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + await to_thread.run_sync( + lambda: process.wait(timeout=5), + abandon_on_cancel=True, + ) + + async def _mark_failed(self, message: str) -> None: + await self._set_status(RunnerFailed(error_message=message)) + await self._send_readiness("failed", message) + + async def _set_status(self, status: RunnerStatus) -> None: + self.status = status + self._status_since = _now_utc_iso() + self._status_since_monotonic = time.monotonic() + self._record_milestone("status_changed", status.__class__.__name__) + await self._send_event( + RunnerStatusUpdated( + runner_id=self.bound_instance.bound_runner_id, + runner_status=status, + ) + ) + + async def _send_event(self, event: Event) -> None: + self._last_event_received_at = _now_utc_iso() + self._last_event_type = event.__class__.__name__ + try: + await self._event_sender.send(event) + except (ClosedResourceError, BrokenResourceError): + logger.warning("LarqlRunner event sender closed") + + def _record_milestone(self, name: str, detail: str | None = None) -> None: + self._milestones.append( + RunnerLifecycleMilestone(at=_now_utc_iso(), name=name, detail=detail) + ) + self._last_progress_at = _now_utc_iso() + + def _task_diagnostics(self, task: Task) -> RunnerTaskDiagnostics: + return RunnerTaskDiagnostics( + task_id=str(task.task_id), + task_kind=task.__class__.__name__, + task_status=str(task.task_status.value), + instance_id=str(task.instance_id), + command_id=None, + runner_id=str(self.bound_instance.bound_runner_id), + model_id=str(self.shard_metadata.model_card.model_id), + ) + + def diagnostics(self) -> RunnerSupervisorDiagnostics: + """Return live read-only diagnostics for this LARQL supervisor.""" + + process = self._process + return RunnerSupervisorDiagnostics( + runner_id=str(self.bound_instance.bound_runner_id), + instance_id=str(self.bound_instance.instance.instance_id), + node_id=str(self.bound_instance.bound_node_id), + model_id=str(self.shard_metadata.model_card.model_id), + device_rank=self.shard_metadata.device_rank, + world_size=self.shard_metadata.world_size, + start_layer=self.shard_metadata.start_layer, + end_layer=self.shard_metadata.end_layer, + n_layers=self.shard_metadata.n_layers, + pid=process.pid if process is not None else None, + process_alive=process is not None and process.poll() is None, + exit_code=process.returncode if process is not None else None, + status_kind=self.status.__class__.__name__, + status_since=self._status_since, + seconds_in_status=time.monotonic() - self._status_since_monotonic, + phase=self._phase, + phase_started_at=self._phase_started_at, + seconds_in_phase=time.monotonic() - self._phase_started_monotonic, + last_progress_at=self._last_progress_at, + active_task_id=None, + active_command_id=None, + phase_detail=self._phase_detail, + last_mlx_memory=None, + flight_recorder=list(self._flight_recorder), + pending_task_ids=[str(task_id) for task_id in self.pending], + in_progress_tasks=[ + self._task_diagnostics(task) for task in self.in_progress.values() + ], + completed_task_count=len(self.completed), + cancelled_task_ids=[str(task_id) for task_id in self.cancelled], + last_task_sent_at=self._last_task_sent_at, + last_event_received_at=self._last_event_received_at, + last_event_type=self._last_event_type, + milestones=list(self._milestones), + ) diff --git a/src/exo/worker/tests/unittests/test_runner/test_larql_supervisor.py b/src/exo/worker/tests/unittests/test_runner/test_larql_supervisor.py new file mode 100644 index 000000000..7af4443fc --- /dev/null +++ b/src/exo/worker/tests/unittests/test_runner/test_larql_supervisor.py @@ -0,0 +1,337 @@ +import io +import subprocess +import threading +from collections.abc import Sequence +from pathlib import Path +from typing import cast + +import anyio +import pytest + +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.types.common import NodeId +from exo.shared.types.events import ( + Event, + LarqlRunnerReadinessUpdated, + RunnerStatusUpdated, + TaskAcknowledged, + TaskStatusUpdated, +) +from exo.shared.types.memory import Memory +from exo.shared.types.tasks import LoadModel, Shutdown, TaskStatus +from exo.shared.types.worker.instances import BoundInstance, InstanceId +from exo.shared.types.worker.larql import LarqlExpertRange +from exo.shared.types.worker.runners import ( + RunnerFailed, + RunnerId, + RunnerReady, + RunnerShutdown, + RunnerShuttingDown, +) +from exo.shared.types.worker.shards import LarqlShardMetadata, ShardMetadata +from exo.utils.channels import channel +from exo.worker.runner.larql_supervisor import ( + LarqlRunnerSupervisor, + build_larql_serve_command, +) +from exo.worker.tests.unittests.conftest import get_mlx_ring_instance + + +class _FakeProcess: + pid = 12345 + + def __init__(self) -> None: + self.returncode: int | None = None + self.stdout = io.StringIO("") + self.stderr = io.StringIO("") + self._exit_event = threading.Event() + + def poll(self) -> int | None: + return self.returncode + + def wait(self, timeout: float | None = None) -> int: + self._exit_event.wait(timeout) + if self.returncode is None: + self.returncode = 0 + return self.returncode + + def terminate(self) -> None: + self.returncode = -15 + self._exit_event.set() + + def kill(self) -> None: + self.returncode = -9 + self._exit_event.set() + + def exit(self, returncode: int) -> None: + self.returncode = returncode + self._exit_event.set() + + +class _HungProcess(_FakeProcess): + def __init__(self) -> None: + super().__init__() + self.terminate_called = False + self.kill_called = False + + def wait(self, timeout: float | None = None) -> int: + if self.returncode is None: + raise subprocess.TimeoutExpired("larql", timeout or 0) + return self.returncode + + def terminate(self) -> None: + self.terminate_called = True + + def kill(self) -> None: + self.kill_called = True + self.returncode = -9 + self._exit_event.set() + + +def _larql_shard() -> LarqlShardMetadata: + return LarqlShardMetadata( + model_card=ModelCard( + model_id=ModelId("skulk/gemma-4-26b-a4b-expert-server-q4-k-vindex"), + storage_size=Memory.from_mb(512), + n_layers=46, + hidden_size=4096, + supports_tensor=False, + tasks=[ModelTask.TextGeneration], + ), + device_rank=0, + world_size=1, + start_layer=4, + end_layer=12, + n_layers=46, + vindex_uri="hf://skulk/gemma-4-26b-a4b-expert-server-q4-k-vindex", + preset="expert-server", + local_vindex_path="/tmp/gemma-vindex", + server_port=49152, + expert_range=LarqlExpertRange(start_expert=0, end_expert=8), + ) + + +def _bound_instance(shard: ShardMetadata) -> BoundInstance: + runner_id = RunnerId("runner-a") + node_id = NodeId("node-a") + instance = get_mlx_ring_instance( + instance_id=InstanceId("instance-a"), + model_id=shard.model_card.model_id, + node_to_runner={node_id: runner_id}, + runner_to_shard={runner_id: shard}, + ) + return BoundInstance( + instance=instance, + bound_runner_id=runner_id, + bound_node_id=node_id, + ) + + +def test_build_larql_serve_command_includes_slice_arguments() -> None: + command = build_larql_serve_command( + _larql_shard(), + vindex_path=Path("/tmp/gemma-vindex"), + port=49152, + ) + + assert command == ( + "larql", + "serve", + "/tmp/gemma-vindex", + "--host", + "127.0.0.1", + "--port", + "49152", + "--ffn-only", + "--layers", + "4-11", + "--preset", + "expert-server", + "--experts", + "0-7", + ) + + +def test_build_larql_serve_command_rejects_empty_layer_range() -> None: + shard = _larql_shard().model_copy(update={"end_layer": 4}) + + with pytest.raises(ValueError, match="layers range must be non-empty"): + build_larql_serve_command( + shard, + vindex_path=Path("/tmp/gemma-vindex"), + port=49152, + ) + + +@pytest.mark.asyncio +async def test_larql_supervisor_load_model_starts_process_and_emits_readiness() -> None: + shard = _larql_shard() + event_sender, event_receiver = channel[Event]() + commands: list[tuple[str, ...]] = [] + process = _FakeProcess() + + def process_factory(command: Sequence[str]) -> subprocess.Popen[str]: + commands.append(tuple(str(part) for part in command)) + return cast(subprocess.Popen[str], cast(object, process)) + + supervisor = LarqlRunnerSupervisor.create( + bound_instance=_bound_instance(shard), + event_sender=event_sender, + process_factory=process_factory, + ) + + async def health_check() -> bool: + return True + + supervisor._health_check = health_check # pyright: ignore[reportPrivateUsage] + + async with anyio.create_task_group() as task_group: + task_group.start_soon(supervisor.run) + await supervisor.start_task( + LoadModel(instance_id=supervisor.bound_instance.instance.instance_id) + ) + task_group.cancel_scope.cancel() + + observed = [ + await event_receiver.receive(), + await event_receiver.receive(), + await event_receiver.receive(), + await event_receiver.receive(), + ] + + assert commands + assert isinstance(observed[0], TaskAcknowledged) + assert isinstance(observed[1], RunnerStatusUpdated) + assert isinstance(observed[2], RunnerStatusUpdated) + assert isinstance(observed[2].runner_status, RunnerReady) + assert isinstance(observed[3], LarqlRunnerReadinessUpdated) + assert observed[3].readiness.status == "ready" + + completion = await event_receiver.receive() + assert isinstance(completion, TaskStatusUpdated) + assert completion.task_status == TaskStatus.Complete + + +@pytest.mark.asyncio +async def test_larql_supervisor_marks_failed_when_ready_child_exits() -> None: + shard = _larql_shard() + event_sender, event_receiver = channel[Event]() + process = _FakeProcess() + + def process_factory(_command: Sequence[str]) -> subprocess.Popen[str]: + return cast(subprocess.Popen[str], cast(object, process)) + + supervisor = LarqlRunnerSupervisor.create( + bound_instance=_bound_instance(shard), + event_sender=event_sender, + process_factory=process_factory, + ) + + async def health_check() -> bool: + return True + + supervisor._health_check = health_check # pyright: ignore[reportPrivateUsage] + failed_status: Event | None = None + failed_readiness: Event | None = None + + async with anyio.create_task_group() as task_group: + task_group.start_soon(supervisor.run) + await supervisor.start_task( + LoadModel(instance_id=supervisor.bound_instance.instance.instance_id) + ) + for _ in range(5): + await event_receiver.receive() + + process.exit(42) + + with anyio.fail_after(1): + failed_status = await event_receiver.receive() + failed_readiness = await event_receiver.receive() + + task_group.cancel_scope.cancel() + + assert isinstance(failed_status, RunnerStatusUpdated) + assert isinstance(failed_status.runner_status, RunnerFailed) + assert failed_status.runner_status.error_message == ( + "LARQL child exited unexpectedly with exit code 42" + ) + assert isinstance(failed_readiness, LarqlRunnerReadinessUpdated) + assert failed_readiness.readiness.status == "failed" + assert failed_readiness.readiness.error_message == ( + "LARQL child exited unexpectedly with exit code 42" + ) + + +@pytest.mark.asyncio +async def test_larql_supervisor_shutdown_completes_without_failure() -> None: + shard = _larql_shard() + event_sender, event_receiver = channel[Event]() + supervisor = LarqlRunnerSupervisor.create( + bound_instance=_bound_instance(shard), + event_sender=event_sender, + ) + supervisor._process = cast( # pyright: ignore[reportPrivateUsage] + subprocess.Popen[str], + cast(object, _FakeProcess()), + ) + + await supervisor.start_task( + Shutdown( + instance_id=supervisor.bound_instance.instance.instance_id, + runner_id=supervisor.bound_instance.bound_runner_id, + ) + ) + + observed = [ + await event_receiver.receive(), + await event_receiver.receive(), + await event_receiver.receive(), + await event_receiver.receive(), + ] + + assert isinstance(observed[0], TaskAcknowledged) + assert isinstance(observed[1], RunnerStatusUpdated) + assert isinstance(observed[1].runner_status, RunnerShuttingDown) + assert isinstance(observed[2], TaskStatusUpdated) + assert observed[2].task_status == TaskStatus.Complete + assert isinstance(observed[3], RunnerStatusUpdated) + assert isinstance(observed[3].runner_status, RunnerShutdown) + + +@pytest.mark.asyncio +async def test_larql_supervisor_shutdown_kills_child_after_wait_timeout() -> None: + """Hung child waits fall through to kill during supervisor shutdown.""" + + shard = _larql_shard() + event_sender, event_receiver = channel[Event]() + process = _HungProcess() + supervisor = LarqlRunnerSupervisor.create( + bound_instance=_bound_instance(shard), + event_sender=event_sender, + ) + supervisor._process = cast( # pyright: ignore[reportPrivateUsage] + subprocess.Popen[str], + cast(object, process), + ) + + await supervisor.start_task( + Shutdown( + instance_id=supervisor.bound_instance.instance.instance_id, + runner_id=supervisor.bound_instance.bound_runner_id, + ) + ) + + observed = [ + await event_receiver.receive(), + await event_receiver.receive(), + await event_receiver.receive(), + await event_receiver.receive(), + ] + + assert process.terminate_called + assert process.kill_called + assert process.returncode == -9 + assert isinstance(observed[1], RunnerStatusUpdated) + assert isinstance(observed[1].runner_status, RunnerShuttingDown) + assert isinstance(observed[2], TaskStatusUpdated) + assert observed[2].task_status == TaskStatus.Complete diff --git a/website/docs/architecture-reference.md b/website/docs/architecture-reference.md index 5564dd8db..ee4805b3f 100644 --- a/website/docs/architecture-reference.md +++ b/website/docs/architecture-reference.md @@ -45,6 +45,15 @@ This file is intentionally dense. If you find a stale fact, fix it inline rather - `src/exo/worker/runner/image_models/runner.py` — image generation - **Communicates via:** `mp.Queue` from worker (incoming tasks); `mp.Queue` to worker (outgoing events); `mlx.distributed` collectives with peer runners +### LarqlRunner + +- **Status:** internal Phase 2 supervision/readiness path implemented; no placement flow creates it yet +- **Role:** worker-managed runner subtype that supervises an upstream `larql serve` child process for vindex-backed FFN / expert slices +- **Lives in:** `src/exo/worker/runner/larql_supervisor.py` +- **Decision records:** `docs/adr/0001-larql-runner-type.md`, `docs/adr/0002-head-mlx-cold-larql.md`, `docs/adr/0003-vindex-provenance.md` +- **Gating:** Phase 4 implementation is blocked until the Phase 3 MLX FFN delegation spike confirms ADR-B +- **Runtime invariant:** existing MLX runners remain the head path; the MLX head never loads a vindex + ### Router (libp2p) - **Role:** transport for all inter-node communication @@ -83,6 +92,7 @@ This file is intentionally dense. If you find a stale fact, fix it inline rather - **Custom cards:** `SKULK_CUSTOM_MODEL_CARDS_DIR` (default `SKULK_DATA_HOME/custom_model_cards`) as TOML - **Built-in cards:** `resources/inference_model_cards/` as TOML - **Optional model store:** shared host with rsync-style staging — `src/exo/store/` +- **Planned vindex source:** Skulk consumes `hf://...` vindex directories produced by the separate `skulk-vindex-publisher` repo; Skulk does not extract vindexes in-tree ## Pubsub topics diff --git a/website/docs/architecture.md b/website/docs/architecture.md index 8c728c0c8..44fcb79e7 100644 --- a/website/docs/architecture.md +++ b/website/docs/architecture.md @@ -19,6 +19,7 @@ The design choices that shape almost everything else: - **libp2p pub/sub for transport.** Topics carry commands, events, election messages, and connection updates between nodes. - **MLX as the inference backend.** Pipeline-parallel and tensor-parallel sharding strategies sit on top of `mlx.distributed`'s ring or jaccl/RDMA backends. - **Subprocess isolation for runners.** Each model instance runs in its own `mp.Process` with its own MLX/Metal context, so a crash or hang in one runner can't bring down the rest of the node. +- **LARQL slice mode is planned, not active.** The accepted planning ADRs keep MLX as the head runtime. Phase 2 adds internal `LarqlRunner` supervision/readiness for cold FFN/expert slices, but placement remains gated by the Phase 3 MLX delegation spike. ## The shape of a node @@ -59,6 +60,8 @@ Each subsystem has its own concern: - **API** is a FastAPI app that exposes inference endpoints in four wire formats (OpenAI Chat Completions, OpenAI Responses, Anthropic Messages, Ollama) and Skulk-native control endpoints (placements, diagnostics, traces, config). It also serves the dashboard build at `/`. - **Storage** is a collection of on-disk responsibilities: the event log (msgpack + zstd), the model cache directory, custom model cards (per-user TOML files), and the optional shared model store. +The LARQL roadmap adds a runner subtype, `LarqlRunner`, that remains worker-managed but supervises an upstream `larql serve` process instead of loading an MLX model directly. Phase 2 implements this internal supervision/readiness path, but no current placement flow creates LARQL runners. The design is captured in `docs/adr/0001-larql-runner-type.md` and is intentionally additive to the current MLX runner path. + ## The shape of a cluster ```mermaid @@ -160,6 +163,25 @@ A snapshot-bootstrap rollout has one operational rule: once a master starts comp Inference happens entirely inside the runner subprocess. Skulk wraps MLX (and the upstream mlx-lm model implementations) in a layer that handles distributed coordination, family-specific behavior, and operator-controlled knobs. +### Planned LARQL slice mode + +The Phase 1 LARQL ADRs define a future second placement mode. The selected head +node remains an MLX runner and owns the hot path: embeddings, attention, norms, +router, and locally assigned layers. Phase 2 adds the internal `LarqlRunner` +supervisor that can start and readiness-check a vindex-backed LARQL HTTP +server, but cold-tier placement remains future work. + +Important constraints: + +- Existing MLX placement is unchanged for models that fit on the head node. +- The MLX head never loads a vindex. +- Skulk consumes HuggingFace-hosted vindexes; extraction and publication belong + to the separate `skulk-vindex-publisher` repo. +- Vindexes are directory-shaped model-store artifacts with replay-safe + readiness state separate from ordinary MLX runner status. +- Phase 4 slice placement must not start until the Phase 3 spike proves MLX can + delegate per-layer FFN work and continue generation with acceptable overhead. + ### Pipeline parallelism For models too large for a single device, Skulk splits the layer stack across ranks. Each rank holds a contiguous range of layers (`start_layer` to `end_layer`). Layers communicate via `mlx.distributed.send` / `recv_like` over the `ring` backend (sockets) or `jaccl` (RDMA, when available).