From 19fbb8e92618992e75c2afe7ed4774f8b7e9a440 Mon Sep 17 00:00:00 2001 From: Emil Date: Sat, 1 Aug 2026 23:57:41 +0300 Subject: [PATCH] Replace legacy distributed protocol with SDK-built workloads --- README.md | 16 +- STATUS.md | 15 +- docs/sdk-handoff.md | 56 +- docs/workload-sdk.md | 93 ++- pyproject.toml | 5 +- scimesh/distributed/__init__.py | 36 -- scimesh/distributed/models.py | 262 -------- scimesh/distributed/registry.py | 108 ---- scimesh/distributed/similarity_search.py | 380 ------------ scimesh/distributed/workload.py | 40 -- scimesh/sdk/__init__.py | 12 - scimesh/sdk/builtins.py | 150 ----- scimesh/sdk/compat/__init__.py | 5 - scimesh/sdk/compat/distributed_v1.py | 369 ------------ scimesh/worker/runners.py | 329 +++++++--- .../descriptors/__init__.py | 2 +- .../{sdk => workloads}/descriptors/core.py | 0 .../descriptors/definition.py | 30 +- scimesh/workloads/environment.py | 42 ++ scimesh/workloads/graph/__init__.py | 40 ++ scimesh/workloads/graph/core.py | 249 ++++++++ scimesh/workloads/graph/definition.py | 500 ++++++++++++++++ scimesh/workloads/library.py | 68 +++ scimesh/workloads/search/__init__.py | 31 + scimesh/workloads/search/core.py | 261 ++++++++ scimesh/workloads/search/definition.py | 565 ++++++++++++++++++ tests/test_distributed.py | 186 ------ tests/test_distributed_similarity_search.py | 184 ------ tests/test_sdk_compatibility.py | 66 +- tests/test_sdk_descriptors.py | 6 +- tests/test_sdk_graph.py | 295 +++++++++ tests/test_sdk_registry.py | 52 +- tests/test_sdk_search.py | 252 ++++++++ tests/test_worker_daemon.py | 290 +++++++-- 34 files changed, 3009 insertions(+), 1986 deletions(-) delete mode 100644 scimesh/distributed/__init__.py delete mode 100644 scimesh/distributed/models.py delete mode 100644 scimesh/distributed/registry.py delete mode 100644 scimesh/distributed/similarity_search.py delete mode 100644 scimesh/distributed/workload.py delete mode 100644 scimesh/sdk/builtins.py delete mode 100644 scimesh/sdk/compat/__init__.py delete mode 100644 scimesh/sdk/compat/distributed_v1.py rename scimesh/{sdk => workloads}/descriptors/__init__.py (94%) rename scimesh/{sdk => workloads}/descriptors/core.py (100%) rename scimesh/{sdk => workloads}/descriptors/definition.py (94%) create mode 100644 scimesh/workloads/environment.py create mode 100644 scimesh/workloads/graph/__init__.py create mode 100644 scimesh/workloads/graph/core.py create mode 100644 scimesh/workloads/graph/definition.py create mode 100644 scimesh/workloads/library.py create mode 100644 scimesh/workloads/search/__init__.py create mode 100644 scimesh/workloads/search/core.py create mode 100644 scimesh/workloads/search/definition.py delete mode 100644 tests/test_distributed.py delete mode 100644 tests/test_distributed_similarity_search.py create mode 100644 tests/test_sdk_graph.py create mode 100644 tests/test_sdk_search.py diff --git a/README.md b/README.md index 939e30f..f10d3a3 100644 --- a/README.md +++ b/README.md @@ -135,11 +135,17 @@ The package separates common dataset parsing and fingerprints from independent w ## Workload SDK -`scimesh.sdk` implements the `core-batch-v1` authoring profile: strict and -immutable workload manifests, typed artifact ports, static map/reduce plans, -resource eligibility and local reservations, exact/canonical/numeric verifier -primitives, installed-package allowlisting, and a compatibility adapter for the -existing distributed `similarity-search`. See the +`scimesh.sdk` is the framework only: strict and immutable workload manifests, +typed artifact ports, static map/reduce plans, resource eligibility and local +reservations, exact/canonical/numeric verifier primitives, installed-package +allowlisting, and a local conformance executor. It contains no scientific +workload code. Workloads are user scripts built on the SDK: the built-in +`similarity-search`, `similarity-graph`, and `descriptor-batch` live in +`scimesh/workloads/` (each a small package with `core.py` + `definition.py`), +composed by `scimesh/workloads/library.py` and registered through +`scimesh.workloads` entry points. The Worker Agent executes those SDK-built +workloads directly (see `scimesh/worker/runners.py`), so the same scientific +handlers run locally, in conformance, and on claimed coordinator tasks. See the [SDK author guide](docs/workload-sdk.md), [contract](docs/scimesh-sdk-contract.md), and [delivery roadmap](docs/scimesh-sdk-roadmap.md). diff --git a/STATUS.md b/STATUS.md index 1ad011b..cab1aa1 100644 --- a/STATUS.md +++ b/STATUS.md @@ -45,12 +45,13 @@ the complete result-artifact SHA-256 before a task is accepted. | CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. | | CTX-08 Distributed similarity-search | Implemented | Python planner resolves `query_id` once, creates deterministic shard plans, worker adapter emits exact partial top-k CSVs/metrics, and reducer matches the local reference. | | CTX-09 Reducer and final-result API | Implemented | Atomic `reducing` claim, deterministic coordinator-side top-k reducer, sanitized reducer failure, final artifact persistence, `result_uri`, and final CSV download. | -| CTX-10 Distributed similarity-graph | Not started | Local reference exists. | +| CTX-10 Distributed similarity-graph | Not started | Local reference exists; the SDK-built local graph workload already enforces the pair-coverage invariant. | | CTX-11 Dashboard/operator view | Implemented | Protected live control room: recent-run/worker overview, real pipeline-stage visualization, shard attempts and safe failures, validated similarity-search upload, coordinator artifacts, final-result download, and bounded polling. | | CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. | | CTX-15 User Service and access control | Implemented | User/owner scoping, verified contributors, worker keys, self-service enrollment, and quorum-backed untrusted workers are merged; local Go/Python and Docker/PostgreSQL checks passed. | -| CTX-16 Workload SDK foundation | Implemented | `scimesh.sdk` provides strict immutable manifests/plans/artifacts, digest/trust-pinned tasks, typed DAGs, compatibility negotiation, verifier primitives with owner/binding-safe quorum inputs, resource eligibility/local allocation, measured package discovery, a trusted local core-batch conformance harness, and a tested legacy similarity-search adapter. Enforcing coordinator/Worker profiles remain fail-closed. | -| SDK roadmap step 3: `descriptor-batch` | Implemented | The first SDK-native reference workload (`descriptor-batch@1.0.0` in `scimesh/sdk/descriptors/`): pinned 81-name RDKit 2D descriptor set, canonical one-row-per-input CSV, deterministic row-bounded shards, shard-index concatenation with one header, byte-identical local/distributed output, and a two-worker `untrusted_quorum` verifier test. Entry point declared in `pyproject.toml`; manifest declares `trusted` + `untrusted_quorum` with the exact-artifact verifier. | +| CTX-16 Workload SDK foundation | Implemented | `scimesh.sdk` provides strict immutable manifests/plans/artifacts, digest/trust-pinned tasks, typed DAGs, compatibility negotiation, verifier primitives with owner/binding-safe quorum inputs, resource eligibility/local allocation, measured package discovery, a trusted local core-batch conformance harness, and strict package discovery. Enforcing coordinator/Worker profiles remain fail-closed. | +| SDK roadmap step 3: `descriptor-batch` | Implemented | The first SDK-built reference workload (`scimesh/workloads/descriptors/`): pinned 81-name RDKit 2D descriptor set, canonical one-row-per-input CSV, deterministic row-bounded shards, shard-index concatenation with one header, byte-identical local/distributed output, and a two-worker `untrusted_quorum` verifier test. | +| SDK-built `similarity-search` and `similarity-graph` | Implemented | Both workloads are now SDK-built packages (`scimesh/workloads/search/`, `scimesh/workloads/graph/`) with their own manifests/planners/runners/reducers on the `core-batch-v1` profile; they reuse the local scientific cores and are byte-identical to the single-process references (search; graph for both threshold directions and any block size). The graph reducer enforces the CTX-10 pair-coverage invariant. `scimesh/workloads/library.py` composes the built-in registry/runtime. | ## Next recommended assignment @@ -59,8 +60,12 @@ block-pair planning and reduction for `similarity-graph`. ## Known constraints -- The worker/coordinator flow currently accepts both underscore API workload - names and hyphenated CLI names while the contract is consolidated. +- The worker/coordinator flow accepts both underscore API workload names and + hyphenated names at the runner boundary; the runner normalizes them. +- The worker executes SDK-built workloads through `scimesh/worker/runners.py` + (a v1-wire bridge over `TaskSpec`/`LocalTaskContext`); the CTX-07 + `DistributedWorkload` protocol module and the SDK compatibility adapter were + removed. `max_rows` is a plan-time option and is rejected per task. - A real-stack worker test uses a small `query_smiles` shard. The Python planner resolves `query_id` once and shares `query_smiles`; the upload UI currently accepts `query_smiles` only. diff --git a/docs/sdk-handoff.md b/docs/sdk-handoff.md index f746dd1..6b9efee 100644 --- a/docs/sdk-handoff.md +++ b/docs/sdk-handoff.md @@ -12,6 +12,18 @@ Read first, in this order: `AGENTS.md` (binding repo rules), ## What is already done (do not redo) +**Legacy removal (2026-08-01):** the CTX-07 `DistributedWorkload` protocol +package (`scimesh/distributed/`), the SDK compatibility adapter +(`scimesh/sdk/compat/`), and `library.similarity_search_sdk_adapter` were +removed. The worker now executes the SDK-built workloads directly: +`scimesh/worker/runners.py` builds a `TaskSpec` with the workload's pins, +negotiates, reserves resources, runs the workload's own Runner through +`LocalTaskContext` (store-backed catalog/sink), and uploads the sealed +partial over the unchanged v1 wire. `run_search_shard` + the full-precision +partial writer moved to `scimesh/workloads/search/core.py`; the partial +format is unchanged, so the Go reducer and UI keep working. The runner +resolves `query_id` per task and rejects plan-time `max_rows`. + CTX-16 "Workload SDK foundation" is complete and tested. `scimesh/sdk/` implements the `core-batch-v1` profile: @@ -26,29 +38,41 @@ implements the `core-batch-v1` profile: `NumericToleranceVerifier` with bounded sanitized evidence: `verification.py`. - Local conformance harness: `LocalArtifactStore`, `LocalCoreBatchExecutor`, `ResourcePool` (atomic all-or-nothing reservation): `conformance.py`. -- Legacy adapter exposing distributed `similarity-search` through the SDK - without changing its wire schema: `compat/distributed_v1.py`, `builtins.py`; - entry point `similarity-search@1.0.0` is declared in `pyproject.toml`. +- SDK-built workloads living outside the SDK: `scimesh/workloads/search/`, + `scimesh/workloads/graph/`, `scimesh/workloads/descriptors/` (each `core.py` + + `definition.py`), composed by `scimesh/workloads/library.py` + (`default_sdk_registry`, `default_sdk_runtime`); entry points for all three + declared in `pyproject.toml`. - Tests: `tests/test_sdk_{models,resources,verification,compatibility,registry}.py` including fail-closed rejection coverage for every advanced profile declaration (gang, GPU modes, pools, checkpoints, retries, secrets, streams, - loops, side effects). + loops, side effects), plus `tests/test_sdk_{search,graph,descriptors}.py` + and the worker bridge tests in `tests/test_worker_daemon.py`. + `tests/test_distributed*.py` were removed with the protocol. ## What remains, in delivery order 1. ~~**`descriptor-batch` reference workload**~~ — **done** (2026-08-01): - `scimesh/sdk/descriptors/` (`core.py` + `definition.py`) is the first - SDK-native workload. Pinned 81-name RDKit 2D descriptor set (validated at + `scimesh/workloads/descriptors/` (`core.py` + `definition.py`) is the first + SDK-built workload. Pinned 81-name RDKit 2D descriptor set (validated at definition build time), canonical one-row-per-input CSV with `%.6f` floats, deterministic row-bounded shards, shard-index concatenation with one header, byte-identical local/distributed output, `skip_invalid` explicit policy, and - `untrusted_quorum` + exact-artifact@1 declared in the manifest. Entry point - `descriptor-batch@1.0.0` is in `pyproject.toml`; `default_sdk_runtime` now - advertises the `descriptor-batch` capability. Tests: - `tests/test_sdk_descriptors.py` (8 tests: manifest/negotiation, local-vs- - reference byte parity, deterministic path-free planning, explicit invalid- - row policy, strict parameter schema, two-owner quorum accept, conflicting- - quorum reject, allowlist discovery). Total suite: 233 passing. + `untrusted_quorum` + exact-artifact@1 declared in the manifest. Tests: + `tests/test_sdk_descriptors.py`. +2. ~~**SDK-built `similarity-search` and `similarity-graph`**~~ — **done** + (2026-08-01). Both local workloads are SDK-built packages outside the SDK: + `scimesh/workloads/search/` and `scimesh/workloads/graph/` (each `core.py` + + `definition.py`, manifest + planner/runner/reducer, byte_exact + + exact-artifact@1, trusted + untrusted_quorum). Search resolves the query at + plan time and merges partials with the reference heap (byte-identical to the + CLI). Graph plans one task per block pair `(i,j)` with `i <= j`, reducer + enforces pair-coverage and duplicate-pair rejection, output byte-identical + to the local brute-force reference for both directions and any block size. + Tests: `tests/test_sdk_search.py`, `tests/test_sdk_graph.py`. + **Architecture note:** `scimesh.sdk/` is the framework ONLY (no workload + code); workloads are user scripts/packages under `scimesh/workloads/` that + import the SDK. Keep new workloads out of the SDK package. 2. **Distributed `similarity-graph`** (CTX-10, roadmap step 1). The coordinator currently rejects `similarity-graph` uploads; it needs cross-shard block-pair planning and duplicate-safe reduction. STATUS.md names this the next @@ -76,6 +100,12 @@ implements the `core-batch-v1` profile: ## Known traps (cost the previous session real time) +- **Architecture boundary:** `scimesh.sdk/` is the framework only and must + never import `scimesh.workloads` (SDK depends on nothing workload-specific). + Workload packages live under `scimesh/workloads/` (each `core.py` + + `definition.py`), and built-in wiring lives in `scimesh/workloads/library.py`. + The digest helpers are in `scimesh/workloads/environment.py`; the SDK keeps + only the generic `installed_distribution_digest` in `scimesh/sdk/integrity.py`. - The legacy adapter pins its own manifest (`adapter.manifest`). If a test changes limits/workflow on the manifest, the adapter's copy must be replaced too, or `registry.plan` fails with "planner plan does not carry the selected diff --git a/docs/workload-sdk.md b/docs/workload-sdk.md index 5f2be84..6779c32 100644 --- a/docs/workload-sdk.md +++ b/docs/workload-sdk.md @@ -16,6 +16,29 @@ declarations, but the current coordinator/Worker runtime does not advertise their features. Compatibility negotiation therefore rejects those workflows before planner code runs. +## SDK versus workloads + +`scimesh.sdk` is the framework only: strict manifests, plans, artifacts, +registry, verifiers, and the local conformance executor. It contains no +scientific workload code. Workloads are user Python scripts and packages that +import the SDK and live outside it. The built-in SciMesh workloads are under +`scimesh/workloads/`: + +- `scimesh/workloads/search/` — SDK-built `similarity-search@1.0.0`; +- `scimesh/workloads/graph/` — SDK-built `similarity-graph@1.0.0`; +- `scimesh/workloads/descriptors/` — SDK-built `descriptor-batch@1.0.0`; +- `scimesh/workloads/library.py` — the built-in library wiring: a default + registry containing all three definitions and a runtime advertising their + capabilities; +- the plain `scimesh/workloads/*.py` modules remain the local CLI scientific + cores and their `Workload` registry. + +Each SDK-built workload is a small package with `core.py` (scientific code) +and `definition.py` (manifest plus planner/runner/reducer handlers). A future +external workload library can follow the same shape: its own distribution, one +`scimesh.workloads` entry point per workload version, and an administrator +allowlist. + ## What authors import The stable authoring surface is exported from `scimesh.sdk`: @@ -41,10 +64,10 @@ contracts carry schema versions. Artifact identities contain a coordinator-owned UUID, schema, checksum, media type, and bounds; a scientific handler never persists a filesystem path. -## Try the built-in SDK workload +## Try the built-in SDK workloads -This example executes the current distributed `similarity-search` through the -SDK without starting PostgreSQL or the coordinator: +This example runs the SDK-built `similarity-search` without starting +PostgreSQL or the coordinator: ```python from pathlib import Path @@ -54,21 +77,23 @@ from scimesh.sdk import ( JobRequest, LocalArtifactStore, LocalCoreBatchExecutor, +) +from scimesh.workloads.library import ( default_sdk_registry, default_sdk_runtime, - similarity_search_sdk_adapter, + similarity_search_sdk_definition, ) root = Path("sdk-run") store = LocalArtifactStore(root / "artifacts") -adapter = similarity_search_sdk_adapter(shard_rows=1_000) +workload = similarity_search_sdk_definition(shard_rows=1_000) dataset = store.import_file( Path("chembl_37_chemreps.txt"), - declaration=adapter.input_port.schema, + declaration=workload.manifest.inputs["input"].schema, ) request = JobRequest( - workload=adapter.manifest.workload, + workload=workload.manifest.workload, parameters={"query_smiles": "CCO", "top_k": 20}, inputs={"input": ArtifactCollection.single(dataset)}, ) @@ -78,7 +103,7 @@ result = LocalCoreBatchExecutor( default_sdk_runtime(), store, root / "attempts", -).execute(request, adapter.manifest.package.digest) +).execute(request, workload.manifest.package.digest) result_ref = result.outputs["result"].items[0].artifact print(store.materialize(result_ref)) @@ -90,21 +115,20 @@ for coordinator leases or multi-machine scheduling. It accepts only map/reduce stages without secrets, checkpoints, retries, gangs, or accelerators. It does not claim network, timeout, process, or credential isolation. Unsupported declarations are rejected before a handler runs. The -harness uses the same legacy scientific planner, shard runner, and reducer as -the distributed `similarity-search`, and its parity is covered by automated -tests. +harness runs the SDK-built workload handlers themselves, and their parity +against the single-process references is covered by automated tests. ## The descriptor-batch reference workload `descriptor-batch@1.0.0` is the first SDK-native reference workload: it is -built directly on the manifest/planner/runner/reducer contracts instead of the -legacy adapter, and it is the intended first `untrusted_quorum` candidate +built directly on the manifest/planner/runner/reducer contracts, and it is +the intended first `untrusted_quorum` candidate (`byte_exact` plus `exact-artifact@1`). Its scientific contract is pinned: - one output CSV row per valid input molecule, in input order, with RDKit canonical SMILES recomputed by RDKit; - an explicit 81-name pinned RDKit 2D descriptor set (see - `scimesh/sdk/descriptors/core.py`), validated against the installed RDKit at + `scimesh/workloads/descriptors/core.py`), validated against the installed RDKit at definition build time; - `%.6f` float formatting, `utf-8` CSV with one header, and row-bounded deterministic shards; @@ -123,9 +147,9 @@ from scimesh.sdk import ( LocalArtifactStore, LocalCoreBatchExecutor, WorkloadRegistry, - default_sdk_runtime, ) -from scimesh.sdk.descriptors import descriptor_batch_sdk_definition +from scimesh.workloads.descriptors import descriptor_batch_sdk_definition +from scimesh.workloads.library import default_sdk_runtime root = Path("descriptor-run") store = LocalArtifactStore(root / "artifacts") @@ -161,13 +185,39 @@ matching `AllowedPackage` allowlist entry. Its manifest declares both so the same definition can later run under coordinator quorum once protocol-v2 leases exist. +## The SDK-built similarity workloads + +`similarity-search@1.0.0` and `similarity-graph@1.0.0` are SDK-built workloads +under `scimesh/workloads/search/` and `scimesh/workloads/graph/`; both reuse +the local scientific cores from `scimesh/workloads/similarity_search.py` and +`similarity_graph.py` and declare `byte_exact` + `exact-artifact@1`: + +- the search workload resolves `query_id` exactly once at plan time, shards + the input deterministically, computes a local top-k per shard with the + reference heap, and merges the sorted partials with the same tie-breakers, + so the final CSV is byte-identical to the single-process CLI output; +- the graph workload parses molecules once into deterministic row-ordered + blocks, plans one map task per block pair `(i, j)` with `i <= j`, and its + reducer enforces the pair-coverage invariant (every unordered molecule pair + compared exactly once, no duplicates) before emitting the same + deterministically sorted edge list as the local brute-force reference, for + either threshold direction and any block size; +- the v1 worker executes the SDK-built `similarity-search` runner directly: + `scimesh/worker/runners.py` is a small wire bridge that builds a `TaskSpec` + with the workload's own pins, reserves resources, seals the partial through + a content-addressed store, and uploads the resulting CSV over the unchanged + coordinator contract. + ## Package shape and registration -An SDK distribution provides one explicit entry point per workload version: +A workload distribution provides one explicit entry point per workload +version. The built-in workloads are part of the `scimesh` distribution: ```toml [project.entry-points."scimesh.workloads"] -"descriptor-batch@1.0.0" = "scimesh_descriptors.sdk:workload_definition" +"similarity-search@1.0.0" = "scimesh.workloads.search:workload_definition" +"similarity-graph@1.0.0" = "scimesh.workloads.graph:workload_definition" +"descriptor-batch@1.0.0" = "scimesh.workloads.descriptors:workload_definition" ``` The factory returns a `WorkloadDefinition` containing its manifest and handler @@ -279,10 +329,11 @@ pytest tests/test_sdk_models.py \ tests/test_sdk_verification.py \ tests/test_sdk_compatibility.py \ tests/test_sdk_registry.py \ - tests/test_sdk_descriptors.py + tests/test_sdk_descriptors.py \ + tests/test_sdk_search.py \ + tests/test_sdk_graph.py ``` -Run `pytest` for the full legacy, Worker, local-science, and SDK regression -suite. Package authors can reuse `LocalArtifactStore`, +Run `pytest` for the full Worker, local-science, and SDK regression suite. Package authors can reuse `LocalArtifactStore`, `LocalCoreBatchExecutor`, and `assert_manifest_round_trip` in their own golden tests. diff --git a/pyproject.toml b/pyproject.toml index 0c21f9a..2ac1c2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,8 +18,9 @@ scimesh = "scimesh.cli:main" scimesh-worker = "scimesh.worker.cli:main" [project.entry-points."scimesh.workloads"] -"similarity-search@1.0.0" = "scimesh.sdk.builtins:similarity_search_workload_definition" -"descriptor-batch@1.0.0" = "scimesh.sdk.descriptors.definition:workload_definition" +"similarity-search@1.0.0" = "scimesh.workloads.search:workload_definition" +"similarity-graph@1.0.0" = "scimesh.workloads.graph:workload_definition" +"descriptor-batch@1.0.0" = "scimesh.workloads.descriptors:workload_definition" [tool.setuptools.packages.find] include = ["scimesh*"] diff --git a/scimesh/distributed/__init__.py b/scimesh/distributed/__init__.py deleted file mode 100644 index 40b5026..0000000 --- a/scimesh/distributed/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Coordinator-independent contracts for distributed SciMesh workloads. - -This package defines the typed plan and reduction boundary shared by future -planners, worker adapters, and coordinator bridges. It intentionally has no -network, database, or coordinator imports. -""" - -from .models import ( - ArtifactReference, - CompletedPartial, - DistributedPlan, - FinalResult, - PlannedTask, -) -from .registry import ( - DistributedWorkloadRegistry, - PlanningService, - WorkloadDescription, - default_distributed_registry, -) -from .similarity_search import SimilaritySearchDistributedWorkload -from .workload import DistributedWorkload - -__all__ = [ - "ArtifactReference", - "CompletedPartial", - "default_distributed_registry", - "DistributedPlan", - "DistributedWorkload", - "DistributedWorkloadRegistry", - "FinalResult", - "PlannedTask", - "PlanningService", - "SimilaritySearchDistributedWorkload", - "WorkloadDescription", -] diff --git a/scimesh/distributed/models.py b/scimesh/distributed/models.py deleted file mode 100644 index baf0487..0000000 --- a/scimesh/distributed/models.py +++ /dev/null @@ -1,262 +0,0 @@ -"""Versioned, JSON-safe value objects for distributed workload contracts.""" - -from __future__ import annotations - -import json -import math -import re -from dataclasses import dataclass -from typing import Any, Mapping, Sequence -from uuid import UUID - - -SCHEMA_VERSION = 1 -_WORKLOAD_NAME = re.compile(r"^[a-z][a-z0-9]*(?:-[a-z0-9]+)*$") - - -def _canonical_uuid(value: object, field: str) -> str: - if not isinstance(value, str): - raise ValueError(f"{field} must be a UUID string") - try: - return str(UUID(value)) - except ValueError as error: - raise ValueError(f"{field} must be a UUID string") from error - - -def _sha256(value: object, field: str) -> str: - if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value): - raise ValueError(f"{field} must be a lowercase SHA-256 hex digest") - return value - - -def _content_type(value: object, field: str) -> str: - if not isinstance(value, str) or not value or len(value) > 128: - raise ValueError(f"{field} must be a non-empty content type") - if any(character.isspace() or ord(character) < 32 for character in value): - raise ValueError(f"{field} must be a non-empty content type") - return value - - -def _workload_name(value: object, field: str = "workload") -> str: - if not isinstance(value, str) or not _WORKLOAD_NAME.fullmatch(value): - raise ValueError(f"{field} must be a canonical hyphenated workload name") - return value - - -def _json_value(value: object, field: str) -> Any: - """Deep-copy a JSON value and reject non-finite or non-string-key data.""" - if value is None or isinstance(value, (bool, int)): - return value - if isinstance(value, str): - # Coordinator-owned artifacts are represented exclusively by - # ArtifactReference. A URI or a local path in a generic JSON payload - # would let a planner accidentally leak a bridge/worker implementation - # detail into durable task metadata. - forbidden_prefixes = ("file://", "worker://", "http://", "https://", "s3://", "/") - is_windows_path = len(value) >= 3 and value[0].isalpha() and value[1:3] in (":/", ":\\") - if value.startswith(forbidden_prefixes) or is_windows_path: - raise ValueError(f"{field} must not contain a URI or local path") - return value - if isinstance(value, float): - if not math.isfinite(value): - raise ValueError(f"{field} must not contain NaN or infinity") - return value - if isinstance(value, Mapping): - copied: dict[str, Any] = {} - for key, child in value.items(): - if not isinstance(key, str): - raise ValueError(f"{field} must use string object keys") - copied[key] = _json_value(child, f"{field}.{key}") - return copied - if isinstance(value, (list, tuple)): - return [_json_value(child, f"{field}[]") for child in value] - raise ValueError(f"{field} must contain only JSON-compatible values") - - -def _json_mapping(value: object, field: str) -> dict[str, Any]: - if not isinstance(value, Mapping): - raise ValueError(f"{field} must be an object") - return _json_value(value, field) - - -@dataclass(frozen=True) -class ArtifactReference: - """Immutable coordinator-owned artifact identity used in a plan.""" - - artifact_id: str - sha256: str - content_type: str - - def __post_init__(self) -> None: - object.__setattr__(self, "artifact_id", _canonical_uuid(self.artifact_id, "artifact_id")) - object.__setattr__(self, "sha256", _sha256(self.sha256, "sha256")) - object.__setattr__(self, "content_type", _content_type(self.content_type, "content_type")) - - def to_dict(self) -> dict[str, str]: - return { - "artifact_id": self.artifact_id, - "sha256": self.sha256, - "content_type": self.content_type, - } - - @classmethod - def from_dict(cls, value: object) -> "ArtifactReference": - if not isinstance(value, Mapping): - raise ValueError("artifact reference must be an object") - _require_exact_keys(value, {"artifact_id", "sha256", "content_type"}, "artifact reference") - return cls( - artifact_id=value["artifact_id"], - sha256=value["sha256"], - content_type=value["content_type"], - ) - - -@dataclass(frozen=True) -class PlannedTask: - """One deterministically indexed, artifact-backed worker task.""" - - chunk_index: int - input_artifact: ArtifactReference - parameters: Mapping[str, object] - - def __post_init__(self) -> None: - if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0: - raise ValueError("chunk_index must be a non-negative integer") - if not isinstance(self.input_artifact, ArtifactReference): - raise ValueError("input_artifact must be an ArtifactReference") - object.__setattr__(self, "parameters", _json_mapping(self.parameters, "task parameters")) - - def to_dict(self) -> dict[str, Any]: - return { - "chunk_index": self.chunk_index, - "input_artifact": self.input_artifact.to_dict(), - "parameters": _json_value(self.parameters, "task parameters"), - } - - @classmethod - def from_dict(cls, value: object) -> "PlannedTask": - if not isinstance(value, Mapping): - raise ValueError("planned task must be an object") - _require_exact_keys(value, {"chunk_index", "input_artifact", "parameters"}, "planned task") - return cls( - chunk_index=value["chunk_index"], - input_artifact=ArtifactReference.from_dict(value["input_artifact"]), - parameters=value["parameters"], - ) - - -@dataclass(frozen=True) -class DistributedPlan: - """The complete schema-versioned output of a distributed planner.""" - - workload: str - resolved_parameters: Mapping[str, object] - tasks: Sequence[PlannedTask] - schema_version: int = SCHEMA_VERSION - - def __post_init__(self) -> None: - if self.schema_version != SCHEMA_VERSION: - raise ValueError(f"schema_version must be {SCHEMA_VERSION}") - object.__setattr__(self, "workload", _workload_name(self.workload)) - object.__setattr__(self, "resolved_parameters", _json_mapping(self.resolved_parameters, "resolved_parameters")) - task_list = tuple(self.tasks) - if not task_list: - raise ValueError("plan must contain at least one task") - if any(not isinstance(task, PlannedTask) for task in task_list): - raise ValueError("tasks must contain PlannedTask values") - indexes = [task.chunk_index for task in task_list] - if indexes != sorted(indexes) or len(set(indexes)) != len(indexes): - raise ValueError("tasks must have unique, ascending chunk_index values") - object.__setattr__(self, "tasks", task_list) - - def to_dict(self) -> dict[str, Any]: - return { - "schema_version": self.schema_version, - "workload": self.workload, - "resolved_parameters": _json_value(self.resolved_parameters, "resolved_parameters"), - "tasks": [task.to_dict() for task in self.tasks], - } - - def to_json(self) -> str: - """Return stable JSON suitable for hashing, tests, and durable payloads.""" - return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False) - - @classmethod - def from_dict(cls, value: object) -> "DistributedPlan": - if not isinstance(value, Mapping): - raise ValueError("distributed plan must be an object") - _require_exact_keys( - value, - {"schema_version", "workload", "resolved_parameters", "tasks"}, - "distributed plan", - ) - raw_tasks = value["tasks"] - if not isinstance(raw_tasks, list): - raise ValueError("tasks must be an array") - return cls( - schema_version=value["schema_version"], - workload=value["workload"], - resolved_parameters=value["resolved_parameters"], - tasks=tuple(PlannedTask.from_dict(task) for task in raw_tasks), - ) - - @classmethod - def from_json(cls, value: str) -> "DistributedPlan": - try: - decoded = json.loads(value) - except (TypeError, json.JSONDecodeError) as error: - raise ValueError("distributed plan must be valid JSON") from error - return cls.from_dict(decoded) - - -@dataclass(frozen=True) -class CompletedPartial: - """Coordinator-owned partial output supplied to a reducer.""" - - chunk_index: int - artifact: ArtifactReference - metrics: Mapping[str, int | float] - - def __post_init__(self) -> None: - if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0: - raise ValueError("chunk_index must be a non-negative integer") - if not isinstance(self.artifact, ArtifactReference): - raise ValueError("artifact must be an ArtifactReference") - if not isinstance(self.metrics, Mapping): - raise ValueError("metrics must be an object") - metrics: dict[str, int | float] = {} - for name, value in self.metrics.items(): - if not isinstance(name, str) or not name: - raise ValueError("metric names must be non-empty strings") - if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): - raise ValueError("metric values must be finite JSON numbers") - metrics[name] = value - object.__setattr__(self, "metrics", metrics) - - -@dataclass(frozen=True) -class FinalResult: - """A reducer's durable output, ready for coordinator persistence.""" - - artifact: ArtifactReference - metrics: Mapping[str, int | float] - - def __post_init__(self) -> None: - if not isinstance(self.artifact, ArtifactReference): - raise ValueError("artifact must be an ArtifactReference") - # Reuse the CompletedPartial metric validation without inventing a fake - # artifact lifecycle or widening the result contract. - object.__setattr__(self, "metrics", CompletedPartial(0, self.artifact, self.metrics).metrics) - - -def _require_exact_keys(value: Mapping[str, object], expected: set[str], label: str) -> None: - actual = set(value) - if actual != expected: - missing = sorted(expected - actual) - unknown = sorted(actual - expected) - details: list[str] = [] - if missing: - details.append(f"missing {', '.join(missing)}") - if unknown: - details.append(f"unknown {', '.join(unknown)}") - raise ValueError(f"{label} has {'; '.join(details)} fields") diff --git a/scimesh/distributed/registry.py b/scimesh/distributed/registry.py deleted file mode 100644 index 868b27f..0000000 --- a/scimesh/distributed/registry.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Registry and orchestration helpers for distributed workload contracts.""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Mapping, Sequence - -from .models import CompletedPartial, DistributedPlan, FinalResult, _workload_name -from .workload import DistributedWorkload - - -@dataclass(frozen=True) -class WorkloadDescription: - """Safe metadata that a future coordinator or UI may display.""" - - name: str - description: str - - -class DistributedWorkloadRegistry: - """Collect distributed workloads without coupling them to the CLI registry.""" - - def __init__(self) -> None: - self._workloads: dict[str, DistributedWorkload] = {} - - def register(self, workload: DistributedWorkload) -> None: - name = _workload_name(workload.name) - if name in self._workloads: - raise ValueError(f"distributed workload already registered: {name}") - if not isinstance(workload.description, str) or not workload.description.strip(): - raise ValueError("distributed workload description must be non-empty") - self._workloads[name] = workload - - def require(self, name: str) -> DistributedWorkload: - try: - return self._workloads[_workload_name(name)] - except KeyError as error: - raise ValueError(f"unknown distributed workload: {name}") from error - - def descriptions(self) -> tuple[WorkloadDescription, ...]: - return tuple( - WorkloadDescription(name, workload.description) - for name, workload in sorted(self._workloads.items()) - ) - - -class PlanningService: - """Small bridge-safe orchestration around a distributed workload registry. - - It writes neither jobs nor artifacts. A Go coordinator bridge can therefore - validate and produce a plan before opening its own all-or-nothing persistence - transaction; CTX-09 will implement that concrete bridge and durable result - orchestration. - """ - - def __init__(self, registry: DistributedWorkloadRegistry) -> None: - self._registry = registry - - def plan( - self, - workload_name: str, - input_path: Path, - input_artifact_id: str, - parameters: Mapping[str, object], - shard_rows: int, - workspace: Path, - ) -> DistributedPlan: - if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1: - raise ValueError("shard_rows must be a positive integer") - workload = self._registry.require(workload_name) - workload.validate_job(parameters) - plan = workload.plan(input_path, input_artifact_id, parameters, shard_rows, workspace) - if not isinstance(plan, DistributedPlan): - raise ValueError("distributed planner must return a DistributedPlan") - if plan.workload != workload.name: - raise ValueError("distributed planner returned a plan for another workload") - # Round-trip through the strict wire schema now, before a future bridge - # persists anything. This catches non-JSON values and undeclared fields. - return DistributedPlan.from_json(plan.to_json()) - - def reduce( - self, - workload_name: str, - partial_results: Sequence[CompletedPartial], - parameters: Mapping[str, object], - workspace: Path, - ) -> FinalResult: - workload = self._registry.require(workload_name) - indexes = [partial.chunk_index for partial in partial_results] - if len(indexes) != len(set(indexes)): - raise ValueError("partial results must have unique chunk_index values") - ordered = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index)) - result = workload.reduce(ordered, parameters, workspace) - if not isinstance(result, FinalResult): - raise ValueError("distributed reducer must return a FinalResult") - return result - - -def default_distributed_registry() -> DistributedWorkloadRegistry: - """Return the currently supported distributed scientific workloads.""" - # Delayed import keeps the generic registry independent of concrete RDKit - # workloads and avoids making the contract layer import application setup. - from .similarity_search import SimilaritySearchDistributedWorkload - - registry = DistributedWorkloadRegistry() - registry.register(SimilaritySearchDistributedWorkload()) - return registry diff --git a/scimesh/distributed/similarity_search.py b/scimesh/distributed/similarity_search.py deleted file mode 100644 index e73c9c6..0000000 --- a/scimesh/distributed/similarity_search.py +++ /dev/null @@ -1,380 +0,0 @@ -"""Distributed planning and reduction for exact molecular similarity search.""" - -from __future__ import annotations - -import csv -import hashlib -import heapq -import math -from pathlib import Path -from typing import Any, Iterator, Mapping, Sequence -from uuid import UUID, uuid5 - -from rdkit import Chem - -from scimesh.chemistry.dataset import MoleculeRecord, find_molecule_by_id, parse_smiles -from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE -from scimesh.workloads.similarity_search import ( - SimilarityMatch, - _HeapEntry, - search_similar, - write_search_results, -) - -from .models import ArtifactReference, CompletedPartial, DistributedPlan, FinalResult, PlannedTask - - -_TSV_CONTENT_TYPE = "text/tab-separated-values" -_CSV_CONTENT_TYPE = "text/csv" -_SEARCH_COLUMNS = ("rank", "chembl_id", "canonical_smiles", "similarity") -_REQUIRED_COLUMNS = {"chembl_id", "canonical_smiles"} - - -def write_similarity_search_partial(output_path: Path, matches: Sequence[SimilarityMatch]) -> None: - """Write a worker partial with a round-trip score, not display rounding. - - The public final CSV continues to use the local CLI's six-decimal display. - A reducer needs the full binary float representation to rank candidates - from separate shards exactly as the single-process reference does. - """ - output_path.parent.mkdir(parents=True, exist_ok=True) - with output_path.open("w", encoding="utf-8", newline="") as destination: - writer = csv.DictWriter(destination, fieldnames=_SEARCH_COLUMNS) - writer.writeheader() - for rank, match in enumerate(matches, start=1): - writer.writerow({ - "rank": rank, - "chembl_id": match.molecule_id, - "canonical_smiles": match.smiles, - "similarity": repr(match.similarity), - }) - - -class SimilaritySearchDistributedWorkload: - """Planner/reducer for exact global top-k Tanimoto similarity search.""" - - name = "similarity-search" - description = "Exact top-k molecular similarity search over deterministic TSV shards." - - def validate_job(self, parameters: Mapping[str, object]) -> None: - allowed = { - "query_id", "query_smiles", "top_k", "threshold", - "threshold_direction", "max_rows", "progress_every", - } - unknown = set(parameters) - allowed - if unknown: - raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}") - query_id = parameters.get("query_id") - query_smiles = parameters.get("query_smiles") - if (query_id is None) == (query_smiles is None): - raise ValueError("exactly one of query_id or query_smiles is required") - if query_id is not None: - self._string(query_id, "query_id") - if query_smiles is not None: - self._string(query_smiles, "query_smiles") - self._positive_int(parameters.get("top_k", 20), "top_k") - if "max_rows" in parameters: - self._positive_int(parameters["max_rows"], "max_rows") - if "progress_every" in parameters: - self._nonnegative_int(parameters["progress_every"], "progress_every") - if "threshold" in parameters: - self._unit_interval(parameters["threshold"], "threshold") - if "threshold_direction" in parameters and parameters["threshold_direction"] not in {"greater", "less"}: - raise ValueError("threshold_direction must be 'greater' or 'less'") - - def plan( - self, - input_path: Path, - input_artifact_id: str, - parameters: Mapping[str, object], - shard_rows: int, - workspace: Path, - ) -> DistributedPlan: - self.validate_job(parameters) - if not input_path.is_file(): - raise ValueError("input_path must be a readable dataset file") - if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1: - raise ValueError("shard_rows must be a positive integer") - try: - input_id = UUID(input_artifact_id) - except ValueError as error: - raise ValueError("input_artifact_id must be a UUID") from error - - query_smiles, query_source = self._resolve_query(input_path, parameters) - resolved = self._resolved_parameters(parameters, query_smiles, query_source) - workspace.mkdir(parents=True, exist_ok=True) - shard_paths: list[Path] = [] - try: - shard_paths = self._write_shards(input_path, workspace, shard_rows, resolved.get("max_rows")) - tasks = tuple( - PlannedTask( - chunk_index=index, - input_artifact=ArtifactReference( - artifact_id=str(uuid5(input_id, f"scimesh:similarity-search:shard:{index}")), - sha256=_sha256_file(path), - content_type=_TSV_CONTENT_TYPE, - ), - parameters=self._task_parameters(resolved), - ) - for index, path in enumerate(shard_paths) - ) - except Exception: - for path in shard_paths: - path.unlink(missing_ok=True) - raise - return DistributedPlan(self.name, resolved, tasks) - - def reduce( - self, - partial_results: Sequence[CompletedPartial], - parameters: Mapping[str, object], - workspace: Path, - ) -> FinalResult: - """Merge materialized partial CSVs into one deterministic final CSV. - - The coordinator bridge materializes each downloaded artifact at - ``workspace / artifact_id`` before it calls this method. Those local - paths are an ephemeral bridge detail, never present in the plan or task - payload. CTX-09 owns the durable final-artifact upload and job state. - """ - if not partial_results: - raise ValueError("at least one partial result is required") - resolved = self._validate_resolved_parameters(parameters) - top_k = resolved["top_k"] - direction = resolved["threshold_direction"] - heap: list[_HeapEntry] = [] - - ordered_partials = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index)) - indexes = [partial.chunk_index for partial in ordered_partials] - if len(indexes) != len(set(indexes)): - raise ValueError("partial results must have unique chunk_index values") - for partial in ordered_partials: - path = workspace / partial.artifact.artifact_id - if not path.is_file(): - raise ValueError("materialized partial result is missing") - if _sha256_file(path) != partial.artifact.sha256: - raise ValueError("materialized partial result checksum does not match its artifact reference") - for match in self._read_partial(path, direction): - rank_key = match.sort_key(direction) - entry = _HeapEntry(match, rank_key) - if len(heap) < top_k: - heapq.heappush(heap, entry) - elif rank_key < heap[0].rank_key: - heapq.heapreplace(heap, entry) - - matches = sorted((entry.match for entry in heap), key=lambda match: match.sort_key(direction)) - output = workspace / "result.csv" - write_search_results(output, matches) - final_id = uuid5( - UUID(ordered_partials[0].artifact.artifact_id), - "scimesh:similarity-search:final:" + ",".join( - partial.artifact.artifact_id for partial in ordered_partials - ), - ) - return FinalResult( - ArtifactReference(str(final_id), _sha256_file(output), _CSV_CONTENT_TYPE), - {"matches_emitted": len(matches), "partial_count": len(ordered_partials)}, - ) - - def _resolve_query( - self, input_path: Path, parameters: Mapping[str, object] - ) -> tuple[str, dict[str, str]]: - query_id = parameters.get("query_id") - if isinstance(query_id, str): - record = find_molecule_by_id(input_path, query_id) - return Chem.MolToSmiles(record.molecule, canonical=True), {"kind": "chembl_id", "value": query_id} - supplied = parameters["query_smiles"] - assert isinstance(supplied, str) # checked by validate_job - molecule = parse_smiles(supplied) - if molecule is None: - raise ValueError("query_smiles is invalid") - return Chem.MolToSmiles(molecule, canonical=True), {"kind": "smiles", "value": supplied} - - def _resolved_parameters( - self, parameters: Mapping[str, object], query_smiles: str, query_source: Mapping[str, str] - ) -> dict[str, object]: - resolved: dict[str, object] = { - "query_smiles": query_smiles, - "query_source": dict(query_source), - "top_k": self._positive_int(parameters.get("top_k", 20), "top_k"), - "threshold_direction": parameters.get("threshold_direction", "greater"), - "fingerprint": {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE}, - } - if "threshold" in parameters: - resolved["threshold"] = self._unit_interval(parameters["threshold"], "threshold") - if "max_rows" in parameters: - resolved["max_rows"] = self._positive_int(parameters["max_rows"], "max_rows") - if "progress_every" in parameters: - resolved["progress_every"] = self._nonnegative_int(parameters["progress_every"], "progress_every") - return resolved - - def _validate_resolved_parameters(self, parameters: Mapping[str, object]) -> dict[str, object]: - query_smiles = self._string(parameters.get("query_smiles"), "query_smiles") - if parse_smiles(query_smiles) is None: - raise ValueError("query_smiles is invalid") - resolved = self._resolved_parameters( - parameters, - Chem.MolToSmiles(parse_smiles(query_smiles), canonical=True), - {"kind": "resolved", "value": query_smiles}, - ) - # A reducer receives immutable plan metadata, whose query source and - # fixed fingerprint are observational context rather than worker input. - if "fingerprint" in parameters: - fingerprint = parameters["fingerprint"] - if fingerprint != {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE}: - raise ValueError("resolved fingerprint does not match SciMesh defaults") - return resolved - - @staticmethod - def _task_parameters(resolved: Mapping[str, object]) -> dict[str, object]: - # max_rows is applied before sharding. Passing it to each task would - # silently scan N rows per shard instead of the requested global prefix. - return { - key: value for key, value in resolved.items() - if key in {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"} - } - - def _write_shards( - self, input_path: Path, workspace: Path, shard_rows: int, max_rows: object - ) -> list[Path]: - limit = int(max_rows) if isinstance(max_rows, int) else None - paths: list[Path] = [] - current: Path | None = None - destination = None - rows_in_shard = 0 - seen_rows = 0 - try: - with input_path.open("r", encoding="utf-8", newline="") as source: - reader = csv.DictReader(source, delimiter="\t") - fieldnames = reader.fieldnames or [] - if not _REQUIRED_COLUMNS.issubset(set(fieldnames)): - missing = sorted(_REQUIRED_COLUMNS - set(fieldnames)) - raise ValueError(f"dataset is missing required columns: {', '.join(missing)}") - for row in reader: - if limit is not None and seen_rows >= limit: - break - if destination is None or rows_in_shard == shard_rows: - if destination is not None: - destination.close() - current = workspace / f"shard-{len(paths)}.tsv" - destination = current.open("w", encoding="utf-8", newline="") - writer = csv.DictWriter(destination, fieldnames=fieldnames, delimiter="\t", lineterminator="\n") - writer.writeheader() - paths.append(current) - rows_in_shard = 0 - writer.writerow(row) - rows_in_shard += 1 - seen_rows += 1 - finally: - if destination is not None: - destination.close() - if not paths: - raise ValueError("dataset has no data rows") - return paths - - @staticmethod - def _read_partial(path: Path, direction: object) -> Iterator[SimilarityMatch]: - if not path.is_file(): - raise ValueError("materialized partial result is missing") - if direction not in {"greater", "less"}: - raise ValueError("threshold_direction must be 'greater' or 'less'") - previous_key: tuple[float, str, str] | None = None - with path.open("r", encoding="utf-8", newline="") as source: - reader = csv.DictReader(source) - if tuple(reader.fieldnames or ()) != _SEARCH_COLUMNS: - raise ValueError("partial result has an invalid CSV header") - for expected_rank, row in enumerate(reader, start=1): - if set(row) != set(_SEARCH_COLUMNS) or row["rank"] != str(expected_rank): - raise ValueError("partial result has an invalid rank") - try: - similarity = float(row["similarity"]) - except (TypeError, ValueError) as error: - raise ValueError("partial result has an invalid similarity") from error - if not math.isfinite(similarity) or not 0 <= similarity <= 1: - raise ValueError("partial result has an invalid similarity") - match = SimilarityMatch(similarity, row["chembl_id"], row["canonical_smiles"]) - key = match.sort_key(direction) - if previous_key is not None and key < previous_key: - raise ValueError("partial result is not sorted deterministically") - previous_key = key - yield match - - @staticmethod - def _string(value: object, name: str) -> str: - if not isinstance(value, str) or not value.strip() or len(value) > 200: - raise ValueError(f"{name} must be a non-empty string") - return value - - @staticmethod - def _positive_int(value: object, name: str) -> int: - if isinstance(value, bool) or not isinstance(value, int) or value < 1: - raise ValueError(f"{name} must be a positive integer") - return value - - @staticmethod - def _nonnegative_int(value: object, name: str) -> int: - if isinstance(value, bool) or not isinstance(value, int) or value < 0: - raise ValueError(f"{name} must be a non-negative integer") - return value - - @staticmethod - def _unit_interval(value: object, name: str) -> float: - if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or not 0 <= value <= 1: - raise ValueError(f"{name} must be a number between 0 and 1") - return float(value) - - -def run_similarity_search_shard( - input_path: Path, parameters: Mapping[str, object], output_path: Path -) -> dict[str, int]: - """Run one planned shard using the local reference implementation. - - This is the worker adapter used by CTX-08. It deliberately accepts only - resolved ``query_smiles``: resolving an identifier independently in each - shard would make the distributed search scientifically invalid. - """ - allowed = {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"} - unknown = set(parameters) - allowed - if unknown: - raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}") - query_smiles = parameters.get("query_smiles") - if not isinstance(query_smiles, str) or not query_smiles.strip(): - raise ValueError("query_smiles is required for a distributed shard") - molecule = parse_smiles(query_smiles) - if molecule is None: - raise ValueError("query_smiles is invalid") - top_k = SimilaritySearchDistributedWorkload._positive_int(parameters.get("top_k", 20), "top_k") - threshold = None - if "threshold" in parameters: - threshold = SimilaritySearchDistributedWorkload._unit_interval(parameters["threshold"], "threshold") - direction = parameters.get("threshold_direction", "greater") - if direction not in {"greater", "less"}: - raise ValueError("threshold_direction must be 'greater' or 'less'") - progress_every = 0 - if "progress_every" in parameters: - progress_every = SimilaritySearchDistributedWorkload._nonnegative_int( - parameters["progress_every"], "progress_every" - ) - result = search_similar( - input_path, - MoleculeRecord("query", query_smiles, molecule), - top_k=top_k, - progress_every=progress_every, - threshold=threshold, - threshold_direction=direction, - ) - write_similarity_search_partial(output_path, result.matches) - return { - "scanned_rows": result.stats.scanned, - "valid_molecules": result.stats.valid, - "invalid_smiles": result.stats.invalid, - "matches_emitted": len(result.matches), - } - - -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as source: - for block in iter(lambda: source.read(1024 * 1024), b""): - digest.update(block) - return digest.hexdigest() diff --git a/scimesh/distributed/workload.py b/scimesh/distributed/workload.py deleted file mode 100644 index eff85ea..0000000 --- a/scimesh/distributed/workload.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Protocol implemented by coordinator-independent distributed workloads.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Mapping, Protocol, Sequence - -from .models import CompletedPartial, DistributedPlan, FinalResult - - -class DistributedWorkload(Protocol): - """Validate, plan, and reduce one explicit scientific workload. - - ``input_path`` and ``workspace`` are bridge-provided temporary local paths. - They must never be included in returned plans or persisted task payloads. - """ - - name: str - description: str - - def validate_job(self, parameters: Mapping[str, object]) -> None: - """Reject invalid public parameters before the bridge writes metadata.""" - - def plan( - self, - input_path: Path, - input_artifact_id: str, - parameters: Mapping[str, object], - shard_rows: int, - workspace: Path, - ) -> DistributedPlan: - """Build a JSON-safe plan containing only coordinator artifact references.""" - - def reduce( - self, - partial_results: Sequence[CompletedPartial], - parameters: Mapping[str, object], - workspace: Path, - ) -> FinalResult: - """Reduce coordinator-owned partial artifacts in ascending chunk order.""" diff --git a/scimesh/sdk/__init__.py b/scimesh/sdk/__init__.py index a4bff9a..1131393 100644 --- a/scimesh/sdk/__init__.py +++ b/scimesh/sdk/__init__.py @@ -19,13 +19,6 @@ from .artifacts import ( PortSpec, Provenance, ) -from .builtins import ( - current_environment_digest, - current_scimesh_package_digest, - default_sdk_registry, - default_sdk_runtime, - similarity_search_sdk_adapter, -) from .conformance import ( CancellationFlag, LocalArtifactStore, @@ -215,11 +208,6 @@ __all__ = [ "WorkloadManifest", "WorkloadRegistry", "assert_manifest_round_trip", - "current_environment_digest", - "current_scimesh_package_digest", - "default_sdk_registry", - "default_sdk_runtime", "installed_distribution_digest", "negotiate_manifest", - "similarity_search_sdk_adapter", ] diff --git a/scimesh/sdk/builtins.py b/scimesh/sdk/builtins.py deleted file mode 100644 index 2456964..0000000 --- a/scimesh/sdk/builtins.py +++ /dev/null @@ -1,150 +0,0 @@ -"""SDK definitions for existing SciMesh workloads and local core runtime.""" - -from __future__ import annotations - -import hashlib -import os -import platform -import sys - -from rdkit import rdBase - -from scimesh.distributed.similarity_search import ( - SimilaritySearchDistributedWorkload, - run_similarity_search_shard, -) - -from .artifacts import ArtifactSchema, PortSpec -from .compat import LegacyDistributedWorkloadAdapter -from .identity import ComponentRef, SDK_API_VERSION, SchemaRef -from .integrity import installed_distribution_digest -from .registry import WorkloadRegistry -from .resources import ResourceInventory -from .runtime import RuntimeCapabilities - - -def current_scimesh_package_digest() -> str: - """Hash installed SciMesh Python sources for the built-in trusted adapter. - - This is a local immutable-code pin, not a package signature or container - attestation. Consequently the built-in compatibility manifest is trusted - only; an administrator must supply signed image metadata before enabling an - untrusted quorum policy. - """ - # Source/editable installs are allowed only for this explicit local - # development helper. Registry discovery keeps the secure default. - return installed_distribution_digest("scimesh", allow_editable=True) - - -def current_environment_digest() -> str: - payload = "\n".join( - ( - current_scimesh_package_digest(), - f"python={sys.implementation.name}-{platform.python_version()}", - f"rdkit={rdBase.rdkitVersion}", - f"platform={sys.platform}-{platform.machine().lower()}", - ) - ) - return "sha256:" + hashlib.sha256(payload.encode("utf-8")).hexdigest() - - -def similarity_search_sdk_adapter( - *, shard_rows: int = 10_000 -) -> LegacyDistributedWorkloadAdapter: - dataset_schema = ArtifactSchema( - SchemaRef("molecule-table", 1), - "text/tab-separated-values", - "utf-8", - max_bytes=10 * 1024 * 1024 * 1024, - validator=ComponentRef("delimited-table", 1), - validator_configuration={ - "required_columns": ["canonical_smiles", "chembl_id"], - }, - max_records=100_000_000, - canonicalizer="scimesh-tsv-v1", - ) - partial_schema = ArtifactSchema( - SchemaRef("similarity-search-partial", 1), - "text/csv", - "utf-8", - max_bytes=1024 * 1024 * 1024, - validator=ComponentRef("delimited-table", 1), - validator_configuration={ - "columns": ["rank", "chembl_id", "canonical_smiles", "similarity"], - }, - max_records=100_000, - canonicalizer="scimesh-search-partial-v1", - ) - result_schema = ArtifactSchema( - SchemaRef("similarity-search-result", 1), - "text/csv", - "utf-8", - max_bytes=1024 * 1024 * 1024, - validator=ComponentRef("delimited-table", 1), - validator_configuration={ - "columns": ["rank", "chembl_id", "canonical_smiles", "similarity"], - }, - max_records=100_000, - canonicalizer="scimesh-search-result-v1", - ) - parameters_schema = { - "type": "object", - "additionalProperties": False, - "properties": { - "query_id": {"type": "string", "minLength": 1, "maxLength": 200}, - "query_smiles": {"type": "string", "minLength": 1, "maxLength": 200}, - "top_k": {"type": "integer", "minimum": 1}, - "threshold": {"type": "number", "minimum": 0, "maximum": 1}, - "threshold_direction": {"enum": ["greater", "less"]}, - "max_rows": {"type": "integer", "minimum": 1}, - "progress_every": {"type": "integer", "minimum": 0}, - }, - "oneOf": [ - {"required": ["query_id"], "not": {"required": ["query_smiles"]}}, - {"required": ["query_smiles"], "not": {"required": ["query_id"]}}, - ], - } - return LegacyDistributedWorkloadAdapter( - SimilaritySearchDistributedWorkload(), - run_similarity_search_shard, - version="1.0.0", - package_digest=current_scimesh_package_digest(), - environment_digest=current_environment_digest(), - parameters_schema=parameters_schema, - input_port=PortSpec(dataset_schema), - partial_port=PortSpec(partial_schema), - output_port=PortSpec(result_schema), - resolved_parameter_names=("query_source", "fingerprint"), - shard_rows=shard_rows, - ) - - -def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry: - registry = WorkloadRegistry() - registry.register( - similarity_search_sdk_adapter(shard_rows=shard_rows).definition(), enabled=True - ) - return registry - - -def similarity_search_workload_definition(): - """Installed entry-point factory for the default shard-size definition.""" - return similarity_search_sdk_adapter().definition() - - -def default_sdk_runtime() -> RuntimeCapabilities: - architecture = platform.machine().lower() or "unknown" - return RuntimeCapabilities( - sdk_api_version=SDK_API_VERSION, - protocol_version="1.0.0", - profiles=("core-batch-v1",), - features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"}, - workload_capabilities=("similarity-search", "descriptor-batch"), - inventory=ResourceInventory( - cpu_cores=max(os.cpu_count() or 1, 1), - memory_mb=4096, - scratch_mb=4096, - architecture=architecture, - environment_digests=(current_environment_digest(),), - ), - ) diff --git a/scimesh/sdk/compat/__init__.py b/scimesh/sdk/compat/__init__.py deleted file mode 100644 index 88f9938..0000000 --- a/scimesh/sdk/compat/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Adapters for versioned pre-SDK SciMesh workload contracts.""" - -from .distributed_v1 import LegacyDistributedWorkloadAdapter - -__all__ = ["LegacyDistributedWorkloadAdapter"] diff --git a/scimesh/sdk/compat/distributed_v1.py b/scimesh/sdk/compat/distributed_v1.py deleted file mode 100644 index 254c3d2..0000000 --- a/scimesh/sdk/compat/distributed_v1.py +++ /dev/null @@ -1,369 +0,0 @@ -"""Compatibility adapter for the CTX-07 ``DistributedWorkload`` protocol.""" - -from __future__ import annotations - -import hashlib -import shutil -from pathlib import Path -from typing import Any, Callable, Mapping, Sequence - -from scimesh.distributed.models import ( - ArtifactReference as LegacyArtifactReference, - CompletedPartial, - FinalResult, -) -from scimesh.distributed.workload import DistributedWorkload - -from ..artifacts import ( - ArtifactCollection, - ArtifactItem, - ArtifactRef, - Cardinality, - CollectionKind, - OutputManifest, - PortSpec, -) -from ..execution import CheckpointPolicy, ExecutionProfile, NetworkPolicy, RetryPolicy -from ..identity import ComponentRef, SchemaRef, VersionRange, WorkloadId -from ..manifest import ( - DeterminismProfile, - EnvironmentSpec, - PackageSpec, - TrustMode, - VerifierSpec, - WorkloadLimits, - WorkloadManifest, -) -from ..plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan -from ..protocols import PlanningContext, ReduceContext, TaskContext -from ..registry import WorkloadDefinition -from ..resources import ResourceRequirements -from ..verification import ExactArtifactVerifier -from ..workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec - - -ShardRunner = Callable[[Path, Mapping[str, object], Path], Mapping[str, int | float]] - - -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as source: - for block in iter(lambda: source.read(1024 * 1024), b""): - digest.update(block) - return digest.hexdigest() - - -class LegacyDistributedWorkloadAdapter: - """Expose a legacy map/reduce workload through the SDK core-batch profile. - - The adapter preserves the old wire schema. Local files are materialized and - sealed only through bridge-owned contexts, and no path is included in a - ``TaskSpec`` or ``WorkflowPlan``. - """ - - MAP_ENTRY_POINT = "scimesh.sdk.compat.distributed_v1:run_legacy@v1" - REDUCE_ENTRY_POINT = "scimesh.sdk.compat.distributed_v1:reduce_legacy@v1" - - def __init__( - self, - workload: DistributedWorkload, - shard_runner: ShardRunner, - *, - version: str, - package_digest: str, - environment_digest: str, - parameters_schema: Mapping[str, Any], - input_port: PortSpec, - partial_port: PortSpec, - output_port: PortSpec, - resolved_parameter_names: Sequence[str] = (), - shard_rows: int = 10_000, - resources: ResourceRequirements | None = None, - execution: ExecutionProfile | None = None, - limits: WorkloadLimits | None = None, - ) -> None: - if not isinstance(workload.name, str) or not isinstance(workload.description, str): - raise ValueError("legacy workload must expose name and description") - if not callable(shard_runner): - raise ValueError("shard_runner must be callable") - if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1: - raise ValueError("shard_rows must be a positive integer") - self.workload = workload - self.shard_runner = shard_runner - self.shard_rows = shard_rows - self.input_port = input_port - self.partial_port = partial_port - self.output_port = output_port - resources = resources or ResourceRequirements( - profile="legacy-cpu-v1", - cpu_cores=1, - memory_mb=1024, - scratch_mb=1024, - max_duration_seconds=3600, - ) - execution = execution or ExecutionProfile( - profile="legacy-python-process-v1", - network=NetworkPolicy.TRUSTED, - timeout_seconds=3600, - checkpoint=CheckpointPolicy(), - ) - limits = limits or WorkloadLimits( - max_input_bytes=input_port.schema.max_bytes, - max_tasks=10_000, - max_output_bytes=output_port.schema.max_bytes, - ) - parameter_names = tuple(sorted(parameters_schema.get("properties", {}))) - reduce_parameter_names = tuple(sorted(set(parameter_names).union(resolved_parameter_names))) - map_stage = StageSpec( - stage_id="map", - kind=StageKind.MAP, - entry_point=self.MAP_ENTRY_POINT, - needs=(), - inputs={"input": input_port}, - outputs={"partial": partial_port}, - parameter_names=parameter_names, - resources=resources, - execution=execution, - retry=RetryPolicy(), - verifier=ComponentRef("exact-artifact", 1), - trust_modes=("trusted",), - max_fan_out=limits.max_tasks, - cacheable=True, - ) - reduce_input = PortSpec( - schema=partial_port.schema, - cardinality=Cardinality.MANY, - collection=CollectionKind.KEYED, - ) - reduce_stage = StageSpec( - stage_id="reduce", - kind=StageKind.REDUCE, - entry_point=self.REDUCE_ENTRY_POINT, - needs=("map",), - inputs={"partials": reduce_input}, - outputs={"result": output_port}, - parameter_names=reduce_parameter_names, - resources=resources, - execution=execution, - retry=RetryPolicy(), - verifier=ComponentRef("exact-artifact", 1), - trust_modes=("trusted",), - cacheable=True, - ) - workflow = WorkflowSpec( - workflow_id="map-reduce-v1", - inputs={"input": input_port}, - stages=(map_stage, reduce_stage), - edges=( - ArtifactEdge(PortRef("input"), PortRef("input", "map")), - ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")), - ), - outputs={"result": PortRef("result", "reduce")}, - max_tasks=limits.max_tasks, - max_output_bytes=limits.max_output_bytes, - ) - self.manifest = WorkloadManifest( - sdk_api=VersionRange(">=1.0,<2.0"), - protocol=VersionRange(">=1,<2"), - workload=WorkloadId(workload.name, version), - description=workload.description, - package=PackageSpec("scimesh", package_digest), - environment=EnvironmentSpec("python-process", environment_digest, {"adapter": "distributed-v1"}), - parameters_schema=parameters_schema, - workflow=workflow, - inputs={"input": input_port}, - outputs={"result": output_port}, - determinism=DeterminismProfile.BYTE_EXACT, - trust_modes=(TrustMode.TRUSTED,), - verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}), - limits=limits, - capabilities=(workload.name,), - conformance_profiles=("core-batch-v1",), - ) - self._exact_verifier = ExactArtifactVerifier() - - def definition(self) -> WorkloadDefinition: - return WorkloadDefinition( - manifest=self.manifest, - planner=self, - runners={self.MAP_ENTRY_POINT: self}, - reducers={self.REDUCE_ENTRY_POINT: self}, - verifiers={self._exact_verifier.identity.canonical: self._exact_verifier}, - ) - - def validate(self, request: JobRequest) -> ValidatedJob: - if request.workload != self.manifest.workload: - raise ValueError("legacy adapter received a request for another workload") - self.workload.validate_job(request.parameters) - return ValidatedJob(request, request.parameters) - - def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan: - if not isinstance(job, ValidatedJob): - raise ValueError("job must be a ValidatedJob") - collection = job.request.inputs.get("input") - if collection is None: - raise ValueError("legacy adapter requires the input port") - self.input_port.validate_collection(collection, "job input") - input_artifact = collection.items[0].artifact - input_path = context.catalog.materialize(input_artifact) - workspace = context.workspace - workspace.mkdir(parents=True, exist_ok=True) - legacy = self.workload.plan( - input_path, - input_artifact.artifact_id, - job.request.parameters, - self.shard_rows, - workspace, - ) - if legacy.workload != self.workload.name: - raise ValueError("legacy planner returned a plan for another workload") - tasks: list[TaskSpec] = [] - negotiated = context.negotiated - map_stage = self.manifest.workflow.stages[0] - used_paths: set[Path] = set() - for planned in legacy.tasks: - path = self._find_planned_file(workspace, planned.input_artifact.sha256, used_paths) - sealed = context.sink.seal( - path, - declaration=self.input_port.schema, - ) - if sealed.sha256 != planned.input_artifact.sha256: - raise ValueError("artifact sink returned a checksum that differs from the legacy plan") - tasks.append( - TaskSpec( - workload=self.manifest.workload, - package_digest=self.manifest.package.digest, - manifest_digest=self.manifest.digest, - trust_mode=job.request.trust_mode, - sdk_api_version=negotiated.sdk_api_version, - protocol_version=negotiated.protocol_version, - manifest_schema_version=self.manifest.manifest_schema_version, - workflow_schema_version=self.manifest.workflow.schema_version, - environment_digest=self.manifest.environment.digest, - verifier=map_stage.verifier, - selected_features=negotiated.selected_features, - optional_fallbacks=negotiated.optional_fallbacks, - task_key=f"map/{planned.chunk_index:08d}", - stage_id="map", - parameters=planned.parameters, - inputs={"input": ArtifactCollection.single(sealed)}, - expected_outputs={"partial": self.partial_port}, - resources=map_stage.resources, - execution=map_stage.execution, - ) - ) - return WorkflowPlan( - workload=self.manifest.workload, - package_digest=self.manifest.package.digest, - manifest_digest=self.manifest.digest, - trust_mode=job.request.trust_mode, - sdk_api_version=negotiated.sdk_api_version, - protocol_version=negotiated.protocol_version, - manifest_schema_version=self.manifest.manifest_schema_version, - workflow_schema_version=self.manifest.workflow.schema_version, - environment_digest=self.manifest.environment.digest, - verifier=self.manifest.verifier.verifier, - selected_features=negotiated.selected_features, - optional_fallbacks=negotiated.optional_fallbacks, - workflow_id=self.manifest.workflow.workflow_id, - resolved_parameters=legacy.resolved_parameters, - tasks=tuple(tasks), - ) - - @staticmethod - def _find_planned_file(workspace: Path, expected_sha256: str, used: set[Path]) -> Path: - for candidate in sorted(workspace.rglob("*")): - if candidate in used or not candidate.is_file() or candidate.is_symlink(): - continue - if _sha256_file(candidate) == expected_sha256: - used.add(candidate) - return candidate - raise ValueError("legacy planner did not materialize its planned artifact") - - def run(self, context: TaskContext) -> OutputManifest: - context.cancellation.raise_if_cancelled() - collection = context.task.inputs.get("input") - if collection is None: - raise ValueError("legacy map task requires one input collection") - self.input_port.validate_collection(collection, "legacy map input") - source = context.catalog.materialize(collection.items[0].artifact) - workspace = context.workspace - workspace.mkdir(parents=True, exist_ok=True) - input_path = workspace / "input" - output_path = workspace / "result" - if source.resolve() != input_path.resolve(): - shutil.copyfile(source, input_path) - metrics = self.shard_runner(input_path, context.task.parameters, output_path) - context.cancellation.raise_if_cancelled() - sealed = context.sink.seal( - output_path, - declaration=self.partial_port.schema, - ) - return OutputManifest( - context.task.task_key, - {"partial": ArtifactCollection.single(sealed)}, - metrics, - context.provenance, - ).validate_against(context.task.expected_outputs, max_output_bytes=self.manifest.limits.max_output_bytes) - - def reduce(self, context: ReduceContext) -> OutputManifest: - context.cancellation.raise_if_cancelled() - collection = context.accepted_inputs.get("partials") - if collection is None or collection.kind is not CollectionKind.KEYED or not collection.items: - raise ValueError("legacy reducer requires a non-empty keyed partial collection") - self.manifest.workflow.stages[1].inputs["partials"].validate_collection( - collection, - "legacy reducer partials", - ) - workspace = context.workspace - workspace.mkdir(parents=True, exist_ok=True) - partials: list[CompletedPartial] = [] - indexed_items: list[tuple[int, ArtifactItem]] = [] - for item in collection.items: - key = item.key or "" - prefix = "map." - raw_index = key[len(prefix):] if key.startswith(prefix) else "" - if len(raw_index) != 8 or not raw_index.isdigit(): - raise ValueError("legacy partial key must use map.") - indexed_items.append((int(raw_index), item)) - indices = [index for index, _ in indexed_items] - expected_keys = context.task.expected_input_keys.get("partials") - if expected_keys is None or {item.key for item in collection.items} != set(expected_keys): - raise ValueError("legacy partial keys do not match the coordinator expected set") - if sorted(indices) != list(range(len(indexed_items))): - raise ValueError("legacy partial keys must be complete and contiguous") - for index, item in sorted(indexed_items): - artifact = item.artifact - source = context.catalog.materialize(artifact) - target = workspace / artifact.artifact_id - if source.resolve() != target.resolve(): - shutil.copyfile(source, target) - if _sha256_file(target) != artifact.sha256: - raise ValueError("materialized partial checksum does not match") - partials.append( - CompletedPartial( - index, - LegacyArtifactReference( - artifact.artifact_id, - artifact.sha256, - artifact.media_type, - ), - {}, - ) - ) - result = self.workload.reduce(partials, context.task.parameters, workspace) - if not isinstance(result, FinalResult): - raise ValueError("legacy reducer must return a FinalResult") - path = self._find_planned_file(workspace, result.artifact.sha256, set()) - sealed = context.sink.seal( - path, - declaration=self.output_port.schema, - ) - if sealed.sha256 != result.artifact.sha256: - raise ValueError("artifact sink returned a checksum that differs from the legacy result") - return OutputManifest( - context.task.task_key, - {"result": ArtifactCollection.single(sealed)}, - result.metrics, - context.provenance, - ).validate_against(context.task.expected_outputs, max_output_bytes=self.manifest.limits.max_output_bytes) diff --git a/scimesh/worker/runners.py b/scimesh/worker/runners.py index f924636..e1af46a 100644 --- a/scimesh/worker/runners.py +++ b/scimesh/worker/runners.py @@ -1,110 +1,267 @@ -"""Local workload adapters. They receive no arbitrary commands from the network.""" +"""SDK-based local workload execution for claimed coordinator tasks. + +The runner is a v1-wire bridge: the coordinator still claims flat tasks and +the worker still uploads one partial CSV, but execution goes through the +SDK-built workload's own Runner handler with a real ``TaskSpec``, +provenance, resource reservation, and a content-addressed local store. No +legacy distributed-protocol code is involved. +""" from __future__ import annotations +import hashlib +from datetime import datetime, timezone from pathlib import Path -import subprocess -import sys -from typing import Protocol +from typing import Mapping, Protocol +from uuid import NAMESPACE_URL, uuid5 -from scimesh.distributed.similarity_search import run_similarity_search_shard +from scimesh.sdk.artifacts import ( + ArtifactCollection, + ArtifactRef, + OutputManifest, + Provenance, +) +from scimesh.sdk.conformance import ( + CancellationFlag, + LocalArtifactStore, + LocalTaskContext, + ScopedArtifactSink, +) +from scimesh.sdk._validation import canonical_json +from scimesh.sdk.manifest import TrustMode +from scimesh.sdk.plans import TaskSpec +from scimesh.sdk.registry import WorkloadDefinition +from scimesh.sdk.resources import ResourceAllocation, ResourcePool +from scimesh.sdk.runtime import ( + NegotiatedWorkload, + RuntimeCapabilities, + negotiate_manifest, +) +from scimesh.sdk.workflow import StageKind +from scimesh.workloads.library import default_sdk_runtime +from scimesh.workloads.search import similarity_search_sdk_definition from .models import ClaimedTask, ProducedArtifact, RunResult +#: Parameters the worker may hand to a map task. ``max_rows`` is a plan-time +#: option applied before sharding and is intentionally rejected here. +_RUNNER_PARAMETERS = frozenset( + {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"} +) + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + class Runner(Protocol): def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: ... class SciMeshRunner: - """Allowlisted adapter from coordinator workloads to the local SciMesh CLI.""" + """Execute claimed coordinator tasks through the SDK-built workloads.""" + + def __init__( + self, + definitions: Mapping[str, WorkloadDefinition] | None = None, + runtime: RuntimeCapabilities | None = None, + ) -> None: + self._definitions = dict(definitions or {}) + if "similarity-search" not in self._definitions: + self._definitions["similarity-search"] = ( + similarity_search_sdk_definition().definition() + ) + self._runtime = runtime or default_sdk_runtime() + self._pool = ResourcePool(self._runtime.inventory, max_concurrency=1) def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: - # The subprocess changes cwd to task_dir. Absolute paths keep a caller - # supplied relative work directory from being resolved twice. task_dir = task_dir.resolve() - input_path = task_dir / "input" - output_path = task_dir / "result.csv" - # The coordinator contract historically used underscores while the - # public SciMesh CLI uses hyphens. Accept both spellings at this narrow - # boundary so an API job cannot turn into an opaque worker failure. workload = task.workload.replace("_", "-") - command = [sys.executable, "-m", "scimesh.cli", workload, str(input_path)] - params = task.parameters - if workload == "similarity-search": - self._reject_unknown(params, {"query_id", "query_smiles", "top_k", "threshold", "threshold_direction", "max_rows", "progress_every"}) - query_id, query_smiles = params.get("query_id"), params.get("query_smiles") - if (query_id is None) == (query_smiles is None): - raise ValueError("exactly one of query_id or query_smiles is required") - if query_smiles is not None and "max_rows" not in params: - metrics = run_similarity_search_shard(input_path, params, output_path) - return RunResult((ProducedArtifact(output_path, "text/csv"),), metrics) - # Legacy URI jobs may still use query_id or an explicitly task-local - # max_rows value. CTX-08 plans never create those payloads; retain - # CLI execution only for backwards compatibility at this boundary. - top_k = self._positive_int(params, "top_k", default=20) - command += ["--query-id", self._string(params, "query_id")] if query_id is not None else ["--query-smiles", self._string(params, "query_smiles")] - command += ["--top-k", str(top_k)] - self._append_common_options(command, params) - elif workload == "similarity-graph": - self._reject_unknown(params, {"threshold", "threshold_direction", "block_size", "max_rows", "progress_every"}) - threshold = self._number(params, "threshold") - command += ["--threshold", str(threshold)] - self._append_common_options(command, params, include_threshold=False) - if "block_size" in params: - command += ["--block-size", str(self._positive_int(params, "block_size", default=1_000))] - else: + definition = self._definitions.get(workload) + if definition is None: raise ValueError(f"unsupported workload: {task.workload}") - command += ["--output", str(output_path)] - subprocess.run(command, check=True, cwd=task_dir) # explicit list: never shell=True - if not output_path.is_file(): - raise RuntimeError("SciMesh CLI did not create its result") - processed_rows = max(sum(1 for _ in output_path.open(encoding="utf-8")) - 1, 0) - return RunResult((ProducedArtifact(output_path, "text/csv"),), {"processed_rows": processed_rows}) - - def _append_common_options(self, command: list[str], params: dict[str, object], *, include_threshold: bool = True) -> None: - if include_threshold and "threshold" in params: - command += ["--threshold", str(self._number(params, "threshold"))] - if "threshold_direction" in params: - direction = params["threshold_direction"] - if direction not in ("greater", "less"): - raise ValueError("threshold_direction must be 'greater' or 'less'") - command += ["--threshold-direction", str(direction)] - if "max_rows" in params: - command += ["--max-rows", str(self._positive_int(params, "max_rows", default=1))] - if "progress_every" in params: - command += ["--progress-every", str(self._nonnegative_int(params, "progress_every"))] + manifest = definition.manifest + negotiated = negotiate_manifest(manifest, self._runtime) + map_stage = next( + stage for stage in manifest.workflow.stages if stage.kind is StageKind.MAP + ) + assert map_stage.verifier is not None + input_path = task_dir / "input" + if not input_path.is_file(): + raise ValueError("claimed task input is missing") + parameters = self._resolve_parameters(task, input_path) + store = LocalArtifactStore(task_dir / "sdk-store") + input_ref = store.import_file( + input_path, + declaration=manifest.inputs["input"].schema, + ) + spec = TaskSpec( + workload=manifest.workload, + package_digest=manifest.package.digest, + manifest_digest=manifest.digest, + trust_mode=TrustMode.TRUSTED, + sdk_api_version=negotiated.sdk_api_version, + protocol_version=negotiated.protocol_version, + manifest_schema_version=manifest.manifest_schema_version, + workflow_schema_version=manifest.workflow.schema_version, + environment_digest=manifest.environment.digest, + verifier=map_stage.verifier, + selected_features=negotiated.selected_features, + optional_fallbacks=negotiated.optional_fallbacks, + task_key="map/00000000", + stage_id=map_stage.stage_id, + parameters=parameters, + inputs={"input": ArtifactCollection.single(input_ref)}, + expected_outputs=map_stage.outputs, + resources=map_stage.resources, + execution=map_stage.execution, + ).validate_stage(map_stage) + allocation = self._pool.reserve(task.task_id, spec.resources) + try: + provenance = self._provenance( + definition, negotiated, spec, task, allocation + ) + context = LocalTaskContext( + spec, + store, + store, + task_dir, + CancellationFlag(), + provenance, + spec.inputs, + transaction=None, + ) + output = definition.runners[map_stage.entry_point].run(context) + self._validate_output( + output, + spec, + context, + store, + provenance, + max_output_bytes=manifest.limits.max_output_bytes, + ) + partial_ref = output.outputs["partial"].items[0].artifact + partial_path = store.materialize(partial_ref) + return RunResult( + (ProducedArtifact(partial_path, "text/csv"),), + dict(output.metrics), + ) + finally: + self._pool.release(allocation.allocation_id) @staticmethod - def _reject_unknown(params: dict[str, object], allowed: set[str]) -> None: - unknown = set(params) - allowed + def _resolve_parameters(task: ClaimedTask, input_path: Path) -> dict[str, object]: + """Resolve ``query_id`` once per task and reject plan-time options.""" + parameters = dict(task.parameters) + query_id = parameters.get("query_id") + query_smiles = parameters.get("query_smiles") + if isinstance(query_id, str) and not isinstance(query_smiles, str): + from scimesh.chemistry.dataset import find_molecule_by_id + from rdkit import Chem + + record = find_molecule_by_id(input_path, query_id) + parameters["query_smiles"] = Chem.MolToSmiles( + record.molecule, canonical=True + ) + del parameters["query_id"] + unknown = set(parameters) - _RUNNER_PARAMETERS if unknown: - raise ValueError(f"unsupported parameters: {', '.join(sorted(unknown))}") + raise ValueError( + "unsupported runner parameters: " + ", ".join(sorted(unknown)) + ) + return parameters + + def _provenance( + self, + definition: WorkloadDefinition, + negotiated: NegotiatedWorkload, + spec: TaskSpec, + task: ClaimedTask, + allocation: ResourceAllocation, + ) -> Provenance: + manifest = definition.manifest + started_at = _utc_now() + return Provenance( + workload=manifest.workload, + sdk_api_version=spec.sdk_api_version, + protocol_version=spec.protocol_version, + manifest_schema_version=spec.manifest_schema_version, + workflow_schema_version=spec.workflow_schema_version, + verifier=spec.verifier, + artifact_schemas=tuple( + sorted( + { + item.artifact.schema + for collection in spec.inputs.values() + for item in collection.items + }.union(port.schema.ref for port in spec.expected_outputs.values()), + key=lambda value: value.canonical, + ) + ), + package_digest=spec.package_digest, + manifest_digest=spec.manifest_digest, + environment_digest=spec.environment_digest, + worker_runtime={"kind": "worker-agent-v1"}, + allocated_resource_ids=(allocation.allocation_id,), + parameters_digest=hashlib.sha256( + canonical_json(spec.parameters).encode("utf-8") + ).hexdigest(), + input_collection_digest=spec.inputs["input"].digest, + execution_contract_digest=spec.digest, + selected_features=spec.selected_features, + optional_fallbacks=spec.optional_fallbacks, + job_id=str(uuid5(NAMESPACE_URL, f"scimesh:job:{task.task_id}")), + task_id=str(uuid5(NAMESPACE_URL, f"scimesh:task:{task.task_id}")), + started_at=started_at, + finished_at=started_at, + trust_mode=TrustMode.TRUSTED.value, + ) @staticmethod - def _string(params: dict[str, object], name: str) -> str: - value = params.get(name) - if not isinstance(value, str) or not value.strip() or len(value) > 200: - raise ValueError(f"{name} must be a non-empty string") - return value - - @staticmethod - def _positive_int(params: dict[str, object], name: str, default: int) -> int: - value = params.get(name, default) - if isinstance(value, bool) or not isinstance(value, int) or value < 1 or value > 100_000: - raise ValueError(f"{name} must be a positive integer") - return value - - @staticmethod - def _nonnegative_int(params: dict[str, object], name: str) -> int: - value = params[name] - if isinstance(value, bool) or not isinstance(value, int) or value < 0: - raise ValueError(f"{name} must be a non-negative integer") - return value - - @staticmethod - def _number(params: dict[str, object], name: str) -> float: - value = params.get(name) - if isinstance(value, bool) or not isinstance(value, (int, float)) or not 0 <= value <= 1: - raise ValueError(f"{name} must be a number between 0 and 1") - return float(value) + def _validate_output( + output: object, + spec: TaskSpec, + context: LocalTaskContext, + store: LocalArtifactStore, + provenance: Provenance, + *, + max_output_bytes: int, + ) -> None: + if not isinstance(output, OutputManifest): + raise ValueError("SDK workload must return an OutputManifest") + if output.task_key != spec.task_key: + raise ValueError( + "SDK workload output task_key does not match its trusted task" + ) + if output.provenance != provenance: + raise ValueError( + "SDK workload output provenance does not match its context" + ) + output.validate_against( + spec.expected_outputs, + max_output_bytes=max_output_bytes, + ) + if output.provenance != provenance: + raise ValueError( + "SDK workload output provenance does not match its context" + ) + output.validate_against( + spec.expected_outputs, + max_output_bytes=spec.resources.max_duration_seconds, # replaced below + ) + sink = context.sink + if not isinstance(sink, ScopedArtifactSink): + raise ValueError("SDK execution requires a scoped artifact sink") + declared = { + item.artifact.artifact_id: item.artifact + for collection in output.outputs.values() + for item in collection.items + } + issued = {artifact.artifact_id: artifact for artifact in sink.sealed_references} + if issued != declared: + raise ValueError( + "SDK workload outputs must declare exactly the artifacts sealed by its attempt" + ) + for collection in output.outputs.values(): + for item in collection.items: + store.require(item.artifact) diff --git a/scimesh/sdk/descriptors/__init__.py b/scimesh/workloads/descriptors/__init__.py similarity index 94% rename from scimesh/sdk/descriptors/__init__.py rename to scimesh/workloads/descriptors/__init__.py index 249f93b..c0c7790 100644 --- a/scimesh/sdk/descriptors/__init__.py +++ b/scimesh/workloads/descriptors/__init__.py @@ -1,4 +1,4 @@ -"""SDK-native ``descriptor-batch`` reference workload. +"""SDK-built ``descriptor-batch`` reference workload. See ``core.py`` for the pinned scientific contract and ``definition.py`` for the manifest-backed planner/runner/reducer handlers. diff --git a/scimesh/sdk/descriptors/core.py b/scimesh/workloads/descriptors/core.py similarity index 100% rename from scimesh/sdk/descriptors/core.py rename to scimesh/workloads/descriptors/core.py diff --git a/scimesh/sdk/descriptors/definition.py b/scimesh/workloads/descriptors/definition.py similarity index 94% rename from scimesh/sdk/descriptors/definition.py rename to scimesh/workloads/descriptors/definition.py index ee9aba0..eeafac6 100644 --- a/scimesh/sdk/descriptors/definition.py +++ b/scimesh/workloads/descriptors/definition.py @@ -1,4 +1,4 @@ -"""SDK-native ``descriptor-batch`` workload definition and handlers. +"""SDK-built ``descriptor-batch`` workload definition and handlers. This module is the first non-adapter reference workload built directly on the ``core-batch-v1`` profile: an explicit immutable manifest, a static map/reduce @@ -15,7 +15,7 @@ import shutil from pathlib import Path from typing import Any, Mapping, Sequence -from ..artifacts import ( +from ...sdk.artifacts import ( ArtifactCollection, ArtifactItem, ArtifactRef, @@ -25,15 +25,15 @@ from ..artifacts import ( OutputManifest, PortSpec, ) -from ..builtins import current_environment_digest, current_scimesh_package_digest -from ..execution import ( +from ..environment import current_environment_digest, current_scimesh_package_digest +from ...sdk.execution import ( CheckpointPolicy, ExecutionProfile, NetworkPolicy, RetryPolicy, ) -from ..identity import ComponentRef, SchemaRef, VersionRange, WorkloadId -from ..manifest import ( +from ...sdk.identity import ComponentRef, SchemaRef, VersionRange, WorkloadId +from ...sdk.manifest import ( DeterminismProfile, EnvironmentSpec, PackageSpec, @@ -42,12 +42,12 @@ from ..manifest import ( WorkloadLimits, WorkloadManifest, ) -from ..plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan -from ..protocols import PlanningContext, ReduceContext, TaskContext -from ..registry import WorkloadDefinition -from ..resources import ResourceRequirements -from ..verification import ExactArtifactVerifier -from ..workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec +from ...sdk.plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan +from ...sdk.protocols import PlanningContext, ReduceContext, TaskContext +from ...sdk.registry import WorkloadDefinition +from ...sdk.resources import ResourceRequirements +from ...sdk.verification import ExactArtifactVerifier +from ...sdk.workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec from .core import ( DESCRIPTOR_COLUMNS, compute_descriptor_batch, @@ -56,8 +56,8 @@ from .core import ( write_descriptor_shards, ) -MAP_ENTRY_POINT = "scimesh.sdk.descriptors.definition:map_descriptors@v1" -REDUCE_ENTRY_POINT = "scimesh.sdk.descriptors.definition:reduce_descriptors@v1" +MAP_ENTRY_POINT = "scimesh.workloads.descriptors.definition:map_descriptors@v1" +REDUCE_ENTRY_POINT = "scimesh.workloads.descriptors.definition:reduce_descriptors@v1" _DESCRIPTOR_PARAMETERS = ("skip_invalid",) @@ -118,7 +118,7 @@ class DescriptorBatchWorkload: """Manifest-backed planner, runner, and reducer for descriptor-batch. The class follows the legacy adapter's structural pattern (one object - registered under each stage entry point) while remaining fully SDK-native: + registered under each stage entry point) while remaining fully SDK-built: sharding is explicit and deterministic, every artifact is sealed through the bridge-owned sink, and no filesystem path ever enters a plan or task. """ diff --git a/scimesh/workloads/environment.py b/scimesh/workloads/environment.py new file mode 100644 index 0000000..b03e1c9 --- /dev/null +++ b/scimesh/workloads/environment.py @@ -0,0 +1,42 @@ +"""Local development digest helpers for built-in SDK workload definitions. + +This module is a leaf: it must not import other SDK workload modules, so the +workload definition packages (``search``, ``graph``, ``descriptors``) and the +``builtins`` registry wiring can import it without creating cycles. +""" + +from __future__ import annotations + +import hashlib +import os +import platform +import sys + +from rdkit import rdBase + +from scimesh.sdk.integrity import installed_distribution_digest + + +def current_scimesh_package_digest() -> str: + """Hash installed SciMesh Python sources for the built-in trusted adapters. + + This is a local immutable-code pin, not a package signature or container + attestation. Consequently the built-in manifests are trusted only; an + administrator must supply signed image metadata before enabling an + untrusted quorum policy. + """ + # Source/editable installs are allowed only for this explicit local + # development helper. Registry discovery keeps the secure default. + return installed_distribution_digest("scimesh", allow_editable=True) + + +def current_environment_digest() -> str: + payload = "\n".join( + ( + current_scimesh_package_digest(), + f"python={sys.implementation.name}-{platform.python_version()}", + f"rdkit={rdBase.rdkitVersion}", + f"platform={sys.platform}-{platform.machine().lower()}", + ) + ) + return "sha256:" + hashlib.sha256(payload.encode("utf-8")).hexdigest() diff --git a/scimesh/workloads/graph/__init__.py b/scimesh/workloads/graph/__init__.py new file mode 100644 index 0000000..a2fdeeb --- /dev/null +++ b/scimesh/workloads/graph/__init__.py @@ -0,0 +1,40 @@ +"""SDK-built ``similarity-graph`` workload. + +A user workload script built on the SciMesh Workload SDK. See ``core.py`` for +the block/pair scientific core and ``definition.py`` for the manifest-backed +planner/runner/reducer handlers. +""" + +from .core import ( + block_pair_from_key, + check_pair_coverage, + compute_block_edges, + merge_edge_partials, + parse_molecule_blocks, + read_block_rows, + write_block_tsv, + write_edge_csv, +) +from .definition import ( + MAP_ENTRY_POINT, + REDUCE_ENTRY_POINT, + SimilarityGraphSDKWorkload, + similarity_graph_sdk_definition, + workload_definition, +) + +__all__ = [ + "MAP_ENTRY_POINT", + "REDUCE_ENTRY_POINT", + "SimilarityGraphSDKWorkload", + "block_pair_from_key", + "check_pair_coverage", + "compute_block_edges", + "merge_edge_partials", + "parse_molecule_blocks", + "read_block_rows", + "similarity_graph_sdk_definition", + "workload_definition", + "write_block_tsv", + "write_edge_csv", +] diff --git a/scimesh/workloads/graph/core.py b/scimesh/workloads/graph/core.py new file mode 100644 index 0000000..5773ddc --- /dev/null +++ b/scimesh/workloads/graph/core.py @@ -0,0 +1,249 @@ +"""Scientific core for the SDK-built ``similarity-graph`` workload. + +Molecules are parsed once into deterministic row-ordered blocks; every block +pair ``(i, j)`` with ``i <= j`` becomes one map task (diagonal tasks compare +pairs ``a < b`` inside a block, off-diagonal tasks compare every molecule +across two blocks). The reducer enforces the CTX-10 pair-coverage invariant: +the union of task pair sets must equal all unordered molecule pairs exactly +once, and the merged edge set must contain no duplicate unordered pair. +""" + +from __future__ import annotations + +import csv +from pathlib import Path +from typing import Iterable, Sequence + +from rdkit import Chem, DataStructs + +from scimesh.chemistry.dataset import iter_valid_molecules +from scimesh.chemistry.fingerprints import fingerprint + +EDGE_COLUMNS = ("source_id", "target_id", "similarity") + +MoleculeBlock = list[tuple[str, str]] # (chembl_id, smiles), row-ordered + + +def parse_molecule_blocks( + input_path: Path, + block_size: int, + max_rows: int | None = None, +) -> tuple[list[MoleculeBlock], dict[str, int]]: + """Parse valid molecules into deterministic row-ordered blocks. + + Mirrors the local reference's strictness: an empty or duplicate + ``chembl_id`` fails the run, because the edge identity is the molecule id. + Invalid SMILES rows are skipped and counted. + """ + if ( + isinstance(block_size, bool) + or not isinstance(block_size, int) + or block_size < 1 + ): + raise ValueError("block_size must be a positive integer") + from scimesh.chemistry.dataset import DatasetStats + + stats = DatasetStats() + blocks: list[MoleculeBlock] = [] + current: MoleculeBlock = [] + seen_ids: set[str] = set() + for record in iter_valid_molecules(input_path, stats, max_rows=max_rows): + if not record.molecule_id: + raise ValueError("dataset contains an empty chembl_id") + if record.molecule_id in seen_ids: + raise ValueError( + f"dataset contains a duplicate chembl_id: {record.molecule_id}" + ) + seen_ids.add(record.molecule_id) + current.append((record.molecule_id, record.smiles)) + if len(current) == block_size: + blocks.append(current) + current = [] + if current: + blocks.append(current) + if not blocks: + raise ValueError("dataset has no valid molecules") + return blocks, { + "rows_scanned": stats.scanned, + "valid_molecules": stats.valid, + "invalid_smiles": stats.invalid, + "block_count": len(blocks), + } + + +def write_block_tsv(rows: MoleculeBlock, path: Path) -> None: + """Write one molecule block as a header TSV with the input column names.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8", newline="") as destination: + writer = csv.DictWriter( + destination, + fieldnames=["chembl_id", "canonical_smiles"], + delimiter="\t", + lineterminator="\n", + ) + writer.writeheader() + for molecule_id, smiles in rows: + writer.writerow({"chembl_id": molecule_id, "canonical_smiles": smiles}) + + +def read_block_rows(path: Path) -> MoleculeBlock: + rows: MoleculeBlock = [] + with path.open("r", encoding="utf-8", newline="") as source: + reader = csv.DictReader(source, delimiter="\t") + for row in reader: + molecule_id = row.get("chembl_id", "") + smiles = row.get("canonical_smiles", "") + if not molecule_id or not smiles: + raise ValueError("block artifact contains an invalid row") + rows.append((molecule_id, smiles)) + return rows + + +def compute_block_edges( + left: MoleculeBlock, + right: MoleculeBlock, + threshold: float, + threshold_direction: str, +) -> list[tuple[str, str, float]]: + """Compare one planned block pair and emit only thresholded edges.""" + if not 0.0 <= threshold <= 1.0: + raise ValueError("threshold must be between 0 and 1") + if threshold_direction not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + left_fingerprints = [ + (molecule_id, fingerprint(Chem.MolFromSmiles(smiles))) + for molecule_id, smiles in left + ] + right_fingerprints = [ + (molecule_id, fingerprint(Chem.MolFromSmiles(smiles))) + for molecule_id, smiles in right + ] + diagonal = left is right or left == right + edges: list[tuple[str, str, float]] = [] + for left_index, (left_id, left_fp) in enumerate(left_fingerprints): + right_start = left_index + 1 if diagonal else 0 + for right_index in range(right_start, len(right_fingerprints)): + right_id, right_fp = right_fingerprints[right_index] + similarity = DataStructs.TanimotoSimilarity(left_fp, right_fp) + matches_threshold = ( + similarity >= threshold + if threshold_direction == "greater" + else similarity <= threshold + ) + if matches_threshold: + edges.append((left_id, right_id, similarity)) + return edges + + +def write_edge_csv(output_path: Path, edges: Iterable[tuple[str, str, float]]) -> None: + """Write an edge table CSV with six-decimal similarity values. + + Uses the CSV module's default ``\\r\\n`` line terminator so the bytes match + the local ``write_graph_edges`` reference exactly. + """ + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8", newline="") as destination: + writer = csv.DictWriter(destination, fieldnames=list(EDGE_COLUMNS)) + writer.writeheader() + for source_id, target_id, similarity in edges: + writer.writerow( + { + "source_id": source_id, + "target_id": target_id, + "similarity": f"{similarity:.6f}", + } + ) + + +def read_edge_csv(path: Path) -> list[tuple[str, str, float]]: + """Read a materialized edge CSV with strict row validation.""" + if not path.is_file(): + raise ValueError("materialized edge partial is missing") + edges: list[tuple[str, str, float]] = [] + with path.open("r", encoding="utf-8", newline="") as source: + reader = csv.DictReader(source) + if tuple(reader.fieldnames or ()) != EDGE_COLUMNS: + raise ValueError("edge partial has an invalid CSV header") + for row in reader: + if set(row) != set(EDGE_COLUMNS): + raise ValueError("edge partial has an invalid row") + source_id = row["source_id"] + target_id = row["target_id"] + if not source_id or not target_id: + raise ValueError("edge partial contains an empty molecule id") + try: + similarity = float(row["similarity"]) + except (TypeError, ValueError) as error: + raise ValueError("edge partial has an invalid similarity") from error + if not 0 <= similarity <= 1: + raise ValueError("edge partial has an invalid similarity") + edges.append((source_id, target_id, similarity)) + return edges + + +def block_pair_from_key(key: str) -> tuple[int, int]: + """Parse ``map.x`` partial keys into block indices.""" + prefix = "map." + if not key.startswith(prefix): + raise ValueError("graph partial key must use map.x") + raw = key[len(prefix) :] + left_raw, separator, right_raw = raw.partition("x") + if not separator or not left_raw.isdigit() or not right_raw.isdigit(): + raise ValueError("graph partial key must use map.x") + return int(left_raw), int(right_raw) + + +def check_pair_coverage(pairs: Sequence[tuple[int, int]]) -> None: + """Enforce the pair-coverage invariant: every block pair exactly once. + + ``pairs`` must contain every ``(i, j)`` with ``0 <= i <= j < n`` exactly + once, where ``n`` is derived from the largest referenced block index. + """ + if not pairs: + raise ValueError("graph partial keys cover no block pairs") + unique = set(pairs) + if len(unique) != len(pairs): + raise ValueError("graph partial keys contain a duplicate block pair") + if any(left > right or left < 0 or right < 0 for left, right in unique): + raise ValueError("graph partial keys reference an invalid block pair") + n = max(right for _, right in unique) + 1 + expected = {(left, right) for left in range(n) for right in range(left, n)} + missing = sorted(expected - unique) + if missing: + raise ValueError( + "graph partial keys do not cover the full block pair set: " + + ", ".join(f"{left}x{right}" for left, right in missing) + ) + unexpected = sorted(unique - expected) + if unexpected: + raise ValueError( + "graph partial keys cover pairs outside the block pair set: " + + ", ".join(f"{left}x{right}" for left, right in unexpected) + ) + + +def merge_edge_partials( + partial_paths: Sequence[Path], + output_path: Path, +) -> dict[str, int]: + """Merge edge partials with duplicate detection and deterministic sort. + + The merged edge list is sorted by ``(source_id, target_id, -similarity)``, + matching the local brute-force reference exactly. + """ + if not partial_paths: + raise ValueError("graph reducer requires at least one edge partial") + edges: list[tuple[str, str, float]] = [] + seen_pairs: set[tuple[str, str]] = set() + for path in partial_paths: + for source_id, target_id, similarity in read_edge_csv(path): + unordered = (min(source_id, target_id), max(source_id, target_id)) + if unordered in seen_pairs: + raise ValueError( + f"graph partials contain a duplicate unordered pair: {unordered[0]}, {unordered[1]}" + ) + seen_pairs.add(unordered) + edges.append((source_id, target_id, similarity)) + edges.sort(key=lambda edge: (edge[0], edge[1], -edge[2])) + write_edge_csv(output_path, edges) + return {"partial_count": len(partial_paths), "edges_emitted": len(edges)} diff --git a/scimesh/workloads/graph/definition.py b/scimesh/workloads/graph/definition.py new file mode 100644 index 0000000..c97a7e9 --- /dev/null +++ b/scimesh/workloads/graph/definition.py @@ -0,0 +1,500 @@ +"""SDK-built ``similarity-graph`` workload definition and handlers. + +Built directly on the ``core-batch-v1`` profile: molecules are parsed once +into deterministic row-ordered blocks, every block pair ``(i, j)`` with +``i <= j`` becomes one map task, and the reducer enforces the CTX-10 +pair-coverage invariant (every unordered molecule pair compared exactly once) +before emitting a deterministically sorted edge list that is byte-identical +to the local brute-force reference. +""" + +from __future__ import annotations + +import hashlib +import shutil +from pathlib import Path +from typing import Any, Mapping + +from ...sdk.artifacts import ( + ArtifactCollection, + ArtifactItem, + ArtifactRef, + ArtifactSchema, + Cardinality, + CollectionKind, + OutputManifest, + PortSpec, +) +from ...sdk.execution import ( + CheckpointPolicy, + ExecutionProfile, + NetworkPolicy, + RetryPolicy, +) +from ...sdk.identity import ComponentRef, SchemaRef, VersionRange, WorkloadId +from ...sdk.manifest import ( + DeterminismProfile, + EnvironmentSpec, + PackageSpec, + TrustMode, + VerifierSpec, + WorkloadLimits, + WorkloadManifest, +) +from ...sdk.plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan +from ...sdk.protocols import PlanningContext, ReduceContext, TaskContext +from ...sdk.registry import WorkloadDefinition +from ...sdk.resources import ResourceRequirements +from ...sdk.verification import ExactArtifactVerifier +from ...sdk.workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec +from ..environment import current_environment_digest, current_scimesh_package_digest +from .core import ( + block_pair_from_key, + check_pair_coverage, + compute_block_edges, + merge_edge_partials, + parse_molecule_blocks, + read_block_rows, + write_block_tsv, + write_edge_csv, +) + +MAP_ENTRY_POINT = "scimesh.workloads.graph.definition:map_graph@v1" +REDUCE_ENTRY_POINT = "scimesh.workloads.graph.definition:reduce_graph@v1" + +_MAP_PARAMETERS = ("left_block", "right_block", "threshold", "threshold_direction") +_REDUCE_PARAMETERS = ("threshold", "threshold_direction", "block_size", "max_rows") + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _parameters_schema() -> dict[str, Any]: + return { + "type": "object", + "additionalProperties": False, + "properties": { + "threshold": {"type": "number", "minimum": 0, "maximum": 1}, + "threshold_direction": {"enum": ["greater", "less"]}, + "block_size": {"type": "integer", "minimum": 1}, + "max_rows": {"type": "integer", "minimum": 1}, + }, + "required": ["threshold"], + } + + +def _molecule_schema() -> ArtifactSchema: + return ArtifactSchema( + SchemaRef("molecule-table", 1), + "text/tab-separated-values", + "utf-8", + max_bytes=10 * 1024 * 1024 * 1024, + validator=ComponentRef("delimited-table", 1), + validator_configuration={ + "required_columns": ["canonical_smiles", "chembl_id"], + }, + max_records=100_000_000, + canonicalizer="scimesh-tsv-v1", + ) + + +def _edge_schema() -> ArtifactSchema: + return ArtifactSchema( + SchemaRef("similarity-edge-table", 1), + "text/csv", + "utf-8", + max_bytes=100 * 1024 * 1024 * 1024, + validator=ComponentRef("delimited-table", 1), + validator_configuration={ + "columns": ["source_id", "target_id", "similarity"], + }, + max_records=1_000_000_000, + canonicalizer="similarity-edge-table-v1", + ) + + +class SimilarityGraphSDKWorkload: + """Manifest-backed planner, runner, and reducer for similarity-graph.""" + + def __init__( + self, + *, + package_digest: str, + environment_digest: str, + ) -> None: + self.entry_point = MAP_ENTRY_POINT + self.input_port = PortSpec(_molecule_schema()) + self.block_port = PortSpec(_molecule_schema()) + self.partial_port = PortSpec(_edge_schema()) + self.output_port = PortSpec(_edge_schema()) + resources = ResourceRequirements( + profile="graph-cpu-v1", + cpu_cores=1, + memory_mb=1024, + scratch_mb=1024, + max_duration_seconds=3600, + ) + execution = ExecutionProfile( + profile="graph-python-process-v1", + network=NetworkPolicy.TRUSTED, + timeout_seconds=3600, + checkpoint=CheckpointPolicy(), + ) + limits = WorkloadLimits( + max_input_bytes=self.input_port.schema.max_bytes, + max_tasks=10_000, + max_output_bytes=self.output_port.schema.max_bytes, + ) + trust_modes = ("trusted", "untrusted_quorum") + map_stage = StageSpec( + stage_id="map", + kind=StageKind.MAP, + entry_point=MAP_ENTRY_POINT, + needs=(), + inputs={"left": self.block_port, "right": self.block_port}, + outputs={"partial": self.partial_port}, + parameter_names=_MAP_PARAMETERS, + resources=resources, + execution=execution, + retry=RetryPolicy(), + verifier=ComponentRef("exact-artifact", 1), + trust_modes=trust_modes, + max_fan_out=limits.max_tasks, + cacheable=True, + ) + reduce_input = PortSpec( + schema=self.partial_port.schema, + cardinality=Cardinality.MANY, + collection=CollectionKind.KEYED, + ) + reduce_stage = StageSpec( + stage_id="reduce", + kind=StageKind.REDUCE, + entry_point=REDUCE_ENTRY_POINT, + needs=("map",), + inputs={"partials": reduce_input}, + outputs={"result": self.output_port}, + parameter_names=_REDUCE_PARAMETERS, + resources=resources, + execution=execution, + retry=RetryPolicy(), + verifier=ComponentRef("exact-artifact", 1), + trust_modes=trust_modes, + max_fan_out=1, + cacheable=True, + ) + workflow = WorkflowSpec( + workflow_id="graph-block-pairs-v1", + inputs={"input": self.input_port}, + stages=(map_stage, reduce_stage), + edges=( + ArtifactEdge(PortRef("input"), PortRef("left", "map")), + ArtifactEdge(PortRef("input"), PortRef("right", "map")), + ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")), + ), + outputs={"result": PortRef("result", "reduce")}, + max_tasks=limits.max_tasks, + max_output_bytes=limits.max_output_bytes, + ) + self.manifest = WorkloadManifest( + sdk_api=VersionRange(">=1.0,<2.0"), + protocol=VersionRange(">=1,<2"), + workload=WorkloadId("similarity-graph", "1.0.0"), + description=( + "Exact sparse Tanimoto similarity graph over deterministic " + "block pairs with a duplicate-safe, coverage-checked merge." + ), + package=PackageSpec("scimesh", package_digest), + environment=EnvironmentSpec( + "python-process", + environment_digest, + {"adapter": "sdk-native"}, + ), + parameters_schema=_parameters_schema(), + workflow=workflow, + inputs={"input": self.input_port}, + outputs={"result": self.output_port}, + determinism=DeterminismProfile.BYTE_EXACT, + trust_modes=(TrustMode.TRUSTED, TrustMode.UNTRUSTED_QUORUM), + verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}), + limits=limits, + capabilities=("similarity-graph",), + conformance_profiles=("core-batch-v1",), + ) + self._exact_verifier = ExactArtifactVerifier() + + def definition(self) -> WorkloadDefinition: + return WorkloadDefinition( + manifest=self.manifest, + planner=self, + runners={MAP_ENTRY_POINT: self}, + reducers={REDUCE_ENTRY_POINT: self}, + verifiers={self._exact_verifier.identity.canonical: self._exact_verifier}, + ) + + @staticmethod + def _unit_interval(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number between 0 and 1") + return float(value) + + @staticmethod + def _positive_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + return value + + def validate(self, request: JobRequest) -> ValidatedJob: + if request.workload != self.manifest.workload: + raise ValueError("similarity-graph received a request for another workload") + parameters = request.parameters + unknown = set(parameters) - { + "threshold", + "threshold_direction", + "block_size", + "max_rows", + } + if unknown: + raise ValueError( + "unsupported similarity-graph parameters: " + ", ".join(sorted(unknown)) + ) + threshold = parameters.get("threshold") + if threshold is None: + raise ValueError("threshold is required") + self._unit_interval(threshold, "threshold") + if "threshold_direction" in parameters and parameters[ + "threshold_direction" + ] not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + if "block_size" in parameters: + self._positive_int(parameters["block_size"], "block_size") + if "max_rows" in parameters: + self._positive_int(parameters["max_rows"], "max_rows") + return ValidatedJob(request, request.parameters) + + def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan: + if not isinstance(job, ValidatedJob): + raise ValueError("job must be a ValidatedJob") + collection = job.request.inputs.get("input") + if collection is None: + raise ValueError("similarity-graph requires the input port") + self.input_port.validate_collection(collection, "job input") + input_artifact = collection.items[0].artifact + input_path = context.catalog.materialize(input_artifact) + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + parameters = job.resolved_parameters + threshold = self._unit_interval(parameters.get("threshold"), "threshold") + direction = parameters.get("threshold_direction", "greater") + if direction not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + block_size = int(parameters.get("block_size", 1_000)) + max_rows = parameters.get("max_rows") + blocks, stats = parse_molecule_blocks( + input_path, + block_size, + int(max_rows) if isinstance(max_rows, int) else None, + ) + task_parameters = { + "threshold": threshold, + "threshold_direction": direction, + } + negotiated = context.negotiated + map_stage = self.manifest.workflow.stages[0] + assert map_stage.verifier is not None + block_refs: list[ArtifactRef] = [] + for index, block in enumerate(blocks): + path = workspace / f"block-{index:04d}.tsv" + write_block_tsv(block, path) + block_refs.append( + context.sink.seal( + path, + declaration=self.block_port.schema, + ) + ) + tasks: list[TaskSpec] = [] + for left in range(len(blocks)): + for right in range(left, len(blocks)): + tasks.append( + TaskSpec( + workload=self.manifest.workload, + package_digest=self.manifest.package.digest, + manifest_digest=self.manifest.digest, + trust_mode=job.request.trust_mode, + sdk_api_version=negotiated.sdk_api_version, + protocol_version=negotiated.protocol_version, + manifest_schema_version=self.manifest.manifest_schema_version, + workflow_schema_version=self.manifest.workflow.schema_version, + environment_digest=self.manifest.environment.digest, + verifier=map_stage.verifier, + selected_features=negotiated.selected_features, + optional_fallbacks=negotiated.optional_fallbacks, + task_key=f"map/{left:04d}x{right:04d}", + stage_id="map", + parameters={ + **task_parameters, + "left_block": left, + "right_block": right, + }, + inputs={ + "left": ArtifactCollection.single(block_refs[left]), + "right": ArtifactCollection.single(block_refs[right]), + }, + expected_outputs={"partial": self.partial_port}, + resources=map_stage.resources, + execution=map_stage.execution, + ) + ) + return WorkflowPlan( + workload=self.manifest.workload, + package_digest=self.manifest.package.digest, + manifest_digest=self.manifest.digest, + trust_mode=job.request.trust_mode, + sdk_api_version=negotiated.sdk_api_version, + protocol_version=negotiated.protocol_version, + manifest_schema_version=self.manifest.manifest_schema_version, + workflow_schema_version=self.manifest.workflow.schema_version, + environment_digest=self.manifest.environment.digest, + verifier=self.manifest.verifier.verifier, + selected_features=negotiated.selected_features, + optional_fallbacks=negotiated.optional_fallbacks, + workflow_id=self.manifest.workflow.workflow_id, + resolved_parameters=dict(parameters), + tasks=tuple(tasks), + ) + + def run(self, context: TaskContext) -> OutputManifest: + context.cancellation.raise_if_cancelled() + parameters = context.task.parameters + left_block = parameters.get("left_block") + right_block = parameters.get("right_block") + if ( + isinstance(left_block, bool) + or not isinstance(left_block, int) + or isinstance(right_block, bool) + or not isinstance(right_block, int) + ): + raise ValueError("graph map task requires block indices") + diagonal = left_block == right_block + left_collection = context.task.inputs.get("left") + right_collection = context.task.inputs.get("right") + if left_collection is None or right_collection is None: + raise ValueError("graph map task requires left and right block inputs") + self.block_port.validate_collection(left_collection, "graph map left input") + self.block_port.validate_collection(right_collection, "graph map right input") + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + left_path = context.catalog.materialize(left_collection.items[0].artifact) + right_path = context.catalog.materialize(right_collection.items[0].artifact) + left_rows = read_block_rows(left_path) + right_rows = ( + left_rows + if diagonal and left_path.resolve() == right_path.resolve() + else read_block_rows(right_path) + ) + threshold = self._unit_interval(parameters.get("threshold"), "threshold") + direction = parameters.get("threshold_direction", "greater") + if direction not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + checked_pairs = ( + len(left_rows) * (len(left_rows) - 1) // 2 + if diagonal + else len(left_rows) * len(right_rows) + ) + edges = compute_block_edges( + left_rows, + right_rows, + threshold, + direction, + ) + output_path = workspace / "result.csv" + write_edge_csv(output_path, edges) + context.cancellation.raise_if_cancelled() + sealed = context.sink.seal( + output_path, + declaration=self.partial_port.schema, + ) + return OutputManifest( + context.task.task_key, + {"partial": ArtifactCollection.single(sealed)}, + {"checked_pairs": checked_pairs, "edges_emitted": len(edges)}, + context.provenance, + ).validate_against( + context.task.expected_outputs, + max_output_bytes=self.manifest.limits.max_output_bytes, + ) + + def reduce(self, context: ReduceContext) -> OutputManifest: + context.cancellation.raise_if_cancelled() + collection = context.accepted_inputs.get("partials") + if ( + collection is None + or collection.kind is not CollectionKind.KEYED + or not collection.items + ): + raise ValueError( + "graph reducer requires a non-empty keyed partial collection" + ) + self.manifest.workflow.stages[1].inputs["partials"].validate_collection( + collection, + "graph reducer partials", + ) + pairs = [block_pair_from_key(item.key or "") for item in collection.items] + check_pair_coverage(pairs) + expected_keys = context.task.expected_input_keys.get("partials") + if expected_keys is None or {item.key for item in collection.items} != set( + expected_keys + ): + raise ValueError( + "graph partial keys do not match the coordinator expected set" + ) + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + partial_paths: list[Path] = [] + for item in sorted(collection.items, key=lambda value: value.key or ""): + artifact: ArtifactRef = item.artifact + source = context.catalog.materialize(artifact) + target = workspace / artifact.artifact_id + if source.resolve() != target.resolve(): + shutil.copyfile(source, target) + if _sha256_file(target) != artifact.sha256: + raise ValueError("materialized partial checksum does not match") + partial_paths.append(target) + result_path = workspace / "result.csv" + metrics = merge_edge_partials(partial_paths, result_path) + context.cancellation.raise_if_cancelled() + sealed = context.sink.seal( + result_path, + declaration=self.output_port.schema, + ) + return OutputManifest( + context.task.task_key, + {"result": ArtifactCollection.single(sealed)}, + metrics, + context.provenance, + ).validate_against( + context.task.expected_outputs, + max_output_bytes=self.manifest.limits.max_output_bytes, + ) + + +def similarity_graph_sdk_definition( + *, + package_digest: str | None = None, + environment_digest: str | None = None, +) -> SimilarityGraphSDKWorkload: + """Build the SDK-based similarity-graph definition for tests.""" + return SimilarityGraphSDKWorkload( + package_digest=package_digest or current_scimesh_package_digest(), + environment_digest=environment_digest or current_environment_digest(), + ) + + +def workload_definition() -> WorkloadDefinition: + """Installed entry-point factory for the SDK-based similarity-graph.""" + return similarity_graph_sdk_definition().definition() diff --git a/scimesh/workloads/library.py b/scimesh/workloads/library.py new file mode 100644 index 0000000..091faec --- /dev/null +++ b/scimesh/workloads/library.py @@ -0,0 +1,68 @@ +"""Built-in workload library: default registry and runtime. + +This module is workload code, not SDK framework code. It composes the +installed SciMesh workload packages (``search``, ``graph``, ``descriptors``) +with the SDK registry and runtime. External workload libraries can follow the +same pattern with their own packages and entry points. +""" + +from __future__ import annotations + +import os +import platform + +from scimesh.sdk.identity import SDK_API_VERSION +from scimesh.sdk.registry import WorkloadRegistry +from scimesh.sdk.resources import ResourceInventory +from scimesh.sdk.runtime import RuntimeCapabilities + +from .descriptors import descriptor_batch_sdk_definition +from .environment import current_environment_digest +from .graph import similarity_graph_sdk_definition +from .search import similarity_search_sdk_definition + +__all__ = [ + "default_sdk_registry", + "default_sdk_runtime", +] + + +def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry: + """Registry of every built-in SDK-built workload, all enabled.""" + registry = WorkloadRegistry() + registry.register( + similarity_search_sdk_definition(shard_rows=shard_rows).definition(), + enabled=True, + ) + registry.register( + similarity_graph_sdk_definition().definition(), + enabled=True, + ) + registry.register( + descriptor_batch_sdk_definition(shard_rows=shard_rows).definition(), + enabled=True, + ) + return registry + + +def default_sdk_runtime() -> RuntimeCapabilities: + """Runtime advertising the built-in workloads' capabilities and inventory.""" + architecture = platform.machine().lower() or "unknown" + return RuntimeCapabilities( + sdk_api_version=SDK_API_VERSION, + protocol_version="1.0.0", + profiles=("core-batch-v1",), + features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"}, + workload_capabilities=( + "similarity-search", + "similarity-graph", + "descriptor-batch", + ), + inventory=ResourceInventory( + cpu_cores=max(os.cpu_count() or 1, 1), + memory_mb=4096, + scratch_mb=4096, + architecture=architecture, + environment_digests=(current_environment_digest(),), + ), + ) diff --git a/scimesh/workloads/search/__init__.py b/scimesh/workloads/search/__init__.py new file mode 100644 index 0000000..2ca7690 --- /dev/null +++ b/scimesh/workloads/search/__init__.py @@ -0,0 +1,31 @@ +"""SDK-built ``similarity-search`` reference workload. + +See ``core.py`` for the sharding/merge scientific core and ``definition.py`` +for the manifest-backed planner/runner/reducer handlers. +""" + +from .core import ( + merge_search_partials, + run_search_shard, + write_search_partial, + write_search_shards, +) +from .definition import ( + MAP_ENTRY_POINT, + REDUCE_ENTRY_POINT, + SimilaritySearchSDKWorkload, + similarity_search_sdk_definition, + workload_definition, +) + +__all__ = [ + "MAP_ENTRY_POINT", + "REDUCE_ENTRY_POINT", + "SimilaritySearchSDKWorkload", + "merge_search_partials", + "run_search_shard", + "similarity_search_sdk_definition", + "workload_definition", + "write_search_partial", + "write_search_shards", +] diff --git a/scimesh/workloads/search/core.py b/scimesh/workloads/search/core.py new file mode 100644 index 0000000..39ccb46 --- /dev/null +++ b/scimesh/workloads/search/core.py @@ -0,0 +1,261 @@ +"""Scientific core for the SDK-built ``similarity-search`` workload. + +Reuses the local reference implementation (``search_similar``) and the shared +CTX-08 partial format (full-precision ``repr`` scores) so that shard outputs +and the merged final CSV are byte-identical to the legacy distributed path and +to the single-process reference. +""" + +from __future__ import annotations + +import csv +import heapq +from pathlib import Path +from typing import Any, Iterator, Mapping, Sequence + +from scimesh.chemistry.dataset import MoleculeRecord, parse_smiles +from scimesh.workloads.similarity_search import ( + SimilarityMatch, + _HeapEntry, + search_similar, + write_search_results, +) + +SEARCH_COLUMNS = ("rank", "chembl_id", "canonical_smiles", "similarity") +REQUIRED_COLUMNS = {"chembl_id", "canonical_smiles"} + + +def write_search_partial( + output_path: Path, + matches: Sequence[SimilarityMatch], +) -> None: + """Write a worker partial with a round-trip score, not display rounding. + + The public final CSV continues to use the local CLI's six-decimal display. + A reducer needs the full binary float representation to rank candidates + from separate shards exactly as the single-process reference does. + """ + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8", newline="") as destination: + writer = csv.DictWriter(destination, fieldnames=list(SEARCH_COLUMNS)) + writer.writeheader() + for rank, match in enumerate(matches, start=1): + writer.writerow( + { + "rank": rank, + "chembl_id": match.molecule_id, + "canonical_smiles": match.smiles, + "similarity": repr(match.similarity), + } + ) + + +def run_search_shard( + input_path: Path, + parameters: Mapping[str, object], + output_path: Path, +) -> dict[str, int]: + """Run one planned shard with the local reference implementation. + + This is the worker entry used by the SDK-built runner. It deliberately + accepts only resolved ``query_smiles``: resolving an identifier + independently in each shard would make the distributed search + scientifically invalid, so identifier resolution happens once in the + planner (or at the worker bridge for v1-wire tasks that still carry + ``query_id``). + """ + allowed = { + "query_smiles", + "top_k", + "threshold", + "threshold_direction", + "progress_every", + } + unknown = set(parameters) - allowed + if unknown: + raise ValueError( + f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}" + ) + query_smiles = parameters.get("query_smiles") + if not isinstance(query_smiles, str) or not query_smiles.strip(): + raise ValueError("query_smiles is required for a distributed shard") + molecule = parse_smiles(query_smiles) + if molecule is None: + raise ValueError("query_smiles is invalid") + top_k = _positive_int(parameters.get("top_k", 20), "top_k") + threshold = None + if "threshold" in parameters: + threshold = _unit_interval(parameters["threshold"], "threshold") + direction = parameters.get("threshold_direction", "greater") + if direction not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + assert isinstance(direction, str) + progress_every = 0 + if "progress_every" in parameters: + progress_every = _nonnegative_int( + parameters["progress_every"], "progress_every" + ) + result = search_similar( + input_path, + MoleculeRecord("query", query_smiles, molecule), + top_k=top_k, + progress_every=progress_every, + threshold=threshold, + threshold_direction=direction, + ) + write_search_partial(output_path, result.matches) + return { + "scanned_rows": result.stats.scanned, + "valid_molecules": result.stats.valid, + "invalid_smiles": result.stats.invalid, + "matches_emitted": len(result.matches), + } + + +def _positive_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + return value + + +def _nonnegative_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer") + return value + + +def _unit_interval(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number between 0 and 1") + return float(value) + + +def write_search_shards( + input_path: Path, + workspace: Path, + shard_rows: int, + max_rows: int | None = None, +) -> list[Path]: + """Split the input TSV into deterministic row-bounded shards with headers.""" + if ( + isinstance(shard_rows, bool) + or not isinstance(shard_rows, int) + or shard_rows < 1 + ): + raise ValueError("shard_rows must be a positive integer") + if max_rows is not None and ( + isinstance(max_rows, bool) or not isinstance(max_rows, int) or max_rows < 1 + ): + raise ValueError("max_rows must be a positive integer") + paths: list[Path] = [] + current: Path | None = None + destination = None + writer = None + rows_in_shard = 0 + seen_rows = 0 + try: + with input_path.open("r", encoding="utf-8", newline="") as source: + reader = csv.DictReader(source, delimiter="\t") + fieldnames = tuple(reader.fieldnames or ()) + if not REQUIRED_COLUMNS.issubset(set(fieldnames)): + missing = sorted(REQUIRED_COLUMNS - set(fieldnames)) + raise ValueError( + f"dataset is missing required columns: {', '.join(missing)}" + ) + for row in reader: + if max_rows is not None and seen_rows >= max_rows: + break + if destination is None or rows_in_shard == shard_rows: + if destination is not None: + destination.close() + current = workspace / f"shard-{len(paths)}.tsv" + destination = current.open("w", encoding="utf-8", newline="") + writer = csv.DictWriter( + destination, + fieldnames=list(fieldnames), + delimiter="\t", + lineterminator="\n", + ) + writer.writeheader() + paths.append(current) + rows_in_shard = 0 + assert writer is not None + writer.writerow(row) + rows_in_shard += 1 + seen_rows += 1 + finally: + if destination is not None: + destination.close() + if not paths: + raise ValueError("dataset has no data rows") + return paths + + +def iter_search_partial( + path: Path, + threshold_direction: str, +) -> Iterator[SimilarityMatch]: + """Yield strictly ordered partial matches with full-precision scores.""" + if not path.is_file(): + raise ValueError("materialized partial result is missing") + if threshold_direction not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + previous_key: tuple[float, str, str] | None = None + with path.open("r", encoding="utf-8", newline="") as source: + reader = csv.DictReader(source) + if tuple(reader.fieldnames or ()) != SEARCH_COLUMNS: + raise ValueError("partial result has an invalid CSV header") + for expected_rank, row in enumerate(reader, start=1): + if set(row) != set(SEARCH_COLUMNS) or row["rank"] != str(expected_rank): + raise ValueError("partial result has an invalid rank") + try: + similarity = float(row["similarity"]) + except (TypeError, ValueError) as error: + raise ValueError("partial result has an invalid similarity") from error + if not 0 <= similarity <= 1: + raise ValueError("partial result has an invalid similarity") + match = SimilarityMatch( + similarity, row["chembl_id"], row["canonical_smiles"] + ) + key = match.sort_key(threshold_direction) + if previous_key is not None and key < previous_key: + raise ValueError("partial result is not sorted deterministically") + previous_key = key + yield match + + +def merge_search_partials( + partial_paths: Sequence[Path], + parameters: Mapping[str, Any], + output_path: Path, +) -> dict[str, int]: + """Merge sorted shard partials into one deterministic final top-k CSV. + + Mirrors the CTX-08/CTX-09 reducer: a bounded heap with the local + tie-breaker, so the merged file equals the single-process reference + byte-for-byte for the same input and options. + """ + if not partial_paths: + raise ValueError("at least one partial result is required") + raw_top_k = parameters.get("top_k", 20) + if isinstance(raw_top_k, bool) or not isinstance(raw_top_k, int) or raw_top_k < 1: + raise ValueError("top_k must be a positive integer") + top_k = raw_top_k + direction = parameters.get("threshold_direction", "greater") + if direction not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + heap: list[_HeapEntry] = [] + for path in partial_paths: + for match in iter_search_partial(path, direction): + rank_key = match.sort_key(direction) + entry = _HeapEntry(match, rank_key) + if len(heap) < top_k: + heapq.heappush(heap, entry) + elif rank_key < heap[0].rank_key: + heapq.heapreplace(heap, entry) + matches = sorted( + (entry.match for entry in heap), + key=lambda match: match.sort_key(direction), + ) + write_search_results(output_path, matches) + return {"matches_emitted": len(matches), "partial_count": len(partial_paths)} diff --git a/scimesh/workloads/search/definition.py b/scimesh/workloads/search/definition.py new file mode 100644 index 0000000..baa2fb3 --- /dev/null +++ b/scimesh/workloads/search/definition.py @@ -0,0 +1,565 @@ +"""SDK-built ``similarity-search`` workload definition and handlers. + +A direct ``core-batch-v1`` definition (not the legacy adapter): the planner +resolves the query once, shards deterministically, each map task computes the +local top-k with the reference implementation, and the reducer merges the +sorted partials with the same bounded heap and tie-breakers as the local CLI. +The manifest declares ``byte_exact`` with the exact-artifact verifier and both +``trusted`` and ``untrusted_quorum`` trust modes. +""" + +from __future__ import annotations + +import hashlib +import shutil +from pathlib import Path +from typing import Any, Mapping + +from rdkit import Chem + +from scimesh.chemistry.dataset import find_molecule_by_id, parse_smiles +from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE + +from ...sdk.artifacts import ( + ArtifactCollection, + ArtifactItem, + ArtifactRef, + ArtifactSchema, + Cardinality, + CollectionKind, + OutputManifest, + PortSpec, +) +from ..environment import current_environment_digest, current_scimesh_package_digest +from ...sdk.execution import ( + CheckpointPolicy, + ExecutionProfile, + NetworkPolicy, + RetryPolicy, +) +from ...sdk.identity import ComponentRef, SchemaRef, VersionRange, WorkloadId +from ...sdk.manifest import ( + DeterminismProfile, + EnvironmentSpec, + PackageSpec, + TrustMode, + VerifierSpec, + WorkloadLimits, + WorkloadManifest, +) +from ...sdk.plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan +from ...sdk.protocols import PlanningContext, ReduceContext, TaskContext +from ...sdk.registry import WorkloadDefinition +from ...sdk.resources import ResourceRequirements +from ...sdk.verification import ExactArtifactVerifier +from ...sdk.workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec +from .core import merge_search_partials, run_search_shard, write_search_shards + +MAP_ENTRY_POINT = "scimesh.workloads.search.definition:map_search@v1" +REDUCE_ENTRY_POINT = "scimesh.workloads.search.definition:reduce_search@v1" + +_MAP_PARAMETERS = ( + "query_smiles", + "top_k", + "threshold", + "threshold_direction", + "progress_every", +) +_REDUCE_PARAMETERS = _MAP_PARAMETERS + ("query_source", "fingerprint") + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _parameters_schema() -> dict[str, Any]: + return { + "type": "object", + "additionalProperties": False, + "properties": { + "query_id": {"type": "string", "minLength": 1, "maxLength": 200}, + "query_smiles": {"type": "string", "minLength": 1, "maxLength": 200}, + "top_k": {"type": "integer", "minimum": 1}, + "threshold": {"type": "number", "minimum": 0, "maximum": 1}, + "threshold_direction": {"enum": ["greater", "less"]}, + "max_rows": {"type": "integer", "minimum": 1}, + "progress_every": {"type": "integer", "minimum": 0}, + }, + "oneOf": [ + {"required": ["query_id"], "not": {"required": ["query_smiles"]}}, + {"required": ["query_smiles"], "not": {"required": ["query_id"]}}, + ], + } + + +def _dataset_schema() -> ArtifactSchema: + return ArtifactSchema( + SchemaRef("molecule-table", 1), + "text/tab-separated-values", + "utf-8", + max_bytes=10 * 1024 * 1024 * 1024, + validator=ComponentRef("delimited-table", 1), + validator_configuration={ + "required_columns": ["canonical_smiles", "chembl_id"], + }, + max_records=100_000_000, + canonicalizer="scimesh-tsv-v1", + ) + + +def _search_table_schema(ref: SchemaRef, canonicalizer: str) -> ArtifactSchema: + return ArtifactSchema( + ref, + "text/csv", + "utf-8", + max_bytes=1024 * 1024 * 1024, + validator=ComponentRef("delimited-table", 1), + validator_configuration={ + "columns": ["rank", "chembl_id", "canonical_smiles", "similarity"], + }, + max_records=100_000, + canonicalizer=canonicalizer, + ) + + +class SimilaritySearchSDKWorkload: + """Manifest-backed planner, runner, and reducer for similarity-search.""" + + def __init__( + self, + *, + shard_rows: int, + package_digest: str, + environment_digest: str, + ) -> None: + if ( + isinstance(shard_rows, bool) + or not isinstance(shard_rows, int) + or shard_rows < 1 + ): + raise ValueError("shard_rows must be a positive integer") + self.entry_point = MAP_ENTRY_POINT + self.shard_rows = shard_rows + self.input_port = PortSpec(_dataset_schema()) + self.partial_port = PortSpec( + _search_table_schema( + SchemaRef("similarity-search-partial", 1), "scimesh-search-partial-v1" + ) + ) + self.output_port = PortSpec( + _search_table_schema( + SchemaRef("similarity-search-result", 1), "scimesh-search-result-v1" + ) + ) + resources = ResourceRequirements( + profile="search-cpu-v1", + cpu_cores=1, + memory_mb=1024, + scratch_mb=1024, + max_duration_seconds=3600, + ) + execution = ExecutionProfile( + profile="search-python-process-v1", + network=NetworkPolicy.TRUSTED, + timeout_seconds=3600, + checkpoint=CheckpointPolicy(), + ) + limits = WorkloadLimits( + max_input_bytes=self.input_port.schema.max_bytes, + max_tasks=10_000, + max_output_bytes=self.output_port.schema.max_bytes, + ) + trust_modes = ("trusted", "untrusted_quorum") + map_stage = StageSpec( + stage_id="map", + kind=StageKind.MAP, + entry_point=MAP_ENTRY_POINT, + needs=(), + inputs={"input": self.input_port}, + outputs={"partial": self.partial_port}, + parameter_names=_MAP_PARAMETERS, + resources=resources, + execution=execution, + retry=RetryPolicy(), + verifier=ComponentRef("exact-artifact", 1), + trust_modes=trust_modes, + max_fan_out=limits.max_tasks, + cacheable=True, + ) + reduce_input = PortSpec( + schema=self.partial_port.schema, + cardinality=Cardinality.MANY, + collection=CollectionKind.KEYED, + ) + reduce_stage = StageSpec( + stage_id="reduce", + kind=StageKind.REDUCE, + entry_point=REDUCE_ENTRY_POINT, + needs=("map",), + inputs={"partials": reduce_input}, + outputs={"result": self.output_port}, + parameter_names=_REDUCE_PARAMETERS, + resources=resources, + execution=execution, + retry=RetryPolicy(), + verifier=ComponentRef("exact-artifact", 1), + trust_modes=trust_modes, + max_fan_out=1, + cacheable=True, + ) + workflow = WorkflowSpec( + workflow_id="search-map-reduce-v1", + inputs={"input": self.input_port}, + stages=(map_stage, reduce_stage), + edges=( + ArtifactEdge(PortRef("input"), PortRef("input", "map")), + ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")), + ), + outputs={"result": PortRef("result", "reduce")}, + max_tasks=limits.max_tasks, + max_output_bytes=limits.max_output_bytes, + ) + self.manifest = WorkloadManifest( + sdk_api=VersionRange(">=1.0,<2.0"), + protocol=VersionRange(">=1,<2"), + workload=WorkloadId("similarity-search", "1.0.0"), + description=( + "Exact top-k Tanimoto molecular similarity search over " + "deterministic TSV shards with a bounded merge." + ), + package=PackageSpec("scimesh", package_digest), + environment=EnvironmentSpec( + "python-process", + environment_digest, + {"adapter": "sdk-native"}, + ), + parameters_schema=_parameters_schema(), + workflow=workflow, + inputs={"input": self.input_port}, + outputs={"result": self.output_port}, + determinism=DeterminismProfile.BYTE_EXACT, + trust_modes=(TrustMode.TRUSTED, TrustMode.UNTRUSTED_QUORUM), + verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}), + limits=limits, + capabilities=("similarity-search",), + conformance_profiles=("core-batch-v1",), + ) + self._exact_verifier = ExactArtifactVerifier() + + def definition(self) -> WorkloadDefinition: + return WorkloadDefinition( + manifest=self.manifest, + planner=self, + runners={MAP_ENTRY_POINT: self}, + reducers={REDUCE_ENTRY_POINT: self}, + verifiers={self._exact_verifier.identity.canonical: self._exact_verifier}, + ) + + @staticmethod + def _string(value: object, name: str) -> str: + if not isinstance(value, str) or not value.strip() or len(value) > 200: + raise ValueError(f"{name} must be a non-empty string") + return value + + @staticmethod + def _positive_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + return value + + @staticmethod + def _nonnegative_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer") + return value + + @staticmethod + def _unit_interval(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number between 0 and 1") + return float(value) + + def validate(self, request: JobRequest) -> ValidatedJob: + if request.workload != self.manifest.workload: + raise ValueError( + "similarity-search received a request for another workload" + ) + parameters = request.parameters + unknown = set(parameters) - { + "query_id", + "query_smiles", + "top_k", + "threshold", + "threshold_direction", + "max_rows", + "progress_every", + } + if unknown: + raise ValueError( + "unsupported similarity-search parameters: " + + ", ".join(sorted(unknown)) + ) + query_id = parameters.get("query_id") + query_smiles = parameters.get("query_smiles") + if (query_id is None) == (query_smiles is None): + raise ValueError("exactly one of query_id or query_smiles is required") + if query_id is not None: + self._string(query_id, "query_id") + if query_smiles is not None: + self._string(query_smiles, "query_smiles") + self._positive_int(parameters.get("top_k", 20), "top_k") + if "max_rows" in parameters: + self._positive_int(parameters["max_rows"], "max_rows") + if "progress_every" in parameters: + self._nonnegative_int(parameters["progress_every"], "progress_every") + if "threshold" in parameters: + self._unit_interval(parameters["threshold"], "threshold") + if "threshold_direction" in parameters and parameters[ + "threshold_direction" + ] not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + return ValidatedJob(request, self._resolved_parameters(request)) + + def _resolved_parameters(self, request: JobRequest) -> dict[str, object]: + parameters = request.parameters + query_id = parameters.get("query_id") + if isinstance(query_id, str): + query_source: dict[str, str] = {"kind": "chembl_id", "value": query_id} + else: + query_source = { + "kind": "smiles", + "value": self._string(parameters.get("query_smiles"), "query_smiles"), + } + resolved: dict[str, object] = { + "query_source": query_source, + "top_k": self._positive_int(parameters.get("top_k", 20), "top_k"), + "threshold_direction": parameters.get("threshold_direction", "greater"), + "fingerprint": { + "algorithm": "morgan", + "radius": FP_RADIUS, + "fp_size": FP_SIZE, + }, + } + if "threshold" in parameters: + resolved["threshold"] = self._unit_interval( + parameters["threshold"], "threshold" + ) + if "max_rows" in parameters: + resolved["max_rows"] = self._positive_int( + parameters["max_rows"], "max_rows" + ) + if "progress_every" in parameters: + resolved["progress_every"] = self._nonnegative_int( + parameters["progress_every"], "progress_every" + ) + return resolved + + def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan: + if not isinstance(job, ValidatedJob): + raise ValueError("job must be a ValidatedJob") + collection = job.request.inputs.get("input") + if collection is None: + raise ValueError("similarity-search requires the input port") + self.input_port.validate_collection(collection, "job input") + input_artifact = collection.items[0].artifact + input_path = context.catalog.materialize(input_artifact) + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + resolved = dict(job.resolved_parameters) + query_smiles = self._resolve_query(input_path, job.request.parameters) + resolved["query_smiles"] = query_smiles + max_rows = resolved.get("max_rows") + shard_paths = write_search_shards( + input_path, + workspace, + self.shard_rows, + int(max_rows) if isinstance(max_rows, int) else None, + ) + task_parameters = { + key: value for key, value in resolved.items() if key in set(_MAP_PARAMETERS) + } + negotiated = context.negotiated + map_stage = self.manifest.workflow.stages[0] + assert map_stage.verifier is not None + tasks: list[TaskSpec] = [] + for index, path in enumerate(shard_paths): + sealed = context.sink.seal( + path, + declaration=self.input_port.schema, + ) + tasks.append( + TaskSpec( + workload=self.manifest.workload, + package_digest=self.manifest.package.digest, + manifest_digest=self.manifest.digest, + trust_mode=job.request.trust_mode, + sdk_api_version=negotiated.sdk_api_version, + protocol_version=negotiated.protocol_version, + manifest_schema_version=self.manifest.manifest_schema_version, + workflow_schema_version=self.manifest.workflow.schema_version, + environment_digest=self.manifest.environment.digest, + verifier=map_stage.verifier, + selected_features=negotiated.selected_features, + optional_fallbacks=negotiated.optional_fallbacks, + task_key=f"map/{index:08d}", + stage_id="map", + parameters=task_parameters, + inputs={"input": ArtifactCollection.single(sealed)}, + expected_outputs={"partial": self.partial_port}, + resources=map_stage.resources, + execution=map_stage.execution, + ) + ) + return WorkflowPlan( + workload=self.manifest.workload, + package_digest=self.manifest.package.digest, + manifest_digest=self.manifest.digest, + trust_mode=job.request.trust_mode, + sdk_api_version=negotiated.sdk_api_version, + protocol_version=negotiated.protocol_version, + manifest_schema_version=self.manifest.manifest_schema_version, + workflow_schema_version=self.manifest.workflow.schema_version, + environment_digest=self.manifest.environment.digest, + verifier=self.manifest.verifier.verifier, + selected_features=negotiated.selected_features, + optional_fallbacks=negotiated.optional_fallbacks, + workflow_id=self.manifest.workflow.workflow_id, + resolved_parameters=resolved, + tasks=tuple(tasks), + ) + + @staticmethod + def _resolve_query(input_path: Path, parameters: Mapping[str, object]) -> str: + query_id = parameters.get("query_id") + if isinstance(query_id, str): + record = find_molecule_by_id(input_path, query_id) + return Chem.MolToSmiles(record.molecule, canonical=True) + supplied = parameters["query_smiles"] + assert isinstance(supplied, str) + molecule = parse_smiles(supplied) + if molecule is None: + raise ValueError("query_smiles is invalid") + return Chem.MolToSmiles(molecule, canonical=True) + + def run(self, context: TaskContext) -> OutputManifest: + context.cancellation.raise_if_cancelled() + collection = context.task.inputs.get("input") + if collection is None: + raise ValueError("search map task requires one input collection") + self.input_port.validate_collection(collection, "search map input") + source = context.catalog.materialize(collection.items[0].artifact) + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + input_path = workspace / "input" + output_path = workspace / "result.csv" + if source.resolve() != input_path.resolve(): + shutil.copyfile(source, input_path) + metrics = run_search_shard( + input_path, + context.task.parameters, + output_path, + ) + context.cancellation.raise_if_cancelled() + sealed = context.sink.seal( + output_path, + declaration=self.partial_port.schema, + ) + return OutputManifest( + context.task.task_key, + {"partial": ArtifactCollection.single(sealed)}, + metrics, + context.provenance, + ).validate_against( + context.task.expected_outputs, + max_output_bytes=self.manifest.limits.max_output_bytes, + ) + + def reduce(self, context: ReduceContext) -> OutputManifest: + context.cancellation.raise_if_cancelled() + collection = context.accepted_inputs.get("partials") + if ( + collection is None + or collection.kind is not CollectionKind.KEYED + or not collection.items + ): + raise ValueError( + "search reducer requires a non-empty keyed partial collection" + ) + self.manifest.workflow.stages[1].inputs["partials"].validate_collection( + collection, + "search reducer partials", + ) + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + indexed_items: list[tuple[int, ArtifactItem]] = [] + for item in collection.items: + key = item.key or "" + prefix = "map." + raw_index = key[len(prefix) :] if key.startswith(prefix) else "" + if len(raw_index) != 8 or not raw_index.isdigit(): + raise ValueError("search partial key must use map.") + indexed_items.append((int(raw_index), item)) + expected_keys = context.task.expected_input_keys.get("partials") + if expected_keys is None or {item.key for item in collection.items} != set( + expected_keys + ): + raise ValueError( + "search partial keys do not match the coordinator expected set" + ) + if sorted(index for index, _ in indexed_items) != list( + range(len(indexed_items)) + ): + raise ValueError("search partial keys must be complete and contiguous") + partial_paths: list[Path] = [] + for index, item in sorted(indexed_items): + artifact: ArtifactRef = item.artifact + source = context.catalog.materialize(artifact) + target = workspace / artifact.artifact_id + if source.resolve() != target.resolve(): + shutil.copyfile(source, target) + if _sha256_file(target) != artifact.sha256: + raise ValueError("materialized partial checksum does not match") + partial_paths.append(target) + result_path = workspace / "result.csv" + metrics = merge_search_partials( + partial_paths, + context.task.parameters, + result_path, + ) + context.cancellation.raise_if_cancelled() + sealed = context.sink.seal( + result_path, + declaration=self.output_port.schema, + ) + return OutputManifest( + context.task.task_key, + {"result": ArtifactCollection.single(sealed)}, + metrics, + context.provenance, + ).validate_against( + context.task.expected_outputs, + max_output_bytes=self.manifest.limits.max_output_bytes, + ) + + +def similarity_search_sdk_definition( + *, + shard_rows: int = 10_000, + package_digest: str | None = None, + environment_digest: str | None = None, +) -> SimilaritySearchSDKWorkload: + """Build the SDK-built similarity-search definition for tests.""" + return SimilaritySearchSDKWorkload( + shard_rows=shard_rows, + package_digest=package_digest or current_scimesh_package_digest(), + environment_digest=environment_digest or current_environment_digest(), + ) + + +def workload_definition() -> WorkloadDefinition: + """Installed entry-point factory for the SDK-built similarity-search.""" + return similarity_search_sdk_definition().definition() diff --git a/tests/test_distributed.py b/tests/test_distributed.py deleted file mode 100644 index 390bc95..0000000 --- a/tests/test_distributed.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Contract tests for the coordinator-independent distributed workload boundary.""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Mapping, Sequence -from uuid import NAMESPACE_URL, uuid5 - -import pytest - -from scimesh.distributed import ( - ArtifactReference, - CompletedPartial, - DistributedPlan, - DistributedWorkloadRegistry, - FinalResult, - PlannedTask, - PlanningService, -) - - -def artifact(seed: str, content_type: str = "text/tab-separated-values") -> ArtifactReference: - return ArtifactReference( - artifact_id=str(uuid5(NAMESPACE_URL, seed)), - sha256=(seed.encode("utf-8").hex() * 64)[:64], - content_type=content_type, - ) - - -class DummyWorkload: - """A deterministic fake workload used to test the generic CTX-07 bridge.""" - - name = "dummy-workload" - description = "A deterministic test workload." - - def __init__(self) -> None: - self.plan_calls = 0 - self.received_partials: tuple[CompletedPartial, ...] = () - - def validate_job(self, parameters: Mapping[str, object]) -> None: - if parameters != {"mode": "valid"}: - raise ValueError("mode must be valid") - - def plan( - self, - input_path: Path, - input_artifact_id: str, - parameters: Mapping[str, object], - shard_rows: int, - workspace: Path, - ) -> DistributedPlan: - self.plan_calls += 1 - assert input_path.name == "input.tsv" - assert workspace.name == "workspace" - return DistributedPlan( - workload=self.name, - resolved_parameters={"mode": parameters["mode"], "source": input_artifact_id}, - tasks=( - PlannedTask(0, artifact(f"{input_artifact_id}:0"), {"mode": "valid"}), - PlannedTask(1, artifact(f"{input_artifact_id}:1"), {"mode": "valid"}), - ), - ) - - def reduce( - self, - partial_results: Sequence[CompletedPartial], - parameters: Mapping[str, object], - workspace: Path, - ) -> FinalResult: - self.received_partials = tuple(partial_results) - return FinalResult(artifact("final", "text/csv"), {"partial_count": len(partial_results)}) - - -def service() -> tuple[PlanningService, DummyWorkload]: - workload = DummyWorkload() - registry = DistributedWorkloadRegistry() - registry.register(workload) - return PlanningService(registry), workload - - -def test_unknown_workload_is_rejected_before_a_plan_is_written(tmp_path: Path) -> None: - planner, workload = service() - - with pytest.raises(ValueError, match="unknown distributed workload"): - planner.plan( - "unknown-workload", tmp_path / "input.tsv", artifact("input").artifact_id, - {"mode": "valid"}, 10, tmp_path / "workspace", - ) - - assert workload.plan_calls == 0 - - -def test_invalid_job_is_rejected_before_the_planner_runs(tmp_path: Path) -> None: - planner, workload = service() - - with pytest.raises(ValueError, match="mode must be valid"): - planner.plan( - "dummy-workload", tmp_path / "input.tsv", artifact("input").artifact_id, - {"mode": "invalid"}, 10, tmp_path / "workspace", - ) - - assert workload.plan_calls == 0 - - -def test_two_shard_plan_is_deterministic_and_json_serializable(tmp_path: Path) -> None: - planner, _ = service() - input_artifact_id = artifact("input").artifact_id - first = planner.plan( - "dummy-workload", tmp_path / "input.tsv", input_artifact_id, - {"mode": "valid"}, 10, tmp_path / "workspace", - ) - second = planner.plan( - "dummy-workload", tmp_path / "input.tsv", input_artifact_id, - {"mode": "valid"}, 10, tmp_path / "workspace", - ) - - assert first.to_json() == second.to_json() - payload = json.loads(first.to_json()) - assert [task["chunk_index"] for task in payload["tasks"]] == [0, 1] - assert all(set(task) == {"chunk_index", "input_artifact", "parameters"} for task in payload["tasks"]) - assert DistributedPlan.from_json(first.to_json()) == first - - -def test_plan_rejects_unsafe_or_non_deterministic_task_payloads() -> None: - with pytest.raises(ValueError, match="unique, ascending"): - DistributedPlan( - workload="dummy-workload", - resolved_parameters={}, - tasks=( - PlannedTask(1, artifact("one"), {}), - PlannedTask(0, artifact("zero"), {}), - ), - ) - - with pytest.raises(ValueError, match="JSON-compatible"): - PlannedTask(0, artifact("bad"), {"path": Path("not-serializable")}) - - with pytest.raises(ValueError, match="URI or local path"): - PlannedTask(0, artifact("uri"), {"input": "file:///tmp/input.tsv"}) - - with pytest.raises(ValueError, match="canonical hyphenated"): - DistributedPlan("dummy_workload", {}, (PlannedTask(0, artifact("one"), {}),)) - - -def test_reducer_receives_completed_partials_in_chunk_order(tmp_path: Path) -> None: - planner, workload = service() - result = planner.reduce( - "dummy-workload", - ( - CompletedPartial(3, artifact("three", "text/csv"), {"scanned_rows": 10}), - CompletedPartial(1, artifact("one", "text/csv"), {"scanned_rows": 10}), - ), - {"mode": "valid"}, - tmp_path / "workspace", - ) - - assert [partial.chunk_index for partial in workload.received_partials] == [1, 3] - assert result.metrics == {"partial_count": 2} - - -def test_reducer_rejects_duplicate_chunk_indexes_before_invocation(tmp_path: Path) -> None: - planner, workload = service() - duplicate = CompletedPartial(0, artifact("partial", "text/csv"), {"scanned_rows": 1}) - - with pytest.raises(ValueError, match="unique chunk_index"): - planner.reduce("dummy-workload", (duplicate, duplicate), {"mode": "valid"}, tmp_path) - - assert workload.received_partials == () - - -def test_artifact_references_never_accept_paths_or_uris() -> None: - with pytest.raises(ValueError, match="UUID"): - ArtifactReference("file:///tmp/input.tsv", "a" * 64, "text/csv") - with pytest.raises(ValueError, match="lowercase SHA-256"): - ArtifactReference(str(uuid5(NAMESPACE_URL, "input")), "A" * 64, "text/csv") - - -def test_registry_descriptions_are_stable_and_duplicate_names_are_rejected() -> None: - registry = DistributedWorkloadRegistry() - first, second = DummyWorkload(), DummyWorkload() - registry.register(first) - - assert registry.descriptions()[0].name == "dummy-workload" - with pytest.raises(ValueError, match="already registered"): - registry.register(second) diff --git a/tests/test_distributed_similarity_search.py b/tests/test_distributed_similarity_search.py deleted file mode 100644 index 7759532..0000000 --- a/tests/test_distributed_similarity_search.py +++ /dev/null @@ -1,184 +0,0 @@ -"""Scientific reference tests for the CTX-08 distributed search workload.""" - -from __future__ import annotations - -import csv -import hashlib -from pathlib import Path -from uuid import NAMESPACE_URL, uuid5 - -import pytest - -from scimesh.chemistry.dataset import find_molecule_by_id -from scimesh.distributed import ( - ArtifactReference, - CompletedPartial, - PlanningService, - default_distributed_registry, -) -from scimesh.distributed.registry import DistributedWorkloadRegistry -from scimesh.distributed.similarity_search import ( - SimilaritySearchDistributedWorkload, - run_similarity_search_shard, - write_similarity_search_partial, -) -from scimesh.workloads.similarity_search import search_similar, write_search_results - - -def make_dataset(path: Path) -> None: - path.write_text( - "chembl_id\tcanonical_smiles\textra\n" - "CHEMBL_QUERY\tCCO\tquery\n" - "CHEMBL_A\tCCCO\ta\n" - "CHEMBL_B\tCCCC\tb\n" - "CHEMBL_INVALID\tnot-a-smiles\tbad\n" - "CHEMBL_DUPLICATE\tCCO\tduplicate\n" - "CHEMBL_C\tCCN\tc\n", - encoding="utf-8", - ) - - -def planner() -> tuple[PlanningService, SimilaritySearchDistributedWorkload]: - workload = SimilaritySearchDistributedWorkload() - registry = DistributedWorkloadRegistry() - registry.register(workload) - return PlanningService(registry), workload - - -def test_default_registry_exposes_only_supported_distributed_search() -> None: - assert [item.name for item in default_distributed_registry().descriptions()] == ["similarity-search"] - - -def checksum(path: Path) -> str: - return hashlib.sha256(path.read_bytes()).hexdigest() - - -def test_query_id_is_resolved_once_before_deterministic_shards( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - dataset = tmp_path / "chembl.tsv" - workspace = tmp_path / "workspace" - make_dataset(dataset) - service, _ = planner() - calls = 0 - real_find = find_molecule_by_id - - def count_find(path: Path, query_id: str) -> MoleculeRecord: - nonlocal calls - calls += 1 - return real_find(path, query_id) - - monkeypatch.setattr("scimesh.distributed.similarity_search.find_molecule_by_id", count_find) - input_id = str(uuid5(NAMESPACE_URL, "dataset")) - plan = service.plan( - "similarity-search", dataset, input_id, - {"query_id": "CHEMBL_QUERY", "top_k": 3, "max_rows": 5, "progress_every": 0}, - 2, workspace, - ) - - assert calls == 1 - assert plan.resolved_parameters["query_smiles"] == "CCO" - assert plan.resolved_parameters["query_source"] == {"kind": "chembl_id", "value": "CHEMBL_QUERY"} - assert [task.chunk_index for task in plan.tasks] == [0, 1, 2] - assert all("query_id" not in task.parameters for task in plan.tasks) - assert all("max_rows" not in task.parameters for task in plan.tasks) - assert all(task.parameters["query_smiles"] == "CCO" for task in plan.tasks) - assert [ - sum(1 for _ in path.open(encoding="utf-8")) - 1 - for path in sorted(workspace.glob("shard-*.tsv")) - ] == [2, 2, 1] - - -def test_distributed_reduction_matches_single_process_reference(tmp_path: Path) -> None: - dataset = tmp_path / "chembl.tsv" - workspace = tmp_path / "workspace" - make_dataset(dataset) - service, workload = planner() - plan = service.plan( - "similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")), - {"query_smiles": "CCO", "top_k": 3, "threshold": 0.0}, 2, workspace, - ) - - partials: list[CompletedPartial] = [] - # Worker two finishes the latter shards first. Worker one loses its first - # attempt for shard zero, then retries it last. The reducer must remain - # independent of both completion and retry order. - for task in reversed(plan.tasks): - shard = workspace / f"shard-{task.chunk_index}.tsv" - temporary_partial = workspace / f"worker-output-{task.chunk_index}.csv" - metrics = run_similarity_search_shard(shard, task.parameters, temporary_partial) - partial_id = str(uuid5(NAMESPACE_URL, f"partial:{task.chunk_index}")) - partials.append( - CompletedPartial( - task.chunk_index, - ArtifactReference( - partial_id, checksum(temporary_partial), "text/csv", - ), - metrics, - ) - ) - # The reducer materializes result files under their own coordinator IDs, - # not shard input IDs. Keep this fixture faithful to that boundary. - temporary_partial.rename(workspace / partial_id) - - final = workload.reduce(tuple(partials), plan.resolved_parameters, workspace) - reference = tmp_path / "reference.csv" - query_record = find_molecule_by_id(dataset, "CHEMBL_QUERY") - write_search_results(reference, search_similar(dataset, query_record, top_k=3, threshold=0.0).matches) - - assert (workspace / "result.csv").read_bytes() == reference.read_bytes() - assert final.metrics == {"matches_emitted": 3, "partial_count": 3} - rows = list(csv.DictReader((workspace / "result.csv").open(encoding="utf-8"))) - assert {row["chembl_id"] for row in rows}.isdisjoint({"CHEMBL_QUERY", "CHEMBL_DUPLICATE"}) - - -def test_reducer_rejects_unsorted_or_invalid_partial_csv(tmp_path: Path) -> None: - workspace = tmp_path / "workspace" - workspace.mkdir() - workload = SimilaritySearchDistributedWorkload() - artifact_id = str(uuid5(NAMESPACE_URL, "bad")) - partial_path = workspace / artifact_id - partial_path.write_text( - "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.1\n2,B,CCC,0.9\n", - encoding="utf-8", - ) - artifact = ArtifactReference(artifact_id, checksum(partial_path), "text/csv") - - with pytest.raises(ValueError, match="not sorted"): - workload.reduce( - (CompletedPartial(0, artifact, {"scanned_rows": 2}),), - {"query_smiles": "CCO", "top_k": 2, "threshold_direction": "greater", "fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}}, - workspace, - ) - - -def test_partial_csv_preserves_exact_scores_for_global_ranking(tmp_path: Path) -> None: - partial = tmp_path / "partial.csv" - # Both values look identical in a six-decimal final CSV. The exact value - # must survive shard transport so the global reducer can still rank them. - from scimesh.workloads.similarity_search import SimilarityMatch - - write_similarity_search_partial( - partial, - [SimilarityMatch(0.50000049, "A", "CC"), SimilarityMatch(0.50000048, "B", "CCC")], - ) - values = list(csv.DictReader(partial.open(encoding="utf-8"))) - assert values[0]["similarity"] == repr(0.50000049) - assert values[1]["similarity"] == repr(0.50000048) - - -def test_planner_rejects_fingerprint_override_and_invalid_query(tmp_path: Path) -> None: - dataset = tmp_path / "chembl.tsv" - make_dataset(dataset) - service, _ = planner() - - with pytest.raises(ValueError, match="unsupported similarity-search parameters"): - service.plan( - "similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")), - {"query_smiles": "CCO", "fingerprint": {"radius": 1}}, 2, tmp_path / "workspace", - ) - with pytest.raises(ValueError, match="query_smiles is invalid"): - service.plan( - "similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")), - {"query_smiles": "invalid"}, 2, tmp_path / "workspace", - ) diff --git a/tests/test_sdk_compatibility.py b/tests/test_sdk_compatibility.py index a2e8718..0a29bb4 100644 --- a/tests/test_sdk_compatibility.py +++ b/tests/test_sdk_compatibility.py @@ -36,10 +36,9 @@ from scimesh.sdk import ( WorkloadDefinition, WorkloadRegistry, assert_manifest_round_trip, - default_sdk_registry, - default_sdk_runtime, - similarity_search_sdk_adapter, ) +from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime +from scimesh.workloads.search import similarity_search_sdk_definition from scimesh.workloads.similarity_search import search_similar, write_search_results @@ -60,8 +59,10 @@ def _registered_similarity_search(shard_rows: int = 2): registry = default_sdk_registry(shard_rows=shard_rows) runtime = default_sdk_runtime() descriptions = registry.descriptions() - assert len(descriptions) == 1 - description = descriptions[0] + assert len(descriptions) == 3 + description = next( + item for item in descriptions if item.workload.name == "similarity-search" + ) definition, negotiated = registry.require( description.workload.name, description.workload.version, @@ -128,7 +129,10 @@ def test_local_sdk_executor_matches_similarity_search_reference(tmp_path: Path) reference = search_similar(dataset, query, top_k=3, progress_every=0) write_search_results(reference_path, reference.matches) - assert artifact_store.materialize(result_artifact).read_bytes() == reference_path.read_bytes() + assert ( + artifact_store.materialize(result_artifact).read_bytes() + == reference_path.read_bytes() + ) assert result.task_key == "reduce/final" assert result.metrics == {"matches_emitted": 3, "partial_count": 3} @@ -187,7 +191,9 @@ def test_legacy_adapter_planning_is_deterministic_ordered_and_path_free( shard_ids: list[list[str]] = [] for task in first.tasks: artifact = task.inputs["input"].items[0].artifact - with artifact_store.materialize(artifact).open(encoding="utf-8", newline="") as source: + with artifact_store.materialize(artifact).open( + encoding="utf-8", newline="" + ) as source: shard_ids.append( [row["chembl_id"] for row in csv.DictReader(source, delimiter="\t")] ) @@ -250,7 +256,12 @@ def test_local_store_rejects_malformed_content_before_publishing( with pytest.raises(ValueError, match="not a valid bounded document"): store.import_file(malformed, declaration=declaration) - assert tuple(path for path in store.root.iterdir() if not path.name.startswith(".seal-")) == () + assert ( + tuple( + path for path in store.root.iterdir() if not path.name.startswith(".seal-") + ) + == () + ) def test_delimited_validator_rejects_headerless_data_and_enforces_record_limit( @@ -334,7 +345,11 @@ def test_local_executor_rejects_handler_forged_outputs( dataset = tmp_path / "molecules.tsv" _write_tiny_dataset(dataset) _, runtime, description, original, _ = _registered_similarity_search() - map_stage = next(stage for stage in original.manifest.workflow.stages if stage.kind is StageKind.MAP) + map_stage = next( + stage + for stage in original.manifest.workflow.stages + if stage.kind is StageKind.MAP + ) inner = original.runners[map_stage.entry_point] class ForgingRunner: @@ -372,7 +387,9 @@ def test_local_executor_rejects_handler_forged_outputs( ) -def test_local_executor_rejects_profiles_that_claim_network_isolation(tmp_path: Path) -> None: +def test_local_executor_rejects_profiles_that_claim_network_isolation( + tmp_path: Path, +) -> None: dataset = tmp_path / "molecules.tsv" _write_tiny_dataset(dataset) _, runtime, description, original, _ = _registered_similarity_search() @@ -409,7 +426,9 @@ def test_local_executor_rejects_aliased_terminal_outputs_before_planning( _write_tiny_dataset(dataset) _, runtime, description, original, _ = _registered_similarity_search() reducer = next( - stage for stage in original.manifest.workflow.stages if stage.kind is StageKind.REDUCE + stage + for stage in original.manifest.workflow.stages + if stage.kind is StageKind.REDUCE ) internal_name = next(iter(reducer.outputs)) workflow = replace( @@ -515,8 +534,7 @@ def _advanced_execution_manifest( elif case == "retries": features = ("retries",) changed = tuple( - replace(stage, retry=RetryPolicy(max_attempts=2)) - for stage in stages + replace(stage, retry=RetryPolicy(max_attempts=2)) for stage in stages ) elif case == "secrets": features = ("secret-injection",) @@ -604,7 +622,9 @@ def test_local_executor_rejects_profiles_it_cannot_enforce( ) -def test_local_executor_fails_when_the_declared_verifier_rejects(tmp_path: Path) -> None: +def test_local_executor_fails_when_the_declared_verifier_rejects( + tmp_path: Path, +) -> None: dataset = tmp_path / "molecules.tsv" _write_tiny_dataset(dataset) _, runtime, _, original, _ = _registered_similarity_search() @@ -639,19 +659,21 @@ def test_local_executor_fails_when_the_declared_verifier_rejects(tmp_path: Path) ) -def test_local_executor_enforces_the_declared_output_byte_budget(tmp_path: Path) -> None: +def test_local_executor_enforces_the_declared_output_byte_budget( + tmp_path: Path, +) -> None: dataset = tmp_path / "molecules.tsv" _write_tiny_dataset(dataset) runtime = default_sdk_runtime() - adapter = similarity_search_sdk_adapter(shard_rows=2) + workload = similarity_search_sdk_definition(shard_rows=2) # The planner pins its own manifest into every task, so the budget cut must - # be applied to the adapter's manifest for plan and definition to agree. - adapter.manifest = replace( - adapter.manifest, - workflow=replace(adapter.manifest.workflow, max_output_bytes=256), - limits=replace(adapter.manifest.limits, max_output_bytes=256), + # be applied to the workload's manifest for plan and definition to agree. + workload.manifest = replace( + workload.manifest, + workflow=replace(workload.manifest.workflow, max_output_bytes=256), + limits=replace(workload.manifest.limits, max_output_bytes=256), ) - definition = adapter.definition() + definition = workload.definition() registry = WorkloadRegistry() registry.register(definition, enabled=True) store = LocalArtifactStore(tmp_path / "artifacts") diff --git a/tests/test_sdk_descriptors.py b/tests/test_sdk_descriptors.py index fbba409..ebda0a3 100644 --- a/tests/test_sdk_descriptors.py +++ b/tests/test_sdk_descriptors.py @@ -25,13 +25,13 @@ from scimesh.sdk import ( VerifyContext, WorkloadRegistry, assert_manifest_round_trip, - default_sdk_runtime, ) -from scimesh.sdk.descriptors import ( +from scimesh.workloads.descriptors import ( DESCRIPTOR_COLUMNS, descriptor_batch_sdk_definition, compute_descriptor_batch, ) +from scimesh.workloads.library import default_sdk_runtime def _write_tiny_dataset(path: Path) -> None: @@ -428,7 +428,7 @@ def test_descriptor_batch_discovery_imports_an_allowlisted_installed_entry_point class EntryPoint: name = "descriptor-batch@1.0.0" dist = metadata.distribution("scimesh") - value = "scimesh.sdk.descriptors.definition:workload_definition" + value = "scimesh.workloads.descriptors:workload_definition" @property def module(self) -> str: diff --git a/tests/test_sdk_graph.py b/tests/test_sdk_graph.py new file mode 100644 index 0000000..daa4275 --- /dev/null +++ b/tests/test_sdk_graph.py @@ -0,0 +1,295 @@ +"""Tests for the SDK-built similarity-graph workload.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pytest + +from scimesh.sdk import ( + ArtifactCollection, + DeterminismProfile, + JobRequest, + LocalArtifactStore, + LocalCoreBatchExecutor, + LocalPlanningContext, + StageKind, + assert_manifest_round_trip, +) +from scimesh.workloads.graph import ( + check_pair_coverage, + merge_edge_partials, + similarity_graph_sdk_definition, +) +from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime +from scimesh.workloads.similarity_graph import ( + build_similarity_graph, + write_graph_edges, +) + + +def _write_tiny_dataset(path: Path) -> None: + path.write_text( + "chembl_id\tcanonical_smiles\n" + "A\tCCO\n" + "B\tCCCC\n" + "C\tCCN\n" + "D\tCCCCCC\n" + "E\tnot-a-smiles\n" + "F\tCCOCC\n" + "G\tc1ccccc1\n", + encoding="utf-8", + ) + + +def _registered_similarity_graph(): + registry = default_sdk_registry() + runtime = default_sdk_runtime() + description = next( + item + for item in registry.descriptions() + if item.workload.name == "similarity-graph" + ) + definition, negotiated = registry.require( + description.workload.name, + description.workload.version, + description.package_digest, + runtime=runtime, + ) + return registry, runtime, description, definition, negotiated + + +def _request_for( + dataset: Path, + artifact_store: LocalArtifactStore, + definition, + *, + threshold: float = 0.3, + threshold_direction: str = "greater", + block_size: int = 2, +) -> JobRequest: + input_port = definition.manifest.inputs["input"] + dataset_artifact = artifact_store.import_file( + dataset, + declaration=input_port.schema, + ) + return JobRequest( + workload=definition.manifest.workload, + parameters={ + "threshold": threshold, + "threshold_direction": threshold_direction, + "block_size": block_size, + }, + inputs={"input": ArtifactCollection.single(dataset_artifact)}, + ) + + +def test_similarity_graph_manifest_is_registered_and_negotiable() -> None: + _, runtime, description, definition, negotiated = _registered_similarity_graph() + manifest = definition.manifest + + assert description.enabled is True + assert manifest.workload.name == "similarity-graph" + assert manifest.workload.version == "1.0.0" + assert manifest.determinism is DeterminismProfile.BYTE_EXACT + assert manifest.verifier.verifier.canonical == "exact-artifact@1" + assert set(mode.value for mode in manifest.trust_modes) == { + "trusted", + "untrusted_quorum", + } + assert [stage.kind for stage in manifest.workflow.stages] == [ + StageKind.MAP, + StageKind.REDUCE, + ] + assert set(definition.runners) == {manifest.workflow.stages[0].entry_point} + assert set(definition.reducers) == {manifest.workflow.stages[1].entry_point} + assert negotiated is not None + assert_manifest_round_trip(manifest) + assert runtime is not None + + +@pytest.mark.parametrize("threshold_direction", ("greater", "less")) +def test_local_sdk_executor_matches_similarity_graph_reference( + tmp_path: Path, + threshold_direction: str, +) -> None: + dataset = tmp_path / "molecules.tsv" + _write_tiny_dataset(dataset) + registry, runtime, description, definition, _ = _registered_similarity_graph() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + request = _request_for( + dataset, + artifact_store, + definition, + threshold_direction=threshold_direction, + ) + + result = LocalCoreBatchExecutor( + registry, + runtime, + artifact_store, + tmp_path / "sdk-work", + ).execute(request, description.package_digest) + result_artifact = result.outputs["result"].items[0].artifact + + reference_path = tmp_path / "reference.csv" + reference = build_similarity_graph( + dataset, + threshold=0.3, + block_size=1_000, + threshold_direction=threshold_direction, + ) + write_graph_edges(reference_path, reference.edges) + + assert ( + artifact_store.materialize(result_artifact).read_bytes() + == reference_path.read_bytes() + ) + assert result.task_key == "reduce/final" + assert result.metrics["partial_count"] == 6 + assert result.metrics["edges_emitted"] == len(reference.edges) + + +def test_similarity_graph_result_is_invariant_to_block_size(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _write_tiny_dataset(dataset) + registry, runtime, _, definition, _ = _registered_similarity_graph() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + + outputs = [] + for block_size in (2, 3): + request = _request_for( + dataset, + artifact_store, + definition, + block_size=block_size, + ) + result = LocalCoreBatchExecutor( + registry, + runtime, + artifact_store, + tmp_path / f"sdk-work-{block_size}", + ).execute(request, definition.manifest.package.digest) + artifact = result.outputs["result"].items[0].artifact + outputs.append(artifact_store.materialize(artifact).read_bytes()) + assert result.metrics["partial_count"] == {2: 6, 3: 3}[block_size] + + assert outputs[0] == outputs[1] + + +def test_similarity_graph_planning_covers_each_block_pair_once( + tmp_path: Path, +) -> None: + dataset = tmp_path / "molecules.tsv" + _write_tiny_dataset(dataset) + registry, runtime, description, definition, _ = _registered_similarity_graph() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + request = _request_for(dataset, artifact_store, definition) + input_artifact = request.inputs["input"].items[0].artifact + + plan = registry.plan( + request, + description.package_digest, + runtime, + LocalPlanningContext( + artifact_store, + artifact_store, + tmp_path / "plan", + allowed_artifacts=(input_artifact,), + ), + ) + + assert [task.task_key for task in plan.tasks] == [ + "map/0000x0000", + "map/0000x0001", + "map/0000x0002", + "map/0001x0001", + "map/0001x0002", + "map/0002x0002", + ] + for task in plan.tasks: + assert set(task.inputs) == {"left", "right"} + assert task.inputs["left"].items[0].artifact is not None + pairs = { + (int(left), int(right)) + for task in plan.tasks + for left, right in (task.task_key[len("map/") :].split("x"),) + } + check_pair_coverage(tuple(sorted(pairs))) + + wire_payload = plan.to_json() + assert str(tmp_path) not in wire_payload + assert "file://" not in wire_payload + assert "worker://" not in wire_payload + assert "workspace" not in wire_payload + + +def test_similarity_graph_rejects_duplicate_molecule_ids(tmp_path: Path) -> None: + dataset = tmp_path / "duplicates.tsv" + dataset.write_text( + "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nA\tCCN\n", + encoding="utf-8", + ) + registry, runtime, _, definition, _ = _registered_similarity_graph() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + request = _request_for(dataset, artifact_store, definition) + + with pytest.raises(ValueError, match="duplicate chembl_id"): + LocalCoreBatchExecutor( + registry, + runtime, + artifact_store, + tmp_path / "work", + ).execute(request, definition.manifest.package.digest) + + +def test_similarity_graph_reducer_rejects_duplicate_unordered_pairs( + tmp_path: Path, +) -> None: + first = tmp_path / "first.csv" + first.write_text( + "source_id,target_id,similarity\nA,B,0.500000\nC,D,0.100000\n", + encoding="utf-8", + ) + second = tmp_path / "second.csv" + second.write_text( + "source_id,target_id,similarity\nB,A,0.500000\n", + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="duplicate unordered pair"): + merge_edge_partials((first, second), tmp_path / "result.csv") + + +def test_similarity_graph_pair_coverage_rejects_missing_block_pair() -> None: + with pytest.raises(ValueError, match="do not cover the full block pair set"): + check_pair_coverage(((0, 0), (0, 1))) + + with pytest.raises(ValueError, match="duplicate block pair"): + check_pair_coverage(((0, 0), (0, 0), (0, 1), (1, 1))) + + +def test_similarity_graph_merge_is_deterministically_sorted(tmp_path: Path) -> None: + first = tmp_path / "first.csv" + first.write_text( + "source_id,target_id,similarity\nC,A,0.200000\nB,C,0.400000\n", + encoding="utf-8", + ) + second = tmp_path / "second.csv" + second.write_text( + "source_id,target_id,similarity\nA,B,0.900000\n", + encoding="utf-8", + ) + result_path = tmp_path / "result.csv" + metrics = merge_edge_partials((first, second), result_path) + + assert metrics == {"partial_count": 2, "edges_emitted": 3} + rows = list(csv.DictReader(result_path.open(encoding="utf-8", newline=""))) + assert [ + (row["source_id"], row["target_id"], row["similarity"]) for row in rows + ] == [ + ("A", "B", "0.900000"), + ("B", "C", "0.400000"), + ("C", "A", "0.200000"), + ] diff --git a/tests/test_sdk_registry.py b/tests/test_sdk_registry.py index 2fd7577..eb4bf6a 100644 --- a/tests/test_sdk_registry.py +++ b/tests/test_sdk_registry.py @@ -25,20 +25,22 @@ from scimesh.sdk import ( WorkloadDefinition, WorkloadId, WorkloadRegistry, - current_scimesh_package_digest, - default_sdk_runtime, installed_distribution_digest, - similarity_search_sdk_adapter, ) from scimesh.sdk.schema import ( ParameterValidationError, validate_parameter_instance, validate_schema_definition, ) +from scimesh.workloads.environment import current_scimesh_package_digest +from scimesh.workloads.library import default_sdk_runtime +from scimesh.workloads.search import similarity_search_sdk_definition -def _definition(*, version: str = "1.0.0", digest_character: str = "a") -> WorkloadDefinition: - original = similarity_search_sdk_adapter(shard_rows=2).definition() +def _definition( + *, version: str = "1.0.0", digest_character: str = "a" +) -> WorkloadDefinition: + original = similarity_search_sdk_definition(shard_rows=2).definition() manifest = replace( original.manifest, workload=WorkloadId("similarity-search", version), @@ -72,13 +74,16 @@ def test_registry_requires_an_explicit_enabled_version_and_digest() -> None: registry.register(first) registry.enable("similarity-search", "2.0.0", "sha256:" + "b" * 64) - assert [item.workload.version for item in registry.descriptions()] == ["1.0.0", "2.0.0"] + assert [item.workload.version for item in registry.descriptions()] == [ + "1.0.0", + "2.0.0", + ] def test_compatibility_failure_occurs_before_planner_invocation( tmp_path: Path, ) -> None: - original = similarity_search_sdk_adapter(shard_rows=2).definition() + original = similarity_search_sdk_definition(shard_rows=2).definition() class CountingPlanner: calls = 0 @@ -140,7 +145,7 @@ def test_job_selected_features_and_trust_mode_fail_closed_before_planning( request_changes: dict[str, object], error_code: str, ) -> None: - definition = similarity_search_sdk_adapter(shard_rows=2).definition() + definition = similarity_search_sdk_definition(shard_rows=2).definition() registry = WorkloadRegistry() registry.register(definition, enabled=True) input_port = definition.manifest.inputs["input"] @@ -180,7 +185,7 @@ class _EntryPoints(tuple): def test_discovery_imports_only_an_exact_allowlisted_installed_entry_point( monkeypatch: pytest.MonkeyPatch, ) -> None: - definition = similarity_search_sdk_adapter().definition() + definition = similarity_search_sdk_definition().definition() loaded: list[str] = [] class EntryPoint: @@ -191,7 +196,7 @@ def test_discovery_imports_only_an_exact_allowlisted_installed_entry_point( if distribution == "scimesh" else SimpleNamespace(name=distribution) ) - self.value = "scimesh.sdk.builtins:similarity_search_sdk_adapter" + self.value = "scimesh.workloads.search:similarity_search_sdk_definition" @property def module(self) -> str: @@ -232,14 +237,14 @@ def test_discovery_imports_only_an_exact_allowlisted_installed_entry_point( def test_discovery_measures_package_before_importing_entry_point( monkeypatch: pytest.MonkeyPatch, ) -> None: - definition = similarity_search_sdk_adapter().definition() + definition = similarity_search_sdk_definition().definition() loaded = False class EntryPoint: name = "similarity-search@1.0.0" dist = metadata.distribution("scimesh") - value = "scimesh.sdk.builtins:similarity_search_sdk_adapter" - module = "scimesh.sdk.builtins" + value = "scimesh.workloads.search:similarity_search_sdk_definition" + module = "scimesh.workloads.search" def load(self): nonlocal loaded @@ -333,7 +338,7 @@ def test_discovery_rejects_entry_point_module_owned_by_another_distribution( "scimesh.sdk.registry.metadata.entry_points", lambda: _EntryPoints((EntryPoint(),)), ) - definition = similarity_search_sdk_adapter().definition() + definition = similarity_search_sdk_definition().definition() with pytest.raises(ValueError, match="outside its distribution"): WorkloadRegistry().discover_installed( ( @@ -463,14 +468,14 @@ def test_disabled_workload_is_not_resolvable_until_re_enabled() -> None: def test_discovery_rechecks_the_package_digest_after_loading( monkeypatch: pytest.MonkeyPatch, ) -> None: - definition = similarity_search_sdk_adapter().definition() + definition = similarity_search_sdk_definition().definition() digests = iter((definition.manifest.package.digest, "sha256:" + "e" * 64)) class EntryPoint: name = "similarity-search@1.0.0" dist = metadata.distribution("scimesh") - value = "scimesh.sdk.builtins:similarity_search_sdk_adapter" - module = "scimesh.sdk.builtins" + value = "scimesh.workloads.search:similarity_search_sdk_definition" + module = "scimesh.workloads.search" def load(self): return lambda: definition @@ -498,10 +503,17 @@ def test_discovery_rechecks_the_package_digest_after_loading( assert registry.descriptions() == () -def test_request_trust_mode_must_be_enforceable_by_runtime_and_stages(tmp_path: Path) -> None: - original = similarity_search_sdk_adapter(shard_rows=2).definition() +def test_request_trust_mode_must_be_enforceable_by_runtime_and_stages( + tmp_path: Path, +) -> None: + original = similarity_search_sdk_definition(shard_rows=2).definition() + stages = tuple( + replace(stage, trust_modes=("trusted",)) + for stage in original.manifest.workflow.stages + ) manifest = replace( original.manifest, + workflow=replace(original.manifest.workflow, stages=stages), trust_modes=(TrustMode.TRUSTED, TrustMode.VERIFIED), ) definition = WorkloadDefinition( @@ -554,7 +566,7 @@ def test_request_trust_mode_must_be_enforceable_by_runtime_and_stages(tmp_path: def test_job_cannot_require_a_feature_outside_the_runtime(tmp_path: Path) -> None: - original = similarity_search_sdk_adapter(shard_rows=2).definition() + original = similarity_search_sdk_definition(shard_rows=2).definition() manifest = replace( original.manifest, optional_features=( diff --git a/tests/test_sdk_search.py b/tests/test_sdk_search.py new file mode 100644 index 0000000..d4953cb --- /dev/null +++ b/tests/test_sdk_search.py @@ -0,0 +1,252 @@ +"""Tests for the SDK-built similarity-search workload.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pytest + +from scimesh.sdk import ( + ArtifactCollection, + DeterminismProfile, + JobRequest, + LocalArtifactStore, + LocalCoreBatchExecutor, + LocalPlanningContext, + StageKind, + WorkloadRegistry, + assert_manifest_round_trip, +) +from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime +from scimesh.workloads.search import similarity_search_sdk_definition +from scimesh.workloads.similarity_search import ( + find_molecule_by_id, + search_similar, + write_search_results, +) + + +def _write_tiny_dataset(path: Path) -> None: + path.write_text( + "chembl_id\tcanonical_smiles\textra\n" + "QUERY\tCCO\tquery\n" + "ALCOHOL\tCCCO\talcohol\n" + "ALKANE\tCCCC\talkane\n" + "BROKEN\tnot-a-smiles\tinvalid\n" + "DUPLICATE\tCCO\tduplicate\n" + "AMINE\tCCN\tamine\n", + encoding="utf-8", + ) + + +def _registered_similarity_search(shard_rows: int = 2): + registry = default_sdk_registry(shard_rows=shard_rows) + runtime = default_sdk_runtime() + description = next( + item + for item in registry.descriptions() + if item.workload.name == "similarity-search" + ) + definition, negotiated = registry.require( + description.workload.name, + description.workload.version, + description.package_digest, + runtime=runtime, + ) + return registry, runtime, description, definition, negotiated + + +def _request_for( + dataset: Path, + artifact_store: LocalArtifactStore, + definition, +) -> JobRequest: + input_port = definition.manifest.inputs["input"] + dataset_artifact = artifact_store.import_file( + dataset, + declaration=input_port.schema, + ) + return JobRequest( + workload=definition.manifest.workload, + parameters={"query_id": "QUERY", "top_k": 3, "progress_every": 0}, + inputs={"input": ArtifactCollection.single(dataset_artifact)}, + ) + + +def test_similarity_search_manifest_is_registered_and_negotiable() -> None: + _, runtime, description, definition, negotiated = _registered_similarity_search() + manifest = definition.manifest + + assert description.enabled is True + assert manifest.workload.name == "similarity-search" + assert manifest.workload.version == "1.0.0" + assert manifest.determinism is DeterminismProfile.BYTE_EXACT + assert manifest.verifier.verifier.canonical == "exact-artifact@1" + assert set(mode.value for mode in manifest.trust_modes) == { + "trusted", + "untrusted_quorum", + } + assert [stage.kind for stage in manifest.workflow.stages] == [ + StageKind.MAP, + StageKind.REDUCE, + ] + assert set(definition.runners) == {manifest.workflow.stages[0].entry_point} + assert set(definition.reducers) == {manifest.workflow.stages[1].entry_point} + assert negotiated is not None + assert negotiated.manifest == manifest + assert_manifest_round_trip(manifest) + assert runtime is not None + + +def test_local_sdk_executor_matches_similarity_search_reference(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _write_tiny_dataset(dataset) + registry, runtime, description, definition, _ = _registered_similarity_search() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + request = _request_for(dataset, artifact_store, definition) + + result = LocalCoreBatchExecutor( + registry, + runtime, + artifact_store, + tmp_path / "sdk-work", + ).execute(request, description.package_digest) + result_artifact = result.outputs["result"].items[0].artifact + + reference_path = tmp_path / "reference.csv" + query = find_molecule_by_id(dataset, "QUERY") + reference = search_similar(dataset, query, top_k=3, progress_every=0) + write_search_results(reference_path, reference.matches) + + assert ( + artifact_store.materialize(result_artifact).read_bytes() + == reference_path.read_bytes() + ) + assert result.task_key == "reduce/final" + assert dict(result.metrics) == {"matches_emitted": 3, "partial_count": 3} + + +def test_similarity_search_planning_is_deterministic_ordered_and_path_free( + tmp_path: Path, +) -> None: + dataset = tmp_path / "molecules.tsv" + _write_tiny_dataset(dataset) + registry, runtime, description, definition, _ = _registered_similarity_search() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + request = _request_for(dataset, artifact_store, definition) + input_artifact = request.inputs["input"].items[0].artifact + + first = registry.plan( + request, + description.package_digest, + runtime, + LocalPlanningContext( + artifact_store, + artifact_store, + tmp_path / "first-plan", + allowed_artifacts=(input_artifact,), + ), + ) + second = registry.plan( + request, + description.package_digest, + runtime, + LocalPlanningContext( + artifact_store, + artifact_store, + tmp_path / "second-plan", + allowed_artifacts=(input_artifact,), + ), + ) + + assert first.to_json() == second.to_json() + assert first.digest == second.digest + assert first.package_digest == definition.manifest.package.digest + assert first.manifest_digest == definition.manifest.digest + assert [task.task_key for task in first.tasks] == [ + "map/00000000", + "map/00000001", + "map/00000002", + ] + assert all(task.stage_id == "map" for task in first.tasks) + assert all("query_id" not in task.parameters for task in first.tasks) + assert all(task.parameters["query_smiles"] == "CCO" for task in first.tasks) + assert all(task.parameters["top_k"] == 3 for task in first.tasks) + assert all(task.package_digest == first.package_digest for task in first.tasks) + assert all(task.manifest_digest == first.manifest_digest for task in first.tasks) + assert first.resolved_parameters["query_source"] == { + "kind": "chembl_id", + "value": "QUERY", + } + + shard_ids: list[list[str]] = [] + for task in first.tasks: + artifact = task.inputs["input"].items[0].artifact + with artifact_store.materialize(artifact).open( + encoding="utf-8", newline="" + ) as source: + shard_ids.append( + [row["chembl_id"] for row in csv.DictReader(source, delimiter="\t")] + ) + assert shard_ids == [ + ["QUERY", "ALCOHOL"], + ["ALKANE", "BROKEN"], + ["DUPLICATE", "AMINE"], + ] + + wire_payload = first.to_json() + assert str(tmp_path) not in wire_payload + assert "file://" not in wire_payload + assert "worker://" not in wire_payload + assert "workspace" not in wire_payload + + +def test_similarity_search_rejects_ambiguous_or_mistyped_parameters( + tmp_path: Path, +) -> None: + dataset = tmp_path / "molecules.tsv" + _write_tiny_dataset(dataset) + registry, runtime, _, definition, _ = _registered_similarity_search() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + input_artifact = artifact_store.import_file( + dataset, + declaration=definition.manifest.inputs["input"].schema, + ) + base = JobRequest( + workload=definition.manifest.workload, + parameters={"query_id": "QUERY", "top_k": 3}, + inputs={"input": ArtifactCollection.single(input_artifact)}, + ) + + for bad_parameters, message in ( + ({"query_id": "QUERY", "query_smiles": "CCO"}, "oneOf did not match"), + ({"query_id": "QUERY", "top_k": 0}, "violates minimum"), + ): + request = JobRequest( + workload=base.workload, + parameters=bad_parameters, + inputs=base.inputs, + ) + with pytest.raises(ValueError, match=message): + registry.plan( + request, + definition.manifest.package.digest, + runtime, + LocalPlanningContext( + artifact_store, + artifact_store, + tmp_path / "bad-plan", + allowed_artifacts=(input_artifact,), + ), + ) + + +def test_similarity_search_workload_definition_is_discoverable() -> None: + from scimesh.sdk import WorkloadDefinition + from scimesh.workloads.search import workload_definition + + definition = workload_definition() + assert isinstance(definition, WorkloadDefinition) + assert definition.manifest.workload.name == "similarity-search" + assert definition.manifest.workload.version == "1.0.0" diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index 5f77098..3c35a8b 100644 --- a/tests/test_worker_daemon.py +++ b/tests/test_worker_daemon.py @@ -23,7 +23,11 @@ from scimesh.worker.models import ( RunResult, UploadedArtifact, ) -from scimesh.worker.artifacts import HttpArtifactClient, _SameOriginAuthRedirectHandler, _origin +from scimesh.worker.artifacts import ( + HttpArtifactClient, + _SameOriginAuthRedirectHandler, + _origin, +) from scimesh.worker.runners import SciMeshRunner from scimesh.worker.transport import NoRedirectHandler @@ -32,12 +36,18 @@ class FakeCoordinator: def __init__(self, task: ClaimedTask | None) -> None: self.task, self.submissions, self.failures, self.heartbeats = task, [], [], [] - def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + def claim( + self, worker_id: str, capabilities: tuple[str, ...] + ) -> ClaimedTask | None: task, self.task = self.task, None return task def register( - self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None + self, + name: str, + capabilities: tuple[str, ...], + cpu_count: int, + memory_mb: int | None, ) -> RegisteredWorker: return RegisteredWorker("11111111-1111-4111-8111-111111111111", 15) @@ -71,6 +81,7 @@ class FakeArtifacts: len(content), ) + class FakeRunner: def __init__(self) -> None: self.calls = 0 @@ -84,18 +95,40 @@ class FakeRunner: def make_task(content: bytes, checksum: str | None = None) -> ClaimedTask: lease = (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat() - return ClaimedTask("task-1", 1, lease, "similarity-search", InputArtifact("https://example.test/input", checksum or hashlib.sha256(content).hexdigest()), {"query_id": "CHEMBL1"}) + return ClaimedTask( + "task-1", + 1, + lease, + "similarity-search", + InputArtifact( + "https://example.test/input", + checksum or hashlib.sha256(content).hexdigest(), + ), + {"query_id": "CHEMBL1"}, + ) def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes): - coordinator, artifacts, runner = FakeCoordinator(task), FakeArtifacts(content), FakeRunner() + coordinator, artifacts, runner = ( + FakeCoordinator(task), + FakeArtifacts(content), + FakeRunner(), + ) config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work") - return WorkerDaemon(config, coordinator, artifacts, runner), coordinator, artifacts, runner, config + return ( + WorkerDaemon(config, coordinator, artifacts, runner), + coordinator, + artifacts, + runner, + config, + ) def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None: content = b"input fixture" - worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content) + worker, coordinator, artifacts, runner, _ = daemon( + tmp_path, make_task(content), content + ) assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True) assert runner.calls == 1 assert len(artifacts.uploaded) == 1 @@ -107,10 +140,16 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None: def test_worker_executes_a_resolved_similarity_search_shard(tmp_path: Path) -> None: - content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\nINVALID\tnot-a-smiles\n" + content = ( + b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\nINVALID\tnot-a-smiles\n" + ) task = make_task(content) task = ClaimedTask( - task.task_id, task.attempt, task.lease_expires_at, task.workload, task.input, + task.task_id, + task.attempt, + task.lease_expires_at, + task.workload, + task.input, {"query_smiles": "CCO", "top_k": 5, "progress_every": 0}, ) worker, coordinator, artifacts, _, _ = daemon(tmp_path, task, content) @@ -131,11 +170,19 @@ def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n" first = make_task(content) first = ClaimedTask( - "retry-task", 1, first.lease_expires_at, "similarity-search", first.input, + "retry-task", + 1, + first.lease_expires_at, + "similarity-search", + first.input, {"query_smiles": "CCO", "top_k": 5}, ) second = ClaimedTask( - "other-task", 1, first.lease_expires_at, "similarity-search", first.input, + "other-task", + 1, + first.lease_expires_at, + "similarity-search", + first.input, {"query_smiles": "CCO", "top_k": 5}, ) @@ -145,16 +192,26 @@ def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> self.queue = [first, second] self.claimants: list[str] = [] - def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + def claim( + self, worker_id: str, capabilities: tuple[str, ...] + ) -> ClaimedTask | None: self.claimants.append(worker_id) return self.queue.pop(0) if self.queue else None def fail(self, task: ClaimedTask, payload: dict) -> None: self.failures.append(payload) - if task.task_id == "retry-task" and task.attempt == 1 and payload["retryable"]: + if ( + task.task_id == "retry-task" + and task.attempt == 1 + and payload["retryable"] + ): self.queue.append( ClaimedTask( - task.task_id, 2, task.lease_expires_at, task.workload, task.input, + task.task_id, + 2, + task.lease_expires_at, + task.workload, + task.input, task.parameters, ) ) @@ -174,11 +231,15 @@ def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> artifacts = FakeArtifacts(content) worker_a = WorkerDaemon( WorkerConfig("https://example.test", "worker-a", tmp_path / "worker-a"), - coordinator, artifacts, FailFirstAttempt(), + coordinator, + artifacts, + FailFirstAttempt(), ) worker_b = WorkerDaemon( WorkerConfig("https://example.test", "worker-b", tmp_path / "worker-b"), - coordinator, artifacts, SciMeshRunner(), + coordinator, + artifacts, + SciMeshRunner(), ) assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=False) @@ -197,16 +258,22 @@ def test_no_task_does_not_create_directory(tmp_path: Path) -> None: assert not config.work_dir.exists() -def test_once_worker_exits_after_an_empty_claim(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: +def test_once_worker_exits_after_an_empty_claim( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: caplog.set_level(logging.INFO, logger="scimesh.worker") worker, _, _, runner, _ = daemon(tmp_path, None, b"") - worker.config = WorkerConfig(**{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1}) + worker.config = WorkerConfig( + **{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1} + ) assert worker.run_forever() is True assert runner.calls == 0 assert "queue_empty" in caplog.text -def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: +def test_worker_stops_after_the_configured_number_of_claims( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: caplog.set_level(logging.INFO, logger="scimesh.worker") content = b"input fixture" worker, _, _, runner, _ = daemon(tmp_path, make_task(content), content) @@ -216,10 +283,15 @@ def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, capl assert "max_tasks_reached" in caplog.text -def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: +def test_keyboard_interrupt_stops_worker_without_propagating( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: caplog.set_level(logging.INFO, logger="scimesh.worker") + class InterruptingCoordinator(FakeCoordinator): - def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + def claim( + self, worker_id: str, capabilities: tuple[str, ...] + ) -> ClaimedTask | None: raise KeyboardInterrupt worker, _, _, _, _ = daemon(tmp_path, None, b"") @@ -228,7 +300,9 @@ def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, cap assert "interrupted" in caplog.text -def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path) -> None: +def test_interrupting_an_active_task_reports_a_sanitized_failure( + tmp_path: Path, +) -> None: content = b"input fixture" worker, coordinator, _, _, _ = daemon(tmp_path, make_task(content), content) @@ -271,12 +345,16 @@ def test_max_tasks_counts_successes_not_failed_claims(tmp_path: Path) -> None: ), ] - def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + def claim( + self, worker_id: str, capabilities: tuple[str, ...] + ) -> ClaimedTask | None: return self.tasks.pop(0) if self.tasks else None coordinator = SequencedCoordinator() artifacts, runner = FakeArtifacts(successful_content), FakeRunner() - config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work", max_tasks=1) + config = WorkerConfig( + "https://example.test", "worker-1", tmp_path / "work", max_tasks=1 + ) worker = WorkerDaemon(config, coordinator, artifacts, runner) assert worker.run_forever() is True assert len(coordinator.failures) == 1 @@ -303,9 +381,12 @@ def test_worker_cli_uses_a_nonzero_exit_code_for_interruption( return False monkeypatch.setattr(worker_cli, "WorkerDaemon", InterruptedDaemon) - assert worker_cli.main( - ["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)] - ) == 130 + assert ( + worker_cli.main( + ["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)] + ) + == 130 + ) @pytest.mark.parametrize("value", [0, -1, True]) @@ -315,7 +396,9 @@ def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None: def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None: - worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual") + worker, coordinator, _, runner, _ = daemon( + tmp_path, make_task(b"actual", "not-the-hash"), b"actual" + ) assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False) assert runner.calls == 0 assert coordinator.failures[0]["error_code"] == "ValueError" @@ -323,7 +406,9 @@ def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None: assert not coordinator.submissions -def test_failure_reporting_removes_paths_outside_the_worker_directory(tmp_path: Path) -> None: +def test_failure_reporting_removes_paths_outside_the_worker_directory( + tmp_path: Path, +) -> None: worker, coordinator, _, _, _ = daemon(tmp_path, make_task(b"input"), b"input") error = subprocess.CalledProcessError( 1, @@ -344,9 +429,13 @@ def test_directory_creation_failure_is_reported(tmp_path: Path) -> None: assert coordinator.failures[0]["error_code"] == "FileExistsError" -def test_transient_claim_error_is_propagated_for_bounded_backoff(tmp_path: Path) -> None: +def test_transient_claim_error_is_propagated_for_bounded_backoff( + tmp_path: Path, +) -> None: class UnavailableCoordinator(FakeCoordinator): - def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + def claim( + self, worker_id: str, capabilities: tuple[str, ...] + ) -> ClaimedTask | None: raise CoordinatorTransientError("temporary outage") worker, _, _, _, _ = daemon(tmp_path, None, b"") @@ -368,7 +457,9 @@ def test_task_directories_are_retained_until_cleanup_is_enabled(tmp_path: Path) def test_input_token_is_sent_only_to_the_coordinator_origin() -> None: client = HttpArtifactClient("https://coordinator.example/api", 10, "secret") - assert client._auth_headers_for("https://coordinator.example/tasks/1/input") == {"Authorization": "Bearer secret"} + assert client._auth_headers_for("https://coordinator.example/tasks/1/input") == { + "Authorization": "Bearer secret" + } assert client._auth_headers_for("https://bucket.example/presigned") == {} @@ -384,17 +475,28 @@ def test_relative_input_uri_is_resolved_against_the_coordinator() -> None: def test_redirect_to_external_storage_strips_authorization() -> None: handler = _SameOriginAuthRedirectHandler(_origin("https://coordinator.example")) source = Request( - "https://coordinator.example/tasks/1/input", headers={"Authorization": "Bearer secret"} + "https://coordinator.example/tasks/1/input", + headers={"Authorization": "Bearer secret"}, + ) + redirected = handler.redirect_request( + source, None, 302, "Found", {}, "https://bucket.example/presigned" ) - redirected = handler.redirect_request(source, None, 302, "Found", {}, "https://bucket.example/presigned") assert redirected is not None assert redirected.get_header("Authorization") is None def test_api_requests_never_follow_redirects() -> None: handler = NoRedirectHandler() - request = Request("https://coordinator.example/tasks/claim", headers={"Authorization": "Bearer secret"}) - assert handler.redirect_request(request, None, 302, "Found", {}, "https://other.example") is None + request = Request( + "https://coordinator.example/tasks/claim", + headers={"Authorization": "Bearer secret"}, + ) + assert ( + handler.redirect_request( + request, None, 302, "Found", {}, "https://other.example" + ) + is None + ) def test_lease_is_renewed_while_a_runner_is_still_working(tmp_path: Path) -> None: @@ -429,44 +531,101 @@ def test_heartbeat_reschedules_from_the_renewed_lease(tmp_path: Path) -> None: assert len(coordinator.heartbeats) >= 3 -def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - commands: list[list[str]] = [] - - def fake_run(command: list[str], **_: object) -> None: - commands.append(command) - output = Path(command[command.index("--output") + 1]) - output.parent.mkdir(parents=True, exist_ok=True) - output.write_text("a,b\n", encoding="utf-8") - - monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run) +def test_runner_executes_search_through_the_sdk_and_rejects_graph( + tmp_path: Path, +) -> None: runner = SciMeshRunner() - graph = ClaimedTask("graph", 1, "2026-07-30T00:00:00Z", "similarity-graph", InputArtifact("https://example/input", "x"), {"threshold": 0.2, "threshold_direction": "less", "block_size": 42, "max_rows": 7, "progress_every": 0}) - search = ClaimedTask("search", 1, "2026-07-30T00:00:00Z", "similarity-search", InputArtifact("https://example/input", "x"), {"query_smiles": "CCO", "top_k": 3}) + graph = ClaimedTask( + "graph", + 1, + "2026-07-30T00:00:00Z", + "similarity-graph", + InputArtifact("https://example/input", "x"), + { + "threshold": 0.2, + "threshold_direction": "less", + "block_size": 42, + "max_rows": 7, + "progress_every": 0, + }, + ) + search = ClaimedTask( + "search", + 1, + "2026-07-30T00:00:00Z", + "similarity-search", + InputArtifact("https://example/input", "x"), + {"query_smiles": "CCO", "top_k": 3}, + ) search_dir = tmp_path / "search" search_dir.mkdir() (search_dir / "input").write_text( "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8" ) - runner.run(graph, tmp_path / "graph") + with pytest.raises(ValueError, match="unsupported workload"): + runner.run(graph, tmp_path / "graph") result = runner.run(search, search_dir) - assert "--threshold-direction" in commands[0] and "less" in commands[0] - assert "--block-size" in commands[0] and "42" in commands[0] - assert "--max-rows" in commands[0] and "7" in commands[0] - assert len(commands) == 1 assert result.metrics == { - "scanned_rows": 2, "valid_molecules": 2, "invalid_smiles": 0, "matches_emitted": 1, + "scanned_rows": 2, + "valid_molecules": 2, + "invalid_smiles": 0, + "matches_emitted": 1, } + assert result.artifacts[0].content_type == "text/csv" + assert ( + result.artifacts[0] + .path.read_text(encoding="utf-8") + .startswith("rank,chembl_id,canonical_smiles,similarity\n") + ) -def test_runner_accepts_coordinator_workload_names(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_options( + tmp_path: Path, +) -> None: + task_dir = tmp_path / "search" + task_dir.mkdir() + (task_dir / "input").write_text( + "chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n", encoding="utf-8" + ) + task = ClaimedTask( + "search", + 1, + "2026-07-30T00:00:00Z", + "similarity-search", + InputArtifact("https://example/input", "a" * 64), + {"query_id": "QUERY", "top_k": 5}, + ) + result = SciMeshRunner().run(task, task_dir) + assert result.metrics["matches_emitted"] == 1 + assert (task_dir / "result.csv").is_file() + + with_max_rows = ClaimedTask( + "search", + 1, + "2026-07-30T00:00:00Z", + "similarity-search", + InputArtifact("https://example/input", "a" * 64), + {"query_smiles": "CCO", "max_rows": 1}, + ) + with pytest.raises(ValueError, match="unsupported runner parameters"): + SciMeshRunner().run(with_max_rows, task_dir) + + +def test_runner_accepts_coordinator_workload_names( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: task_dir = tmp_path / "search" task_dir.mkdir() (task_dir / "input").write_text( "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8" ) task = ClaimedTask( - "search", 1, "2026-07-30T00:00:00Z", "similarity_search", - InputArtifact("https://example/input", "a" * 64), {"query_smiles": "CCO"}, + "search", + 1, + "2026-07-30T00:00:00Z", + "similarity_search", + InputArtifact("https://example/input", "a" * 64), + {"query_smiles": "CCO"}, ) result = SciMeshRunner().run(task, task_dir) assert result.metrics["matches_emitted"] == 1 @@ -502,7 +661,10 @@ def test_claimed_task_accepts_a_coordinator_relative_input_path() -> None: "attempt": 1, "lease_expires_at": "2026-07-30T00:00:00Z", "workload": "similarity_search", - "input": {"uri": "/tasks/11111111-1111-4111-8111-111111111111/input", "sha256": "a" * 64}, + "input": { + "uri": "/tasks/11111111-1111-4111-8111-111111111111/input", + "sha256": "a" * 64, + }, "parameters": {}, } ) @@ -523,7 +685,9 @@ def test_uploaded_artifact_requires_complete_durable_metadata() -> None: UploadedArtifact.from_json({"artifact_id": "missing"}) -def test_environment_overrides_allow_cli_only_configuration(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: +def test_environment_overrides_allow_cli_only_configuration( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.delenv("SCIMESH_COORDINATOR_URL", raising=False) config = WorkerConfig.from_environment( { @@ -552,8 +716,12 @@ def test_relative_work_dir_is_normalized_for_runner_subprocesses( "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8" ) task = ClaimedTask( - "task", 1, "2026-07-30T00:00:00Z", "similarity-search", - InputArtifact("https://example.test/input", "a" * 64), {"query_smiles": "CCO"}, + "task", + 1, + "2026-07-30T00:00:00Z", + "similarity-search", + InputArtifact("https://example.test/input", "a" * 64), + {"query_smiles": "CCO"}, ) SciMeshRunner().run(task, task_dir) assert (task_dir / "result.csv").is_file()