Replace legacy distributed protocol with SDK-built workloads
This commit is contained in:
Emil
@@ -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).
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
+43
-13
@@ -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
|
||||
|
||||
+72
-21
@@ -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.
|
||||
|
||||
+3
-2
@@ -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*"]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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."""
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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(),),
|
||||
),
|
||||
)
|
||||
@@ -1,5 +0,0 @@
|
||||
"""Adapters for versioned pre-SDK SciMesh workload contracts."""
|
||||
|
||||
from .distributed_v1 import LegacyDistributedWorkloadAdapter
|
||||
|
||||
__all__ = ["LegacyDistributedWorkloadAdapter"]
|
||||
@@ -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.<eight-digit-index>")
|
||||
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)
|
||||
+243
-86
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
+15
-15
@@ -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.
|
||||
"""
|
||||
@@ -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()
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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.<i>x<j>`` partial keys into block indices."""
|
||||
prefix = "map."
|
||||
if not key.startswith(prefix):
|
||||
raise ValueError("graph partial key must use map.<left>x<right>")
|
||||
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.<left>x<right>")
|
||||
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)}
|
||||
@@ -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()
|
||||
@@ -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(),),
|
||||
),
|
||||
)
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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)}
|
||||
@@ -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.<eight-digit-index>")
|
||||
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()
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"),
|
||||
]
|
||||
+32
-20
@@ -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=(
|
||||
|
||||
@@ -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"
|
||||
+229
-61
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user