diff --git a/.opencode/session-goal.md b/.opencode/session-goal.md new file mode 100644 index 0000000..e374e3d --- /dev/null +++ b/.opencode/session-goal.md @@ -0,0 +1,26 @@ +COMPLETED +# Session Goal + +адаптируй код под sdk, где это необходимо. при надобности доработай SDK. главное чтобы workloads можно было дописывать не трогая остальной код программы, так как он в будущем будет закрытым. Workloads - в первую очередь пользовательские скрипты, поэтому sdk должен полностью покрывать необходимый функционал. + +## Plan + +1. SDK: добавить высокоуровневый каркас MapReduceWorkload (scimesh/sdk/batch.py) — манифест/стейджи/definition собираются автоматически, планировщик/раннер/редуктор — общий скелет с хуками (partition_input, compute_shard, parse/validate_partial_keys, reduce_partials, domain_validate). +2. Рефакторинг: descriptor-batch, similarity-search, similarity-graph переписать на базовый класс (поведение/байты не меняются — парность покрыта тестами). +3. Worker (закрываемый код): обобщить SciMeshRunner — загрузка ворклоадов из конфига/дискавери (allowlist), инвентарь из конфига воркера, fail-closed для неподдерживаемых форм; конфиг: SCIMESH_CAPABILITIES, SCIMESH_WORKLOAD_ALLOWLIST. +4. CLI: добавить общий `scimesh workload list|run` (generic SDK-инструмент, без workload-специфичной логики) — пользовательские скрипты можно запускать локально без правки остального кода. +5. Тесты: test_sdk_batch.py (каркас + хуки + fail-closed), тесты воркера на не-search ворклоаде, CLI-тесты; регрессия парности. +6. Документация: workload-sdk.md (авторский гайд на базе MapReduceWorkload), handoff, STATUS. +7. Полный прогон pytest, финальная верификация. + +## Progress + +- [x] `scimesh/sdk/batch.py`: `MapReduceWorkload` — identity/parameters/ports + 3 научных хука; сборка манифеста, map/reduce стейджей, workflow, pinned handlers, exact-artifact verifier; хуки: domain_validate, resolved_parameters, resolved_parameters_for_plan, plan_tasks, parse/validate_partial_keys, map_stage_inputs; экспортирован из scimesh.sdk. +- [x] descriptor-batch, similarity-search, similarity-graph переписаны на MapReduceWorkload; парность с локальными reference сохранена (тесты byte-identical зелёные). `query_id`-резолюция переехала в run_search_shard (ворклоад сам валидирует параметры). +- [x] Worker обобщён: SciMeshRunner принимает definitions+inventory+runtime, `for_worker(config)` грузит ворклоады через allowlist-дискавери (entry points) или built-in fallback; fail-closed для map-стейджей не по v1-контракту (single input); параметры таски проходят насквозь, валидация в ворклоаде. Конфиг: SCIMESH_CAPABILITIES, SCIMESH_WORKLOAD_ALLOWLIST (JSON {distribution,name,version,digest}); парсер вынесен в SDK (`workload_allowlist_from_json`). +- [x] CLI: `scimesh workload list|run` (generic; SCIMESH_WORKLOAD_ALLOWLIST поддерживается; runtime строится из discovered-ворклоадов); зарегистрирован как ворклоад-модуль. +- [x] `default_sdk_registry(allowlist=...)` и `default_sdk_runtime(workload_capabilities=..., environment_digests=...)` в library. +- [x] Тесты: test_sdk_batch.py (5), test_cli_workload.py (6, включая end-to-end allowlisted custom workload), worker: generic execution (descriptor-batch), v1-contract rejection (graph), for_worker discovery, config parsing. +- [x] Документация: workload-sdk.md (раздел "Authoring a workload" + worker/CLI), handoff, STATUS, README. +- [x] Финальная верификация: 249 passed; scimesh workload list/run работают; scimesh.sdk не импортирует workloads (grep чист). +- Изменения НЕ закоммичены (по AGENTS.md коммит только по явной просьбе). diff --git a/README.md b/README.md index f10d3a3..7a3fbbd 100644 --- a/README.md +++ b/README.md @@ -145,7 +145,9 @@ workload code. Workloads are user scripts built on the SDK: the built-in 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 +handlers run locally, in conformance, and on claimed coordinator tasks. +`scimesh workload list` and `scimesh workload run` run any SDK workload from +the command line. See the [SDK author guide](docs/workload-sdk.md), [contract](docs/scimesh-sdk-contract.md), and [delivery roadmap](docs/scimesh-sdk-roadmap.md). diff --git a/STATUS.md b/STATUS.md index cab1aa1..38a3bab 100644 --- a/STATUS.md +++ b/STATUS.md @@ -51,7 +51,8 @@ the complete result-artifact SHA-256 before a task is accepted. | 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 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. | +| SDK-built `similarity-search` and `similarity-graph` | Implemented | Both workloads are SDK-built packages (`scimesh/workloads/search/`, `scimesh/workloads/graph/`) built on the `MapReduceWorkload` authoring scaffold (`scimesh/sdk/batch.py`); 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. | +| SDK authoring scaffold | Implemented | `MapReduceWorkload` (exported from `scimesh.sdk`) assembles manifest, map/reduce stages, workflow, and digest-pinned handlers from three scientific hooks (partition/compute/merge), with overridable hooks for domain validation, plan-time resolution, custom task planning, and partial-key policy. The generic `scimesh workload list|run` CLI and the worker's allowlist-driven loading (`SCIMESH_WORKLOAD_ALLOWLIST`, `SCIMESH_CAPABILITIES`) let new workloads run without touching other code. | ## Next recommended assignment @@ -63,9 +64,10 @@ block-pair planning and reduction for `similarity-graph`. - 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 workload-generic v1-wire bridge over `TaskSpec`/`LocalTaskContext`); + `query_id` resolution and parameter validation live in the workload itself. + `max_rows` is a plan-time option and is rejected per task by the stage + projection. - A real-stack worker test uses a small `query_smiles` shard. The Python planner resolves `query_id` once and shares `query_smiles`; the upload UI currently accepts `query_smiles` only. diff --git a/docs/sdk-handoff.md b/docs/sdk-handoff.md index 6b9efee..6d91b22 100644 --- a/docs/sdk-handoff.md +++ b/docs/sdk-handoff.md @@ -24,6 +24,24 @@ 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`. +**Authoring scaffold (2026-08-01):** `scimesh/sdk/batch.py` adds +`MapReduceWorkload` — the primary authoring surface for `core-batch-v1`. A +subclass declares identity/parameters/ports and three scientific hooks +(`partition_input`, `compute_shard`, `reduce_partials`); the SDK assembles the +manifest, map/reduce stages, workflow, digest-pinned handlers, and the +exact-artifact verifier. Overridable hooks: `domain_validate`, +`resolved_parameters`, `resolved_parameters_for_plan`, `plan_tasks`, +`parse_partial_key`/`validate_partial_keys`, `map_stage_inputs` (multi-input +map stages share the external input schema). All three built-in workloads are +refactored onto it. Generic `scimesh workload list|run` CLI added (no +workload-specific logic). The worker loads workloads generically: +`SCIMESH_WORKLOAD_ALLOWLIST` (JSON `{distribution, name, version, digest}`, +discovery via entry points) or built-in fallback; `SCIMESH_CAPABILITIES` +overrides advertised capabilities; workloads with multi-input map stages are +rejected by the v1 bridge. `query_id` resolution moved into the search +workload's `run_search_shard`; the worker passes task parameters through and +the workload validates them. + CTX-16 "Workload SDK foundation" is complete and tested. `scimesh/sdk/` implements the `core-batch-v1` profile: diff --git a/docs/workload-sdk.md b/docs/workload-sdk.md index 6779c32..68be303 100644 --- a/docs/workload-sdk.md +++ b/docs/workload-sdk.md @@ -55,7 +55,12 @@ The stable authoring surface is exported from `scimesh.sdk`: protocols; - `OutputManifest` and `Provenance` describe sealed durable results; - `WorkloadRegistry` resolves an exact name, version, package digest, runtime, - environment, and feature set. It never selects an implicit latest version. + environment, and feature set. It never selects an implicit latest version; +- `MapReduceWorkload` is the primary authoring scaffold for `core-batch-v1`: + a subclass declares its identity, parameter schema, artifact ports, and + three scientific hooks (partition, compute, merge), and the SDK assembles + the manifest, map/reduce stages, workflow, digest-pinned handlers, and the + exact-artifact verifier. See "Authoring a workload" below. Persisted manifests, requests, plans, tasks, expansions, outputs, candidates, decisions, and failures are frozen, recursively immutable, JSON-safe, @@ -202,11 +207,126 @@ the local scientific cores from `scimesh/workloads/similarity_search.py` and 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. +- the v1 worker executes SDK-built workloads directly: + `scimesh/worker/runners.py` is a workload-generic wire bridge that builds a + `TaskSpec` with the workload's own pins, negotiates against a runtime + derived from the loaded definitions, reserves resources, seals the partial + through a content-addressed store, and uploads the resulting CSV over the + unchanged coordinator contract. The worker loads workloads from + `SCIMESH_WORKLOAD_ALLOWLIST` (a JSON array of + `{distribution, name, version, digest}` entries matched against installed + `scimesh.workloads` entry points) or falls back to the built-in + `similarity-search`; advertised capabilities come from + `SCIMESH_CAPABILITIES`. Workloads whose map stage needs more than one input + port are rejected with a clear message until the coordinator contract + supports them. + +## Authoring a workload + +A workload is a user script that imports the SDK. For the standard +`core-batch-v1` shape (one input dataset, shards, partials, one merged result) +subclass `MapReduceWorkload` and implement the three scientific hooks; the +framework provides everything else: + +```python +from pathlib import Path +from typing import Any, Mapping, Sequence + +from scimesh.sdk import ( + ArtifactSchema, + ComponentRef, + MapReduceWorkload, + PortSpec, + SchemaRef, + WorkloadId, +) + +class CountRowsWorkload(MapReduceWorkload): + workload_id = WorkloadId("count-rows", "1.0.0") + description = "Count TSV data rows per shard and concatenate the counts." + parameters_schema = { + "type": "object", + "additionalProperties": False, + "properties": {"prefix": {"type": "string", "minLength": 1, "maxLength": 50}}, + } + input_port = PortSpec(ArtifactSchema( + SchemaRef("molecule-table", 1), "text/tab-separated-values", "utf-8", + max_bytes=10**9, validator=ComponentRef("delimited-table", 1), + validator_configuration={"required_columns": ["canonical_smiles", "chembl_id"]}, + )) + partial_port = output_port = PortSpec(ArtifactSchema( + SchemaRef("count-table", 1), "text/csv", "utf-8", + max_bytes=10**9, validator=ComponentRef("delimited-table", 1), + validator_configuration={"columns": ["id", "rows"]}, + )) + map_parameter_names = ("prefix",) + + def partition_input(self, input_path, parameters, workspace): # -> list[Path] + ... # deterministic shard files, one per map task + + def compute_shard(self, inputs, parameters, output_path): # -> Mapping[str, int|float] + ... # one map task; inputs maps each map port to a materialized file + + def reduce_partials(self, partial_paths, parameters, output_path): # -> Mapping[str, int|float] + ... # deterministic merge of the accepted partials +``` + +The base class then provides `validate`, `plan`, `run`, `reduce`, and +`definition()`; the registry, negotiation, resource reservation, verification, +and the local conformance executor treat the result like any other workload: + +```python +from scimesh.sdk import ( + ArtifactCollection, + JobRequest, + LocalArtifactStore, + LocalCoreBatchExecutor, + WorkloadRegistry, +) +from scimesh.workloads.library import default_sdk_runtime + +workload = CountRowsWorkload(package_digest=..., environment_digest=...) +registry = WorkloadRegistry() +registry.register(workload.definition(), enabled=True) + +store = LocalArtifactStore(Path("artifacts")) +artifact = store.import_file(Path("tiny.tsv"), declaration=workload.manifest.inputs["input"].schema) +request = JobRequest(workload=workload.manifest.workload, parameters={"prefix": "x"}, + inputs={"input": ArtifactCollection.single(artifact)}) +result = LocalCoreBatchExecutor(registry, default_sdk_runtime(), store, Path("work")) \ + .execute(request, workload.manifest.package.digest) +``` + +Hooks you can override beyond the three scientific ones: + +- `domain_validate(parameters)` — extra job-parameter validation (the JSON + schema already ran); +- `resolved_parameters(request)` / `resolved_parameters_for_plan(job, input_path, resolved)` + — values persisted into the plan (for example one-time query resolution); +- `plan_tasks(...)` — custom task construction (the graph workload uses this + to plan one task per block pair with two block inputs); +- `parse_partial_key(key)` / `validate_partial_keys(parsed)` — partial-key + policy (default: `map.`, contiguous; the graph workload + parses `map.x` and enforces the pair-coverage invariant); +- `map_stage_inputs` — a map stage with more than one input port (each extra + port must share the external input schema). + +Anything outside this model uses the lower-level SDK value objects directly. +Authoring rules: keep the scientific core callable without a coordinator, +inline a strict JSON parameter schema, declare artifact schemas with bounds, +return only sink-sealed artifacts, and select a verifier compatible with +determinism and trust. + +To run a workload from the command line without writing any program code: + +```bash +scimesh workload list +scimesh workload run count-rows --input tiny.tsv --params '{"prefix": "x"}' -o result.csv +``` + +`scimesh workload` is a generic SDK tool; it contains no workload-specific +logic, so new workloads do not require changes to the CLI or any other part of +the program. ## Package shape and registration diff --git a/scimesh/sdk/__init__.py b/scimesh/sdk/__init__.py index 1131393..17e234e 100644 --- a/scimesh/sdk/__init__.py +++ b/scimesh/sdk/__init__.py @@ -19,6 +19,7 @@ from .artifacts import ( PortSpec, Provenance, ) +from .batch import MapReduceWorkload from .conformance import ( CancellationFlag, LocalArtifactStore, @@ -77,6 +78,7 @@ from .registry import ( WorkloadDefinition, WorkloadDescription, WorkloadRegistry, + workload_allowlist_from_json, ) from .resources import ( AcceleratorDevice, @@ -155,6 +157,7 @@ __all__ = [ "LocalTaskContext", "LoopSpec", "MANIFEST_SCHEMA_VERSION", + "MapReduceWorkload", "NegotiatedWorkload", "NetworkPolicy", "NumericTolerance", @@ -207,6 +210,7 @@ __all__ = [ "WorkloadLimits", "WorkloadManifest", "WorkloadRegistry", + "workload_allowlist_from_json", "assert_manifest_round_trip", "installed_distribution_digest", "negotiate_manifest", diff --git a/scimesh/sdk/batch.py b/scimesh/sdk/batch.py new file mode 100644 index 0000000..fed5a7c --- /dev/null +++ b/scimesh/sdk/batch.py @@ -0,0 +1,554 @@ +"""High-level core-batch-v1 authoring scaffold for map/reduce workloads. + +``MapReduceWorkload`` is the SDK's primary authoring surface for the +``core-batch-v1`` profile: a subclass declares its identity, parameter +schema, artifact ports, and three scientific hooks (partitioning, per-shard +computation, partial merging), and the base class assembles the immutable +manifest, the map/reduce stages, the workflow DAG, the digest-pinned +planner/runner/reducer handlers, and the exact-artifact verifier. + +The base class deliberately supports only the static byte-exact map/reduce +shape with a single external input and one output port. Workloads that need +a different DAG (for example block-pair tasks with two map inputs) override +the ``map_stage_inputs``, ``plan_tasks``, ``parse_partial_key``, and +``validate_partial_keys`` hooks; anything outside the model must use the +lower-level SDK value objects directly. +""" + +from __future__ import annotations + +import hashlib +import shutil +from pathlib import Path +from typing import Any, Mapping, Sequence + +from .artifacts import ( + ArtifactCollection, + ArtifactItem, + ArtifactRef, + ArtifactSchema, + Cardinality, + CollectionKind, + OutputManifest, + PortSpec, +) +from .execution import ( + CheckpointPolicy, + ExecutionProfile, + NetworkPolicy, + RetryPolicy, +) +from .identity import ComponentRef, 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 + +_EXACT_ARTIFACT = ComponentRef("exact-artifact", 1) +_EXACT_VERIFIER = ExactArtifactVerifier() + + +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 _default_entry_point(module: str, kind: str) -> str: + return f"{module}:{kind}@v1" + + +class MapReduceWorkload: + """Base class for static byte-exact map/reduce workloads. + + Required class attributes: + + - ``workload_id``: ``WorkloadId`` identity (name + version); + - ``parameters_schema``: strict JSON object schema (the registry validates + ``additionalProperties: false`` for the top level); + - ``input_port``, ``partial_port``, ``output_port``: typed ``PortSpec`` + values for the external input, one map partial, and the final result. + + Scientific hooks (override as needed): + + - ``domain_validate(parameters)``: extra job-parameter validation; + - ``resolved_parameters(request)``: values carried into the plan; + - ``partition_input(input_path, parameters, workspace)``: deterministic + shard splitting; returns one TSV/CSV file per map task; + - ``plan_tasks(shard_paths, resolved, job, negotiated, map_stage, context)``: + task construction (default: one task per shard, ``map/``); + - ``compute_shard(inputs, parameters, output_path)``: one map task; + ``inputs`` maps every ``map_stage_inputs`` port to a materialized file; + - ``parse_partial_key(key)`` and ``validate_partial_keys(parsed)``: + partial-key policy for the reducer; + - ``reduce_partials(partial_paths, parameters, output_path)``: the + deterministic merge. + + Optional class attributes: ``map_stage_inputs`` (default one ``input`` + port), ``map_parameter_names``, ``reduce_parameter_names``, ``capabilities``, + ``trust_modes``, ``workflow_id``, ``limits``, ``resources``, ``execution``, + ``map_entry_point``, ``reduce_entry_point``. + """ + + workload_id: WorkloadId + description: str = "" + parameters_schema: Mapping[str, Any] = {} + input_port: PortSpec + partial_port: PortSpec + output_port: PortSpec + + map_stage_inputs: Mapping[str, PortSpec] | None = None + map_parameter_names: tuple[str, ...] = () + reduce_parameter_names: tuple[str, ...] = () + capabilities: tuple[str, ...] | None = None + trust_modes: tuple[TrustMode, ...] = (TrustMode.TRUSTED, TrustMode.UNTRUSTED_QUORUM) + workflow_id: str | None = None + limits: WorkloadLimits | None = None + resources: ResourceRequirements | None = None + execution: ExecutionProfile | None = None + map_entry_point: str | None = None + reduce_entry_point: str | None = None + + def __init__( + self, + *, + package_digest: str, + environment_digest: str, + ) -> None: + workload_id = self.workload_id + if not isinstance(workload_id, WorkloadId): + raise ValueError("workload_id must be a WorkloadId") + if not isinstance(self.input_port, PortSpec): + raise ValueError("input_port must be a PortSpec") + if not isinstance(self.partial_port, PortSpec): + raise ValueError("partial_port must be a PortSpec") + if not isinstance(self.output_port, PortSpec): + raise ValueError("output_port must be a PortSpec") + if not self.parameters_schema: + raise ValueError("parameters_schema must be provided") + module = type(self).__module__ + self.map_entry_point = self.map_entry_point or _default_entry_point( + module, "map" + ) + self.reduce_entry_point = self.reduce_entry_point or _default_entry_point( + module, "reduce" + ) + self.entry_point = self.map_entry_point + self.map_stage_inputs = dict( + self.map_stage_inputs or {"input": self.input_port} + ) + if set(self.map_stage_inputs) != {"input"} and any( + port.schema != self.input_port.schema + for port in self.map_stage_inputs.values() + ): + raise ValueError( + "additional map inputs must share the external input schema" + ) + self.map_parameter_names = tuple(self.map_parameter_names) + self.reduce_parameter_names = tuple( + self.reduce_parameter_names or self.map_parameter_names + ) + self.capabilities = tuple(self.capabilities or (workload_id.name,)) + if workload_id.name not in self.capabilities: + raise ValueError("capabilities must include the workload name") + name = workload_id.name + resources = self.resources or ResourceRequirements( + profile=f"{name}-cpu-v1", + cpu_cores=1, + memory_mb=1024, + scratch_mb=1024, + max_duration_seconds=3600, + ) + execution = self.execution or ExecutionProfile( + profile=f"{name}-python-process-v1", + network=NetworkPolicy.TRUSTED, + timeout_seconds=3600, + checkpoint=CheckpointPolicy(), + ) + limits = self.limits or WorkloadLimits( + max_input_bytes=self.input_port.schema.max_bytes, + max_tasks=10_000, + max_output_bytes=self.output_port.schema.max_bytes, + ) + trust_values = tuple(mode.value for mode in self.trust_modes) + map_stage = StageSpec( + stage_id="map", + kind=StageKind.MAP, + entry_point=self.map_entry_point, + needs=(), + inputs=self.map_stage_inputs, + outputs={"partial": self.partial_port}, + parameter_names=self.map_parameter_names, + resources=resources, + execution=execution, + retry=RetryPolicy(), + verifier=_EXACT_ARTIFACT, + trust_modes=trust_values, + 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=self.reduce_entry_point, + needs=("map",), + inputs={"partials": reduce_input}, + outputs={"result": self.output_port}, + parameter_names=self.reduce_parameter_names, + resources=resources, + execution=execution, + retry=RetryPolicy(), + verifier=_EXACT_ARTIFACT, + trust_modes=trust_values, + max_fan_out=1, + cacheable=True, + ) + edges = [ + ArtifactEdge(PortRef("input"), PortRef(port, "map")) + for port in self.map_stage_inputs + ] + [ + ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")), + ] + workflow = WorkflowSpec( + workflow_id=self.workflow_id or f"{name}-map-reduce-v1", + inputs={"input": self.input_port}, + stages=(map_stage, reduce_stage), + edges=tuple(edges), + 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=workload_id, + description=self.description, + package=PackageSpec("scimesh", package_digest), + environment=EnvironmentSpec( + "python-process", + environment_digest, + {"adapter": "sdk-native"}, + ), + parameters_schema=self.parameters_schema, + workflow=workflow, + inputs={"input": self.input_port}, + outputs={"result": self.output_port}, + determinism=DeterminismProfile.BYTE_EXACT, + trust_modes=self.trust_modes, + verifier=VerifierSpec(_EXACT_ARTIFACT, {}), + limits=limits, + capabilities=self.capabilities, + conformance_profiles=("core-batch-v1",), + ) + self._exact_verifier = _EXACT_VERIFIER + self._resources = resources + self._execution = execution + self._limits = limits + + # ------------------------------------------------------------------ + # Public assembly + # ------------------------------------------------------------------ + + def definition(self) -> WorkloadDefinition: + map_entry_point = self.map_entry_point + reduce_entry_point = self.reduce_entry_point + assert map_entry_point is not None and reduce_entry_point is not None + return WorkloadDefinition( + manifest=self.manifest, + planner=self, + runners={map_entry_point: self}, + reducers={reduce_entry_point: self}, + verifiers={_EXACT_ARTIFACT.canonical: self._exact_verifier}, + ) + + # ------------------------------------------------------------------ + # Scientific hooks + # ------------------------------------------------------------------ + + def domain_validate(self, parameters: Mapping[str, Any]) -> None: + """Extra job-parameter validation beyond the JSON schema.""" + + def resolved_parameters(self, request: JobRequest) -> dict[str, Any]: + """Values persisted as the plan's resolved parameters.""" + return dict(request.parameters) + + def partition_input( + self, + input_path: Path, + parameters: Mapping[str, Any], + workspace: Path, + ) -> list[Path]: + """Split the materialized input into deterministic shard files.""" + raise NotImplementedError("partition_input must be implemented") + + def plan_tasks( + self, + shard_paths: Sequence[Path], + resolved: Mapping[str, Any], + job: ValidatedJob, + negotiated: Any, + map_stage: StageSpec, + context: PlanningContext, + ) -> list[TaskSpec]: + """Build one map TaskSpec per planned shard.""" + task_parameters = self.task_parameters(resolved) + tasks: list[TaskSpec] = [] + for index, path in enumerate(shard_paths): + sealed = context.sink.seal( + path, + declaration=self.input_port.schema, + ) + tasks.append( + self.task_spec( + map_stage, + job, + negotiated, + f"map/{index:08d}", + task_parameters, + {"input": ArtifactCollection.single(sealed)}, + ) + ) + return tasks + + def task_parameters(self, resolved: Mapping[str, Any]) -> dict[str, Any]: + """Project resolved parameters onto the map stage projection.""" + return { + key: value + for key, value in resolved.items() + if key in set(self.map_parameter_names) + } + + def compute_shard( + self, + inputs: Mapping[str, Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + """Compute one map task and write its partial CSV.""" + raise NotImplementedError("compute_shard must be implemented") + + def parse_partial_key(self, key: str) -> Any: + """Parse one ``map.`` partial key into an orderable identity.""" + prefix = "map." + if not key.startswith(prefix): + raise ValueError("partial key must use map.") + raw = key[len(prefix) :] + if len(raw) != 8 or not raw.isdigit(): + raise ValueError("partial key must use map.") + return int(raw) + + def validate_partial_keys(self, parsed: Sequence[Any]) -> None: + """Enforce the reducer's partial-set invariant (default: contiguous).""" + indices = sorted(int(value) for value in parsed) + if indices != list(range(len(indices))): + raise ValueError("partial keys must be complete and contiguous") + + def reduce_partials( + self, + partial_paths: Sequence[Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + """Merge materialized partials into one deterministic final CSV.""" + raise NotImplementedError("reduce_partials must be implemented") + + # ------------------------------------------------------------------ + # Framework handlers + # ------------------------------------------------------------------ + + def validate(self, request: JobRequest) -> ValidatedJob: + if request.workload != self.manifest.workload: + raise ValueError("workload received a request for another workload") + self.domain_validate(request.parameters) + return ValidatedJob(request, self.resolved_parameters(request)) + + 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("workload requires the input port") + self.input_port.validate_collection(collection, "job input") + input_path = context.catalog.materialize(collection.items[0].artifact) + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + resolved = dict(job.resolved_parameters) + resolved = self.resolved_parameters_for_plan(job, input_path, resolved) + shard_paths = self.partition_input(input_path, resolved, workspace) + negotiated = context.negotiated + map_stage = self.manifest.workflow.stages[0] + tasks = self.plan_tasks( + shard_paths, resolved, job, negotiated, map_stage, context + ) + return self.workflow_plan(job, negotiated, resolved, tasks) + + def resolved_parameters_for_plan( + self, + job: ValidatedJob, + input_path: Path, + resolved: dict[str, Any], + ) -> dict[str, Any]: + """Hook to enrich resolved parameters at plan time (query resolution).""" + return resolved + + def workflow_plan( + self, + job: ValidatedJob, + negotiated: Any, + resolved: Mapping[str, Any], + tasks: Sequence[TaskSpec], + ) -> WorkflowPlan: + 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(resolved), + tasks=tuple(tasks), + ) + + def task_spec( + self, + map_stage: StageSpec, + job: ValidatedJob, + negotiated: Any, + task_key: str, + parameters: Mapping[str, Any], + inputs: Mapping[str, ArtifactCollection], + ) -> TaskSpec: + assert map_stage.verifier is not None + return 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=task_key, + stage_id=map_stage.stage_id, + parameters=parameters, + inputs=inputs, + expected_outputs=map_stage.outputs, + resources=map_stage.resources, + execution=map_stage.execution, + ).validate_stage(map_stage) + + def run(self, context: TaskContext) -> OutputManifest: + context.cancellation.raise_if_cancelled() + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + inputs: dict[str, Path] = {} + for name, port in self.map_stage_inputs.items(): + collection = context.task.inputs.get(name) + if collection is None: + raise ValueError(f"map task requires the {name} input") + port.validate_collection(collection, f"map input {name}") + assert collection.items + inputs[name] = context.catalog.materialize(collection.items[0].artifact) + output_path = workspace / "result.csv" + metrics = self.compute_shard( + inputs, + 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._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: + raise ValueError("reducer requires a non-empty keyed partial collection") + if collection.kind is not CollectionKind.KEYED or not collection.items: + raise ValueError("reducer requires a non-empty keyed partial collection") + self.manifest.workflow.stages[1].inputs["partials"].validate_collection( + collection, + "reducer partials", + ) + parsed = [self.parse_partial_key(item.key or "") for item in collection.items] + self.validate_partial_keys(parsed) + 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("partial keys do not match the coordinator expected set") + workspace = context.workspace + workspace.mkdir(parents=True, exist_ok=True) + partial_paths: list[Path] = [] + for index, item in sorted( + enumerate(collection.items), + key=lambda entry: entry[1].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 = self.reduce_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._limits.max_output_bytes, + ) diff --git a/scimesh/sdk/registry.py b/scimesh/sdk/registry.py index 619e36e..874dfa5 100644 --- a/scimesh/sdk/registry.py +++ b/scimesh/sdk/registry.py @@ -24,8 +24,20 @@ from .identity import ComponentRef, WorkloadId from .integrity import installed_distribution_digest from .manifest import WorkloadManifest from .plans import JobRequest, ValidatedJob, WorkflowPlan -from .protocols import Planner, PlanningContext, PlanningResources, Reducer, Runner, Verifier -from .runtime import CompatibilityError, NegotiatedWorkload, RuntimeCapabilities, negotiate_manifest +from .protocols import ( + Planner, + PlanningContext, + PlanningResources, + Reducer, + Runner, + Verifier, +) +from .runtime import ( + CompatibilityError, + NegotiatedWorkload, + RuntimeCapabilities, + negotiate_manifest, +) from .schema import validate_parameter_instance from .workflow import StageKind @@ -51,11 +63,11 @@ def _validate_entry_point_ownership(entry_point: metadata.EntryPoint) -> None: raise ValueError("workload entry point has an invalid module path") raw_top_level = distribution.read_text("top_level.txt") - declared = { - line.strip() - for line in raw_top_level.splitlines() - if line.strip() - } if raw_top_level is not None else set() + declared = ( + {line.strip() for line in raw_top_level.splitlines() if line.strip()} + if raw_top_level is not None + else set() + ) root_name = parts[0] if root_name not in declared: raise ValueError("workload entry point module is outside its distribution") @@ -81,9 +93,18 @@ def _validate_entry_point_ownership(entry_point: metadata.EntryPoint) -> None: ownership_root = package_root.resolve() candidates = [ *(Path(str(module_base) + suffix) for suffix in machinery.SOURCE_SUFFIXES), - *(Path(str(module_base) + suffix) for suffix in machinery.EXTENSION_SUFFIXES), - *(module_base / ("__init__" + suffix) for suffix in machinery.SOURCE_SUFFIXES), - *(module_base / ("__init__" + suffix) for suffix in machinery.EXTENSION_SUFFIXES), + *( + Path(str(module_base) + suffix) + for suffix in machinery.EXTENSION_SUFFIXES + ), + *( + module_base / ("__init__" + suffix) + for suffix in machinery.SOURCE_SUFFIXES + ), + *( + module_base / ("__init__" + suffix) + for suffix in machinery.EXTENSION_SUFFIXES + ), ] else: if len(parts) != 1: @@ -92,7 +113,10 @@ def _validate_entry_point_ownership(entry_point: metadata.EntryPoint) -> None: module_base = Path(distribution.locate_file(root_name)) candidates = [ *(Path(str(module_base) + suffix) for suffix in machinery.SOURCE_SUFFIXES), - *(Path(str(module_base) + suffix) for suffix in machinery.EXTENSION_SUFFIXES), + *( + Path(str(module_base) + suffix) + for suffix in machinery.EXTENSION_SUFFIXES + ), ] existing = tuple(candidate for candidate in candidates if candidate.is_file()) if len(existing) != 1 or not existing[0].resolve().is_relative_to(ownership_root): @@ -126,7 +150,9 @@ class WorkloadDefinition: for name, handler in values.items(): canonical = require_string(name, f"{field} entry point", max_length=256) if not callable(getattr(handler, method, None)): - raise ValueError(f"definition {field} handler must implement {method}") + raise ValueError( + f"definition {field} handler must implement {method}" + ) copied[canonical] = handler object.__setattr__(self, field, MappingProxyType(copied)) for stage in self.manifest.workflow.stages: @@ -149,23 +175,39 @@ class WorkloadDefinition: ) verifier_key = self.manifest.verifier.verifier.canonical if verifier_key not in self.verifiers: - raise ValueError(f"definition has no installed manifest verifier: {verifier_key}") + raise ValueError( + f"definition has no installed manifest verifier: {verifier_key}" + ) for key, verifier in self.verifiers.items(): try: declared_identity = ComponentRef.from_dict(key) except ValueError as error: - raise ValueError("definition verifier keys must be component identities") from error - if declared_identity.canonical != key or getattr(verifier, "identity", None) != declared_identity: - raise ValueError("definition verifier handler identity does not match its key") + raise ValueError( + "definition verifier keys must be component identities" + ) from error + if ( + declared_identity.canonical != key + or getattr(verifier, "identity", None) != declared_identity + ): + raise ValueError( + "definition verifier handler identity does not match its key" + ) manifest_verifier = self.verifiers[verifier_key] handler_configuration = getattr(manifest_verifier, "configuration", None) if handler_configuration is None: if self.manifest.verifier.configuration: - raise ValueError("manifest verifier configuration is not bound by its handler") + raise ValueError( + "manifest verifier configuration is not bound by its handler" + ) elif dict(handler_configuration) != dict(self.manifest.verifier.configuration): - raise ValueError("manifest verifier configuration does not match its handler") + raise ValueError( + "manifest verifier configuration does not match its handler" + ) for stage in self.manifest.workflow.stages: - if stage.verifier is not None and stage.verifier.canonical not in self.verifiers: + if ( + stage.verifier is not None + and stage.verifier.canonical not in self.verifiers + ): raise ValueError( f"definition has no installed stage verifier: {stage.verifier.canonical}" ) @@ -178,13 +220,54 @@ class AllowedPackage: digest: str def __post_init__(self) -> None: - distribution = require_string(self.distribution, "distribution", max_length=128).lower() + distribution = require_string( + self.distribution, "distribution", max_length=128 + ).lower() if not re.fullmatch(r"[a-z0-9]+(?:[-_.][a-z0-9]+)*", distribution): - raise ValueError("distribution must be a canonical Python distribution name") + raise ValueError( + "distribution must be a canonical Python distribution name" + ) object.__setattr__(self, "distribution", distribution.replace("_", "-")) if not isinstance(self.workload, WorkloadId): raise ValueError("allowed workload must be a WorkloadId") - object.__setattr__(self, "digest", require_sha256(self.digest, "allowed digest", prefixed=True)) + object.__setattr__( + self, "digest", require_sha256(self.digest, "allowed digest", prefixed=True) + ) + + +def workload_allowlist_from_json(value: object) -> tuple[AllowedPackage, ...]: + """Parse a JSON array of ``{distribution, name, version, digest}`` allowlist entries.""" + import json + + if value is None or value == "": + return () + if not isinstance(value, str): + raise ValueError("workload allowlist must be a JSON array") + try: + decoded = json.loads(value) + except (TypeError, json.JSONDecodeError, RecursionError) as error: + raise ValueError("workload allowlist must be valid JSON") from error + if not isinstance(decoded, list): + raise ValueError("workload allowlist must be a JSON array") + entries: list[AllowedPackage] = [] + for item in decoded: + if not isinstance(item, dict) or not { + "distribution", + "name", + "version", + "digest", + }.issubset(item): + raise ValueError( + "workload allowlist entries need distribution, name, version, and digest" + ) + entries.append( + AllowedPackage( + str(item["distribution"]), + WorkloadId(str(item["name"]), str(item["version"])), + str(item["digest"]), + ) + ) + return tuple(entries) @dataclass(frozen=True, slots=True) @@ -223,14 +306,18 @@ class WorkloadRegistry: self._enabled: set[tuple[str, str, str]] = set() self._lock = RLock() - def register(self, definition: WorkloadDefinition, *, enabled: bool = False) -> None: + def register( + self, definition: WorkloadDefinition, *, enabled: bool = False + ) -> None: if not isinstance(definition, WorkloadDefinition): raise ValueError("definition must be a WorkloadDefinition") workload = definition.manifest.workload key = (workload.name, workload.version) with self._lock: if key in self._definitions: - raise ValueError(f"workload version already registered: {workload.name}@{workload.version}") + raise ValueError( + f"workload version already registered: {workload.name}@{workload.version}" + ) self._definitions[key] = definition if enabled: self._enabled.add((*key, definition.manifest.package.digest)) @@ -240,7 +327,9 @@ class WorkloadRegistry: with self._lock: definition = self._registered(name, version) if digest != definition.manifest.package.digest: - raise ValueError("package digest does not match the registered manifest") + raise ValueError( + "package digest does not match the registered manifest" + ) self._enabled.add((definition.manifest.workload.name, version, digest)) def disable(self, name: str, version: str, package_digest: str) -> None: @@ -257,7 +346,9 @@ class WorkloadRegistry: try: return self._definitions[(canonical, version)] except KeyError as error: - raise ValueError(f"unknown workload version: {canonical}@{version}") from error + raise ValueError( + f"unknown workload version: {canonical}@{version}" + ) from error def require( self, @@ -270,10 +361,21 @@ class WorkloadRegistry: digest = require_sha256(package_digest, "package_digest", prefixed=True) with self._lock: definition = self._registered(name, version) - identity = (definition.manifest.workload.name, definition.manifest.workload.version, digest) - if digest != definition.manifest.package.digest or identity not in self._enabled: + identity = ( + definition.manifest.workload.name, + definition.manifest.workload.version, + digest, + ) + if ( + digest != definition.manifest.package.digest + or identity not in self._enabled + ): raise ValueError("workload package digest is not enabled") - negotiated = negotiate_manifest(definition.manifest, runtime) if runtime is not None else None + negotiated = ( + negotiate_manifest(definition.manifest, runtime) + if runtime is not None + else None + ) return definition, negotiated def plan( @@ -293,31 +395,41 @@ class WorkloadRegistry: runtime=runtime, ) assert negotiated is not None - self._validate_request_compatibility(request, definition.manifest, runtime, negotiated) + self._validate_request_compatibility( + request, definition.manifest, runtime, negotiated + ) self._validate_request_shape(request, definition.manifest) validated = definition.planner.validate(request) if not isinstance(validated, ValidatedJob) or validated.request != request: - raise ValueError("planner.validate must return a ValidatedJob for the same request") + raise ValueError( + "planner.validate must return a ValidatedJob for the same request" + ) plan = definition.planner.plan( validated, _NegotiatedPlanningContext(context, negotiated), ) if not isinstance(plan, WorkflowPlan) or plan.workload != request.workload: - raise ValueError("planner.plan must return a WorkflowPlan for the requested workload") + raise ValueError( + "planner.plan must return a WorkflowPlan for the requested workload" + ) if ( plan.package_digest != definition.manifest.package.digest or plan.manifest_digest != definition.manifest.digest or plan.trust_mode is not request.trust_mode or plan.sdk_api_version != runtime.sdk_api_version or plan.protocol_version != runtime.protocol_version - or plan.manifest_schema_version != definition.manifest.manifest_schema_version - or plan.workflow_schema_version != definition.manifest.workflow.schema_version + or plan.manifest_schema_version + != definition.manifest.manifest_schema_version + or plan.workflow_schema_version + != definition.manifest.workflow.schema_version or plan.environment_digest != definition.manifest.environment.digest or plan.verifier != definition.manifest.verifier.verifier or plan.selected_features != negotiated.selected_features or plan.optional_fallbacks != negotiated.optional_fallbacks ): - raise ValueError("planner plan does not carry the selected immutable workload pin") + raise ValueError( + "planner plan does not carry the selected immutable workload pin" + ) plan.validate_workflow(definition.manifest.workflow) self._validate_plan_limits(request, plan, definition.manifest) return WorkflowPlan.from_json(plan.to_json()) @@ -369,7 +481,9 @@ class WorkloadRegistry: ) @staticmethod - def _validate_request_shape(request: JobRequest, manifest: WorkloadManifest) -> None: + def _validate_request_shape( + request: JobRequest, manifest: WorkloadManifest + ) -> None: if set(request.inputs) != set(manifest.inputs): raise ValueError("job input ports do not match the manifest") total_bytes = 0 @@ -380,7 +494,9 @@ class WorkloadRegistry: for item in request.inputs[name].items: existing = artifact_references.get(item.artifact.artifact_id) if existing is not None and existing != item.artifact: - raise ValueError("job reuses an artifact ID with conflicting metadata") + raise ValueError( + "job reuses an artifact ID with conflicting metadata" + ) artifact_references[item.artifact.artifact_id] = item.artifact if total_bytes > manifest.limits.max_input_bytes: raise ValueError("job inputs exceed the manifest byte limit") @@ -388,7 +504,15 @@ class WorkloadRegistry: raise ValueError("job inputs exceed the manifest artifact limit") import json from ._validation import thaw_json - if len(json.dumps(thaw_json(request.parameters), allow_nan=False).encode("utf-8")) > manifest.limits.max_parameter_bytes: + + if ( + len( + json.dumps(thaw_json(request.parameters), allow_nan=False).encode( + "utf-8" + ) + ) + > manifest.limits.max_parameter_bytes + ): raise ValueError("job parameters exceed the manifest byte limit") validate_parameter_instance(request.parameters, manifest.parameters_schema) @@ -408,13 +532,23 @@ class WorkloadRegistry: for item in collection.items: existing = references.get(item.artifact.artifact_id) if existing is not None and existing != item.artifact: - raise ValueError("workflow plan reuses an artifact ID with conflicting metadata") + raise ValueError( + "workflow plan reuses an artifact ID with conflicting metadata" + ) references[item.artifact.artifact_id] = item.artifact - if len(canonical_json(task.parameters).encode("utf-8")) > manifest.limits.max_parameter_bytes: - raise ValueError("planned task parameters exceed the manifest byte limit") + if ( + len(canonical_json(task.parameters).encode("utf-8")) + > manifest.limits.max_parameter_bytes + ): + raise ValueError( + "planned task parameters exceed the manifest byte limit" + ) if len(references) > manifest.limits.max_artifacts: raise ValueError("workflow plan exceeds the manifest artifact limit") - if len(canonical_json(plan.resolved_parameters).encode("utf-8")) > manifest.limits.max_parameter_bytes: + if ( + len(canonical_json(plan.resolved_parameters).encode("utf-8")) + > manifest.limits.max_parameter_bytes + ): raise ValueError("resolved parameters exceed the manifest byte limit") def descriptions(self) -> tuple[WorkloadDescription, ...]: @@ -455,7 +589,10 @@ class WorkloadRegistry: for key, approval in allowed.items(): if _normalized_distribution_name(key[0]) != distribution: continue - if entry_point.name != f"{approval.workload.name}@{approval.workload.version}": + if ( + entry_point.name + != f"{approval.workload.name}@{approval.workload.version}" + ): continue _validate_entry_point_ownership(entry_point) # Import policy is process-global, so installed discovery is @@ -463,9 +600,12 @@ class WorkloadRegistry: # cache prefix prevents pre-existing package pyc files from # being consumed, while dont_write_bytecode keeps the # measured source tree unchanged during both load and factory. - with _DISCOVERY_IMPORT_LOCK, TemporaryDirectory( - prefix="scimesh-discovery-cache-" - ) as cache_prefix: + with ( + _DISCOVERY_IMPORT_LOCK, + TemporaryDirectory( + prefix="scimesh-discovery-cache-" + ) as cache_prefix, + ): measured_before = installed_distribution_digest(entry_point.dist) if measured_before != approval.digest: raise ValueError( @@ -479,44 +619,64 @@ class WorkloadRegistry: loaded = entry_point.load() definition = ( loaded() - if callable(loaded) and not isinstance(loaded, WorkloadDefinition) + if callable(loaded) + and not isinstance(loaded, WorkloadDefinition) else loaded ) finally: sys.pycache_prefix = previous_cache_prefix sys.dont_write_bytecode = previous_bytecode_policy - if installed_distribution_digest(entry_point.dist) != measured_before: + if ( + installed_distribution_digest(entry_point.dist) + != measured_before + ): raise ValueError( "installed package content changed while loading its entry point" ) if not isinstance(definition, WorkloadDefinition): - raise ValueError("workload entry point must provide a WorkloadDefinition") + raise ValueError( + "workload entry point must provide a WorkloadDefinition" + ) if definition.manifest.workload != approval.workload: - raise ValueError("discovered workload identity does not match its allowlist entry") + raise ValueError( + "discovered workload identity does not match its allowlist entry" + ) if definition.manifest.package.distribution != approval.distribution: - raise ValueError("discovered package identity does not match its allowlist entry") + raise ValueError( + "discovered package identity does not match its allowlist entry" + ) if definition.manifest.package.digest != approval.digest: - raise ValueError("discovered package digest does not match its allowlist entry") + raise ValueError( + "discovered package digest does not match its allowlist entry" + ) if key in discovered: - raise ValueError("multiple installed entry points match one allowlist entry") + raise ValueError( + "multiple installed entry points match one allowlist entry" + ) pending.append(definition) discovered.add(key) break missing = sorted(set(allowed) - discovered) if missing: identities = ", ".join(f"{name}@{version}" for _, name, version in missing) - raise ValueError("allowlisted workload entry points were not installed: " + identities) + raise ValueError( + "allowlisted workload entry points were not installed: " + identities + ) pending_keys = [ (definition.manifest.workload.name, definition.manifest.workload.version) for definition in pending ] if len(pending_keys) != len(set(pending_keys)): - raise ValueError("multiple allowlisted distributions provide one workload version") + raise ValueError( + "multiple allowlisted distributions provide one workload version" + ) with self._lock: conflicts = [key for key in pending_keys if key in self._definitions] if conflicts: name, version = conflicts[0] - raise ValueError(f"workload version already registered: {name}@{version}") + raise ValueError( + f"workload version already registered: {name}@{version}" + ) definitions = dict(self._definitions) enabled = set(self._enabled) for key, definition in zip(pending_keys, pending): diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index f8a39ec..7237207 100644 --- a/scimesh/worker/cli.py +++ b/scimesh/worker/cli.py @@ -77,7 +77,9 @@ def main(argv: list[str] | None = None) -> int: config = WorkerConfig.from_environment(overrides) except (TypeError, ValueError) as error: parser.error(str(error)) - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s" + ) # One shared token strategy backs both clients: a worker key (exchanged and # refreshed) or a static bearer token, decided by what the config carries. tokens = provider_from_config( @@ -95,7 +97,7 @@ def main(argv: list[str] | None = None) -> int: HttpArtifactClient( config.coordinator_url, config.request_timeout, token_provider=tokens ), - SciMeshRunner(), + SciMeshRunner.for_worker(config), ).run_forever() return 0 if completed_without_interruption else 130 diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py index b56fba7..dbdecfa 100644 --- a/scimesh/worker/config.py +++ b/scimesh/worker/config.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from dataclasses import dataclass from math import isfinite from pathlib import Path @@ -10,6 +11,8 @@ import socket from typing import Mapping from urllib.parse import urlsplit +from scimesh.sdk.registry import AllowedPackage, workload_allowlist_from_json + def _clean_url(value: object | None) -> str | None: """Normalise an optional URL: drop a blank one, strip a trailing slash.""" @@ -31,6 +34,24 @@ def _positive_number(value: object, name: str, *, allow_zero: bool = False) -> N raise ValueError(f"{name} must be {qualifier}") +def _capabilities(value: object) -> tuple[str, ...]: + """Parse a comma-separated capability list into unique non-empty names.""" + if value is None: + return ("similarity-search", "similarity_search") + if not isinstance(value, str) or not value.strip(): + raise ValueError("capabilities must be a comma-separated list") + names = tuple( + dict.fromkeys(item.strip() for item in value.split(",") if item.strip()) + ) + if not names: + raise ValueError("capabilities cannot be empty") + return names + + +def _workload_allowlist(value: object) -> tuple[AllowedPackage, ...]: + return workload_allowlist_from_json(value) + + @dataclass(frozen=True) class WorkerConfig: coordinator_url: str @@ -60,6 +81,11 @@ class WorkerConfig: "similarity-search", "similarity_search", ) + # Optional allowlist of installed SDK workload packages to execute. When + # empty, the worker runs the built-in similarity-search only. Entries are + # ``{distribution, name, version, digest}`` JSON objects matching the + # installed ``scimesh.workloads`` entry points. + workload_allowlist: tuple[AllowedPackage, ...] = () def __post_init__(self) -> None: parsed = urlsplit(self.coordinator_url) @@ -72,8 +98,14 @@ class WorkerConfig: if us.scheme not in {"http", "https"} or not us.hostname: raise ValueError("userservice_url must be an absolute HTTP(S) URL") if self.worker_key is not None and not self.userservice_url: - raise ValueError("worker_key requires userservice_url (SCIMESH_USERSERVICE_URL)") - if isinstance(self.cpu_count, bool) or not isinstance(self.cpu_count, int) or self.cpu_count < 1: + raise ValueError( + "worker_key requires userservice_url (SCIMESH_USERSERVICE_URL)" + ) + if ( + isinstance(self.cpu_count, bool) + or not isinstance(self.cpu_count, int) + or self.cpu_count < 1 + ): raise ValueError("cpu_count must be positive") if self.worker_id is not None and not isinstance(self.worker_id, str): raise ValueError("worker_id must be a string when set") @@ -87,7 +119,9 @@ class WorkerConfig: _positive_number(self.request_timeout, "request_timeout") _positive_number(self.heartbeat_interval, "heartbeat_interval") if self.cleanup_after_seconds is not None: - _positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True) + _positive_number( + self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True + ) if self.max_tasks is not None: if ( isinstance(self.max_tasks, bool) @@ -99,6 +133,18 @@ class WorkerConfig: raise ValueError("exit_when_idle must be a boolean") if not self.capabilities: raise ValueError("capabilities cannot be empty") + if any( + not isinstance(capability, str) or not capability.strip() + for capability in self.capabilities + ): + raise ValueError("capabilities must contain non-empty names") + if len(self.capabilities) != len(set(self.capabilities)): + raise ValueError("capabilities must be unique") + if any( + not isinstance(package, AllowedPackage) + for package in self.workload_allowlist + ): + raise ValueError("workload_allowlist must contain AllowedPackage values") # Runner subprocesses use a task directory as their cwd. Keep the # configured root absolute so input/output paths remain valid there # even when the CLI received a convenient relative --work-dir value. @@ -111,7 +157,9 @@ class WorkerConfig: """Build config from environment, allowing typed CLI values to override it.""" values = overrides or {} - def value(name: str, environment: str, default: object | None = None) -> object | None: + def value( + name: str, environment: str, default: object | None = None + ) -> object | None: override = values.get(name) return override if override is not None else os.getenv(environment, default) @@ -122,20 +170,37 @@ class WorkerConfig: cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1) memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB") max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS") + capabilities = value("capabilities", "SCIMESH_CAPABILITIES") + allowlist = value("workload_allowlist", "SCIMESH_WORKLOAD_ALLOWLIST") + worker_id = value("worker_id", "SCIMESH_WORKER_ID") + work_dir = value("work_dir", "SCIMESH_WORK_DIR", "./scimesh-worker-data") + worker_name = value("worker_name", "SCIMESH_WORKER_NAME", socket.gethostname()) + poll_interval = value("poll_interval", "SCIMESH_POLL_INTERVAL", "2") + request_timeout = value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30") + heartbeat_interval = value( + "heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15" + ) + bearer_token = value("bearer_token", "SCIMESH_BEARER_TOKEN") + worker_key = value("worker_key", "SCIMESH_WORKER_KEY") + userservice_url = _clean_url( + value("userservice_url", "SCIMESH_USERSERVICE_URL") + ) return cls( coordinator_url=url.rstrip("/"), - worker_id=value("worker_id", "SCIMESH_WORKER_ID"), - work_dir=Path(value("work_dir", "SCIMESH_WORK_DIR", "./scimesh-worker-data")), - worker_name=str(value("worker_name", "SCIMESH_WORKER_NAME", socket.gethostname())), + worker_id=str(worker_id) if worker_id is not None else None, + work_dir=Path(str(work_dir)), + worker_name=str(worker_name), cpu_count=int(cpu_count), memory_mb=int(memory_mb) if memory_mb is not None else None, - poll_interval=float(value("poll_interval", "SCIMESH_POLL_INTERVAL", "2")), - request_timeout=float(value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30")), - heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")), - bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"), - worker_key=value("worker_key", "SCIMESH_WORKER_KEY"), - userservice_url=_clean_url(value("userservice_url", "SCIMESH_USERSERVICE_URL")), + poll_interval=float(poll_interval), + request_timeout=float(request_timeout), + heartbeat_interval=float(heartbeat_interval), + bearer_token=str(bearer_token) if bearer_token is not None else None, + worker_key=str(worker_key) if worker_key is not None else None, + userservice_url=userservice_url, cleanup_after_seconds=float(cleanup) if cleanup else None, max_tasks=int(max_tasks) if max_tasks is not None else None, exit_when_idle=bool(values.get("exit_when_idle", False)), + capabilities=_capabilities(capabilities), + workload_allowlist=_workload_allowlist(allowlist), ) diff --git a/scimesh/worker/runners.py b/scimesh/worker/runners.py index e1af46a..3cc3225 100644 --- a/scimesh/worker/runners.py +++ b/scimesh/worker/runners.py @@ -3,13 +3,20 @@ 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. +provenance, resource reservation, and a content-addressed local store. + +The runner is workload-generic: it loads definitions by name (from an +explicit mapping, built-in defaults, or installed-package discovery through +an administrator allowlist) and executes any workload whose map stage has a +single ``input`` port and a single ``partial`` output. Anything else fails +closed with a clear message, so adding a workload never requires touching +worker code. """ from __future__ import annotations import hashlib +import platform from datetime import datetime, timezone from pathlib import Path from typing import Mapping, Protocol @@ -28,28 +35,21 @@ from scimesh.sdk.conformance import ( ScopedArtifactSink, ) from scimesh.sdk._validation import canonical_json +from scimesh.sdk.identity import SDK_API_VERSION 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.registry import WorkloadDefinition, WorkloadRegistry +from scimesh.sdk.resources import ResourceAllocation, ResourceInventory, 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 .config import WorkerConfig 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") @@ -58,22 +58,93 @@ class Runner(Protocol): def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: ... +def _inventory_for( + definitions: Mapping[str, WorkloadDefinition], + *, + cpu_cores: int, + memory_mb: int, +) -> ResourceInventory: + return ResourceInventory( + cpu_cores=cpu_cores, + memory_mb=memory_mb, + scratch_mb=memory_mb, + architecture=platform.machine().lower() or "unknown", + environment_digests=tuple( + dict.fromkeys( + definition.manifest.environment.digest + for definition in definitions.values() + ) + ), + ) + + +def _runtime_for( + definitions: Mapping[str, WorkloadDefinition], inventory: ResourceInventory +) -> RuntimeCapabilities: + 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=tuple(sorted(definitions)), + inventory=inventory, + ) + + class SciMeshRunner: - """Execute claimed coordinator tasks through the SDK-built workloads.""" + """Execute claimed coordinator tasks through SDK-built workloads.""" def __init__( self, definitions: Mapping[str, WorkloadDefinition] | None = None, + *, + inventory: ResourceInventory | None = None, runtime: RuntimeCapabilities | None = None, ) -> None: self._definitions = dict(definitions or {}) if "similarity-search" not in self._definitions: + from scimesh.workloads.search import similarity_search_sdk_definition + self._definitions["similarity-search"] = ( similarity_search_sdk_definition().definition() ) - self._runtime = runtime or default_sdk_runtime() + self._inventory = inventory or _inventory_for( + self._definitions, + cpu_cores=1, + memory_mb=1024, + ) + self._runtime = runtime or _runtime_for(self._definitions, self._inventory) self._pool = ResourcePool(self._runtime.inventory, max_concurrency=1) + @classmethod + def for_worker(cls, config: WorkerConfig) -> "SciMeshRunner": + """Build a runner for one worker: discover allowlisted workloads or use built-ins.""" + definitions: dict[str, WorkloadDefinition] = {} + if config.workload_allowlist: + registry = WorkloadRegistry() + registry.discover_installed(config.workload_allowlist) + for description in registry.descriptions(): + definition, _ = registry.require( + description.workload.name, + description.workload.version, + description.package_digest, + ) + definitions[description.workload.name] = definition + if not definitions: + raise ValueError("workload_allowlist discovered no workloads") + else: + from scimesh.workloads.search import similarity_search_sdk_definition + + definitions["similarity-search"] = ( + similarity_search_sdk_definition().definition() + ) + inventory = _inventory_for( + definitions, + cpu_cores=config.cpu_count, + memory_mb=config.memory_mb or 1024, + ) + return cls(definitions=definitions, inventory=inventory) + def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: task_dir = task_dir.resolve() workload = task.workload.replace("_", "-") @@ -85,11 +156,14 @@ class SciMeshRunner: map_stage = next( stage for stage in manifest.workflow.stages if stage.kind is StageKind.MAP ) + if set(map_stage.inputs) != {"input"} or len(map_stage.outputs) != 1: + raise ValueError( + f"workload {workload} is not executable through the v1 single-input contract" + ) 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, @@ -110,7 +184,7 @@ class SciMeshRunner: optional_fallbacks=negotiated.optional_fallbacks, task_key="map/00000000", stage_id=map_stage.stage_id, - parameters=parameters, + parameters=task.parameters, inputs={"input": ArtifactCollection.single(input_ref)}, expected_outputs=map_stage.outputs, resources=map_stage.resources, @@ -149,28 +223,6 @@ class SciMeshRunner: finally: self._pool.release(allocation.allocation_id) - @staticmethod - 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( - "unsupported runner parameters: " + ", ".join(sorted(unknown)) - ) - return parameters - def _provenance( self, definition: WorkloadDefinition, @@ -241,14 +293,6 @@ class SciMeshRunner: 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") diff --git a/scimesh/workloads/__init__.py b/scimesh/workloads/__init__.py index e5cc2f1..20cd308 100644 --- a/scimesh/workloads/__init__.py +++ b/scimesh/workloads/__init__.py @@ -6,6 +6,7 @@ from scimesh.core.registry import WorkloadRegistry from scimesh.workloads.help import HelpWorkload from scimesh.workloads.similarity_graph import SimilarityGraphWorkload from scimesh.workloads.similarity_search import SimilaritySearchWorkload +from scimesh.workloads.workload_cli import WorkloadCLI def register_workloads(registry: WorkloadRegistry) -> None: @@ -13,3 +14,4 @@ def register_workloads(registry: WorkloadRegistry) -> None: registry.register(HelpWorkload()) registry.register(SimilaritySearchWorkload()) registry.register(SimilarityGraphWorkload()) + registry.register(WorkloadCLI()) diff --git a/scimesh/workloads/descriptors/definition.py b/scimesh/workloads/descriptors/definition.py index eeafac6..5c90251 100644 --- a/scimesh/workloads/descriptors/definition.py +++ b/scimesh/workloads/descriptors/definition.py @@ -1,53 +1,23 @@ """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 -workflow, a row-bounded planner, pinned descriptor computation, deterministic -shard concatenation, and the exact-artifact verifier. It is the intended -first ``untrusted_quorum`` candidate: ``byte_exact`` determinism with whole -file SHA-256 agreement from distinct owners. +A thin subclass of ``MapReduceWorkload``: the SDK assembles the manifest, the +map/reduce stages, the workflow, and the digest-pinned handlers; this module +only declares the scientific contract (pinned descriptors, canonical CSV, +row-bounded shards, header-preserving concatenation) and the three hooks that +partition, compute, and merge. """ from __future__ import annotations -import hashlib -import shutil from pathlib import Path from typing import Any, Mapping, Sequence -from ...sdk.artifacts import ( - ArtifactCollection, - ArtifactItem, - ArtifactRef, - ArtifactSchema, - Cardinality, - CollectionKind, - OutputManifest, - PortSpec, -) +from scimesh.sdk.artifacts import ArtifactSchema, ComponentRef, PortSpec +from scimesh.sdk.batch import MapReduceWorkload +from scimesh.sdk.identity import SchemaRef, WorkloadId +from scimesh.sdk.registry import WorkloadDefinition + 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 ( DESCRIPTOR_COLUMNS, compute_descriptor_batch, @@ -59,16 +29,6 @@ from .core import ( MAP_ENTRY_POINT = "scimesh.workloads.descriptors.definition:map_descriptors@v1" REDUCE_ENTRY_POINT = "scimesh.workloads.descriptors.definition:reduce_descriptors@v1" -_DESCRIPTOR_PARAMETERS = ("skip_invalid",) - - -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 { @@ -114,14 +74,22 @@ def _descriptor_schema() -> ArtifactSchema: ) -class DescriptorBatchWorkload: - """Manifest-backed planner, runner, and reducer for descriptor-batch. +class DescriptorBatchWorkload(MapReduceWorkload): + """Pinned RDKit 2D descriptor computation, one canonical row per input.""" - The class follows the legacy adapter's structural pattern (one object - 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. - """ + workload_id = WorkloadId("descriptor-batch", "1.0.0") + description = ( + "Compute a pinned set of RDKit 2D descriptors, one canonical " + "CSV row per input molecule, in deterministic input order." + ) + parameters_schema = _parameters_schema() + input_port = PortSpec(_input_schema()) + partial_port = PortSpec(_descriptor_schema()) + output_port = PortSpec(_descriptor_schema()) + map_parameter_names = ("skip_invalid",) + reduce_parameter_names = ("skip_invalid",) + map_entry_point = MAP_ENTRY_POINT + reduce_entry_point = REDUCE_ENTRY_POINT def __init__( self, @@ -137,291 +105,51 @@ class DescriptorBatchWorkload: ): raise ValueError("shard_rows must be a positive integer") validate_descriptor_names() - self.entry_point = MAP_ENTRY_POINT self.shard_rows = shard_rows - self.input_port = PortSpec(_input_schema()) - self.partial_port = PortSpec(_descriptor_schema()) - self.output_port = PortSpec(_descriptor_schema()) - resources = ResourceRequirements( - profile="descriptor-cpu-v1", - cpu_cores=1, - memory_mb=1024, - scratch_mb=1024, - max_duration_seconds=3600, + super().__init__( + package_digest=package_digest, + environment_digest=environment_digest, ) - execution = ExecutionProfile( - profile="descriptor-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=_DESCRIPTOR_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=_DESCRIPTOR_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="descriptor-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("descriptor-batch", "1.0.0"), - description=( - "Compute a pinned set of RDKit 2D descriptors, one canonical " - "CSV row per input molecule, in deterministic input order." - ), - 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=("descriptor-batch",), - 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}, + def domain_validate(self, parameters: Mapping[str, Any]) -> None: + value = parameters.get("skip_invalid", True) + if not isinstance(value, bool): + raise ValueError("skip_invalid must be a boolean") + + def partition_input( + self, + input_path: Path, + parameters: Mapping[str, Any], + workspace: Path, + ) -> list[Path]: + return write_descriptor_shards(input_path, workspace, self.shard_rows) + + def compute_shard( + self, + inputs: Mapping[str, Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + return compute_descriptor_batch( + inputs["input"], + output_path, + skip_invalid=self.domain_validate_check(parameters), ) @staticmethod - def _skip_invalid(parameters: Mapping[str, Any]) -> bool: + def domain_validate_check(parameters: Mapping[str, Any]) -> bool: value = parameters.get("skip_invalid", True) if not isinstance(value, bool): raise ValueError("skip_invalid must be a boolean") return value - def validate(self, request: JobRequest) -> ValidatedJob: - if request.workload != self.manifest.workload: - raise ValueError("descriptor-batch received a request for another workload") - self._skip_invalid(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("descriptor-batch 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) - shard_paths = write_descriptor_shards( - input_path, - workspace, - self.shard_rows, - ) - 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=job.resolved_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=job.resolved_parameters, - tasks=tuple(tasks), - ) - - def run(self, context: TaskContext) -> OutputManifest: - context.cancellation.raise_if_cancelled() - collection = context.task.inputs.get("input") - if collection is None: - raise ValueError("descriptor map task requires one input collection") - self.input_port.validate_collection(collection, "descriptor 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 = compute_descriptor_batch( - input_path, - output_path, - skip_invalid=self._skip_invalid(context.task.parameters), - ) - 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( - "descriptor reducer requires a non-empty keyed partial collection" - ) - self.manifest.workflow.stages[1].inputs["partials"].validate_collection( - collection, - "descriptor 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( - "descriptor partial key must use map." - ) - indexed_items.append((int(raw_index), item)) - expected_keys = context.task.expected_input_keys.get("partials") - if expected_keys is None or {item.key for item in collection.items} != set( - expected_keys - ): - raise ValueError( - "descriptor partial keys do not match the coordinator expected set" - ) - if sorted(index for index, _ in indexed_items) != list( - range(len(indexed_items)) - ): - raise ValueError("descriptor 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 = concatenate_descriptor_shards(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 reduce_partials( + self, + partial_paths: Sequence[Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + return concatenate_descriptor_shards(partial_paths, output_path) def descriptor_batch_sdk_definition( @@ -430,7 +158,7 @@ def descriptor_batch_sdk_definition( package_digest: str | None = None, environment_digest: str | None = None, ) -> DescriptorBatchWorkload: - """Build the default local descriptor-batch definition for tests.""" + """Build the default descriptor-batch definition for tests.""" return DescriptorBatchWorkload( shard_rows=shard_rows, package_digest=package_digest or current_scimesh_package_digest(), @@ -439,5 +167,5 @@ def descriptor_batch_sdk_definition( def workload_definition() -> WorkloadDefinition: - """Installed entry-point factory for the default descriptor-batch definition.""" + """Installed entry-point factory for the descriptor-batch workload.""" return descriptor_batch_sdk_definition().definition() diff --git a/scimesh/workloads/graph/definition.py b/scimesh/workloads/graph/definition.py index c97a7e9..8929b53 100644 --- a/scimesh/workloads/graph/definition.py +++ b/scimesh/workloads/graph/definition.py @@ -1,52 +1,30 @@ """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. +A ``MapReduceWorkload`` subclass with two non-default hooks: ``plan_tasks`` +builds one task per block pair ``(i, j)`` with ``i <= j`` (each task receives +two block inputs), and the partial-key hooks parse ``map.x`` keys and +enforce the CTX-10 pair-coverage invariant. The final edge list 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 typing import Any, Mapping, Sequence -from ...sdk.artifacts import ( +from scimesh.sdk.artifacts import ( ArtifactCollection, - ArtifactItem, - ArtifactRef, ArtifactSchema, - Cardinality, - CollectionKind, - OutputManifest, + ComponentRef, 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 scimesh.sdk.batch import MapReduceWorkload +from scimesh.sdk.identity import SchemaRef, WorkloadId +from scimesh.sdk.plans import TaskSpec, ValidatedJob +from scimesh.sdk.protocols import PlanningContext +from scimesh.sdk.registry import WorkloadDefinition +from scimesh.sdk.workflow import StageSpec + from ..environment import current_environment_digest, current_scimesh_package_digest from .core import ( block_pair_from_key, @@ -63,15 +41,6 @@ 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]: @@ -118,141 +87,32 @@ def _edge_schema() -> ArtifactSchema: ) -class SimilarityGraphSDKWorkload: - """Manifest-backed planner, runner, and reducer for similarity-graph.""" +class SimilarityGraphSDKWorkload(MapReduceWorkload): + """Exact sparse Tanimoto graph over deterministic block pairs.""" - 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() + workload_id = WorkloadId("similarity-graph", "1.0.0") + description = ( + "Exact sparse Tanimoto similarity graph over deterministic " + "block pairs with a duplicate-safe, coverage-checked merge." + ) + parameters_schema = _parameters_schema() + input_port = PortSpec(_molecule_schema()) + block_port = PortSpec(_molecule_schema()) + partial_port = PortSpec(_edge_schema()) + output_port = PortSpec(_edge_schema()) + map_stage_inputs = {"left": block_port, "right": block_port} + map_parameter_names = _MAP_PARAMETERS + reduce_parameter_names = ( + "threshold", + "threshold_direction", + "block_size", + "max_rows", + ) + workflow_id = "graph-block-pairs-v1" + map_entry_point = MAP_ENTRY_POINT + reduce_entry_point = REDUCE_ENTRY_POINT - 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 + def domain_validate(self, parameters: Mapping[str, Any]) -> None: unknown = set(parameters) - { "threshold", "threshold_direction", @@ -275,102 +135,87 @@ class SimilarityGraphSDKWorkload: 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'") + @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 partition_input( + self, + input_path: Path, + parameters: Mapping[str, Any], + workspace: Path, + ) -> list[Path]: block_size = int(parameters.get("block_size", 1_000)) max_rows = parameters.get("max_rows") - blocks, stats = parse_molecule_blocks( + 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] = [] + paths: list[Path] = [] 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, - ) + paths.append(path) + return paths + + def plan_tasks( + self, + shard_paths: Sequence[Path], + resolved: Mapping[str, Any], + job: ValidatedJob, + negotiated: Any, + map_stage: StageSpec, + context: PlanningContext, + ) -> list[TaskSpec]: + task_parameters = { + "threshold": self._unit_interval(resolved.get("threshold"), "threshold"), + "threshold_direction": resolved.get("threshold_direction", "greater"), + } + block_refs = [ + context.sink.seal( + path, + declaration=self.input_port.schema, ) + for path in shard_paths + ] tasks: list[TaskSpec] = [] - for left in range(len(blocks)): - for right in range(left, len(blocks)): + for left in range(len(block_refs)): + for right in range(left, len(block_refs)): 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={ + self.task_spec( + map_stage, + job, + negotiated, + f"map/{left:04d}x{right:04d}", + { **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), - ) + return tasks - def run(self, context: TaskContext) -> OutputManifest: - context.cancellation.raise_if_cancelled() - parameters = context.task.parameters + def compute_shard( + self, + inputs: Mapping[str, Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: left_block = parameters.get("left_block") right_block = parameters.get("right_block") if ( @@ -381,106 +226,32 @@ class SimilarityGraphSDKWorkload: ): 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'") + left = read_block_rows(inputs["left"]) + right = left if diagonal else read_block_rows(inputs["right"]) checked_pairs = ( - len(left_rows) * (len(left_rows) - 1) // 2 - if diagonal - else len(left_rows) * len(right_rows) + len(left) * (len(left) - 1) // 2 if diagonal else len(left) * len(right) ) - edges = compute_block_edges( - left_rows, - right_rows, - threshold, - direction, - ) - output_path = workspace / "result.csv" + edges = compute_block_edges(left, right, threshold, direction) 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, - ) + return {"checked_pairs": checked_pairs, "edges_emitted": len(edges)} - 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 parse_partial_key(self, key: str) -> Any: + return block_pair_from_key(key) + + def validate_partial_keys(self, parsed: Sequence[Any]) -> None: + check_pair_coverage(tuple(parsed)) + + def reduce_partials( + self, + partial_paths: Sequence[Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + return merge_edge_partials(partial_paths, output_path) def similarity_graph_sdk_definition( diff --git a/scimesh/workloads/library.py b/scimesh/workloads/library.py index 091faec..b16f638 100644 --- a/scimesh/workloads/library.py +++ b/scimesh/workloads/library.py @@ -27,9 +27,20 @@ __all__ = [ ] -def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry: - """Registry of every built-in SDK-built workload, all enabled.""" +def default_sdk_registry( + *, + shard_rows: int = 10_000, + allowlist: tuple[AllowedPackage, ...] | None = None, +) -> WorkloadRegistry: + """Registry of every built-in SDK-built workload, all enabled. + + When ``allowlist`` is provided, installed workloads are discovered through + the ``scimesh.workloads`` entry points instead of the built-ins. + """ registry = WorkloadRegistry() + if allowlist: + registry.discover_installed(allowlist) + return registry registry.register( similarity_search_sdk_definition(shard_rows=shard_rows).definition(), enabled=True, @@ -45,8 +56,17 @@ def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry: return registry -def default_sdk_runtime() -> RuntimeCapabilities: - """Runtime advertising the built-in workloads' capabilities and inventory.""" +def default_sdk_runtime( + *, + workload_capabilities: tuple[str, ...] | None = None, + environment_digests: tuple[str, ...] | None = None, +) -> RuntimeCapabilities: + """Runtime advertising the built-in workloads' capabilities and inventory. + + ``workload_capabilities`` and ``environment_digests`` override the built-in + defaults, for example when a registry was populated from an allowlist of + installed user workloads instead of the built-ins. + """ architecture = platform.machine().lower() or "unknown" return RuntimeCapabilities( sdk_api_version=SDK_API_VERSION, @@ -54,15 +74,14 @@ def default_sdk_runtime() -> RuntimeCapabilities: profiles=("core-batch-v1",), features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"}, workload_capabilities=( - "similarity-search", - "similarity-graph", - "descriptor-batch", + workload_capabilities + or ("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(),), + environment_digests=environment_digests or (current_environment_digest(),), ), ) diff --git a/scimesh/workloads/search/core.py b/scimesh/workloads/search/core.py index 39ccb46..cde15b2 100644 --- a/scimesh/workloads/search/core.py +++ b/scimesh/workloads/search/core.py @@ -57,14 +57,12 @@ def run_search_shard( ) -> 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``). + Accepts either a resolved ``query_smiles`` or a raw ``query_id`` that is + resolved against the shard (worker tasks on the v1 wire may still carry + the identifier; the SDK planner always resolves it once at plan time). """ allowed = { + "query_id", "query_smiles", "top_k", "threshold", @@ -77,6 +75,14 @@ def run_search_shard( f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}" ) query_smiles = parameters.get("query_smiles") + query_id = parameters.get("query_id") + if isinstance(query_id, str) and not isinstance(query_smiles, str): + from rdkit import Chem + + from scimesh.chemistry.dataset import find_molecule_by_id + + record = find_molecule_by_id(input_path, query_id) + query_smiles = Chem.MolToSmiles(record.molecule, canonical=True) 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) diff --git a/scimesh/workloads/search/definition.py b/scimesh/workloads/search/definition.py index baa2fb3..4130dee 100644 --- a/scimesh/workloads/search/definition.py +++ b/scimesh/workloads/search/definition.py @@ -1,79 +1,41 @@ """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. +A thin subclass of ``MapReduceWorkload``: the SDK assembles the manifest, +stages, workflow, and digest-pinned handlers. This module declares the search +scientific contract (plan-time query resolution, deterministic sharding, local +top-k per shard, bounded heap merge) and the hooks that partition, compute, +and merge. """ from __future__ import annotations -import hashlib -import shutil from pathlib import Path -from typing import Any, Mapping +from typing import Any, Mapping, Sequence 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 scimesh.sdk.artifacts import ArtifactSchema, ComponentRef, PortSpec +from scimesh.sdk.batch import MapReduceWorkload +from scimesh.sdk.identity import SchemaRef, WorkloadId +from scimesh.sdk.plans import JobRequest, ValidatedJob +from scimesh.sdk.registry import WorkloadDefinition -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_id", "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]: @@ -126,8 +88,30 @@ def _search_table_schema(ref: SchemaRef, canonicalizer: str) -> ArtifactSchema: ) -class SimilaritySearchSDKWorkload: - """Manifest-backed planner, runner, and reducer for similarity-search.""" +class SimilaritySearchSDKWorkload(MapReduceWorkload): + """Exact top-k Tanimoto search over deterministic shards with a bounded merge.""" + + workload_id = WorkloadId("similarity-search", "1.0.0") + description = ( + "Exact top-k Tanimoto molecular similarity search over " + "deterministic TSV shards with a bounded merge." + ) + parameters_schema = _parameters_schema() + input_port = PortSpec(_dataset_schema()) + partial_port = PortSpec( + _search_table_schema( + SchemaRef("similarity-search-partial", 1), "scimesh-search-partial-v1" + ) + ) + output_port = PortSpec( + _search_table_schema( + SchemaRef("similarity-search-result", 1), "scimesh-search-result-v1" + ) + ) + map_parameter_names = _MAP_PARAMETERS + reduce_parameter_names = _MAP_PARAMETERS + ("query_source", "fingerprint") + map_entry_point = MAP_ENTRY_POINT + reduce_entry_point = REDUCE_ENTRY_POINT def __init__( self, @@ -142,122 +126,15 @@ class SimilaritySearchSDKWorkload: 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" - ) + super().__init__( + package_digest=package_digest, + environment_digest=environment_digest, ) - 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}, - ) + # ------------------------------------------------------------------ + # Scientific hooks + # ------------------------------------------------------------------ @staticmethod def _string(value: object, name: str) -> str: @@ -283,12 +160,7 @@ class SimilaritySearchSDKWorkload: 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 + def domain_validate(self, parameters: Mapping[str, Any]) -> None: unknown = set(parameters) - { "query_id", "query_smiles", @@ -322,9 +194,8 @@ class SimilaritySearchSDKWorkload: "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]: + def resolved_parameters(self, request: JobRequest) -> dict[str, Any]: parameters = request.parameters query_id = parameters.get("query_id") if isinstance(query_id, str): @@ -334,7 +205,7 @@ class SimilaritySearchSDKWorkload: "kind": "smiles", "value": self._string(parameters.get("query_smiles"), "query_smiles"), } - resolved: dict[str, object] = { + resolved: dict[str, Any] = { "query_source": query_source, "top_k": self._positive_int(parameters.get("top_k", 20), "top_k"), "threshold_direction": parameters.get("threshold_direction", "greater"), @@ -358,192 +229,54 @@ class SimilaritySearchSDKWorkload: ) 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( + def resolved_parameters_for_plan( + self, + job: ValidatedJob, + input_path: Path, + resolved: dict[str, Any], + ) -> dict[str, Any]: + query_id = job.request.parameters.get("query_id") + if isinstance(query_id, str): + record = find_molecule_by_id(input_path, query_id) + resolved["query_smiles"] = Chem.MolToSmiles(record.molecule, canonical=True) + else: + supplied = job.request.parameters["query_smiles"] + assert isinstance(supplied, str) + molecule = parse_smiles(supplied) + if molecule is None: + raise ValueError("query_smiles is invalid") + resolved["query_smiles"] = Chem.MolToSmiles(molecule, canonical=True) + return resolved + + def partition_input( + self, + input_path: Path, + parameters: Mapping[str, Any], + workspace: Path, + ) -> list[Path]: + max_rows = parameters.get("max_rows") + return 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 compute_shard( + self, + inputs: Mapping[str, Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + return run_search_shard(inputs["input"], parameters, output_path) - def run(self, context: TaskContext) -> OutputManifest: - context.cancellation.raise_if_cancelled() - collection = context.task.inputs.get("input") - if collection is None: - raise ValueError("search map task requires one input collection") - self.input_port.validate_collection(collection, "search map input") - source = context.catalog.materialize(collection.items[0].artifact) - workspace = context.workspace - workspace.mkdir(parents=True, exist_ok=True) - input_path = workspace / "input" - output_path = workspace / "result.csv" - if source.resolve() != input_path.resolve(): - shutil.copyfile(source, input_path) - metrics = run_search_shard( - input_path, - context.task.parameters, - output_path, - ) - context.cancellation.raise_if_cancelled() - sealed = context.sink.seal( - output_path, - declaration=self.partial_port.schema, - ) - return OutputManifest( - context.task.task_key, - {"partial": ArtifactCollection.single(sealed)}, - metrics, - context.provenance, - ).validate_against( - context.task.expected_outputs, - max_output_bytes=self.manifest.limits.max_output_bytes, - ) - - def reduce(self, context: ReduceContext) -> OutputManifest: - context.cancellation.raise_if_cancelled() - collection = context.accepted_inputs.get("partials") - if ( - collection is None - or collection.kind is not CollectionKind.KEYED - or not collection.items - ): - raise ValueError( - "search reducer requires a non-empty keyed partial collection" - ) - self.manifest.workflow.stages[1].inputs["partials"].validate_collection( - collection, - "search reducer partials", - ) - workspace = context.workspace - workspace.mkdir(parents=True, exist_ok=True) - indexed_items: list[tuple[int, ArtifactItem]] = [] - for item in collection.items: - key = item.key or "" - prefix = "map." - raw_index = key[len(prefix) :] if key.startswith(prefix) else "" - if len(raw_index) != 8 or not raw_index.isdigit(): - raise ValueError("search partial key must use map.") - indexed_items.append((int(raw_index), item)) - expected_keys = context.task.expected_input_keys.get("partials") - if expected_keys is None or {item.key for item in collection.items} != set( - expected_keys - ): - raise ValueError( - "search partial keys do not match the coordinator expected set" - ) - if sorted(index for index, _ in indexed_items) != list( - range(len(indexed_items)) - ): - raise ValueError("search partial keys must be complete and contiguous") - partial_paths: list[Path] = [] - for index, item in sorted(indexed_items): - artifact: ArtifactRef = item.artifact - source = context.catalog.materialize(artifact) - target = workspace / artifact.artifact_id - if source.resolve() != target.resolve(): - shutil.copyfile(source, target) - if _sha256_file(target) != artifact.sha256: - raise ValueError("materialized partial checksum does not match") - partial_paths.append(target) - result_path = workspace / "result.csv" - metrics = merge_search_partials( - partial_paths, - context.task.parameters, - result_path, - ) - context.cancellation.raise_if_cancelled() - sealed = context.sink.seal( - result_path, - declaration=self.output_port.schema, - ) - return OutputManifest( - context.task.task_key, - {"result": ArtifactCollection.single(sealed)}, - metrics, - context.provenance, - ).validate_against( - context.task.expected_outputs, - max_output_bytes=self.manifest.limits.max_output_bytes, - ) + def reduce_partials( + self, + partial_paths: Sequence[Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + return merge_search_partials(partial_paths, parameters, output_path) def similarity_search_sdk_definition( diff --git a/scimesh/workloads/workload_cli.py b/scimesh/workloads/workload_cli.py new file mode 100644 index 0000000..488bee4 --- /dev/null +++ b/scimesh/workloads/workload_cli.py @@ -0,0 +1,167 @@ +"""Generic SDK workload runner CLI: ``scimesh workload list|run``. + +This is a generic SDK tool, not a workload. It lists the enabled SDK-built +workloads and executes any of them locally through ``LocalCoreBatchExecutor``, +so a user-written workload package can be inspected and verified without +touching any other part of the program. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import tempfile +from pathlib import Path + +from scimesh.sdk import ( + ArtifactCollection, + JobRequest, + LocalArtifactStore, + LocalCoreBatchExecutor, +) +from scimesh.sdk.registry import workload_allowlist_from_json +from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime + + +class WorkloadCLI: + """Inspect and run SDK-built workloads from the command line.""" + + name = "workload" + help = "List and run SDK-built workloads locally." + + def configure_parser(self, parser: argparse.ArgumentParser) -> None: + subparsers = parser.add_subparsers(dest="workload_command", required=True) + + list_parser = subparsers.add_parser( + "list", help="List installed and enabled SDK workloads." + ) + list_parser.set_defaults(workload_handler=self.list_workloads) + + run_parser = subparsers.add_parser( + "run", help="Run one SDK workload locally against an input file." + ) + run_parser.add_argument( + "name", help="Workload name, for example descriptor-batch" + ) + run_parser.add_argument( + "--version", help="Exact workload version (default: the enabled one)" + ) + run_parser.add_argument( + "--input", required=True, type=Path, help="Input dataset file" + ) + run_parser.add_argument( + "--params", default="{}", help="Job parameters as a JSON object" + ) + run_parser.add_argument( + "--shard-rows", + type=int, + default=10_000, + help="Rows per planned shard for workloads that shard by rows", + ) + run_parser.add_argument( + "-o", + "--output", + type=Path, + default=Path("workload_result.csv"), + help="Output path for the final artifact", + ) + run_parser.add_argument( + "--work-dir", + type=Path, + help="Temporary working directory (default: a fresh temporary directory)", + ) + run_parser.set_defaults(workload_handler=self.run_workload) + + def run(self, args: argparse.Namespace) -> int: + handler = getattr(args, "workload_handler", None) + if handler is None: + raise ValueError("select a workload subcommand: list or run") + return handler(args) + + @staticmethod + def _registry(args: argparse.Namespace): + import os + + allowlist = workload_allowlist_from_json( + os.getenv("SCIMESH_WORKLOAD_ALLOWLIST") + ) + return default_sdk_registry( + shard_rows=getattr(args, "shard_rows", 10_000), + allowlist=allowlist, + ) + + def list_workloads(self, args: argparse.Namespace) -> int: + registry = self._registry(args) + descriptions = registry.descriptions() + if not descriptions: + print("No SDK workloads are installed or enabled.") + return 0 + width = max(len(item.workload.name) for item in descriptions) + for item in sorted(descriptions, key=lambda value: value.workload.name): + digest = item.package_digest.removeprefix("sha256:")[:12] + state = "enabled" if item.enabled else "disabled" + print( + f"{item.workload.name:<{width}} {item.workload.version} " + f"{item.description} [{state} {digest}]" + ) + return 0 + + def run_workload(self, args: argparse.Namespace) -> int: + registry = self._registry(args) + descriptions = registry.descriptions() + try: + parameters = json.loads(args.params) + except (TypeError, json.JSONDecodeError, RecursionError) as error: + raise ValueError("--params must be a valid JSON object") from error + if not isinstance(parameters, dict): + raise ValueError("--params must be a JSON object") + description = next( + ( + item + for item in descriptions + if item.workload.name == args.name + and (args.version is None or item.workload.version == args.version) + ), + None, + ) + if description is None: + raise ValueError(f"unknown or disabled SDK workload: {args.name}") + definition, _ = registry.require( + description.workload.name, + description.workload.version, + description.package_digest, + ) + runtime = default_sdk_runtime( + workload_capabilities=tuple(item.workload.name for item in descriptions), + environment_digests=(definition.manifest.environment.digest,), + ) + if not args.input.is_file(): + raise ValueError(f"input file does not exist: {args.input}") + with tempfile.TemporaryDirectory(prefix="scimesh-workload-") as temporary: + root = Path(temporary) + store = LocalArtifactStore(root / "artifacts") + artifact = store.import_file( + args.input, + declaration=definition.manifest.inputs["input"].schema, + ) + request = JobRequest( + workload=definition.manifest.workload, + parameters=parameters, + inputs={"input": ArtifactCollection.single(artifact)}, + ) + result = LocalCoreBatchExecutor( + registry, + runtime, + store, + args.work_dir or root / "attempts", + ).execute(request, description.package_digest) + result_artifact = result.outputs["result"].items[0].artifact + source = store.materialize(result_artifact) + args.output.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source, args.output) + print( + f"Saved {description.workload.name} result to {args.output} " + f"(metrics: {dict(result.metrics)})" + ) + return 0 diff --git a/tests/test_cli_workload.py b/tests/test_cli_workload.py new file mode 100644 index 0000000..9284fd2 --- /dev/null +++ b/tests/test_cli_workload.py @@ -0,0 +1,229 @@ +"""Tests for the generic ``scimesh workload`` CLI.""" + +from __future__ import annotations + +from pathlib import Path + +from scimesh.cli import main + + +def test_workload_cli_lists_sdk_workloads(capsys: object) -> None: + assert main(["workload", "list"]) == 0 + output = capsys.readouterr().out + assert "descriptor-batch" in output + assert "similarity-graph" in output + assert "similarity-search" in output + assert "enabled" in output + + +def test_workload_cli_runs_descriptor_batch(tmp_path: Path, capsys: object) -> None: + dataset = tmp_path / "molecules.tsv" + dataset.write_text( + "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nC\tCCN\n", + encoding="utf-8", + ) + output = tmp_path / "descriptors.csv" + + code = main( + [ + "workload", + "run", + "descriptor-batch", + "--input", + str(dataset), + "--params", + '{"skip_invalid": true}', + "--shard-rows", + "2", + "-o", + str(output), + ] + ) + + assert code == 0 + lines = output.read_text(encoding="utf-8").splitlines() + assert lines[0].startswith("chembl_id,canonical_smiles,ExactMolWt") + assert len(lines) == 4 + assert "rows_emitted" in capsys.readouterr().out + + +def test_workload_cli_runs_similarity_search(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + dataset.write_text( + "chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n", + encoding="utf-8", + ) + output = tmp_path / "search.csv" + + code = main( + [ + "workload", + "run", + "similarity-search", + "--input", + str(dataset), + "--params", + '{"query_smiles": "CCO", "top_k": 5, "progress_every": 0}', + "--shard-rows", + "2", + "-o", + str(output), + ] + ) + + assert code == 0 + lines = output.read_text(encoding="utf-8").splitlines() + assert lines[0] == "rank,chembl_id,canonical_smiles,similarity" + assert len(lines) == 2 + + +def test_workload_cli_rejects_unknown_or_missing_workload(tmp_path: Path) -> None: + import pytest + + dataset = tmp_path / "molecules.tsv" + dataset.write_text("chembl_id\tcanonical_smiles\nA\tCCO\n", encoding="utf-8") + assert main(["workload", "run", "no-such-workload", "--input", str(dataset)]) == 1 + with pytest.raises(SystemExit): + main(["workload", "run", "descriptor-batch"]) + + +def test_workload_cli_rejects_invalid_params_json(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + dataset.write_text("chembl_id\tcanonical_smiles\nA\tCCO\n", encoding="utf-8") + assert ( + main( + [ + "workload", + "run", + "descriptor-batch", + "--input", + str(dataset), + "--params", + "{broken", + ] + ) + == 1 + ) + + +def test_workload_cli_runs_an_allowlisted_custom_workload( + tmp_path: Path, monkeypatch: object, capsys: object +) -> None: + import csv + + from scimesh.sdk import ( + ArtifactSchema, + ComponentRef, + MapReduceWorkload, + PortSpec, + SchemaRef, + WorkloadId, + ) + from scimesh.sdk.registry import WorkloadRegistry + from scimesh.workloads.environment import ( + current_environment_digest, + current_scimesh_package_digest, + ) + + class CountRowsWorkload(MapReduceWorkload): + workload_id = WorkloadId("count-rows", "1.0.0") + description = "Count TSV data rows per shard." + parameters_schema = { + "type": "object", + "additionalProperties": False, + "properties": {}, + } + input_port = PortSpec( + ArtifactSchema( + SchemaRef("molecule-table", 1), + "text/tab-separated-values", + "utf-8", + 10**9, + ComponentRef("delimited-table", 1), + validator_configuration={ + "required_columns": ["canonical_smiles", "chembl_id"] + }, + max_records=10**8, + ) + ) + partial_port = output_port = PortSpec( + ArtifactSchema( + SchemaRef("count-table", 1), + "text/csv", + "utf-8", + 10**9, + ComponentRef("delimited-table", 1), + validator_configuration={"columns": ["id", "rows"]}, + max_records=10**8, + ) + ) + + def partition_input(self, input_path, parameters, workspace): + paths = [] + with input_path.open(encoding="utf-8", newline="") as source: + for index, row in enumerate(csv.DictReader(source, delimiter="\t")): + path = workspace / f"shard-{index}.tsv" + path.write_text( + "chembl_id\tcanonical_smiles\n" + + row["chembl_id"] + + "\t" + + row["canonical_smiles"] + + "\n", + encoding="utf-8", + ) + paths.append(path) + return paths + + def compute_shard(self, inputs, parameters, output_path): + lines = inputs["input"].read_text(encoding="utf-8").splitlines() + rows = max(len(lines) - 1, 0) + output_path.write_text( + "id,rows\nshard," + str(rows) + "\n", encoding="utf-8" + ) + return {"rows": rows} + + def reduce_partials(self, partial_paths, parameters, output_path): + with output_path.open("w", encoding="utf-8") as destination: + destination.write("id,rows\n") + total = 0 + for partial in partial_paths: + rows = partial.read_text(encoding="utf-8").splitlines()[1:] + destination.write("".join(row + "\n" for row in rows)) + total += len(rows) + return {"rows_total": total, "partial_count": len(partial_paths)} + + definition = CountRowsWorkload( + package_digest=current_scimesh_package_digest(), + environment_digest=current_environment_digest(), + ).definition() + + def fake_discover(self: WorkloadRegistry, allowlist) -> None: + self.register(definition, enabled=True) + + monkeypatch.setattr(WorkloadRegistry, "discover_installed", fake_discover) + monkeypatch.setenv( + "SCIMESH_WORKLOAD_ALLOWLIST", + '[{"distribution": "scimesh", "name": "count-rows", "version": "1.0.0", ' + '"digest": "' + current_scimesh_package_digest() + '"}]', + ) + + dataset = tmp_path / "molecules.tsv" + dataset.write_text( + "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n", encoding="utf-8" + ) + output = tmp_path / "counts.csv" + assert ( + main( + [ + "workload", + "run", + "count-rows", + "--input", + str(dataset), + "-o", + str(output), + ] + ) + == 0 + ) + assert output.read_text(encoding="utf-8") == "id,rows\nshard,1\nshard,1\n" diff --git a/tests/test_sdk_batch.py b/tests/test_sdk_batch.py new file mode 100644 index 0000000..2bc6a74 --- /dev/null +++ b/tests/test_sdk_batch.py @@ -0,0 +1,305 @@ +"""Tests for the MapReduceWorkload authoring scaffold.""" + +from __future__ import annotations + +import csv +from dataclasses import replace +from pathlib import Path + +import pytest + +from scimesh.sdk import ( + ArtifactCollection, + ArtifactSchema, + ComponentRef, + DeterminismProfile, + JobRequest, + LocalArtifactStore, + LocalCoreBatchExecutor, + LocalPlanningContext, + MapReduceWorkload, + PortSpec, + SchemaRef, + StageKind, + TrustMode, + WorkloadId, + WorkloadRegistry, + assert_manifest_round_trip, +) +from scimesh.workloads.library import default_sdk_runtime +from scimesh.workloads.environment import ( + current_environment_digest, + current_scimesh_package_digest, +) + + +def _molecule_port() -> PortSpec: + return PortSpec( + ArtifactSchema( + SchemaRef("molecule-table", 1), + "text/tab-separated-values", + "utf-8", + 10**9, + ComponentRef("delimited-table", 1), + validator_configuration={ + "required_columns": ["canonical_smiles", "chembl_id"] + }, + max_records=10**8, + ) + ) + + +def _count_port() -> PortSpec: + return PortSpec( + ArtifactSchema( + SchemaRef("count-table", 1), + "text/csv", + "utf-8", + 10**9, + ComponentRef("delimited-table", 1), + validator_configuration={"columns": ["id", "rows"]}, + max_records=10**8, + ) + ) + + +class CountRowsWorkload(MapReduceWorkload): + """A minimal author-written workload: three scientific hooks only.""" + + workload_id = WorkloadId("count-rows", "1.0.0") + description = "Count TSV data rows per shard and concatenate the counts." + parameters_schema = { + "type": "object", + "additionalProperties": False, + "properties": {"prefix": {"type": "string", "minLength": 1, "maxLength": 50}}, + } + input_port = _molecule_port() + partial_port = _count_port() + output_port = _count_port() + map_parameter_names = ("prefix",) + reduce_parameter_names = ("prefix",) + + def partition_input(self, input_path, parameters, workspace): + paths = [] + with input_path.open(encoding="utf-8", newline="") as source: + for index, row in enumerate(csv.DictReader(source, delimiter="\t")): + path = workspace / f"shard-{index}.tsv" + path.write_text( + "chembl_id\tcanonical_smiles\n" + + row["chembl_id"] + + "\t" + + row["canonical_smiles"] + + "\n", + encoding="utf-8", + ) + paths.append(path) + return paths + + def compute_shard(self, inputs, parameters, output_path): + lines = inputs["input"].read_text(encoding="utf-8").splitlines() + rows = max(len(lines) - 1, 0) + output_path.write_text( + "id,rows\n" + parameters.get("prefix", "shard") + "," + str(rows) + "\n", + encoding="utf-8", + ) + return {"rows": rows} # type: ignore[return-value] + + def reduce_partials(self, partial_paths, parameters, output_path): + total = 0 + with output_path.open("w", encoding="utf-8") as destination: + destination.write("id,rows\n") + for partial in partial_paths: + for index, line in enumerate( + partial.read_text(encoding="utf-8").splitlines() + ): + if index == 0: + continue + destination.write(line + "\n") + total += int(line.split(",")[1]) + return {"rows_total": total, "partial_count": len(partial_paths)} # type: ignore[return-value] + + +def _registered_count_rows(): + workload = CountRowsWorkload( + package_digest=current_scimesh_package_digest(), + environment_digest=current_environment_digest(), + ) + registry = WorkloadRegistry() + registry.register(workload.definition(), enabled=True) + runtime = replace( + default_sdk_runtime(), + workload_capabilities=( + *default_sdk_runtime().workload_capabilities, + "count-rows", + ), + ) + return workload, registry, runtime + + +def _write_dataset(path: Path) -> None: + path.write_text( + "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nC\tCCN\n", + encoding="utf-8", + ) + + +def _request(workload, store, dataset) -> JobRequest: + artifact = store.import_file( + dataset, + declaration=workload.manifest.inputs["input"].schema, + ) + return JobRequest( + workload=workload.manifest.workload, + parameters={"prefix": "x"}, + inputs={"input": ArtifactCollection.single(artifact)}, + ) + + +def test_map_reduce_scaffold_assembles_the_manifest_and_runs(tmp_path: Path) -> None: + workload, registry, runtime = _registered_count_rows() + manifest = workload.manifest + + assert manifest.workload.name == "count-rows" + 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(workload.definition().runners) == { + manifest.workflow.stages[0].entry_point + } + assert set(workload.definition().reducers) == { + manifest.workflow.stages[1].entry_point + } + assert_manifest_round_trip(manifest) + + dataset = tmp_path / "molecules.tsv" + _write_dataset(dataset) + store = LocalArtifactStore(tmp_path / "artifacts") + result = LocalCoreBatchExecutor( + registry, + runtime, + store, + tmp_path / "work", + ).execute(_request(workload, store, dataset), workload.manifest.package.digest) + + assert result.task_key == "reduce/final" + assert dict(result.metrics) == {"rows_total": 3, "partial_count": 3} + text = store.materialize(result.outputs["result"].items[0].artifact).read_text( + encoding="utf-8" + ) + assert text == "id,rows\nx,1\nx,1\nx,1\n" + + +def test_map_reduce_scaffold_derives_pinned_plans_and_parameters( + tmp_path: Path, +) -> None: + workload, registry, runtime = _registered_count_rows() + dataset = tmp_path / "molecules.tsv" + _write_dataset(dataset) + store = LocalArtifactStore(tmp_path / "artifacts") + request = _request(workload, store, dataset) + input_artifact = request.inputs["input"].items[0].artifact + + plan = registry.plan( + request, + workload.manifest.package.digest, + runtime, + LocalPlanningContext( + store, + store, + tmp_path / "plan", + allowed_artifacts=(input_artifact,), + ), + ) + assert [task.task_key for task in plan.tasks] == [ + "map/00000000", + "map/00000001", + "map/00000002", + ] + assert all(task.parameters == {"prefix": "x"} for task in plan.tasks) + assert all(task.package_digest == plan.package_digest for task in plan.tasks) + assert all(task.manifest_digest == plan.manifest_digest for task in plan.tasks) + assert all(task.trust_mode is TrustMode.TRUSTED for task in plan.tasks) + + +def test_map_reduce_scaffold_requires_scientific_hooks() -> None: + class MissingHooksWorkload(MapReduceWorkload): + workload_id = WorkloadId("missing-hooks", "1.0.0") + description = "A workload that forgets its scientific hooks." + parameters_schema = { + "type": "object", + "additionalProperties": False, + "properties": {}, + } + input_port = _molecule_port() + partial_port = _count_port() + output_port = _count_port() + + workload = MissingHooksWorkload( + package_digest=current_scimesh_package_digest(), + environment_digest=current_environment_digest(), + ) + with pytest.raises(NotImplementedError, match="partition_input"): + workload.partition_input(Path("input"), {}, Path("workspace")) + with pytest.raises(NotImplementedError, match="compute_shard"): + workload.compute_shard({}, {}, Path("output")) + with pytest.raises(NotImplementedError, match="reduce_partials"): + workload.reduce_partials([], {}, Path("output")) + + +def test_map_reduce_scaffold_default_partial_keys_are_contiguous() -> None: + workload = CountRowsWorkload( + package_digest=current_scimesh_package_digest(), + environment_digest=current_environment_digest(), + ) + assert workload.parse_partial_key("map.00000000") == 0 + assert workload.parse_partial_key("map.00000002") == 2 + with pytest.raises(ValueError, match="eight-digit-index"): + workload.parse_partial_key("map.0") + workload.validate_partial_keys((0, 1, 2)) + with pytest.raises(ValueError, match="complete and contiguous"): + workload.validate_partial_keys((0, 2)) + with pytest.raises(ValueError, match="complete and contiguous"): + workload.validate_partial_keys((0, 0, 1)) + + +def test_map_reduce_scaffold_rejects_domain_invalid_parameters(tmp_path: Path) -> None: + workload, registry, runtime = _registered_count_rows() + dataset = tmp_path / "molecules.tsv" + _write_dataset(dataset) + store = LocalArtifactStore(tmp_path / "artifacts") + artifact = store.import_file( + dataset, + declaration=workload.manifest.inputs["input"].schema, + ) + + class StrictCountRows(CountRowsWorkload): + def domain_validate(self, parameters): + if "prefix" not in parameters: + raise ValueError("prefix is required") + + strict = StrictCountRows( + package_digest=current_scimesh_package_digest(), + environment_digest=current_environment_digest(), + ) + registry2 = WorkloadRegistry() + registry2.register(strict.definition(), enabled=True) + request = JobRequest( + workload=strict.manifest.workload, + parameters={}, + inputs={"input": ArtifactCollection.single(artifact)}, + ) + with pytest.raises(ValueError, match="prefix is required"): + registry2.plan( + request, + strict.manifest.package.digest, + runtime, + LocalPlanningContext(store, store, tmp_path / "plan"), + ) diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index 3c35a8b..fb292ae 100644 --- a/tests/test_worker_daemon.py +++ b/tests/test_worker_daemon.py @@ -10,7 +10,7 @@ from urllib.request import Request import pytest -from scimesh.worker.config import WorkerConfig +from scimesh.worker.config import WorkerConfig, _workload_allowlist from scimesh.worker import cli as worker_cli from scimesh.worker.cli import build_parser from scimesh.worker.coordinator import CoordinatorTransientError @@ -579,7 +579,7 @@ def test_runner_executes_search_through_the_sdk_and_rejects_graph( ) -def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_options( +def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_parameters( tmp_path: Path, ) -> None: task_dir = tmp_path / "search" @@ -607,7 +607,7 @@ def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_options( InputArtifact("https://example/input", "a" * 64), {"query_smiles": "CCO", "max_rows": 1}, ) - with pytest.raises(ValueError, match="unsupported runner parameters"): + with pytest.raises(ValueError, match="outside the stage projection"): SciMeshRunner().run(with_max_rows, task_dir) @@ -732,3 +732,111 @@ def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None: worker._register_worker() assert worker.worker_id == "11111111-1111-4111-8111-111111111111" assert worker.config.heartbeat_interval == 15 + + +def test_runner_executes_an_arbitrary_sdk_workload(tmp_path: Path) -> None: + from scimesh.workloads.descriptors import descriptor_batch_sdk_definition + + content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n" + task = ClaimedTask( + "task-1", + 1, + (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(), + "descriptor-batch", + InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()), + {"skip_invalid": True}, + ) + task_dir = tmp_path / "task-1" / "1" + task_dir.mkdir(parents=True) + (task_dir / "input").write_bytes(content) + runner = SciMeshRunner( + definitions={ + "descriptor-batch": descriptor_batch_sdk_definition().definition() + } + ) + + result = runner.run(task, task_dir) + header = result.artifacts[0].path.read_text(encoding="utf-8").splitlines()[0] + assert header.startswith("chembl_id,canonical_smiles,ExactMolWt") + assert result.metrics["rows_emitted"] == 2 + + +def test_runner_rejects_workloads_outside_the_v1_single_input_contract( + tmp_path: Path, +) -> None: + from scimesh.workloads.graph import similarity_graph_sdk_definition + + content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n" + task = ClaimedTask( + "graph-task", + 1, + (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(), + "similarity-graph", + InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()), + {"threshold": 0.5}, + ) + task_dir = tmp_path / "graph" + task_dir.mkdir(parents=True) + (task_dir / "input").write_bytes(content) + runner = SciMeshRunner( + definitions={ + "similarity-graph": similarity_graph_sdk_definition().definition() + } + ) + + with pytest.raises(ValueError, match="v1 single-input contract"): + runner.run(task, task_dir) + + +def test_runner_for_worker_discovers_allowlisted_installed_workloads( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from scimesh.sdk.registry import WorkloadRegistry + from scimesh.workloads.search import similarity_search_sdk_definition + + definition = similarity_search_sdk_definition().definition() + + def fake_discover(self: WorkloadRegistry, allowlist) -> None: + assert len(allowlist) == 1 + self.register(definition, enabled=True) + + monkeypatch.setattr(WorkloadRegistry, "discover_installed", fake_discover) + allowlist = _workload_allowlist( + '[{"distribution": "scimesh", "name": "similarity-search", ' + '"version": "1.0.0", "digest": "sha256:' + "a" * 64 + '"}]' + ) + config = WorkerConfig( + "https://example.test", "worker-1", tmp_path / "work", + capabilities=("similarity-search",), + workload_allowlist=allowlist, + ) + runner = SciMeshRunner.for_worker(config) + assert set(runner._definitions) == {"similarity-search"} + + +def test_worker_config_parses_capabilities_and_workload_allowlist( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("SCIMESH_CAPABILITIES", "similarity-search,descriptor-batch") + config = WorkerConfig.from_environment( + {"coordinator_url": "https://example.test", "work_dir": tmp_path} + ) + assert config.capabilities == ("similarity-search", "descriptor-batch") + assert config.workload_allowlist == () + + monkeypatch.setenv( + "SCIMESH_WORKLOAD_ALLOWLIST", + '[{"distribution": "scimesh", "name": "descriptor-batch", ' + '"version": "1.0.0", "digest": "sha256:' + "b" * 64 + '"}]', + ) + config = WorkerConfig.from_environment( + {"coordinator_url": "https://example.test", "work_dir": tmp_path} + ) + assert len(config.workload_allowlist) == 1 + assert config.workload_allowlist[0].workload.name == "descriptor-batch" + + monkeypatch.setenv("SCIMESH_WORKLOAD_ALLOWLIST", "not-json") + with pytest.raises(ValueError, match="valid JSON"): + WorkerConfig.from_environment( + {"coordinator_url": "https://example.test", "work_dir": tmp_path} + )