Add MapReduceWorkload scaffold and generic workload execution

This commit is contained in:
Emil
2026-08-02 01:02:24 +03:00
parent 19fbb8e926
commit bc76f386e5
21 changed files with 2241 additions and 1176 deletions
+26
View File
@@ -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 коммит только по явной просьбе).
+3 -1
View File
@@ -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 composed by `scimesh/workloads/library.py` and registered through
`scimesh.workloads` entry points. The Worker Agent executes those SDK-built `scimesh.workloads` entry points. The Worker Agent executes those SDK-built
workloads directly (see `scimesh/worker/runners.py`), so the same scientific 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), [SDK author guide](docs/workload-sdk.md), [contract](docs/scimesh-sdk-contract.md),
and [delivery roadmap](docs/scimesh-sdk-roadmap.md). and [delivery roadmap](docs/scimesh-sdk-roadmap.md).
+6 -4
View File
@@ -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-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. | | 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 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 ## 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 - The worker/coordinator flow accepts both underscore API workload names and
hyphenated names at the runner boundary; the runner normalizes them. hyphenated names at the runner boundary; the runner normalizes them.
- The worker executes SDK-built workloads through `scimesh/worker/runners.py` - The worker executes SDK-built workloads through `scimesh/worker/runners.py`
(a v1-wire bridge over `TaskSpec`/`LocalTaskContext`); the CTX-07 (a workload-generic v1-wire bridge over `TaskSpec`/`LocalTaskContext`);
`DistributedWorkload` protocol module and the SDK compatibility adapter were `query_id` resolution and parameter validation live in the workload itself.
removed. `max_rows` is a plan-time option and is rejected per task. `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 - 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 planner resolves `query_id` once and shares `query_smiles`; the upload UI
currently accepts `query_smiles` only. currently accepts `query_smiles` only.
+18
View File
@@ -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 format is unchanged, so the Go reducer and UI keep working. The runner
resolves `query_id` per task and rejects plan-time `max_rows`. 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/` CTX-16 "Workload SDK foundation" is complete and tested. `scimesh/sdk/`
implements the `core-batch-v1` profile: implements the `core-batch-v1` profile:
+126 -6
View File
@@ -55,7 +55,12 @@ The stable authoring surface is exported from `scimesh.sdk`:
protocols; protocols;
- `OutputManifest` and `Provenance` describe sealed durable results; - `OutputManifest` and `Provenance` describe sealed durable results;
- `WorkloadRegistry` resolves an exact name, version, package digest, runtime, - `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, Persisted manifests, requests, plans, tasks, expansions, outputs, candidates,
decisions, and failures are frozen, recursively immutable, JSON-safe, 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 compared exactly once, no duplicates) before emitting the same
deterministically sorted edge list as the local brute-force reference, for deterministically sorted edge list as the local brute-force reference, for
either threshold direction and any block size; either threshold direction and any block size;
- the v1 worker executes the SDK-built `similarity-search` runner directly: - the v1 worker executes SDK-built workloads directly:
`scimesh/worker/runners.py` is a small wire bridge that builds a `TaskSpec` `scimesh/worker/runners.py` is a workload-generic wire bridge that builds a
with the workload's own pins, reserves resources, seals the partial through `TaskSpec` with the workload's own pins, negotiates against a runtime
a content-addressed store, and uploads the resulting CSV over the unchanged derived from the loaded definitions, reserves resources, seals the partial
coordinator contract. 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.<eight-digit-index>`, contiguous; the graph workload
parses `map.<i>x<j>` 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 ## Package shape and registration
+4
View File
@@ -19,6 +19,7 @@ from .artifacts import (
PortSpec, PortSpec,
Provenance, Provenance,
) )
from .batch import MapReduceWorkload
from .conformance import ( from .conformance import (
CancellationFlag, CancellationFlag,
LocalArtifactStore, LocalArtifactStore,
@@ -77,6 +78,7 @@ from .registry import (
WorkloadDefinition, WorkloadDefinition,
WorkloadDescription, WorkloadDescription,
WorkloadRegistry, WorkloadRegistry,
workload_allowlist_from_json,
) )
from .resources import ( from .resources import (
AcceleratorDevice, AcceleratorDevice,
@@ -155,6 +157,7 @@ __all__ = [
"LocalTaskContext", "LocalTaskContext",
"LoopSpec", "LoopSpec",
"MANIFEST_SCHEMA_VERSION", "MANIFEST_SCHEMA_VERSION",
"MapReduceWorkload",
"NegotiatedWorkload", "NegotiatedWorkload",
"NetworkPolicy", "NetworkPolicy",
"NumericTolerance", "NumericTolerance",
@@ -207,6 +210,7 @@ __all__ = [
"WorkloadLimits", "WorkloadLimits",
"WorkloadManifest", "WorkloadManifest",
"WorkloadRegistry", "WorkloadRegistry",
"workload_allowlist_from_json",
"assert_manifest_round_trip", "assert_manifest_round_trip",
"installed_distribution_digest", "installed_distribution_digest",
"negotiate_manifest", "negotiate_manifest",
+554
View File
@@ -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/<index>``);
- ``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.<key>`` partial key into an orderable identity."""
prefix = "map."
if not key.startswith(prefix):
raise ValueError("partial key must use map.<eight-digit-index>")
raw = key[len(prefix) :]
if len(raw) != 8 or not raw.isdigit():
raise ValueError("partial key must use map.<eight-digit-index>")
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,
)
+216 -56
View File
@@ -24,8 +24,20 @@ from .identity import ComponentRef, WorkloadId
from .integrity import installed_distribution_digest from .integrity import installed_distribution_digest
from .manifest import WorkloadManifest from .manifest import WorkloadManifest
from .plans import JobRequest, ValidatedJob, WorkflowPlan from .plans import JobRequest, ValidatedJob, WorkflowPlan
from .protocols import Planner, PlanningContext, PlanningResources, Reducer, Runner, Verifier from .protocols import (
from .runtime import CompatibilityError, NegotiatedWorkload, RuntimeCapabilities, negotiate_manifest Planner,
PlanningContext,
PlanningResources,
Reducer,
Runner,
Verifier,
)
from .runtime import (
CompatibilityError,
NegotiatedWorkload,
RuntimeCapabilities,
negotiate_manifest,
)
from .schema import validate_parameter_instance from .schema import validate_parameter_instance
from .workflow import StageKind 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") raise ValueError("workload entry point has an invalid module path")
raw_top_level = distribution.read_text("top_level.txt") raw_top_level = distribution.read_text("top_level.txt")
declared = { declared = (
line.strip() {line.strip() for line in raw_top_level.splitlines() if line.strip()}
for line in raw_top_level.splitlines() if raw_top_level is not None
if line.strip() else set()
} if raw_top_level is not None else set() )
root_name = parts[0] root_name = parts[0]
if root_name not in declared: if root_name not in declared:
raise ValueError("workload entry point module is outside its distribution") 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() ownership_root = package_root.resolve()
candidates = [ candidates = [
*(Path(str(module_base) + suffix) for suffix in machinery.SOURCE_SUFFIXES), *(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), Path(str(module_base) + suffix)
*(module_base / ("__init__" + suffix) for suffix in machinery.EXTENSION_SUFFIXES), 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: else:
if len(parts) != 1: 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)) module_base = Path(distribution.locate_file(root_name))
candidates = [ candidates = [
*(Path(str(module_base) + suffix) for suffix in machinery.SOURCE_SUFFIXES), *(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()) 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): 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(): for name, handler in values.items():
canonical = require_string(name, f"{field} entry point", max_length=256) canonical = require_string(name, f"{field} entry point", max_length=256)
if not callable(getattr(handler, method, None)): 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 copied[canonical] = handler
object.__setattr__(self, field, MappingProxyType(copied)) object.__setattr__(self, field, MappingProxyType(copied))
for stage in self.manifest.workflow.stages: for stage in self.manifest.workflow.stages:
@@ -149,23 +175,39 @@ class WorkloadDefinition:
) )
verifier_key = self.manifest.verifier.verifier.canonical verifier_key = self.manifest.verifier.verifier.canonical
if verifier_key not in self.verifiers: 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(): for key, verifier in self.verifiers.items():
try: try:
declared_identity = ComponentRef.from_dict(key) declared_identity = ComponentRef.from_dict(key)
except ValueError as error: except ValueError as error:
raise ValueError("definition verifier keys must be component identities") from error raise ValueError(
if declared_identity.canonical != key or getattr(verifier, "identity", None) != declared_identity: "definition verifier keys must be component identities"
raise ValueError("definition verifier handler identity does not match its key") ) 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] manifest_verifier = self.verifiers[verifier_key]
handler_configuration = getattr(manifest_verifier, "configuration", None) handler_configuration = getattr(manifest_verifier, "configuration", None)
if handler_configuration is None: if handler_configuration is None:
if self.manifest.verifier.configuration: 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): 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: 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( raise ValueError(
f"definition has no installed stage verifier: {stage.verifier.canonical}" f"definition has no installed stage verifier: {stage.verifier.canonical}"
) )
@@ -178,13 +220,54 @@ class AllowedPackage:
digest: str digest: str
def __post_init__(self) -> None: 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): 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("_", "-")) object.__setattr__(self, "distribution", distribution.replace("_", "-"))
if not isinstance(self.workload, WorkloadId): if not isinstance(self.workload, WorkloadId):
raise ValueError("allowed workload must be a 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) @dataclass(frozen=True, slots=True)
@@ -223,14 +306,18 @@ class WorkloadRegistry:
self._enabled: set[tuple[str, str, str]] = set() self._enabled: set[tuple[str, str, str]] = set()
self._lock = RLock() 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): if not isinstance(definition, WorkloadDefinition):
raise ValueError("definition must be a WorkloadDefinition") raise ValueError("definition must be a WorkloadDefinition")
workload = definition.manifest.workload workload = definition.manifest.workload
key = (workload.name, workload.version) key = (workload.name, workload.version)
with self._lock: with self._lock:
if key in self._definitions: 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 self._definitions[key] = definition
if enabled: if enabled:
self._enabled.add((*key, definition.manifest.package.digest)) self._enabled.add((*key, definition.manifest.package.digest))
@@ -240,7 +327,9 @@ class WorkloadRegistry:
with self._lock: with self._lock:
definition = self._registered(name, version) definition = self._registered(name, version)
if digest != definition.manifest.package.digest: 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)) self._enabled.add((definition.manifest.workload.name, version, digest))
def disable(self, name: str, version: str, package_digest: str) -> None: def disable(self, name: str, version: str, package_digest: str) -> None:
@@ -257,7 +346,9 @@ class WorkloadRegistry:
try: try:
return self._definitions[(canonical, version)] return self._definitions[(canonical, version)]
except KeyError as error: 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( def require(
self, self,
@@ -270,10 +361,21 @@ class WorkloadRegistry:
digest = require_sha256(package_digest, "package_digest", prefixed=True) digest = require_sha256(package_digest, "package_digest", prefixed=True)
with self._lock: with self._lock:
definition = self._registered(name, version) definition = self._registered(name, version)
identity = (definition.manifest.workload.name, definition.manifest.workload.version, digest) identity = (
if digest != definition.manifest.package.digest or identity not in self._enabled: 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") 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 return definition, negotiated
def plan( def plan(
@@ -293,31 +395,41 @@ class WorkloadRegistry:
runtime=runtime, runtime=runtime,
) )
assert negotiated is not None 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) self._validate_request_shape(request, definition.manifest)
validated = definition.planner.validate(request) validated = definition.planner.validate(request)
if not isinstance(validated, ValidatedJob) or validated.request != 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( plan = definition.planner.plan(
validated, validated,
_NegotiatedPlanningContext(context, negotiated), _NegotiatedPlanningContext(context, negotiated),
) )
if not isinstance(plan, WorkflowPlan) or plan.workload != request.workload: 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 ( if (
plan.package_digest != definition.manifest.package.digest plan.package_digest != definition.manifest.package.digest
or plan.manifest_digest != definition.manifest.digest or plan.manifest_digest != definition.manifest.digest
or plan.trust_mode is not request.trust_mode or plan.trust_mode is not request.trust_mode
or plan.sdk_api_version != runtime.sdk_api_version or plan.sdk_api_version != runtime.sdk_api_version
or plan.protocol_version != runtime.protocol_version or plan.protocol_version != runtime.protocol_version
or plan.manifest_schema_version != definition.manifest.manifest_schema_version or plan.manifest_schema_version
or plan.workflow_schema_version != definition.manifest.workflow.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.environment_digest != definition.manifest.environment.digest
or plan.verifier != definition.manifest.verifier.verifier or plan.verifier != definition.manifest.verifier.verifier
or plan.selected_features != negotiated.selected_features or plan.selected_features != negotiated.selected_features
or plan.optional_fallbacks != negotiated.optional_fallbacks 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) plan.validate_workflow(definition.manifest.workflow)
self._validate_plan_limits(request, plan, definition.manifest) self._validate_plan_limits(request, plan, definition.manifest)
return WorkflowPlan.from_json(plan.to_json()) return WorkflowPlan.from_json(plan.to_json())
@@ -369,7 +481,9 @@ class WorkloadRegistry:
) )
@staticmethod @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): if set(request.inputs) != set(manifest.inputs):
raise ValueError("job input ports do not match the manifest") raise ValueError("job input ports do not match the manifest")
total_bytes = 0 total_bytes = 0
@@ -380,7 +494,9 @@ class WorkloadRegistry:
for item in request.inputs[name].items: for item in request.inputs[name].items:
existing = artifact_references.get(item.artifact.artifact_id) existing = artifact_references.get(item.artifact.artifact_id)
if existing is not None and existing != item.artifact: 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 artifact_references[item.artifact.artifact_id] = item.artifact
if total_bytes > manifest.limits.max_input_bytes: if total_bytes > manifest.limits.max_input_bytes:
raise ValueError("job inputs exceed the manifest byte limit") raise ValueError("job inputs exceed the manifest byte limit")
@@ -388,7 +504,15 @@ class WorkloadRegistry:
raise ValueError("job inputs exceed the manifest artifact limit") raise ValueError("job inputs exceed the manifest artifact limit")
import json import json
from ._validation import thaw_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") raise ValueError("job parameters exceed the manifest byte limit")
validate_parameter_instance(request.parameters, manifest.parameters_schema) validate_parameter_instance(request.parameters, manifest.parameters_schema)
@@ -408,13 +532,23 @@ class WorkloadRegistry:
for item in collection.items: for item in collection.items:
existing = references.get(item.artifact.artifact_id) existing = references.get(item.artifact.artifact_id)
if existing is not None and existing != item.artifact: 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 references[item.artifact.artifact_id] = item.artifact
if len(canonical_json(task.parameters).encode("utf-8")) > manifest.limits.max_parameter_bytes: if (
raise ValueError("planned task parameters exceed the manifest byte limit") 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: if len(references) > manifest.limits.max_artifacts:
raise ValueError("workflow plan exceeds the manifest artifact limit") 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") raise ValueError("resolved parameters exceed the manifest byte limit")
def descriptions(self) -> tuple[WorkloadDescription, ...]: def descriptions(self) -> tuple[WorkloadDescription, ...]:
@@ -455,7 +589,10 @@ class WorkloadRegistry:
for key, approval in allowed.items(): for key, approval in allowed.items():
if _normalized_distribution_name(key[0]) != distribution: if _normalized_distribution_name(key[0]) != distribution:
continue continue
if entry_point.name != f"{approval.workload.name}@{approval.workload.version}": if (
entry_point.name
!= f"{approval.workload.name}@{approval.workload.version}"
):
continue continue
_validate_entry_point_ownership(entry_point) _validate_entry_point_ownership(entry_point)
# Import policy is process-global, so installed discovery is # Import policy is process-global, so installed discovery is
@@ -463,9 +600,12 @@ class WorkloadRegistry:
# cache prefix prevents pre-existing package pyc files from # cache prefix prevents pre-existing package pyc files from
# being consumed, while dont_write_bytecode keeps the # being consumed, while dont_write_bytecode keeps the
# measured source tree unchanged during both load and factory. # measured source tree unchanged during both load and factory.
with _DISCOVERY_IMPORT_LOCK, TemporaryDirectory( with (
prefix="scimesh-discovery-cache-" _DISCOVERY_IMPORT_LOCK,
) as cache_prefix: TemporaryDirectory(
prefix="scimesh-discovery-cache-"
) as cache_prefix,
):
measured_before = installed_distribution_digest(entry_point.dist) measured_before = installed_distribution_digest(entry_point.dist)
if measured_before != approval.digest: if measured_before != approval.digest:
raise ValueError( raise ValueError(
@@ -479,44 +619,64 @@ class WorkloadRegistry:
loaded = entry_point.load() loaded = entry_point.load()
definition = ( definition = (
loaded() loaded()
if callable(loaded) and not isinstance(loaded, WorkloadDefinition) if callable(loaded)
and not isinstance(loaded, WorkloadDefinition)
else loaded else loaded
) )
finally: finally:
sys.pycache_prefix = previous_cache_prefix sys.pycache_prefix = previous_cache_prefix
sys.dont_write_bytecode = previous_bytecode_policy 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( raise ValueError(
"installed package content changed while loading its entry point" "installed package content changed while loading its entry point"
) )
if not isinstance(definition, WorkloadDefinition): 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: 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: 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: 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: 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) pending.append(definition)
discovered.add(key) discovered.add(key)
break break
missing = sorted(set(allowed) - discovered) missing = sorted(set(allowed) - discovered)
if missing: if missing:
identities = ", ".join(f"{name}@{version}" for _, name, version in 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 = [ pending_keys = [
(definition.manifest.workload.name, definition.manifest.workload.version) (definition.manifest.workload.name, definition.manifest.workload.version)
for definition in pending for definition in pending
] ]
if len(pending_keys) != len(set(pending_keys)): 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: with self._lock:
conflicts = [key for key in pending_keys if key in self._definitions] conflicts = [key for key in pending_keys if key in self._definitions]
if conflicts: if conflicts:
name, version = conflicts[0] 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) definitions = dict(self._definitions)
enabled = set(self._enabled) enabled = set(self._enabled)
for key, definition in zip(pending_keys, pending): for key, definition in zip(pending_keys, pending):
+4 -2
View File
@@ -77,7 +77,9 @@ def main(argv: list[str] | None = None) -> int:
config = WorkerConfig.from_environment(overrides) config = WorkerConfig.from_environment(overrides)
except (TypeError, ValueError) as error: except (TypeError, ValueError) as error:
parser.error(str(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 # One shared token strategy backs both clients: a worker key (exchanged and
# refreshed) or a static bearer token, decided by what the config carries. # refreshed) or a static bearer token, decided by what the config carries.
tokens = provider_from_config( tokens = provider_from_config(
@@ -95,7 +97,7 @@ def main(argv: list[str] | None = None) -> int:
HttpArtifactClient( HttpArtifactClient(
config.coordinator_url, config.request_timeout, token_provider=tokens config.coordinator_url, config.request_timeout, token_provider=tokens
), ),
SciMeshRunner(), SciMeshRunner.for_worker(config),
).run_forever() ).run_forever()
return 0 if completed_without_interruption else 130 return 0 if completed_without_interruption else 130
+78 -13
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import json
from dataclasses import dataclass from dataclasses import dataclass
from math import isfinite from math import isfinite
from pathlib import Path from pathlib import Path
@@ -10,6 +11,8 @@ import socket
from typing import Mapping from typing import Mapping
from urllib.parse import urlsplit from urllib.parse import urlsplit
from scimesh.sdk.registry import AllowedPackage, workload_allowlist_from_json
def _clean_url(value: object | None) -> str | None: def _clean_url(value: object | None) -> str | None:
"""Normalise an optional URL: drop a blank one, strip a trailing slash.""" """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}") 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) @dataclass(frozen=True)
class WorkerConfig: class WorkerConfig:
coordinator_url: str coordinator_url: str
@@ -60,6 +81,11 @@ class WorkerConfig:
"similarity-search", "similarity-search",
"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: def __post_init__(self) -> None:
parsed = urlsplit(self.coordinator_url) parsed = urlsplit(self.coordinator_url)
@@ -72,8 +98,14 @@ class WorkerConfig:
if us.scheme not in {"http", "https"} or not us.hostname: if us.scheme not in {"http", "https"} or not us.hostname:
raise ValueError("userservice_url must be an absolute HTTP(S) URL") raise ValueError("userservice_url must be an absolute HTTP(S) URL")
if self.worker_key is not None and not self.userservice_url: if self.worker_key is not None and not self.userservice_url:
raise ValueError("worker_key requires userservice_url (SCIMESH_USERSERVICE_URL)") raise ValueError(
if isinstance(self.cpu_count, bool) or not isinstance(self.cpu_count, int) or self.cpu_count < 1: "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") raise ValueError("cpu_count must be positive")
if self.worker_id is not None and not isinstance(self.worker_id, str): if self.worker_id is not None and not isinstance(self.worker_id, str):
raise ValueError("worker_id must be a string when set") 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.request_timeout, "request_timeout")
_positive_number(self.heartbeat_interval, "heartbeat_interval") _positive_number(self.heartbeat_interval, "heartbeat_interval")
if self.cleanup_after_seconds is not None: 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 self.max_tasks is not None:
if ( if (
isinstance(self.max_tasks, bool) isinstance(self.max_tasks, bool)
@@ -99,6 +133,18 @@ class WorkerConfig:
raise ValueError("exit_when_idle must be a boolean") raise ValueError("exit_when_idle must be a boolean")
if not self.capabilities: if not self.capabilities:
raise ValueError("capabilities cannot be empty") 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 # Runner subprocesses use a task directory as their cwd. Keep the
# configured root absolute so input/output paths remain valid there # configured root absolute so input/output paths remain valid there
# even when the CLI received a convenient relative --work-dir value. # 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.""" """Build config from environment, allowing typed CLI values to override it."""
values = overrides or {} 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) override = values.get(name)
return override if override is not None else os.getenv(environment, default) 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) cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB") memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS") 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( return cls(
coordinator_url=url.rstrip("/"), coordinator_url=url.rstrip("/"),
worker_id=value("worker_id", "SCIMESH_WORKER_ID"), worker_id=str(worker_id) if worker_id is not None else None,
work_dir=Path(value("work_dir", "SCIMESH_WORK_DIR", "./scimesh-worker-data")), work_dir=Path(str(work_dir)),
worker_name=str(value("worker_name", "SCIMESH_WORKER_NAME", socket.gethostname())), worker_name=str(worker_name),
cpu_count=int(cpu_count), cpu_count=int(cpu_count),
memory_mb=int(memory_mb) if memory_mb is not None else None, memory_mb=int(memory_mb) if memory_mb is not None else None,
poll_interval=float(value("poll_interval", "SCIMESH_POLL_INTERVAL", "2")), poll_interval=float(poll_interval),
request_timeout=float(value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30")), request_timeout=float(request_timeout),
heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")), heartbeat_interval=float(heartbeat_interval),
bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"), bearer_token=str(bearer_token) if bearer_token is not None else None,
worker_key=value("worker_key", "SCIMESH_WORKER_KEY"), worker_key=str(worker_key) if worker_key is not None else None,
userservice_url=_clean_url(value("userservice_url", "SCIMESH_USERSERVICE_URL")), userservice_url=userservice_url,
cleanup_after_seconds=float(cleanup) if cleanup else None, cleanup_after_seconds=float(cleanup) if cleanup else None,
max_tasks=int(max_tasks) if max_tasks is not None 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)), exit_when_idle=bool(values.get("exit_when_idle", False)),
capabilities=_capabilities(capabilities),
workload_allowlist=_workload_allowlist(allowlist),
) )
+91 -47
View File
@@ -3,13 +3,20 @@
The runner is a v1-wire bridge: the coordinator still claims flat tasks and 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 the worker still uploads one partial CSV, but execution goes through the
SDK-built workload's own Runner handler with a real ``TaskSpec``, SDK-built workload's own Runner handler with a real ``TaskSpec``,
provenance, resource reservation, and a content-addressed local store. No provenance, resource reservation, and a content-addressed local store.
legacy distributed-protocol code is involved.
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 from __future__ import annotations
import hashlib import hashlib
import platform
from datetime import datetime, timezone from datetime import datetime, timezone
from pathlib import Path from pathlib import Path
from typing import Mapping, Protocol from typing import Mapping, Protocol
@@ -28,28 +35,21 @@ from scimesh.sdk.conformance import (
ScopedArtifactSink, ScopedArtifactSink,
) )
from scimesh.sdk._validation import canonical_json from scimesh.sdk._validation import canonical_json
from scimesh.sdk.identity import SDK_API_VERSION
from scimesh.sdk.manifest import TrustMode from scimesh.sdk.manifest import TrustMode
from scimesh.sdk.plans import TaskSpec from scimesh.sdk.plans import TaskSpec
from scimesh.sdk.registry import WorkloadDefinition from scimesh.sdk.registry import WorkloadDefinition, WorkloadRegistry
from scimesh.sdk.resources import ResourceAllocation, ResourcePool from scimesh.sdk.resources import ResourceAllocation, ResourceInventory, ResourcePool
from scimesh.sdk.runtime import ( from scimesh.sdk.runtime import (
NegotiatedWorkload, NegotiatedWorkload,
RuntimeCapabilities, RuntimeCapabilities,
negotiate_manifest, negotiate_manifest,
) )
from scimesh.sdk.workflow import StageKind 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 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: def _utc_now() -> str:
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") 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 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: class SciMeshRunner:
"""Execute claimed coordinator tasks through the SDK-built workloads.""" """Execute claimed coordinator tasks through SDK-built workloads."""
def __init__( def __init__(
self, self,
definitions: Mapping[str, WorkloadDefinition] | None = None, definitions: Mapping[str, WorkloadDefinition] | None = None,
*,
inventory: ResourceInventory | None = None,
runtime: RuntimeCapabilities | None = None, runtime: RuntimeCapabilities | None = None,
) -> None: ) -> None:
self._definitions = dict(definitions or {}) self._definitions = dict(definitions or {})
if "similarity-search" not in self._definitions: if "similarity-search" not in self._definitions:
from scimesh.workloads.search import similarity_search_sdk_definition
self._definitions["similarity-search"] = ( self._definitions["similarity-search"] = (
similarity_search_sdk_definition().definition() 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) 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: def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
task_dir = task_dir.resolve() task_dir = task_dir.resolve()
workload = task.workload.replace("_", "-") workload = task.workload.replace("_", "-")
@@ -85,11 +156,14 @@ class SciMeshRunner:
map_stage = next( map_stage = next(
stage for stage in manifest.workflow.stages if stage.kind is StageKind.MAP 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 assert map_stage.verifier is not None
input_path = task_dir / "input" input_path = task_dir / "input"
if not input_path.is_file(): if not input_path.is_file():
raise ValueError("claimed task input is missing") raise ValueError("claimed task input is missing")
parameters = self._resolve_parameters(task, input_path)
store = LocalArtifactStore(task_dir / "sdk-store") store = LocalArtifactStore(task_dir / "sdk-store")
input_ref = store.import_file( input_ref = store.import_file(
input_path, input_path,
@@ -110,7 +184,7 @@ class SciMeshRunner:
optional_fallbacks=negotiated.optional_fallbacks, optional_fallbacks=negotiated.optional_fallbacks,
task_key="map/00000000", task_key="map/00000000",
stage_id=map_stage.stage_id, stage_id=map_stage.stage_id,
parameters=parameters, parameters=task.parameters,
inputs={"input": ArtifactCollection.single(input_ref)}, inputs={"input": ArtifactCollection.single(input_ref)},
expected_outputs=map_stage.outputs, expected_outputs=map_stage.outputs,
resources=map_stage.resources, resources=map_stage.resources,
@@ -149,28 +223,6 @@ class SciMeshRunner:
finally: finally:
self._pool.release(allocation.allocation_id) 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( def _provenance(
self, self,
definition: WorkloadDefinition, definition: WorkloadDefinition,
@@ -241,14 +293,6 @@ class SciMeshRunner:
spec.expected_outputs, spec.expected_outputs,
max_output_bytes=max_output_bytes, 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 sink = context.sink
if not isinstance(sink, ScopedArtifactSink): if not isinstance(sink, ScopedArtifactSink):
raise ValueError("SDK execution requires a scoped artifact sink") raise ValueError("SDK execution requires a scoped artifact sink")
+2
View File
@@ -6,6 +6,7 @@ from scimesh.core.registry import WorkloadRegistry
from scimesh.workloads.help import HelpWorkload from scimesh.workloads.help import HelpWorkload
from scimesh.workloads.similarity_graph import SimilarityGraphWorkload from scimesh.workloads.similarity_graph import SimilarityGraphWorkload
from scimesh.workloads.similarity_search import SimilaritySearchWorkload from scimesh.workloads.similarity_search import SimilaritySearchWorkload
from scimesh.workloads.workload_cli import WorkloadCLI
def register_workloads(registry: WorkloadRegistry) -> None: def register_workloads(registry: WorkloadRegistry) -> None:
@@ -13,3 +14,4 @@ def register_workloads(registry: WorkloadRegistry) -> None:
registry.register(HelpWorkload()) registry.register(HelpWorkload())
registry.register(SimilaritySearchWorkload()) registry.register(SimilaritySearchWorkload())
registry.register(SimilarityGraphWorkload()) registry.register(SimilarityGraphWorkload())
registry.register(WorkloadCLI())
+61 -333
View File
@@ -1,53 +1,23 @@
"""SDK-built ``descriptor-batch`` workload definition and handlers. """SDK-built ``descriptor-batch`` workload definition and handlers.
This module is the first non-adapter reference workload built directly on the A thin subclass of ``MapReduceWorkload``: the SDK assembles the manifest, the
``core-batch-v1`` profile: an explicit immutable manifest, a static map/reduce map/reduce stages, the workflow, and the digest-pinned handlers; this module
workflow, a row-bounded planner, pinned descriptor computation, deterministic only declares the scientific contract (pinned descriptors, canonical CSV,
shard concatenation, and the exact-artifact verifier. It is the intended row-bounded shards, header-preserving concatenation) and the three hooks that
first ``untrusted_quorum`` candidate: ``byte_exact`` determinism with whole partition, compute, and merge.
file SHA-256 agreement from distinct owners.
""" """
from __future__ import annotations from __future__ import annotations
import hashlib
import shutil
from pathlib import Path from pathlib import Path
from typing import Any, Mapping, Sequence from typing import Any, Mapping, Sequence
from ...sdk.artifacts import ( from scimesh.sdk.artifacts import ArtifactSchema, ComponentRef, PortSpec
ArtifactCollection, from scimesh.sdk.batch import MapReduceWorkload
ArtifactItem, from scimesh.sdk.identity import SchemaRef, WorkloadId
ArtifactRef, from scimesh.sdk.registry import WorkloadDefinition
ArtifactSchema,
Cardinality,
CollectionKind,
OutputManifest,
PortSpec,
)
from ..environment import current_environment_digest, current_scimesh_package_digest 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 ( from .core import (
DESCRIPTOR_COLUMNS, DESCRIPTOR_COLUMNS,
compute_descriptor_batch, compute_descriptor_batch,
@@ -59,16 +29,6 @@ from .core import (
MAP_ENTRY_POINT = "scimesh.workloads.descriptors.definition:map_descriptors@v1" MAP_ENTRY_POINT = "scimesh.workloads.descriptors.definition:map_descriptors@v1"
REDUCE_ENTRY_POINT = "scimesh.workloads.descriptors.definition:reduce_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]: def _parameters_schema() -> dict[str, Any]:
return { return {
@@ -114,14 +74,22 @@ def _descriptor_schema() -> ArtifactSchema:
) )
class DescriptorBatchWorkload: class DescriptorBatchWorkload(MapReduceWorkload):
"""Manifest-backed planner, runner, and reducer for descriptor-batch. """Pinned RDKit 2D descriptor computation, one canonical row per input."""
The class follows the legacy adapter's structural pattern (one object workload_id = WorkloadId("descriptor-batch", "1.0.0")
registered under each stage entry point) while remaining fully SDK-built: description = (
sharding is explicit and deterministic, every artifact is sealed through "Compute a pinned set of RDKit 2D descriptors, one canonical "
the bridge-owned sink, and no filesystem path ever enters a plan or task. "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__( def __init__(
self, self,
@@ -137,291 +105,51 @@ class DescriptorBatchWorkload:
): ):
raise ValueError("shard_rows must be a positive integer") raise ValueError("shard_rows must be a positive integer")
validate_descriptor_names() validate_descriptor_names()
self.entry_point = MAP_ENTRY_POINT
self.shard_rows = shard_rows self.shard_rows = shard_rows
self.input_port = PortSpec(_input_schema()) super().__init__(
self.partial_port = PortSpec(_descriptor_schema()) package_digest=package_digest,
self.output_port = PortSpec(_descriptor_schema()) environment_digest=environment_digest,
resources = ResourceRequirements(
profile="descriptor-cpu-v1",
cpu_cores=1,
memory_mb=1024,
scratch_mb=1024,
max_duration_seconds=3600,
) )
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: def domain_validate(self, parameters: Mapping[str, Any]) -> None:
return WorkloadDefinition( value = parameters.get("skip_invalid", True)
manifest=self.manifest, if not isinstance(value, bool):
planner=self, raise ValueError("skip_invalid must be a boolean")
runners={MAP_ENTRY_POINT: self},
reducers={REDUCE_ENTRY_POINT: self}, def partition_input(
verifiers={self._exact_verifier.identity.canonical: self._exact_verifier}, 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 @staticmethod
def _skip_invalid(parameters: Mapping[str, Any]) -> bool: def domain_validate_check(parameters: Mapping[str, Any]) -> bool:
value = parameters.get("skip_invalid", True) value = parameters.get("skip_invalid", True)
if not isinstance(value, bool): if not isinstance(value, bool):
raise ValueError("skip_invalid must be a boolean") raise ValueError("skip_invalid must be a boolean")
return value return value
def validate(self, request: JobRequest) -> ValidatedJob: def reduce_partials(
if request.workload != self.manifest.workload: self,
raise ValueError("descriptor-batch received a request for another workload") partial_paths: Sequence[Path],
self._skip_invalid(request.parameters) parameters: Mapping[str, Any],
return ValidatedJob(request, request.parameters) output_path: Path,
) -> Mapping[str, int | float]:
def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan: return concatenate_descriptor_shards(partial_paths, output_path)
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.<eight-digit-index>"
)
indexed_items.append((int(raw_index), item))
expected_keys = context.task.expected_input_keys.get("partials")
if expected_keys is None or {item.key for item in collection.items} != set(
expected_keys
):
raise ValueError(
"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 descriptor_batch_sdk_definition( def descriptor_batch_sdk_definition(
@@ -430,7 +158,7 @@ def descriptor_batch_sdk_definition(
package_digest: str | None = None, package_digest: str | None = None,
environment_digest: str | None = None, environment_digest: str | None = None,
) -> DescriptorBatchWorkload: ) -> DescriptorBatchWorkload:
"""Build the default local descriptor-batch definition for tests.""" """Build the default descriptor-batch definition for tests."""
return DescriptorBatchWorkload( return DescriptorBatchWorkload(
shard_rows=shard_rows, shard_rows=shard_rows,
package_digest=package_digest or current_scimesh_package_digest(), package_digest=package_digest or current_scimesh_package_digest(),
@@ -439,5 +167,5 @@ def descriptor_batch_sdk_definition(
def workload_definition() -> WorkloadDefinition: 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() return descriptor_batch_sdk_definition().definition()
+115 -344
View File
@@ -1,52 +1,30 @@
"""SDK-built ``similarity-graph`` workload definition and handlers. """SDK-built ``similarity-graph`` workload definition and handlers.
Built directly on the ``core-batch-v1`` profile: molecules are parsed once A ``MapReduceWorkload`` subclass with two non-default hooks: ``plan_tasks``
into deterministic row-ordered blocks, every block pair ``(i, j)`` with builds one task per block pair ``(i, j)`` with ``i <= j`` (each task receives
``i <= j`` becomes one map task, and the reducer enforces the CTX-10 two block inputs), and the partial-key hooks parse ``map.<i>x<j>`` keys and
pair-coverage invariant (every unordered molecule pair compared exactly once) enforce the CTX-10 pair-coverage invariant. The final edge list is
before emitting a deterministically sorted edge list that is byte-identical byte-identical to the local brute-force reference.
to the local brute-force reference.
""" """
from __future__ import annotations from __future__ import annotations
import hashlib
import shutil
from pathlib import Path 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, ArtifactCollection,
ArtifactItem,
ArtifactRef,
ArtifactSchema, ArtifactSchema,
Cardinality, ComponentRef,
CollectionKind,
OutputManifest,
PortSpec, PortSpec,
) )
from ...sdk.execution import ( from scimesh.sdk.batch import MapReduceWorkload
CheckpointPolicy, from scimesh.sdk.identity import SchemaRef, WorkloadId
ExecutionProfile, from scimesh.sdk.plans import TaskSpec, ValidatedJob
NetworkPolicy, from scimesh.sdk.protocols import PlanningContext
RetryPolicy, from scimesh.sdk.registry import WorkloadDefinition
) from scimesh.sdk.workflow import StageSpec
from ...sdk.identity import ComponentRef, SchemaRef, VersionRange, WorkloadId
from ...sdk.manifest import (
DeterminismProfile,
EnvironmentSpec,
PackageSpec,
TrustMode,
VerifierSpec,
WorkloadLimits,
WorkloadManifest,
)
from ...sdk.plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan
from ...sdk.protocols import PlanningContext, ReduceContext, TaskContext
from ...sdk.registry import WorkloadDefinition
from ...sdk.resources import ResourceRequirements
from ...sdk.verification import ExactArtifactVerifier
from ...sdk.workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec
from ..environment import current_environment_digest, current_scimesh_package_digest from ..environment import current_environment_digest, current_scimesh_package_digest
from .core import ( from .core import (
block_pair_from_key, 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" REDUCE_ENTRY_POINT = "scimesh.workloads.graph.definition:reduce_graph@v1"
_MAP_PARAMETERS = ("left_block", "right_block", "threshold", "threshold_direction") _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]: def _parameters_schema() -> dict[str, Any]:
@@ -118,141 +87,32 @@ def _edge_schema() -> ArtifactSchema:
) )
class SimilarityGraphSDKWorkload: class SimilarityGraphSDKWorkload(MapReduceWorkload):
"""Manifest-backed planner, runner, and reducer for similarity-graph.""" """Exact sparse Tanimoto graph over deterministic block pairs."""
def __init__( workload_id = WorkloadId("similarity-graph", "1.0.0")
self, description = (
*, "Exact sparse Tanimoto similarity graph over deterministic "
package_digest: str, "block pairs with a duplicate-safe, coverage-checked merge."
environment_digest: str, )
) -> None: parameters_schema = _parameters_schema()
self.entry_point = MAP_ENTRY_POINT input_port = PortSpec(_molecule_schema())
self.input_port = PortSpec(_molecule_schema()) block_port = PortSpec(_molecule_schema())
self.block_port = PortSpec(_molecule_schema()) partial_port = PortSpec(_edge_schema())
self.partial_port = PortSpec(_edge_schema()) output_port = PortSpec(_edge_schema())
self.output_port = PortSpec(_edge_schema()) map_stage_inputs = {"left": block_port, "right": block_port}
resources = ResourceRequirements( map_parameter_names = _MAP_PARAMETERS
profile="graph-cpu-v1", reduce_parameter_names = (
cpu_cores=1, "threshold",
memory_mb=1024, "threshold_direction",
scratch_mb=1024, "block_size",
max_duration_seconds=3600, "max_rows",
) )
execution = ExecutionProfile( workflow_id = "graph-block-pairs-v1"
profile="graph-python-process-v1", map_entry_point = MAP_ENTRY_POINT
network=NetworkPolicy.TRUSTED, reduce_entry_point = REDUCE_ENTRY_POINT
timeout_seconds=3600,
checkpoint=CheckpointPolicy(),
)
limits = WorkloadLimits(
max_input_bytes=self.input_port.schema.max_bytes,
max_tasks=10_000,
max_output_bytes=self.output_port.schema.max_bytes,
)
trust_modes = ("trusted", "untrusted_quorum")
map_stage = StageSpec(
stage_id="map",
kind=StageKind.MAP,
entry_point=MAP_ENTRY_POINT,
needs=(),
inputs={"left": self.block_port, "right": self.block_port},
outputs={"partial": self.partial_port},
parameter_names=_MAP_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=limits.max_tasks,
cacheable=True,
)
reduce_input = PortSpec(
schema=self.partial_port.schema,
cardinality=Cardinality.MANY,
collection=CollectionKind.KEYED,
)
reduce_stage = StageSpec(
stage_id="reduce",
kind=StageKind.REDUCE,
entry_point=REDUCE_ENTRY_POINT,
needs=("map",),
inputs={"partials": reduce_input},
outputs={"result": self.output_port},
parameter_names=_REDUCE_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=1,
cacheable=True,
)
workflow = WorkflowSpec(
workflow_id="graph-block-pairs-v1",
inputs={"input": self.input_port},
stages=(map_stage, reduce_stage),
edges=(
ArtifactEdge(PortRef("input"), PortRef("left", "map")),
ArtifactEdge(PortRef("input"), PortRef("right", "map")),
ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")),
),
outputs={"result": PortRef("result", "reduce")},
max_tasks=limits.max_tasks,
max_output_bytes=limits.max_output_bytes,
)
self.manifest = WorkloadManifest(
sdk_api=VersionRange(">=1.0,<2.0"),
protocol=VersionRange(">=1,<2"),
workload=WorkloadId("similarity-graph", "1.0.0"),
description=(
"Exact sparse Tanimoto similarity graph over deterministic "
"block pairs with a duplicate-safe, coverage-checked merge."
),
package=PackageSpec("scimesh", package_digest),
environment=EnvironmentSpec(
"python-process",
environment_digest,
{"adapter": "sdk-native"},
),
parameters_schema=_parameters_schema(),
workflow=workflow,
inputs={"input": self.input_port},
outputs={"result": self.output_port},
determinism=DeterminismProfile.BYTE_EXACT,
trust_modes=(TrustMode.TRUSTED, TrustMode.UNTRUSTED_QUORUM),
verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}),
limits=limits,
capabilities=("similarity-graph",),
conformance_profiles=("core-batch-v1",),
)
self._exact_verifier = ExactArtifactVerifier()
def definition(self) -> WorkloadDefinition: def domain_validate(self, parameters: Mapping[str, Any]) -> None:
return WorkloadDefinition(
manifest=self.manifest,
planner=self,
runners={MAP_ENTRY_POINT: self},
reducers={REDUCE_ENTRY_POINT: self},
verifiers={self._exact_verifier.identity.canonical: self._exact_verifier},
)
@staticmethod
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
@staticmethod
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
def validate(self, request: JobRequest) -> ValidatedJob:
if request.workload != self.manifest.workload:
raise ValueError("similarity-graph received a request for another workload")
parameters = request.parameters
unknown = set(parameters) - { unknown = set(parameters) - {
"threshold", "threshold",
"threshold_direction", "threshold_direction",
@@ -275,102 +135,87 @@ class SimilarityGraphSDKWorkload:
self._positive_int(parameters["block_size"], "block_size") self._positive_int(parameters["block_size"], "block_size")
if "max_rows" in parameters: if "max_rows" in parameters:
self._positive_int(parameters["max_rows"], "max_rows") self._positive_int(parameters["max_rows"], "max_rows")
return ValidatedJob(request, request.parameters)
def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan: @staticmethod
if not isinstance(job, ValidatedJob): def _unit_interval(value: object, name: str) -> float:
raise ValueError("job must be a ValidatedJob") if isinstance(value, bool) or not isinstance(value, (int, float)):
collection = job.request.inputs.get("input") raise ValueError(f"{name} must be a number between 0 and 1")
if collection is None: return float(value)
raise ValueError("similarity-graph requires the input port")
self.input_port.validate_collection(collection, "job input") @staticmethod
input_artifact = collection.items[0].artifact def _positive_int(value: object, name: str) -> int:
input_path = context.catalog.materialize(input_artifact) if isinstance(value, bool) or not isinstance(value, int) or value < 1:
workspace = context.workspace raise ValueError(f"{name} must be a positive integer")
workspace.mkdir(parents=True, exist_ok=True) return value
parameters = job.resolved_parameters
threshold = self._unit_interval(parameters.get("threshold"), "threshold") def partition_input(
direction = parameters.get("threshold_direction", "greater") self,
if direction not in {"greater", "less"}: input_path: Path,
raise ValueError("threshold_direction must be 'greater' or 'less'") parameters: Mapping[str, Any],
workspace: Path,
) -> list[Path]:
block_size = int(parameters.get("block_size", 1_000)) block_size = int(parameters.get("block_size", 1_000))
max_rows = parameters.get("max_rows") max_rows = parameters.get("max_rows")
blocks, stats = parse_molecule_blocks( blocks, _stats = parse_molecule_blocks(
input_path, input_path,
block_size, block_size,
int(max_rows) if isinstance(max_rows, int) else None, int(max_rows) if isinstance(max_rows, int) else None,
) )
task_parameters = { paths: list[Path] = []
"threshold": threshold,
"threshold_direction": direction,
}
negotiated = context.negotiated
map_stage = self.manifest.workflow.stages[0]
assert map_stage.verifier is not None
block_refs: list[ArtifactRef] = []
for index, block in enumerate(blocks): for index, block in enumerate(blocks):
path = workspace / f"block-{index:04d}.tsv" path = workspace / f"block-{index:04d}.tsv"
write_block_tsv(block, path) write_block_tsv(block, path)
block_refs.append( paths.append(path)
context.sink.seal( return paths
path,
declaration=self.block_port.schema, 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] = [] tasks: list[TaskSpec] = []
for left in range(len(blocks)): for left in range(len(block_refs)):
for right in range(left, len(blocks)): for right in range(left, len(block_refs)):
tasks.append( tasks.append(
TaskSpec( self.task_spec(
workload=self.manifest.workload, map_stage,
package_digest=self.manifest.package.digest, job,
manifest_digest=self.manifest.digest, negotiated,
trust_mode=job.request.trust_mode, f"map/{left:04d}x{right:04d}",
sdk_api_version=negotiated.sdk_api_version, {
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=map_stage.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
task_key=f"map/{left:04d}x{right:04d}",
stage_id="map",
parameters={
**task_parameters, **task_parameters,
"left_block": left, "left_block": left,
"right_block": right, "right_block": right,
}, },
inputs={ {
"left": ArtifactCollection.single(block_refs[left]), "left": ArtifactCollection.single(block_refs[left]),
"right": ArtifactCollection.single(block_refs[right]), "right": ArtifactCollection.single(block_refs[right]),
}, },
expected_outputs={"partial": self.partial_port},
resources=map_stage.resources,
execution=map_stage.execution,
) )
) )
return WorkflowPlan( return tasks
workload=self.manifest.workload,
package_digest=self.manifest.package.digest,
manifest_digest=self.manifest.digest,
trust_mode=job.request.trust_mode,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=self.manifest.verifier.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
workflow_id=self.manifest.workflow.workflow_id,
resolved_parameters=dict(parameters),
tasks=tuple(tasks),
)
def run(self, context: TaskContext) -> OutputManifest: def compute_shard(
context.cancellation.raise_if_cancelled() self,
parameters = context.task.parameters inputs: Mapping[str, Path],
parameters: Mapping[str, Any],
output_path: Path,
) -> Mapping[str, int | float]:
left_block = parameters.get("left_block") left_block = parameters.get("left_block")
right_block = parameters.get("right_block") right_block = parameters.get("right_block")
if ( if (
@@ -381,106 +226,32 @@ class SimilarityGraphSDKWorkload:
): ):
raise ValueError("graph map task requires block indices") raise ValueError("graph map task requires block indices")
diagonal = left_block == right_block 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") threshold = self._unit_interval(parameters.get("threshold"), "threshold")
direction = parameters.get("threshold_direction", "greater") direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}: if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or '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 = ( checked_pairs = (
len(left_rows) * (len(left_rows) - 1) // 2 len(left) * (len(left) - 1) // 2 if diagonal else len(left) * len(right)
if diagonal
else len(left_rows) * len(right_rows)
) )
edges = compute_block_edges( edges = compute_block_edges(left, right, threshold, direction)
left_rows,
right_rows,
threshold,
direction,
)
output_path = workspace / "result.csv"
write_edge_csv(output_path, edges) write_edge_csv(output_path, edges)
context.cancellation.raise_if_cancelled() return {"checked_pairs": checked_pairs, "edges_emitted": len(edges)}
sealed = context.sink.seal(
output_path,
declaration=self.partial_port.schema,
)
return OutputManifest(
context.task.task_key,
{"partial": ArtifactCollection.single(sealed)},
{"checked_pairs": checked_pairs, "edges_emitted": len(edges)},
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def reduce(self, context: ReduceContext) -> OutputManifest: def parse_partial_key(self, key: str) -> Any:
context.cancellation.raise_if_cancelled() return block_pair_from_key(key)
collection = context.accepted_inputs.get("partials")
if ( def validate_partial_keys(self, parsed: Sequence[Any]) -> None:
collection is None check_pair_coverage(tuple(parsed))
or collection.kind is not CollectionKind.KEYED
or not collection.items def reduce_partials(
): self,
raise ValueError( partial_paths: Sequence[Path],
"graph reducer requires a non-empty keyed partial collection" parameters: Mapping[str, Any],
) output_path: Path,
self.manifest.workflow.stages[1].inputs["partials"].validate_collection( ) -> Mapping[str, int | float]:
collection, return merge_edge_partials(partial_paths, output_path)
"graph reducer partials",
)
pairs = [block_pair_from_key(item.key or "") for item in collection.items]
check_pair_coverage(pairs)
expected_keys = context.task.expected_input_keys.get("partials")
if expected_keys is None or {item.key for item in collection.items} != set(
expected_keys
):
raise ValueError(
"graph partial keys do not match the coordinator expected set"
)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
partial_paths: list[Path] = []
for item in sorted(collection.items, key=lambda value: value.key or ""):
artifact: ArtifactRef = item.artifact
source = context.catalog.materialize(artifact)
target = workspace / artifact.artifact_id
if source.resolve() != target.resolve():
shutil.copyfile(source, target)
if _sha256_file(target) != artifact.sha256:
raise ValueError("materialized partial checksum does not match")
partial_paths.append(target)
result_path = workspace / "result.csv"
metrics = merge_edge_partials(partial_paths, result_path)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
result_path,
declaration=self.output_port.schema,
)
return OutputManifest(
context.task.task_key,
{"result": ArtifactCollection.single(sealed)},
metrics,
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def similarity_graph_sdk_definition( def similarity_graph_sdk_definition(
+27 -8
View File
@@ -27,9 +27,20 @@ __all__ = [
] ]
def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry: def default_sdk_registry(
"""Registry of every built-in SDK-built workload, all enabled.""" *,
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() registry = WorkloadRegistry()
if allowlist:
registry.discover_installed(allowlist)
return registry
registry.register( registry.register(
similarity_search_sdk_definition(shard_rows=shard_rows).definition(), similarity_search_sdk_definition(shard_rows=shard_rows).definition(),
enabled=True, enabled=True,
@@ -45,8 +56,17 @@ def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry:
return registry return registry
def default_sdk_runtime() -> RuntimeCapabilities: def default_sdk_runtime(
"""Runtime advertising the built-in workloads' capabilities and inventory.""" *,
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" architecture = platform.machine().lower() or "unknown"
return RuntimeCapabilities( return RuntimeCapabilities(
sdk_api_version=SDK_API_VERSION, sdk_api_version=SDK_API_VERSION,
@@ -54,15 +74,14 @@ def default_sdk_runtime() -> RuntimeCapabilities:
profiles=("core-batch-v1",), profiles=("core-batch-v1",),
features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"}, features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"},
workload_capabilities=( workload_capabilities=(
"similarity-search", workload_capabilities
"similarity-graph", or ("similarity-search", "similarity-graph", "descriptor-batch")
"descriptor-batch",
), ),
inventory=ResourceInventory( inventory=ResourceInventory(
cpu_cores=max(os.cpu_count() or 1, 1), cpu_cores=max(os.cpu_count() or 1, 1),
memory_mb=4096, memory_mb=4096,
scratch_mb=4096, scratch_mb=4096,
architecture=architecture, architecture=architecture,
environment_digests=(current_environment_digest(),), environment_digests=environment_digests or (current_environment_digest(),),
), ),
) )
+12 -6
View File
@@ -57,14 +57,12 @@ def run_search_shard(
) -> dict[str, int]: ) -> dict[str, int]:
"""Run one planned shard with the local reference implementation. """Run one planned shard with the local reference implementation.
This is the worker entry used by the SDK-built runner. It deliberately Accepts either a resolved ``query_smiles`` or a raw ``query_id`` that is
accepts only resolved ``query_smiles``: resolving an identifier resolved against the shard (worker tasks on the v1 wire may still carry
independently in each shard would make the distributed search the identifier; the SDK planner always resolves it once at plan time).
scientifically invalid, so identifier resolution happens once in the
planner (or at the worker bridge for v1-wire tasks that still carry
``query_id``).
""" """
allowed = { allowed = {
"query_id",
"query_smiles", "query_smiles",
"top_k", "top_k",
"threshold", "threshold",
@@ -77,6 +75,14 @@ def run_search_shard(
f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}" f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}"
) )
query_smiles = parameters.get("query_smiles") 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(): if not isinstance(query_smiles, str) or not query_smiles.strip():
raise ValueError("query_smiles is required for a distributed shard") raise ValueError("query_smiles is required for a distributed shard")
molecule = parse_smiles(query_smiles) molecule = parse_smiles(query_smiles)
+86 -353
View File
@@ -1,79 +1,41 @@
"""SDK-built ``similarity-search`` workload definition and handlers. """SDK-built ``similarity-search`` workload definition and handlers.
A direct ``core-batch-v1`` definition (not the legacy adapter): the planner A thin subclass of ``MapReduceWorkload``: the SDK assembles the manifest,
resolves the query once, shards deterministically, each map task computes the stages, workflow, and digest-pinned handlers. This module declares the search
local top-k with the reference implementation, and the reducer merges the scientific contract (plan-time query resolution, deterministic sharding, local
sorted partials with the same bounded heap and tie-breakers as the local CLI. top-k per shard, bounded heap merge) and the hooks that partition, compute,
The manifest declares ``byte_exact`` with the exact-artifact verifier and both and merge.
``trusted`` and ``untrusted_quorum`` trust modes.
""" """
from __future__ import annotations from __future__ import annotations
import hashlib
import shutil
from pathlib import Path from pathlib import Path
from typing import Any, Mapping from typing import Any, Mapping, Sequence
from rdkit import Chem from rdkit import Chem
from scimesh.chemistry.dataset import find_molecule_by_id, parse_smiles from scimesh.chemistry.dataset import find_molecule_by_id, parse_smiles
from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE 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 ..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 from .core import merge_search_partials, run_search_shard, write_search_shards
MAP_ENTRY_POINT = "scimesh.workloads.search.definition:map_search@v1" MAP_ENTRY_POINT = "scimesh.workloads.search.definition:map_search@v1"
REDUCE_ENTRY_POINT = "scimesh.workloads.search.definition:reduce_search@v1" REDUCE_ENTRY_POINT = "scimesh.workloads.search.definition:reduce_search@v1"
_MAP_PARAMETERS = ( _MAP_PARAMETERS = (
"query_id",
"query_smiles", "query_smiles",
"top_k", "top_k",
"threshold", "threshold",
"threshold_direction", "threshold_direction",
"progress_every", "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]: def _parameters_schema() -> dict[str, Any]:
@@ -126,8 +88,30 @@ def _search_table_schema(ref: SchemaRef, canonicalizer: str) -> ArtifactSchema:
) )
class SimilaritySearchSDKWorkload: class SimilaritySearchSDKWorkload(MapReduceWorkload):
"""Manifest-backed planner, runner, and reducer for similarity-search.""" """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__( def __init__(
self, self,
@@ -142,122 +126,15 @@ class SimilaritySearchSDKWorkload:
or shard_rows < 1 or shard_rows < 1
): ):
raise ValueError("shard_rows must be a positive integer") raise ValueError("shard_rows must be a positive integer")
self.entry_point = MAP_ENTRY_POINT
self.shard_rows = shard_rows self.shard_rows = shard_rows
self.input_port = PortSpec(_dataset_schema()) super().__init__(
self.partial_port = PortSpec( package_digest=package_digest,
_search_table_schema( environment_digest=environment_digest,
SchemaRef("similarity-search-partial", 1), "scimesh-search-partial-v1"
)
) )
self.output_port = PortSpec(
_search_table_schema(
SchemaRef("similarity-search-result", 1), "scimesh-search-result-v1"
)
)
resources = ResourceRequirements(
profile="search-cpu-v1",
cpu_cores=1,
memory_mb=1024,
scratch_mb=1024,
max_duration_seconds=3600,
)
execution = ExecutionProfile(
profile="search-python-process-v1",
network=NetworkPolicy.TRUSTED,
timeout_seconds=3600,
checkpoint=CheckpointPolicy(),
)
limits = WorkloadLimits(
max_input_bytes=self.input_port.schema.max_bytes,
max_tasks=10_000,
max_output_bytes=self.output_port.schema.max_bytes,
)
trust_modes = ("trusted", "untrusted_quorum")
map_stage = StageSpec(
stage_id="map",
kind=StageKind.MAP,
entry_point=MAP_ENTRY_POINT,
needs=(),
inputs={"input": self.input_port},
outputs={"partial": self.partial_port},
parameter_names=_MAP_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=limits.max_tasks,
cacheable=True,
)
reduce_input = PortSpec(
schema=self.partial_port.schema,
cardinality=Cardinality.MANY,
collection=CollectionKind.KEYED,
)
reduce_stage = StageSpec(
stage_id="reduce",
kind=StageKind.REDUCE,
entry_point=REDUCE_ENTRY_POINT,
needs=("map",),
inputs={"partials": reduce_input},
outputs={"result": self.output_port},
parameter_names=_REDUCE_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=1,
cacheable=True,
)
workflow = WorkflowSpec(
workflow_id="search-map-reduce-v1",
inputs={"input": self.input_port},
stages=(map_stage, reduce_stage),
edges=(
ArtifactEdge(PortRef("input"), PortRef("input", "map")),
ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")),
),
outputs={"result": PortRef("result", "reduce")},
max_tasks=limits.max_tasks,
max_output_bytes=limits.max_output_bytes,
)
self.manifest = WorkloadManifest(
sdk_api=VersionRange(">=1.0,<2.0"),
protocol=VersionRange(">=1,<2"),
workload=WorkloadId("similarity-search", "1.0.0"),
description=(
"Exact top-k Tanimoto molecular similarity search over "
"deterministic TSV shards with a bounded merge."
),
package=PackageSpec("scimesh", package_digest),
environment=EnvironmentSpec(
"python-process",
environment_digest,
{"adapter": "sdk-native"},
),
parameters_schema=_parameters_schema(),
workflow=workflow,
inputs={"input": self.input_port},
outputs={"result": self.output_port},
determinism=DeterminismProfile.BYTE_EXACT,
trust_modes=(TrustMode.TRUSTED, TrustMode.UNTRUSTED_QUORUM),
verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}),
limits=limits,
capabilities=("similarity-search",),
conformance_profiles=("core-batch-v1",),
)
self._exact_verifier = ExactArtifactVerifier()
def definition(self) -> WorkloadDefinition: # ------------------------------------------------------------------
return WorkloadDefinition( # Scientific hooks
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 @staticmethod
def _string(value: object, name: str) -> str: 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") raise ValueError(f"{name} must be a number between 0 and 1")
return float(value) return float(value)
def validate(self, request: JobRequest) -> ValidatedJob: def domain_validate(self, parameters: Mapping[str, Any]) -> None:
if request.workload != self.manifest.workload:
raise ValueError(
"similarity-search received a request for another workload"
)
parameters = request.parameters
unknown = set(parameters) - { unknown = set(parameters) - {
"query_id", "query_id",
"query_smiles", "query_smiles",
@@ -322,9 +194,8 @@ class SimilaritySearchSDKWorkload:
"threshold_direction" "threshold_direction"
] not in {"greater", "less"}: ] not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or '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 parameters = request.parameters
query_id = parameters.get("query_id") query_id = parameters.get("query_id")
if isinstance(query_id, str): if isinstance(query_id, str):
@@ -334,7 +205,7 @@ class SimilaritySearchSDKWorkload:
"kind": "smiles", "kind": "smiles",
"value": self._string(parameters.get("query_smiles"), "query_smiles"), "value": self._string(parameters.get("query_smiles"), "query_smiles"),
} }
resolved: dict[str, object] = { resolved: dict[str, Any] = {
"query_source": query_source, "query_source": query_source,
"top_k": self._positive_int(parameters.get("top_k", 20), "top_k"), "top_k": self._positive_int(parameters.get("top_k", 20), "top_k"),
"threshold_direction": parameters.get("threshold_direction", "greater"), "threshold_direction": parameters.get("threshold_direction", "greater"),
@@ -358,192 +229,54 @@ class SimilaritySearchSDKWorkload:
) )
return resolved return resolved
def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan: def resolved_parameters_for_plan(
if not isinstance(job, ValidatedJob): self,
raise ValueError("job must be a ValidatedJob") job: ValidatedJob,
collection = job.request.inputs.get("input") input_path: Path,
if collection is None: resolved: dict[str, Any],
raise ValueError("similarity-search requires the input port") ) -> dict[str, Any]:
self.input_port.validate_collection(collection, "job input") query_id = job.request.parameters.get("query_id")
input_artifact = collection.items[0].artifact if isinstance(query_id, str):
input_path = context.catalog.materialize(input_artifact) record = find_molecule_by_id(input_path, query_id)
workspace = context.workspace resolved["query_smiles"] = Chem.MolToSmiles(record.molecule, canonical=True)
workspace.mkdir(parents=True, exist_ok=True) else:
resolved = dict(job.resolved_parameters) supplied = job.request.parameters["query_smiles"]
query_smiles = self._resolve_query(input_path, job.request.parameters) assert isinstance(supplied, str)
resolved["query_smiles"] = query_smiles molecule = parse_smiles(supplied)
max_rows = resolved.get("max_rows") if molecule is None:
shard_paths = write_search_shards( 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, input_path,
workspace, workspace,
self.shard_rows, self.shard_rows,
int(max_rows) if isinstance(max_rows, int) else None, 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 compute_shard(
def _resolve_query(input_path: Path, parameters: Mapping[str, object]) -> str: self,
query_id = parameters.get("query_id") inputs: Mapping[str, Path],
if isinstance(query_id, str): parameters: Mapping[str, Any],
record = find_molecule_by_id(input_path, query_id) output_path: Path,
return Chem.MolToSmiles(record.molecule, canonical=True) ) -> Mapping[str, int | float]:
supplied = parameters["query_smiles"] return run_search_shard(inputs["input"], parameters, output_path)
assert isinstance(supplied, str)
molecule = parse_smiles(supplied)
if molecule is None:
raise ValueError("query_smiles is invalid")
return Chem.MolToSmiles(molecule, canonical=True)
def run(self, context: TaskContext) -> OutputManifest: def reduce_partials(
context.cancellation.raise_if_cancelled() self,
collection = context.task.inputs.get("input") partial_paths: Sequence[Path],
if collection is None: parameters: Mapping[str, Any],
raise ValueError("search map task requires one input collection") output_path: Path,
self.input_port.validate_collection(collection, "search map input") ) -> Mapping[str, int | float]:
source = context.catalog.materialize(collection.items[0].artifact) return merge_search_partials(partial_paths, parameters, output_path)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
input_path = workspace / "input"
output_path = workspace / "result.csv"
if source.resolve() != input_path.resolve():
shutil.copyfile(source, input_path)
metrics = run_search_shard(
input_path,
context.task.parameters,
output_path,
)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
output_path,
declaration=self.partial_port.schema,
)
return OutputManifest(
context.task.task_key,
{"partial": ArtifactCollection.single(sealed)},
metrics,
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def reduce(self, context: ReduceContext) -> OutputManifest:
context.cancellation.raise_if_cancelled()
collection = context.accepted_inputs.get("partials")
if (
collection is None
or collection.kind is not CollectionKind.KEYED
or not collection.items
):
raise ValueError(
"search reducer requires a non-empty keyed partial collection"
)
self.manifest.workflow.stages[1].inputs["partials"].validate_collection(
collection,
"search reducer partials",
)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
indexed_items: list[tuple[int, ArtifactItem]] = []
for item in collection.items:
key = item.key or ""
prefix = "map."
raw_index = key[len(prefix) :] if key.startswith(prefix) else ""
if len(raw_index) != 8 or not raw_index.isdigit():
raise ValueError("search partial key must use map.<eight-digit-index>")
indexed_items.append((int(raw_index), item))
expected_keys = context.task.expected_input_keys.get("partials")
if expected_keys is None or {item.key for item in collection.items} != set(
expected_keys
):
raise ValueError(
"search partial keys do not match the coordinator expected set"
)
if sorted(index for index, _ in indexed_items) != list(
range(len(indexed_items))
):
raise ValueError("search partial keys must be complete and contiguous")
partial_paths: list[Path] = []
for index, item in sorted(indexed_items):
artifact: ArtifactRef = item.artifact
source = context.catalog.materialize(artifact)
target = workspace / artifact.artifact_id
if source.resolve() != target.resolve():
shutil.copyfile(source, target)
if _sha256_file(target) != artifact.sha256:
raise ValueError("materialized partial checksum does not match")
partial_paths.append(target)
result_path = workspace / "result.csv"
metrics = merge_search_partials(
partial_paths,
context.task.parameters,
result_path,
)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
result_path,
declaration=self.output_port.schema,
)
return OutputManifest(
context.task.task_key,
{"result": ArtifactCollection.single(sealed)},
metrics,
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def similarity_search_sdk_definition( def similarity_search_sdk_definition(
+167
View File
@@ -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
+229
View File
@@ -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"
+305
View File
@@ -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"),
)
+111 -3
View File
@@ -10,7 +10,7 @@ from urllib.request import Request
import pytest 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 import cli as worker_cli
from scimesh.worker.cli import build_parser from scimesh.worker.cli import build_parser
from scimesh.worker.coordinator import CoordinatorTransientError 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, tmp_path: Path,
) -> None: ) -> None:
task_dir = tmp_path / "search" 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), InputArtifact("https://example/input", "a" * 64),
{"query_smiles": "CCO", "max_rows": 1}, {"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) 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() worker._register_worker()
assert worker.worker_id == "11111111-1111-4111-8111-111111111111" assert worker.worker_id == "11111111-1111-4111-8111-111111111111"
assert worker.config.heartbeat_interval == 15 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}
)