From 5c5a2af0a1d5aed4c76a788ab1368cb67ba9e28a Mon Sep 17 00:00:00 2001 From: Emil Date: Sun, 2 Aug 2026 01:09:13 +0300 Subject: [PATCH] Add molwt-filter workload with default scaffold hooks --- STATUS.md | 1 + docs/sdk-handoff.md | 2 + docs/workload-sdk.md | 13 +- pyproject.toml | 1 + scimesh/sdk/batch.py | 110 ++++++++++- scimesh/workloads/descriptors/definition.py | 22 +-- scimesh/workloads/library.py | 7 +- scimesh/workloads/molwt_filter/__init__.py | 25 +++ scimesh/workloads/molwt_filter/core.py | 89 +++++++++ scimesh/workloads/molwt_filter/definition.py | 172 ++++++++++++++++ tests/test_sdk_batch.py | 57 +++++- tests/test_sdk_compatibility.py | 2 +- tests/test_sdk_molwt_filter.py | 197 +++++++++++++++++++ 13 files changed, 661 insertions(+), 37 deletions(-) create mode 100644 scimesh/workloads/molwt_filter/__init__.py create mode 100644 scimesh/workloads/molwt_filter/core.py create mode 100644 scimesh/workloads/molwt_filter/definition.py create mode 100644 tests/test_sdk_molwt_filter.py diff --git a/STATUS.md b/STATUS.md index 38a3bab..dcddfa5 100644 --- a/STATUS.md +++ b/STATUS.md @@ -52,6 +52,7 @@ the complete result-artifact SHA-256 before a task is accepted. | CTX-16 Workload SDK foundation | Implemented | `scimesh.sdk` provides strict immutable manifests/plans/artifacts, digest/trust-pinned tasks, typed DAGs, compatibility negotiation, verifier primitives with owner/binding-safe quorum inputs, resource eligibility/local allocation, measured package discovery, a trusted local core-batch conformance harness, and strict package discovery. Enforcing coordinator/Worker profiles remain fail-closed. | | SDK roadmap step 3: `descriptor-batch` | Implemented | The first SDK-built reference workload (`scimesh/workloads/descriptors/`): pinned 81-name RDKit 2D descriptor set, canonical one-row-per-input CSV, deterministic row-bounded shards, shard-index concatenation with one header, byte-identical local/distributed output, and a two-worker `untrusted_quorum` verifier test. | | SDK-built `similarity-search` and `similarity-graph` | Implemented | Both workloads are 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-built `molwt-filter` | Implemented | The minimal authoring example (`scimesh/workloads/molwt_filter/`): filters molecules by exact RDKit molecular weight with only one scientific hook, using the scaffold's new default sharding and concatenation hooks. Registered in the built-in library and as a `scimesh.workloads` entry point. | | 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 diff --git a/docs/sdk-handoff.md b/docs/sdk-handoff.md index 6d91b22..ade3679 100644 --- a/docs/sdk-handoff.md +++ b/docs/sdk-handoff.md @@ -24,6 +24,8 @@ partial writer moved to `scimesh/workloads/search/core.py`; the partial format is unchanged, so the Go reducer and UI keep working. The runner resolves `query_id` per task and rejects plan-time `max_rows`. +**Default hooks + molwt-filter (2026-08-01):** `MapReduceWorkload` now provides default `partition_input` (row-bounded, header-preserving sharding for delimited inputs, `shard_rows` class attr) and default `reduce_partials` (`concatenate_partial_tables`, one header, byte-identical). A new built-in `molwt-filter@1.0.0` (`scimesh/workloads/molwt_filter/`) demonstrates the minimal authoring surface: only `compute_shard` is workload code. descriptor-batch dropped its now-redundant partition/reduce overrides. + **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 diff --git a/docs/workload-sdk.md b/docs/workload-sdk.md index 68be303..8c78113 100644 --- a/docs/workload-sdk.md +++ b/docs/workload-sdk.md @@ -27,6 +27,10 @@ import the SDK and live outside it. The built-in SciMesh workloads are under - `scimesh/workloads/search/` — SDK-built `similarity-search@1.0.0`; - `scimesh/workloads/graph/` — SDK-built `similarity-graph@1.0.0`; - `scimesh/workloads/descriptors/` — SDK-built `descriptor-batch@1.0.0`; +- `scimesh/workloads/molwt_filter/` — SDK-built `molwt-filter@1.0.0`, the + minimal authoring example: it only declares identity, parameters, ports, + and the `compute_shard` hook, using the scaffold's default sharding and + concatenation; - `scimesh/workloads/library.py` — the built-in library wiring: a default registry containing all three definitions and a runtime advertising their capabilities; @@ -272,8 +276,13 @@ class CountRowsWorkload(MapReduceWorkload): ``` 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: +`definition()`. For workloads whose map output is a simple filtered or +transformed table, the scaffold's defaults already cover partitioning +(row-bounded shards that keep the header) and reduction (concatenation with +one header), so only `compute_shard` has to be written — that is exactly what +the built-in `molwt-filter` workload does. The registry, negotiation, resource +reservation, verification, and the local conformance executor treat the +result like any other workload: ```python from scimesh.sdk import ( diff --git a/pyproject.toml b/pyproject.toml index 2ac1c2e..6c7a98f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ scimesh-worker = "scimesh.worker.cli:main" "similarity-search@1.0.0" = "scimesh.workloads.search:workload_definition" "similarity-graph@1.0.0" = "scimesh.workloads.graph:workload_definition" "descriptor-batch@1.0.0" = "scimesh.workloads.descriptors:workload_definition" +"molwt-filter@1.0.0" = "scimesh.workloads.molwt_filter:workload_definition" [tool.setuptools.packages.find] include = ["scimesh*"] diff --git a/scimesh/sdk/batch.py b/scimesh/sdk/batch.py index fed5a7c..8efe31f 100644 --- a/scimesh/sdk/batch.py +++ b/scimesh/sdk/batch.py @@ -71,6 +71,37 @@ def _default_entry_point(module: str, kind: str) -> str: return f"{module}:{kind}@v1" +def concatenate_partial_tables( + partial_paths: Sequence[Path], + output_path: Path, +) -> dict[str, int]: + """Concatenate partial CSV/TSV tables with exactly one shared header. + + Every partial must start with the same header line; the first partial is + copied verbatim and later partials contribute only their data rows, so the + merged file is byte-identical to a single-process run over the same rows. + """ + if not partial_paths: + raise ValueError("reducer requires at least one partial") + output_path.parent.mkdir(parents=True, exist_ok=True) + header: str | None = None + rows_emitted = 0 + with output_path.open("w", encoding="utf-8", newline="") as destination: + for index, partial in enumerate(partial_paths): + with partial.open("r", encoding="utf-8", newline="") as source: + for line_index, line in enumerate(source): + if line_index == 0: + if header is None: + header = line + destination.write(line) + elif line != header: + raise ValueError("partial tables have inconsistent headers") + continue + destination.write(line) + rows_emitted += 1 + return {"partial_count": len(partial_paths), "rows_emitted": rows_emitted} + + class MapReduceWorkload: """Base class for static byte-exact map/reduce workloads. @@ -87,7 +118,9 @@ class MapReduceWorkload: - ``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; + shard splitting; returns one TSV/CSV file per map task. The default + splits a delimited table into ``shard_rows``-bounded shards that keep + the header, using the input schema's media type to pick the delimiter; - ``plan_tasks(shard_paths, resolved, job, negotiated, map_stage, context)``: task construction (default: one task per shard, ``map/``); - ``compute_shard(inputs, parameters, output_path)``: one map task; @@ -95,12 +128,14 @@ class MapReduceWorkload: - ``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. + deterministic merge. The default concatenates partial tables with one + header and counts the emitted data rows. 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``. + ``map_entry_point``, ``reduce_entry_point``, ``shard_rows`` (default 1000, + used only by the default ``partition_input``). """ workload_id: WorkloadId @@ -121,6 +156,7 @@ class MapReduceWorkload: execution: ExecutionProfile | None = None map_entry_point: str | None = None reduce_entry_point: str | None = None + shard_rows: int = 1_000 def __init__( self, @@ -296,8 +332,62 @@ class MapReduceWorkload: parameters: Mapping[str, Any], workspace: Path, ) -> list[Path]: - """Split the materialized input into deterministic shard files.""" - raise NotImplementedError("partition_input must be implemented") + """Split the materialized input into deterministic shard files. + + The default implementation shards a delimited table by rows: every + shard keeps the header and holds at most ``self.shard_rows`` data + rows, in input order. The delimiter follows the input schema's media + type. Workloads that partition differently (block pairs, sampling) + override this hook. + """ + import csv + + if isinstance(self.shard_rows, bool) or not isinstance(self.shard_rows, int) or self.shard_rows < 1: + raise ValueError("shard_rows must be a positive integer") + media_type = self.input_port.schema.media_type + if media_type == "text/tab-separated-values": + delimiter = "\t" + elif media_type == "text/csv": + delimiter = "," + else: + raise ValueError( + "default sharding requires a delimited input media type: " + media_type + ) + workspace.mkdir(parents=True, exist_ok=True) + paths: list[Path] = [] + destination = None + writer = None + rows_in_shard = 0 + try: + with input_path.open("r", encoding="utf-8", newline="") as source: + reader = csv.DictReader(source, delimiter=delimiter) + fieldnames = tuple(reader.fieldnames or ()) + if not fieldnames: + raise ValueError("dataset has no header row") + for row in reader: + if destination is None or rows_in_shard == self.shard_rows: + if destination is not None: + destination.close() + current = workspace / f"shard-{len(paths)}.tsv" + destination = current.open("w", encoding="utf-8", newline="") + writer = csv.DictWriter( + destination, + fieldnames=list(fieldnames), + delimiter=delimiter, + lineterminator="\n", + ) + writer.writeheader() + paths.append(current) + rows_in_shard = 0 + assert writer is not None + writer.writerow(row) + rows_in_shard += 1 + finally: + if destination is not None: + destination.close() + if not paths: + raise ValueError("dataset has no data rows") + return paths def plan_tasks( self, @@ -367,8 +457,14 @@ class MapReduceWorkload: 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") + """Merge materialized partials into one deterministic final CSV. + + The default implementation concatenates partial tables in key order + with exactly one header: the first partial is copied verbatim and + every later partial contributes only its data rows. Workloads that + merge (top-k, edge sets) override this hook. + """ + return concatenate_partial_tables(partial_paths, output_path) # ------------------------------------------------------------------ # Framework handlers diff --git a/scimesh/workloads/descriptors/definition.py b/scimesh/workloads/descriptors/definition.py index 5c90251..eaa57aa 100644 --- a/scimesh/workloads/descriptors/definition.py +++ b/scimesh/workloads/descriptors/definition.py @@ -9,8 +9,7 @@ partition, compute, and merge. from __future__ import annotations -from pathlib import Path -from typing import Any, Mapping, Sequence +from typing import Any, Mapping from scimesh.sdk.artifacts import ArtifactSchema, ComponentRef, PortSpec from scimesh.sdk.batch import MapReduceWorkload @@ -21,9 +20,7 @@ from ..environment import current_environment_digest, current_scimesh_package_di from .core import ( DESCRIPTOR_COLUMNS, compute_descriptor_batch, - concatenate_descriptor_shards, validate_descriptor_names, - write_descriptor_shards, ) MAP_ENTRY_POINT = "scimesh.workloads.descriptors.definition:map_descriptors@v1" @@ -116,14 +113,6 @@ class DescriptorBatchWorkload(MapReduceWorkload): if not isinstance(value, bool): raise ValueError("skip_invalid must be a boolean") - def partition_input( - self, - input_path: Path, - parameters: Mapping[str, Any], - workspace: Path, - ) -> list[Path]: - return write_descriptor_shards(input_path, workspace, self.shard_rows) - def compute_shard( self, inputs: Mapping[str, Path], @@ -143,15 +132,6 @@ class DescriptorBatchWorkload(MapReduceWorkload): raise ValueError("skip_invalid must be a boolean") return value - def reduce_partials( - self, - partial_paths: Sequence[Path], - parameters: Mapping[str, Any], - output_path: Path, - ) -> Mapping[str, int | float]: - return concatenate_descriptor_shards(partial_paths, output_path) - - def descriptor_batch_sdk_definition( *, shard_rows: int = 10_000, diff --git a/scimesh/workloads/library.py b/scimesh/workloads/library.py index b16f638..750ffb1 100644 --- a/scimesh/workloads/library.py +++ b/scimesh/workloads/library.py @@ -19,6 +19,7 @@ from scimesh.sdk.runtime import RuntimeCapabilities from .descriptors import descriptor_batch_sdk_definition from .environment import current_environment_digest from .graph import similarity_graph_sdk_definition +from .molwt_filter import molwt_filter_sdk_definition from .search import similarity_search_sdk_definition __all__ = [ @@ -53,6 +54,10 @@ def default_sdk_registry( descriptor_batch_sdk_definition(shard_rows=shard_rows).definition(), enabled=True, ) + registry.register( + molwt_filter_sdk_definition(shard_rows=shard_rows).definition(), + enabled=True, + ) return registry @@ -75,7 +80,7 @@ def default_sdk_runtime( features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"}, workload_capabilities=( workload_capabilities - or ("similarity-search", "similarity-graph", "descriptor-batch") + or ("similarity-search", "similarity-graph", "descriptor-batch", "molwt-filter") ), inventory=ResourceInventory( cpu_cores=max(os.cpu_count() or 1, 1), diff --git a/scimesh/workloads/molwt_filter/__init__.py b/scimesh/workloads/molwt_filter/__init__.py new file mode 100644 index 0000000..e1fca80 --- /dev/null +++ b/scimesh/workloads/molwt_filter/__init__.py @@ -0,0 +1,25 @@ +"""SDK-built ``molwt-filter`` workload. + +A minimal authoring example built on ``MapReduceWorkload``: the scaffold's +default sharding and concatenation hooks are used unchanged, so the workload +only declares identity, parameters, ports, and the single scientific hook. +""" + +from .core import MOLWT_COLUMNS, filter_molecules_by_molwt +from .definition import ( + MAP_ENTRY_POINT, + REDUCE_ENTRY_POINT, + MolwtFilterWorkload, + molwt_filter_sdk_definition, + workload_definition, +) + +__all__ = [ + "MAP_ENTRY_POINT", + "MOLWT_COLUMNS", + "REDUCE_ENTRY_POINT", + "MolwtFilterWorkload", + "filter_molecules_by_molwt", + "molwt_filter_sdk_definition", + "workload_definition", +] diff --git a/scimesh/workloads/molwt_filter/core.py b/scimesh/workloads/molwt_filter/core.py new file mode 100644 index 0000000..34deea0 --- /dev/null +++ b/scimesh/workloads/molwt_filter/core.py @@ -0,0 +1,89 @@ +"""Scientific core for the ``molwt-filter`` workload. + +Keeps one canonical row per input molecule whose exact RDKit molecular weight +falls inside the requested bounds; rows keep input order, invalid SMILES are +skipped or fail the run, and molecular weights are serialized with fixed +``%.6f`` formatting so the output is byte-identical across workers. +""" + +from __future__ import annotations + +import csv +from pathlib import Path +from typing import Mapping + +from rdkit import Chem +from rdkit.Chem import Descriptors + +from scimesh.chemistry.dataset import iter_rows + +MOLWT_COLUMNS = ("chembl_id", "canonical_smiles", "molwt") + + +def _bound(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number") + return float(value) + + +def filter_molecules_by_molwt( + input_path: Path, + output_path: Path, + *, + min_molwt: object, + max_molwt: object, + skip_invalid: bool, +) -> dict[str, int]: + """Write one CSV row per molecule whose MolWt is within [min, max].""" + minimum = _bound(min_molwt, "min_molwt") if min_molwt is not None else None + maximum = _bound(max_molwt, "max_molwt") if max_molwt is not None else None + if minimum is None and maximum is None: + raise ValueError("at least one of min_molwt or max_molwt is required") + if minimum is not None and minimum < 0: + raise ValueError("min_molwt must be non-negative") + if maximum is not None and maximum < 0: + raise ValueError("max_molwt must be non-negative") + if minimum is not None and maximum is not None and minimum > maximum: + raise ValueError("min_molwt must not exceed max_molwt") + if not isinstance(skip_invalid, bool): + raise ValueError("skip_invalid must be a boolean") + + output_path.parent.mkdir(parents=True, exist_ok=True) + scanned = 0 + invalid = 0 + emitted = 0 + with output_path.open("w", encoding="utf-8", newline="") as destination: + writer = csv.DictWriter( + destination, + fieldnames=list(MOLWT_COLUMNS), + lineterminator="\n", + ) + writer.writeheader() + for row in iter_rows(input_path): + scanned += 1 + smiles = row.get("canonical_smiles", "") + molecule = Chem.MolFromSmiles(smiles) + if molecule is None: + invalid += 1 + if not skip_invalid: + raise ValueError(f"row {scanned} has an invalid canonical_smiles") + continue + molwt = Descriptors.MolWt(molecule) + if minimum is not None and molwt < minimum: + continue + if maximum is not None and molwt > maximum: + continue + canonical = Chem.MolToSmiles(molecule, canonical=True) + writer.writerow( + { + "chembl_id": row.get("chembl_id", ""), + "canonical_smiles": canonical, + "molwt": f"{molwt:.6f}", + } + ) + emitted += 1 + return { + "rows_scanned": scanned, + "invalid_rows": invalid, + "rows_emitted": emitted, + } diff --git a/scimesh/workloads/molwt_filter/definition.py b/scimesh/workloads/molwt_filter/definition.py new file mode 100644 index 0000000..69bb81d --- /dev/null +++ b/scimesh/workloads/molwt_filter/definition.py @@ -0,0 +1,172 @@ +"""SDK-built ``molwt-filter`` workload definition. + +The minimal authoring example: a subclass of ``MapReduceWorkload`` that only +declares identity, parameters, ports, and one scientific hook. The default +hooks of the scaffold provide deterministic row-bounded sharding and +header-preserving concatenation, so nothing else is needed. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping + +from scimesh.sdk.artifacts import ArtifactSchema, ComponentRef, PortSpec +from scimesh.sdk.batch import MapReduceWorkload +from scimesh.sdk.identity import SchemaRef, WorkloadId +from scimesh.sdk.registry import WorkloadDefinition + +from ..environment import current_environment_digest, current_scimesh_package_digest +from .core import MOLWT_COLUMNS, filter_molecules_by_molwt + +MAP_ENTRY_POINT = "scimesh.workloads.molwt_filter.definition:map_molwt_filter@v1" +REDUCE_ENTRY_POINT = "scimesh.workloads.molwt_filter.definition:reduce_molwt_filter@v1" + + +def _parameters_schema() -> dict[str, Any]: + return { + "type": "object", + "additionalProperties": False, + "properties": { + "min_molwt": { + "type": "number", + "minimum": 0, + "description": "Keep molecules with MolWt >= this value", + }, + "max_molwt": { + "type": "number", + "minimum": 0, + "description": "Keep molecules with MolWt <= this value", + }, + "skip_invalid": { + "type": "boolean", + "default": True, + "description": "Skip rows with invalid SMILES instead of failing", + }, + }, + } + + +def _molecule_schema() -> ArtifactSchema: + return ArtifactSchema( + SchemaRef("molecule-table", 1), + "text/tab-separated-values", + "utf-8", + max_bytes=10 * 1024 * 1024 * 1024, + validator=ComponentRef("delimited-table", 1), + validator_configuration={ + "required_columns": ["canonical_smiles", "chembl_id"], + }, + max_records=100_000_000, + canonicalizer="scimesh-tsv-v1", + ) + + +def _filtered_schema() -> ArtifactSchema: + return ArtifactSchema( + SchemaRef("molwt-filtered-table", 1), + "text/csv", + "utf-8", + max_bytes=100 * 1024 * 1024 * 1024, + validator=ComponentRef("delimited-table", 1), + validator_configuration={ + "columns": list(MOLWT_COLUMNS), + }, + max_records=100_000_000, + canonicalizer="molwt-filtered-table-v1", + ) + + +class MolwtFilterWorkload(MapReduceWorkload): + """Keep molecules whose exact RDKit molecular weight is within bounds.""" + + workload_id = WorkloadId("molwt-filter", "1.0.0") + description = ( + "Filter molecules by exact RDKit molecular weight, one canonical " + "CSV row per kept input molecule, in deterministic input order." + ) + parameters_schema = _parameters_schema() + input_port = PortSpec(_molecule_schema()) + partial_port = PortSpec(_filtered_schema()) + output_port = PortSpec(_filtered_schema()) + map_parameter_names = ("min_molwt", "max_molwt", "skip_invalid") + reduce_parameter_names = ("min_molwt", "max_molwt", "skip_invalid") + map_entry_point = MAP_ENTRY_POINT + reduce_entry_point = REDUCE_ENTRY_POINT + + def __init__( + self, + *, + shard_rows: int = 10_000, + package_digest: str, + environment_digest: str, + ) -> None: + if ( + isinstance(shard_rows, bool) + or not isinstance(shard_rows, int) + or shard_rows < 1 + ): + raise ValueError("shard_rows must be a positive integer") + self.shard_rows = shard_rows + super().__init__( + package_digest=package_digest, + environment_digest=environment_digest, + ) + + def domain_validate(self, parameters: Mapping[str, Any]) -> None: + unknown = set(parameters) - {"min_molwt", "max_molwt", "skip_invalid"} + if unknown: + raise ValueError( + "unsupported molwt-filter parameters: " + ", ".join(sorted(unknown)) + ) + minimum = parameters.get("min_molwt") + maximum = parameters.get("max_molwt") + if minimum is None and maximum is None: + raise ValueError("at least one of min_molwt or max_molwt is required") + for value, name in ((minimum, "min_molwt"), (maximum, "max_molwt")): + if value is not None and ( + isinstance(value, bool) or not isinstance(value, (int, float)) + ): + raise ValueError(f"{name} must be a number") + if ( + minimum is not None + and maximum is not None + and float(minimum) > float(maximum) + ): + raise ValueError("min_molwt must not exceed max_molwt") + skip_invalid = parameters.get("skip_invalid", True) + if not isinstance(skip_invalid, bool): + raise ValueError("skip_invalid must be a boolean") + + def compute_shard( + self, + inputs: Mapping[str, Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + return filter_molecules_by_molwt( + inputs["input"], + output_path, + min_molwt=parameters.get("min_molwt"), + max_molwt=parameters.get("max_molwt"), + skip_invalid=parameters.get("skip_invalid", True), + ) + + +def molwt_filter_sdk_definition( + *, + shard_rows: int = 10_000, + package_digest: str | None = None, + environment_digest: str | None = None, +) -> MolwtFilterWorkload: + """Build the molwt-filter definition for tests.""" + return MolwtFilterWorkload( + shard_rows=shard_rows, + package_digest=package_digest or current_scimesh_package_digest(), + environment_digest=environment_digest or current_environment_digest(), + ) + + +def workload_definition() -> WorkloadDefinition: + """Installed entry-point factory for molwt-filter.""" + return molwt_filter_sdk_definition().definition() diff --git a/tests/test_sdk_batch.py b/tests/test_sdk_batch.py index 2bc6a74..dfbe7f3 100644 --- a/tests/test_sdk_batch.py +++ b/tests/test_sdk_batch.py @@ -229,7 +229,7 @@ def test_map_reduce_scaffold_derives_pinned_plans_and_parameters( assert all(task.trust_mode is TrustMode.TRUSTED for task in plan.tasks) -def test_map_reduce_scaffold_requires_scientific_hooks() -> None: +def test_map_reduce_scaffold_requires_compute_hook_only() -> None: class MissingHooksWorkload(MapReduceWorkload): workload_id = WorkloadId("missing-hooks", "1.0.0") description = "A workload that forgets its scientific hooks." @@ -246,12 +246,59 @@ def test_map_reduce_scaffold_requires_scientific_hooks() -> None: 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")) + # reduce has a scaffold default: header-preserving concatenation that + # fails closed on an empty partial set. + with pytest.raises(ValueError, match="at least one partial"): + workload.reduce_partials([], {}, Path("merged")) + + +def test_scaffold_default_sharding_is_row_bounded_and_header_preserving( + tmp_path: Path, +) -> None: + import csv + + class DefaultShardingWorkload(MapReduceWorkload): + workload_id = WorkloadId("default-sharding", "1.0.0") + description = "Uses only the scaffold defaults." + parameters_schema = { + "type": "object", + "additionalProperties": False, + "properties": {}, + } + input_port = _molecule_port() + partial_port = _count_port() + output_port = _count_port() + + def compute_shard(self, inputs, parameters, output_path): + raise AssertionError("not exercised") + + dataset = tmp_path / "input.tsv" + dataset.write_text( + "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nC\tCCN\nD\tCCCCCC\n", + encoding="utf-8", + ) + workload = DefaultShardingWorkload( + package_digest=current_scimesh_package_digest(), + environment_digest=current_environment_digest(), + ) + workload.shard_rows = 2 + workspace = tmp_path / "shards" + shards = workload.partition_input(dataset, {}, workspace) + + assert [path.name for path in shards] == ["shard-0.tsv", "shard-1.tsv"] + for path in shards: + with path.open(encoding="utf-8", newline="") as source: + rows = list(csv.DictReader(source, delimiter="\t")) + assert rows[0]["chembl_id"] in {"A", "C"} + assert len(rows) <= 2 + other = tmp_path / "other.tsv" + other.write_text("different_header\nX\tY\n", encoding="utf-8") + with pytest.raises(ValueError, match="inconsistent headers"): + workload.reduce_partials( + [dataset, other], {}, tmp_path / "merged.csv" + ) def test_map_reduce_scaffold_default_partial_keys_are_contiguous() -> None: diff --git a/tests/test_sdk_compatibility.py b/tests/test_sdk_compatibility.py index 0a29bb4..bd79ba6 100644 --- a/tests/test_sdk_compatibility.py +++ b/tests/test_sdk_compatibility.py @@ -59,7 +59,7 @@ def _registered_similarity_search(shard_rows: int = 2): registry = default_sdk_registry(shard_rows=shard_rows) runtime = default_sdk_runtime() descriptions = registry.descriptions() - assert len(descriptions) == 3 + assert len(descriptions) == 4 description = next( item for item in descriptions if item.workload.name == "similarity-search" ) diff --git a/tests/test_sdk_molwt_filter.py b/tests/test_sdk_molwt_filter.py new file mode 100644 index 0000000..8ab6b4f --- /dev/null +++ b/tests/test_sdk_molwt_filter.py @@ -0,0 +1,197 @@ +"""Tests for the molwt-filter workload and the default scaffold hooks.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pytest + +from scimesh.sdk import ( + ArtifactCollection, + DeterminismProfile, + JobRequest, + LocalArtifactStore, + LocalCoreBatchExecutor, + StageKind, + assert_manifest_round_trip, +) +from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime +from scimesh.workloads.molwt_filter import ( + filter_molecules_by_molwt, + molwt_filter_sdk_definition, +) + + +def _write_dataset(path: Path) -> None: + path.write_text( + "chembl_id\tcanonical_smiles\n" + "WATER\tO\n" + "ETHANOL\tCCO\n" + "PENTANE\tCCCCC\n" + "BROKEN\tnot-a-smiles\n" + "HEXANE\tCCCCCC\n", + encoding="utf-8", + ) + + +def _registered_molwt_filter(shard_rows: int = 2): + registry = default_sdk_registry(shard_rows=shard_rows) + runtime = default_sdk_runtime() + description = next( + item for item in registry.descriptions() if item.workload.name == "molwt-filter" + ) + definition, negotiated = registry.require( + description.workload.name, + description.workload.version, + description.package_digest, + runtime=runtime, + ) + return registry, runtime, description, definition, negotiated + + +def _request(definition, store, dataset, *, parameters) -> JobRequest: + artifact = store.import_file( + dataset, + declaration=definition.manifest.inputs["input"].schema, + ) + return JobRequest( + workload=definition.manifest.workload, + parameters=parameters, + inputs={"input": ArtifactCollection.single(artifact)}, + ) + + +def test_molwt_filter_manifest_is_registered_and_negotiable() -> None: + _, runtime, description, definition, negotiated = _registered_molwt_filter() + manifest = definition.manifest + + assert description.enabled is True + assert manifest.workload.name == "molwt-filter" + assert manifest.workload.version == "1.0.0" + assert manifest.determinism is DeterminismProfile.BYTE_EXACT + assert manifest.verifier.verifier.canonical == "exact-artifact@1" + assert [stage.kind for stage in manifest.workflow.stages] == [ + StageKind.MAP, + StageKind.REDUCE, + ] + assert set(definition.runners) == {manifest.workflow.stages[0].entry_point} + assert set(definition.reducers) == {manifest.workflow.stages[1].entry_point} + assert negotiated is not None + assert_manifest_round_trip(manifest) + assert runtime is not None + + +def test_local_sdk_executor_matches_molwt_filter_reference(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _write_dataset(dataset) + registry, runtime, description, definition, _ = _registered_molwt_filter() + store = LocalArtifactStore(tmp_path / "artifacts") + parameters = {"min_molwt": 40.0, "max_molwt": 90.0} + request = _request(definition, store, dataset, parameters=parameters) + + result = LocalCoreBatchExecutor( + registry, + runtime, + store, + tmp_path / "work", + ).execute(request, description.package_digest) + artifact = result.outputs["result"].items[0].artifact + + reference = tmp_path / "reference.csv" + reference_metrics = filter_molecules_by_molwt( + dataset, + reference, + min_molwt=40.0, + max_molwt=90.0, + skip_invalid=True, + ) + + assert store.materialize(artifact).read_bytes() == reference.read_bytes() + assert result.task_key == "reduce/final" + assert dict(result.metrics) == { + "partial_count": 3, + "rows_emitted": reference_metrics["rows_emitted"], + } + with store.materialize(artifact).open(encoding="utf-8", newline="") as source: + rows = list(csv.DictReader(source)) + assert [row["chembl_id"] for row in rows] == ["ETHANOL", "PENTANE", "HEXANE"] + assert all(row["molwt"].count(".") == 1 for row in rows) + + +def test_molwt_filter_single_bound_and_invalid_row_policy(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _write_dataset(dataset) + registry, runtime, _, definition, _ = _registered_molwt_filter() + store = LocalArtifactStore(tmp_path / "artifacts") + + lower = _request( + definition, + store, + dataset, + parameters={"min_molwt": 72.0}, + ) + result = LocalCoreBatchExecutor( + registry, runtime, store, tmp_path / "work-lower" + ).execute(lower, definition.manifest.package.digest) + rows = list( + csv.DictReader( + store.materialize(result.outputs["result"].items[0].artifact).open( + encoding="utf-8", newline="" + ) + ) + ) + assert [row["chembl_id"] for row in rows] == ["PENTANE", "HEXANE"] + + strict = _request( + definition, + store, + dataset, + parameters={"max_molwt": 100.0, "skip_invalid": False}, + ) + with pytest.raises(ValueError, match="invalid canonical_smiles"): + LocalCoreBatchExecutor( + registry, runtime, store, tmp_path / "work-strict" + ).execute(strict, definition.manifest.package.digest) + + +def test_molwt_filter_rejects_invalid_parameters(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _write_dataset(dataset) + registry, runtime, _, definition, _ = _registered_molwt_filter() + store = LocalArtifactStore(tmp_path / "artifacts") + input_artifact = store.import_file( + dataset, + declaration=definition.manifest.inputs["input"].schema, + ) + + for parameters, message in ( + ({}, "at least one of min_molwt or max_molwt"), + ({"min_molwt": 50, "max_molwt": 10}, "min_molwt must not exceed"), + ({"min_molwt": "heavy"}, "type mismatch"), + ({"min_molwt": 10, "bogus": 1}, "unknown field bogus"), + ): + request = JobRequest( + workload=definition.manifest.workload, + parameters=parameters, + inputs={"input": ArtifactCollection.single(input_artifact)}, + ) + with pytest.raises(ValueError, match=message): + LocalCoreBatchExecutor(registry, runtime, store, tmp_path / "bad").execute( + request, definition.manifest.package.digest + ) + + +def test_molwt_filter_uses_scaffold_default_sharding(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _write_dataset(dataset) + workload = molwt_filter_sdk_definition(shard_rows=2) + workspace = tmp_path / "shards" + shards = workload.partition_input(dataset, {}, workspace) + assert [path.name for path in shards] == [ + "shard-0.tsv", + "shard-1.tsv", + "shard-2.tsv", + ] + with shards[0].open(encoding="utf-8", newline="") as source: + assert len(list(csv.DictReader(source, delimiter="\t"))) == 2