Add MapReduceWorkload scaffold and generic workload execution
This commit is contained in:
@@ -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 коммит только по явной просьбе).
|
||||||
@@ -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).
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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
@@ -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):
|
||||||
|
|||||||
@@ -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
@@ -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
@@ -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")
|
||||||
|
|||||||
@@ -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())
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
@@ -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,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(),),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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"
|
||||||
@@ -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
@@ -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}
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user