From 138cfff63251556debf0184bf80a909a7f19a849 Mon Sep 17 00:00:00 2001 From: kaizencode Date: Sun, 9 Aug 2026 22:32:18 +0200 Subject: [PATCH] feat(semantic): auto-start local encoder --- docs/branch/semantic.md | 12 +- docs/changes.md | 4 +- docs/source/details/semantic-search.md | 39 +-- owilix/cli/remote.py | 6 +- owilix/core/tasks/semantic.py | 47 ++-- owilix/semantic_encoder/app.py | 4 + owilix/semantic_encoder/process.py | 227 ++++++++++++++++++ scripts/setup-semantic-rocm.sh | 2 +- tests/owilix/semantic_encoder/test_process.py | 146 +++++++++++ 9 files changed, 440 insertions(+), 47 deletions(-) create mode 100644 owilix/semantic_encoder/process.py create mode 100644 tests/owilix/semantic_encoder/test_process.py diff --git a/docs/branch/semantic.md b/docs/branch/semantic.md index 5e2d522..30ae719 100644 --- a/docs/branch/semantic.md +++ b/docs/branch/semantic.md @@ -24,14 +24,16 @@ - [x] Build and reuse a memory-mapped vector cache - [x] Add `owi remote search --mode semantic` -- [x] Add the standalone Jina encoder service +- [x] Add the persistent Jina encoder process +- [x] Automatically start and reuse the local encoder from semantic search - [x] Document setup and usage ### Phase 1.3: Verification - [x] Semantic unit and CLI tests pass -- [x] Existing non-HTTP suite passes (802 passed, 2 skipped) -- [x] Real French OWIE smoke test passes (5 hits in 0.40s; 9.2ms scan) +- [x] Existing non-HTTP suite passes (807 passed, 2 skipped) +- [x] Real French OWIE auto-start smoke test passes (3 hits in 5.91s; 5.3ms scan) +- [x] A second search reuses the encoder (1 hit in 0.40s; 5.5ms scan) - [x] Existing HTTP suite passes outside the network/thread sandbox (58 passed) --- @@ -43,7 +45,7 @@ - **Alternatives considered**: Remote scan per query and an external vector database. 2. **Decision**: Run Jina in a separate HTTP encoder service. - - **Reason**: Keeping a persistent process avoids reloading Jina for every query. OWILIX and the encoder now share the root Python 3.12 environment. + - **Reason**: Keeping a persistent process avoids reloading Jina for every query. OWILIX starts it automatically, so users only invoke `owi remote search --mode semantic`; OWILIX and the encoder share the root Python 3.12 environment. - **Alternatives considered**: Loading Transformers directly in OWILIX. 3. **Decision**: Require the encoder model to match the model declared by OWIE. @@ -63,7 +65,7 @@ ## Merge Checklist -- [x] All tests pass (860 passed, 2 skipped across semantic and HTTP extras) +- [x] All tests pass (865 passed, 2 skipped across semantic and HTTP extras) - [x] Documentation updated - [x] `docs/changes.md` updated - [ ] `docs/epics.md` updated if epic complete diff --git a/docs/changes.md b/docs/changes.md index f0c4f39..576c32c 100644 --- a/docs/changes.md +++ b/docs/changes.md @@ -52,8 +52,10 @@ vectors, a reusable memory-mapped cache and a persistent Jina encoder service. - Exact cosine ranking over a selected/newest local language partition, with best-chunk deduplication and source OWI title/URL/passage enrichment. - Strict OWIE schema, embedding model and dimension validation. -- Root `semantic` extra and persistent encoder command using the same Python +- Root `semantic` extra and persistent encoder process using the same Python 3.12 `.venv`; platform-specific Torch/ROCm installation remains explicit. +- Automatic health check, detached startup and readiness wait for the local + encoder; semantic-search users no longer start a separate command manually. - CLI controls for OWIE selection, encoder URL/timeout and cache rebuild, plus synthetic Parquet/CLI coverage and end-to-end validation on the French 2026-08-05 snapshot. diff --git a/docs/source/details/semantic-search.md b/docs/source/details/semantic-search.md index 93dd2d7..5b94d8d 100644 --- a/docs/source/details/semantic-search.md +++ b/docs/source/details/semantic-search.md @@ -10,8 +10,9 @@ related source OWI. - OWILIX requires Python 3.12, discovers OWIE/OWI files, owns the cache and ranks hits. -- `owi-semantic-encoder` uses the same root `.venv`, owns Torch, Transformers - and the GPU, and runs as a separate process so the model remains loaded. +- OWILIX automatically manages a local encoder process in the same root + `.venv`. That process owns Torch, Transformers and the GPU so the model + remains loaded between CLI searches. - The data plane is local. Reading OWIE Parquet directly from the remote store for every query is supported neither by this MVP nor recommended for interactive use. @@ -58,31 +59,31 @@ The OWIE metadata chooses the model. The tested snapshot declares model cannot query those vectors; using it requires regenerating the complete OWIE dataset in that model's vector space. -## 3. Start the persistent encoder +## 3. Search ```bash -.venv/bin/owi-semantic-encoder +owi remote search \ + "Quels sont les risques et les enjeux de l'intelligence artificielle pour la société ?" \ + --mode semantic --language fra --limit 10 ``` -On first use, Transformers downloads the model. Keep the service on localhost -unless an authenticated proxy is placed in front of it. +No encoder command is required. OWILIX checks the local health endpoint, +starts the persistent process when it is absent, waits for Jina to be ready, +and then submits the query. Later searches reuse the running process and the +model already loaded in GPU memory. The first search may take longer while +Transformers downloads the model. -Verify its vector-space identity: +Managed process state and startup diagnostics are stored under: -```bash -curl http://127.0.0.1:8765/health +```text +~/.owi/.cache/semantic/encoder/ +├── encoder.pid +└── encoder.log ``` -Expected fields include `ready: true`, the exact Jina model, dimension `1024` -and the detected ROCm device. - -## 4. Search - -```bash -owi remote search \ - "Quels sont les risques et les enjeux de l'intelligence artificielle pour la société ?" \ - --mode semantic --language fra --limit 10 -``` +Automatic startup is limited to loopback HTTP URLs. A custom remote +`--encoder-url` is used when reachable, but OWILIX never launches processes on +remote hosts. To pin a snapshot and inspect the resolved configuration: diff --git a/owilix/cli/remote.py b/owilix/cli/remote.py index f458ec7..3ae357a 100644 --- a/owilix/cli/remote.py +++ b/owilix/cli/remote.py @@ -383,7 +383,7 @@ def search( None, "--encoder-url", envvar="OWI_SEMANTIC_ENCODER_URL", - help="Semantic encoder base URL (default: http://127.0.0.1:8765)", + help="Semantic encoder URL (the default local encoder starts automatically)", ), encoder_timeout: float = typer.Option( 120.0, @@ -401,8 +401,8 @@ def search( Search the OWI lexical index or a locally pulled OWIE vector dataset. Lexical mode (the default) connects to the remote DuckLake/S3 index. - Semantic mode searches a local OWIE schema 2.0 snapshot and calls a - standalone encoder using the exact embedding model declared by OWIE. + Semantic mode searches a local OWIE schema 2.0 snapshot and automatically + manages a persistent local encoder using the model declared by OWIE. In lexical mode, use --fetch to retrieve full document records from OWI. Semantic mode enriches hits from the locally pulled source OWI directly. diff --git a/owilix/core/tasks/semantic.py b/owilix/core/tasks/semantic.py index fe08460..03c005a 100644 --- a/owilix/core/tasks/semantic.py +++ b/owilix/core/tasks/semantic.py @@ -2,9 +2,10 @@ Semantic search over locally pulled OWIE datasets. The task builds a compact, memory-mapped NumPy cache from OWIE schema 2.0 -Parquet files, asks a separate HTTP service to encode the query with the exact -model declared by the dataset, and performs an exact cosine-similarity scan. -The source OWI dataset is used to enrich hits with URLs, titles and passages. +Parquet files, asks a managed local HTTP process to encode the query with the +exact model declared by the dataset, and performs an exact cosine-similarity +scan. The source OWI dataset is used to enrich hits with URLs, titles and +passages. Keeping query encoding behind HTTP avoids adding Torch and Transformers to the OWILIX process and lets the encoder run in a Python/ROCm environment tailored @@ -34,6 +35,7 @@ import pyarrow.parquet as pq from rich.console import Console from owilix.core.types import CommandResult, ErrorType, ExitCode +from owilix.semantic_encoder.process import EncoderProcessError, ensure_encoder logger = logging.getLogger("owilix") @@ -101,7 +103,8 @@ def semantic_search( limit: Maximum number of distinct documents to return. dataset_id: Optional local OWIE dataset UUID. The newest compatible local snapshot is selected when omitted. - encoder_url: Base URL of the standalone semantic encoder. + encoder_url: Base URL of the encoder. The default local encoder is + started automatically when needed. encoder_timeout: HTTP timeout for model health and encoding requests. rebuild_cache: Rebuild the memory-mapped cache even when valid. explain: Print the selected dataset, model and cache information. @@ -139,6 +142,10 @@ def semantic_search( expected_model=selection.model, expected_dimension=selection.dimension, timeout=encoder_timeout, + state_directory=( + Path(manager.owi_path).expanduser() / ".cache" / "semantic" / "encoder" + ), + console=console, ) search_started = time.perf_counter() @@ -595,22 +602,21 @@ def _encode_query( expected_model: str, expected_dimension: int, timeout: float, + state_directory: Path, + console: Console, ) -> np.ndarray: - """Validate the encoder and request one normalized query vector.""" + """Ensure and validate the encoder, then request one query vector.""" base_url = encoder_url.rstrip("/") try: + readiness = ensure_encoder( + base_url=base_url, + state_directory=state_directory, + model=expected_model, + timeout=timeout, + status_callback=console.print, + ) + _validate_encoder_identity(readiness.health, expected_model, expected_dimension) with httpx.Client(timeout=timeout) as client: - health_response = client.get(f"{base_url}/health") - health_response.raise_for_status() - health = health_response.json() - if health.get("ready") is False: - raise SemanticSearchError( - f"Semantic encoder at {base_url} is not ready: {health.get('detail', 'model is loading')}", - ErrorType.NETWORK, - ExitCode.NETWORK_ERROR, - ) - _validate_encoder_identity(health, expected_model, expected_dimension) - response = client.post( f"{base_url}/embed", json={"texts": [query], "task": "retrieval", "prompt_name": "query"}, @@ -619,10 +625,15 @@ def _encode_query( payload = response.json() except SemanticSearchError: raise + except EncoderProcessError as error: + raise SemanticSearchError( + str(error), + ErrorType.NETWORK, + ExitCode.NETWORK_ERROR, + ) from error except (httpx.HTTPError, OSError) as error: raise SemanticSearchError( - f"Could not reach semantic encoder at {base_url}: {error}. " - "Start it first or pass --encoder-url.", + f"Could not reach semantic encoder at {base_url}: {error}", ErrorType.NETWORK, ExitCode.NETWORK_ERROR, ) from error diff --git a/owilix/semantic_encoder/app.py b/owilix/semantic_encoder/app.py index 02ce7b6..ed6d879 100644 --- a/owilix/semantic_encoder/app.py +++ b/owilix/semantic_encoder/app.py @@ -105,3 +105,7 @@ def main() -> None: port=int(os.environ.get("OWI_SEMANTIC_PORT", "8765")), workers=1, ) + + +if __name__ == "__main__": + main() diff --git a/owilix/semantic_encoder/process.py b/owilix/semantic_encoder/process.py new file mode 100644 index 0000000..ece084f --- /dev/null +++ b/owilix/semantic_encoder/process.py @@ -0,0 +1,227 @@ +"""Lifecycle management for the local persistent semantic encoder.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable +from urllib.parse import urlparse + +import httpx + + +StatusCallback = Callable[[str], None] + + +class EncoderProcessError(RuntimeError): + """A local encoder could not be reused, started, or made ready.""" + + +@dataclass(frozen=True) +class EncoderReadiness: + """Health payload returned once the encoder can accept requests.""" + + health: dict[str, Any] + started: bool + + +def ensure_encoder( + base_url: str, + state_directory: Path, + model: str, + timeout: float, + status_callback: StatusCallback | None = None, +) -> EncoderReadiness: + """Reuse a healthy encoder or start and await a local one. + + Automatic startup is deliberately restricted to loopback HTTP endpoints. + A custom remote encoder remains supported through ``--encoder-url``, but + OWILIX never attempts to launch a process for a remote host. + """ + base_url = base_url.rstrip("/") + deadline = time.monotonic() + timeout + try: + health = _request_health(base_url, min(2.0, timeout)) + except httpx.RequestError as error: + bind = _local_bind(base_url) + if bind is None: + raise EncoderProcessError( + f"Could not reach semantic encoder at {base_url}: {error}. " + "Automatic startup is only available for a local HTTP encoder." + ) from error + + process, log_path, started = _start_local_encoder( + state_directory=state_directory, + host=bind[0], + port=bind[1], + model=model, + ) + if status_callback is not None: + if started: + status_callback( + f"Starting the local semantic encoder at {base_url} " + f"(log: {log_path})" + ) + else: + status_callback( + f"Waiting for the local semantic encoder at {base_url}" + ) + + health = _wait_until_ready( + base_url=base_url, + deadline=deadline, + process=process, + log_path=log_path, + ) + if status_callback is not None: + device = health.get("device") + suffix = f" on {device}" if device else "" + status_callback(f"Local semantic encoder is ready{suffix}") + return EncoderReadiness(health=health, started=started) + + _require_ready(health, base_url) + return EncoderReadiness(health=health, started=False) + + +def _request_health(base_url: str, timeout: float) -> dict[str, Any]: + """Fetch and decode one encoder health response.""" + with httpx.Client(timeout=timeout) as client: + response = client.get(f"{base_url}/health") + response.raise_for_status() + payload = response.json() + if not isinstance(payload, dict): + raise ValueError("semantic encoder health response must be a JSON object") + return payload + + +def _require_ready(health: dict[str, Any], base_url: str) -> None: + """Reject a reachable encoder whose model failed to load.""" + if health.get("ready") is not True: + detail = health.get("detail") or "model is not ready" + raise EncoderProcessError(f"Semantic encoder at {base_url} is not ready: {detail}") + + +def _local_bind(base_url: str) -> tuple[str, int] | None: + """Return a safe local bind target for an auto-startable URL.""" + try: + parsed = urlparse(base_url) + port = parsed.port or 80 + except ValueError: + return None + if ( + parsed.scheme != "http" + or parsed.hostname not in {"127.0.0.1", "localhost", "::1"} + or parsed.username is not None + or parsed.password is not None + or parsed.path not in {"", "/"} + or parsed.params + or parsed.query + or parsed.fragment + ): + return None + return parsed.hostname, port + + +def _start_local_encoder( + state_directory: Path, + host: str, + port: int, + model: str, +) -> tuple[subprocess.Popen[bytes] | None, Path, bool]: + """Start a detached encoder unless a managed process is already alive.""" + state_directory.mkdir(parents=True, exist_ok=True) + pid_path = state_directory / "encoder.pid" + log_path = state_directory / "encoder.log" + if _read_live_pid(pid_path) is not None: + return None, log_path, False + + environment = os.environ.copy() + environment["OWI_SEMANTIC_HOST"] = host + environment["OWI_SEMANTIC_PORT"] = str(port) + environment["OWI_SEMANTIC_MODEL"] = model + command = [sys.executable, "-m", "owilix.semantic_encoder.app"] + with log_path.open("ab") as log_file: + process = subprocess.Popen( + command, + stdin=subprocess.DEVNULL, + stdout=log_file, + stderr=subprocess.STDOUT, + env=environment, + start_new_session=True, + close_fds=True, + ) + pid_path.write_text(f"{process.pid}\n", encoding="utf-8") + return process, log_path, True + + +def _read_live_pid(pid_path: Path) -> int | None: + """Return a live managed PID and discard stale state.""" + try: + pid = int(pid_path.read_text(encoding="utf-8").strip()) + except (FileNotFoundError, ValueError): + pid_path.unlink(missing_ok=True) + return None + if pid <= 0: + pid_path.unlink(missing_ok=True) + return None + try: + os.kill(pid, 0) + except ProcessLookupError: + pid_path.unlink(missing_ok=True) + return None + except PermissionError: + # A process that exists but is owned by another user is still live. + return pid + return pid + + +def _wait_until_ready( + base_url: str, + deadline: float, + process: subprocess.Popen[bytes] | None, + log_path: Path, +) -> dict[str, Any]: + """Poll health until startup completes, the child exits, or time expires.""" + last_error: Exception | None = None + while True: + if process is not None and process.poll() is not None: + detail = _last_log_line(log_path) + suffix = f" Last log message: {detail}" if detail else "" + raise EncoderProcessError( + f"The local semantic encoder exited with code {process.returncode}. " + f"See {log_path}.{suffix}" + ) + + remaining = deadline - time.monotonic() + if remaining <= 0: + detail = f" Last error: {last_error}" if last_error else "" + raise EncoderProcessError( + f"The local semantic encoder did not become ready before the timeout. " + f"See {log_path}.{detail}" + ) + + try: + health = _request_health(base_url, min(1.0, remaining)) + except httpx.RequestError as error: + last_error = error + else: + _require_ready(health, base_url) + return health + time.sleep(min(0.25, max(0.0, deadline - time.monotonic()))) + + +def _last_log_line(log_path: Path) -> str | None: + """Read a bounded diagnostic from a failed encoder process.""" + try: + with log_path.open("rb") as log_file: + log_file.seek(0, os.SEEK_END) + size = log_file.tell() + log_file.seek(max(0, size - 4096)) + lines = log_file.read().decode("utf-8", errors="replace").splitlines() + except OSError: + return None + return next((line.strip() for line in reversed(lines) if line.strip()), None) diff --git a/scripts/setup-semantic-rocm.sh b/scripts/setup-semantic-rocm.sh index 8812620..96aa62d 100755 --- a/scripts/setup-semantic-rocm.sh +++ b/scripts/setup-semantic-rocm.sh @@ -35,4 +35,4 @@ if torch.cuda.is_available(): print("GPU:", torch.cuda.get_device_name(0)) PY -echo "Start the encoder with: .venv/bin/owi-semantic-encoder" +echo "Ready. Run 'owi remote search --mode semantic'; the encoder starts automatically." diff --git a/tests/owilix/semantic_encoder/test_process.py b/tests/owilix/semantic_encoder/test_process.py new file mode 100644 index 0000000..ad1a66e --- /dev/null +++ b/tests/owilix/semantic_encoder/test_process.py @@ -0,0 +1,146 @@ +"""Tests for automatic lifecycle management of the semantic encoder.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from unittest.mock import Mock, patch + +import httpx + +from owilix.semantic_encoder.process import ( + EncoderProcessError, + _local_bind, + _start_local_encoder, + ensure_encoder, +) + + +MODEL = "jinaai/jina-embeddings-v5-text-small" +HEALTH = { + "ready": True, + "model": MODEL, + "dimension": 1024, + "device": "ROCm: test GPU", +} + + +def test_ensure_encoder_reuses_healthy_process(tmp_path: Path): + """A ready encoder is reused without starting another process.""" + with ( + patch( + "owilix.semantic_encoder.process._request_health", + return_value=HEALTH, + ), + patch("owilix.semantic_encoder.process._start_local_encoder") as start, + ): + readiness = ensure_encoder( + "http://127.0.0.1:8765", + tmp_path, + MODEL, + timeout=5.0, + ) + + assert readiness.health == HEALTH + assert readiness.started is False + start.assert_not_called() + + +def test_ensure_encoder_starts_local_process_and_waits(tmp_path: Path): + """Connection refusal triggers a detached local process and readiness poll.""" + request = httpx.Request("GET", "http://127.0.0.1:8765/health") + refused = httpx.ConnectError("connection refused", request=request) + process = Mock() + process.poll.return_value = None + statuses: list[str] = [] + log_path = tmp_path / "encoder.log" + + with ( + patch( + "owilix.semantic_encoder.process._request_health", + side_effect=[refused, HEALTH], + ), + patch( + "owilix.semantic_encoder.process._start_local_encoder", + return_value=(process, log_path, True), + ) as start, + ): + readiness = ensure_encoder( + "http://127.0.0.1:8765", + tmp_path, + MODEL, + timeout=5.0, + status_callback=statuses.append, + ) + + assert readiness.health == HEALTH + assert readiness.started is True + start.assert_called_once_with( + state_directory=tmp_path, + host="127.0.0.1", + port=8765, + model=MODEL, + ) + assert statuses[0].startswith("Starting the local semantic encoder") + assert statuses[-1] == "Local semantic encoder is ready on ROCm: test GPU" + + +def test_ensure_encoder_does_not_start_remote_process(tmp_path: Path): + """An unreachable custom remote endpoint is never treated as a bind target.""" + request = httpx.Request("GET", "https://encoder.example/health") + refused = httpx.ConnectError("connection refused", request=request) + with ( + patch( + "owilix.semantic_encoder.process._request_health", + side_effect=refused, + ), + patch("owilix.semantic_encoder.process._start_local_encoder") as start, + ): + try: + ensure_encoder( + "https://encoder.example", + tmp_path, + MODEL, + timeout=5.0, + ) + except EncoderProcessError as error: + assert "only available for a local HTTP encoder" in str(error) + else: + raise AssertionError("unreachable remote encoder should fail") + + start.assert_not_called() + + +def test_start_local_encoder_uses_current_python_and_dataset_model(tmp_path: Path): + """The managed child inherits this environment and the required vector model.""" + process = Mock(pid=4242) + with ( + patch("owilix.semantic_encoder.process._read_live_pid", return_value=None), + patch("owilix.semantic_encoder.process.subprocess.Popen", return_value=process) as popen, + ): + returned, log_path, started = _start_local_encoder( + tmp_path, + host="127.0.0.1", + port=8765, + model=MODEL, + ) + + assert returned is process + assert log_path == tmp_path / "encoder.log" + assert started is True + assert (tmp_path / "encoder.pid").read_text(encoding="utf-8") == "4242\n" + command = popen.call_args.args[0] + environment = popen.call_args.kwargs["env"] + assert command == [sys.executable, "-m", "owilix.semantic_encoder.app"] + assert environment["OWI_SEMANTIC_MODEL"] == MODEL + assert environment["OWI_SEMANTIC_HOST"] == "127.0.0.1" + assert environment["OWI_SEMANTIC_PORT"] == "8765" + + +def test_local_bind_only_accepts_plain_loopback_http(): + """Auto-start cannot be redirected to a non-local or ambiguous endpoint.""" + assert _local_bind("http://127.0.0.1:9999") == ("127.0.0.1", 9999) + assert _local_bind("http://localhost:8765/") == ("localhost", 8765) + assert _local_bind("https://127.0.0.1:8765") is None + assert _local_bind("http://encoder.example:8765") is None + assert _local_bind("http://127.0.0.1:8765/api") is None -- 2.51.2