diff --git a/coordinator/internal/workloads/workloads.json b/coordinator/internal/workloads/workloads.json index 15f0948..bb187b8 100644 --- a/coordinator/internal/workloads/workloads.json +++ b/coordinator/internal/workloads/workloads.json @@ -597,6 +597,206 @@ "upload_ready": true, "verifier": "exact-artifact@1", "version": "1.0.0" + }, + { + "capabilities": [ + "similarity-search-parallel" + ], + "description": "Exact top-k Tanimoto molecular similarity search over deterministic TSV shards with a bounded merge; each shard is fingerprinted and scored across a thread pool. Output is byte-identical to similarity-search.", + "determinism": "byte_exact", + "enabled": true, + "inputs": { + "input": { + "allow_nested_collections": false, + "canonicalizer": "scimesh-tsv-v1", + "encoding": "utf-8", + "max_bytes": 10737418240, + "max_dimensions": [], + "max_records": 100000000, + "media_type": "text/tab-separated-values", + "privacy_class": "project", + "ref": "molecule-table@1", + "retention_class": "durable", + "streaming": false, + "validator": "delimited-table@1", + "validator_configuration": { + "required_columns": [ + "canonical_smiles", + "chembl_id" + ] + } + } + }, + "name": "similarity-search-parallel", + "outputs": { + "result": { + "allow_nested_collections": false, + "canonicalizer": "scimesh-search-result-v1", + "encoding": "utf-8", + "max_bytes": 1073741824, + "max_dimensions": [], + "max_records": 100000, + "media_type": "text/csv", + "privacy_class": "project", + "ref": "similarity-search-result@1", + "retention_class": "durable", + "streaming": false, + "validator": "delimited-table@1", + "validator_configuration": { + "columns": [ + "rank", + "chembl_id", + "canonical_smiles", + "similarity" + ] + } + } + }, + "parameters_schema": { + "additionalProperties": false, + "oneOf": [ + { + "not": { + "required": [ + "query_smiles" + ] + }, + "required": [ + "query_id" + ] + }, + { + "not": { + "required": [ + "query_id" + ] + }, + "required": [ + "query_smiles" + ] + } + ], + "properties": { + "max_rows": { + "minimum": 1, + "type": "integer" + }, + "progress_every": { + "minimum": 0, + "type": "integer" + }, + "query_id": { + "maxLength": 200, + "minLength": 1, + "type": "string" + }, + "query_smiles": { + "maxLength": 200, + "minLength": 1, + "type": "string" + }, + "threads": { + "description": "Threads used to fingerprint and score one shard (default: CPU count).", + "minimum": 1, + "type": "integer" + }, + "threshold": { + "maximum": 1, + "minimum": 0, + "type": "number" + }, + "threshold_direction": { + "enum": [ + "greater", + "less" + ] + }, + "top_k": { + "minimum": 1, + "type": "integer" + } + }, + "type": "object" + }, + "reduction": "top-k", + "trust_modes": [ + "trusted", + "untrusted_quorum" + ], + "ui_elements": [ + { + "default": null, + "field": "query_id", + "group": "", + "help": "ChEMBL id of the query molecule. Provide exactly one of id or SMILES.", + "label": "Query molecule id", + "options": [], + "order": 1, + "placeholder": "", + "widget": "text" + }, + { + "default": null, + "field": "query_smiles", + "group": "", + "help": "SMILES of the query molecule. Provide exactly one of id or SMILES.", + "label": "Query molecule SMILES", + "options": [], + "order": 2, + "placeholder": "", + "widget": "text" + }, + { + "default": 20, + "field": "top_k", + "group": "", + "help": "Number of most similar molecules to keep per shard (global merge keeps the best of these).", + "label": "Top k", + "options": [], + "order": 3, + "placeholder": "", + "widget": "number" + }, + { + "default": "greater", + "field": "threshold_direction", + "group": "", + "help": "Keep molecules with similarity greater or less than the threshold.", + "label": "Direction", + "options": [ + "greater", + "less" + ], + "order": 4, + "placeholder": "", + "widget": "select" + }, + { + "default": null, + "field": "threshold", + "group": "", + "help": "Optional similarity bound: results are filtered to this direction.", + "label": "Similarity threshold", + "options": [], + "order": 5, + "placeholder": "e.g. 0.8", + "widget": "number" + }, + { + "default": null, + "field": "threads", + "group": "", + "help": "Threads used to fingerprint and score one shard (default: CPU count).", + "label": "Threads per shard", + "options": [], + "order": 6, + "placeholder": "auto", + "widget": "number" + } + ], + "upload_ready": true, + "verifier": "exact-artifact@1", + "version": "1.0.0" } ] } diff --git a/pyproject.toml b/pyproject.toml index 7cbe3d7..2bc0d20 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ scimesh = "scimesh.cli:main" [project.entry-points."scimesh.workloads"] "similarity-search@1.0.0" = "scimesh.workloads.search:workload_definition" +"similarity-search-parallel@1.0.0" = "scimesh.workloads.search_parallel: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" diff --git a/scimesh/workloads/library.py b/scimesh/workloads/library.py index 814739a..e996d88 100644 --- a/scimesh/workloads/library.py +++ b/scimesh/workloads/library.py @@ -21,6 +21,7 @@ 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 +from .search_parallel import similarity_search_parallel_sdk_definition __all__ = [ "default_sdk_registry", @@ -46,6 +47,10 @@ def default_sdk_registry( similarity_search_sdk_definition(shard_rows=shard_rows).definition(), enabled=True, ) + registry.register( + similarity_search_parallel_sdk_definition(shard_rows=shard_rows).definition(), + enabled=True, + ) registry.register( similarity_graph_sdk_definition().definition(), enabled=True, @@ -82,6 +87,7 @@ def default_sdk_runtime( workload_capabilities or ( "similarity-search", + "similarity-search-parallel", "similarity-graph", "descriptor-batch", "molwt-filter", diff --git a/scimesh/workloads/search_parallel/__init__.py b/scimesh/workloads/search_parallel/__init__.py new file mode 100644 index 0000000..9b15e36 --- /dev/null +++ b/scimesh/workloads/search_parallel/__init__.py @@ -0,0 +1,28 @@ +"""SDK-built ``similarity-search-parallel`` workload. + +Same contract as ``similarity-search`` with a per-shard thread pool. See +``core.py`` for the parallel scoring core and ``definition.py`` for the +manifest-backed handlers. +""" + +from .core import ( + run_search_shard_parallel, + search_similar_parallel, + write_search_shards, +) +from .definition import ( + MAP_ENTRY_POINT, + SimilaritySearchParallelSDKWorkload, + similarity_search_parallel_sdk_definition, + workload_definition, +) + +__all__ = [ + "MAP_ENTRY_POINT", + "SimilaritySearchParallelSDKWorkload", + "similarity_search_parallel_sdk_definition", + "workload_definition", + "run_search_shard_parallel", + "search_similar_parallel", + "write_search_shards", +] diff --git a/scimesh/workloads/search_parallel/core.py b/scimesh/workloads/search_parallel/core.py new file mode 100644 index 0000000..07c2e23 --- /dev/null +++ b/scimesh/workloads/search_parallel/core.py @@ -0,0 +1,199 @@ +"""Scientific core for the SDK-built ``similarity-search-parallel`` workload. + +The exact same semantics as ``similarity-search`` — identical partial format, +identical bounded merge, byte-identical output — but the per-molecule +fingerprinting and Tanimoto scoring of one shard run across a thread pool +(``threads`` parameter, default = CPU count). + +Parallelism is confined to the scoring phase: ``ThreadPoolExecutor.map`` keeps +the input row order, so the results are merged exactly like the sequential +reference (same ``_HeapEntry`` logic), which makes the output byte-identical +for every thread count by construction. +""" + +from __future__ import annotations + +import heapq +import os +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Mapping + +from rdkit import Chem +from rdkit.Chem import DataStructs + +from scimesh.chemistry.dataset import MoleculeRecord, parse_smiles +from scimesh.chemistry.fingerprints import fingerprint +from scimesh.workloads.search.core import ( + run_search_shard, + write_search_partial, + write_search_shards, +) +from scimesh.workloads.similarity_search import ( + DatasetStats, + SearchResult, + SimilarityMatch, + _HeapEntry, + iter_valid_molecules, +) + + +def search_similar_parallel( + tsv_path: Path, + query: MoleculeRecord, + top_k: int, + *, + threads: int = 0, + max_rows: int | None = None, + threshold: float | None = None, + threshold_direction: str = "greater", +) -> SearchResult: + """Exact top-k matches with a bounded heap, scored by a thread pool. + + Identical selection and ordering to ``search_similar`` for every thread + count: the merge runs in row order over the parallel-computed scores. + """ + if top_k < 1: + raise ValueError("--top-k must be a positive integer") + if threads < 0: + raise ValueError("threads must be a non-negative integer") + if threshold is not None and not 0.0 <= threshold <= 1.0: + raise ValueError("--threshold must be between 0 and 1") + if threshold_direction not in {"greater", "less"}: + raise ValueError("--threshold-direction must be 'greater' or 'less'") + workers = threads or (os.cpu_count() or 1) + + query_fingerprint = fingerprint(query.molecule) + query_canonical_smiles = Chem.MolToSmiles(query.molecule, canonical=True) + stats = DatasetStats() + records = list(iter_valid_molecules(tsv_path, stats, max_rows=max_rows)) + + def score(record: MoleculeRecord): + candidate_smiles = Chem.MolToSmiles(record.molecule, canonical=True) + if ( + record.molecule_id == query.molecule_id + or candidate_smiles == query_canonical_smiles + ): + return None + similarity = DataStructs.TanimotoSimilarity( + query_fingerprint, fingerprint(record.molecule) + ) + if threshold is not None and ( + similarity < threshold + if threshold_direction == "greater" + else similarity > threshold + ): + return None + return similarity + + # map preserves the input order, so the merge below is exactly the + # sequential reference's merge, just over precomputed scores. + with ThreadPoolExecutor(max_workers=workers) as pool: + scored = pool.map(score, records) + + heap: list[_HeapEntry] = [] + for record, similarity in zip(records, scored): + if similarity is None: + continue + match = SimilarityMatch(similarity, record.molecule_id, record.smiles) + rank_key = match.sort_key(threshold_direction) + entry = _HeapEntry(match, rank_key) + if len(heap) < top_k: + heapq.heappush(heap, entry) + elif rank_key < heap[0].rank_key: + heapq.heapreplace(heap, entry) + + matches = [entry.match for entry in sorted(heap, key=lambda e: e.rank_key)] + return SearchResult(matches=matches, stats=stats) + + +def run_search_shard_parallel( + input_path: Path, + parameters: Mapping[str, object], + output_path: Path, +) -> dict[str, int]: + """Run one planned shard with the parallel scoring core. + + Accepts the same parameters as ``similarity-search`` plus ``threads``. + """ + allowed = { + "query_id", + "query_smiles", + "top_k", + "threshold", + "threshold_direction", + "progress_every", + "threads", + } + unknown = set(parameters) - allowed + if unknown: + raise ValueError( + f"unsupported similarity-search-parallel parameters: {', '.join(sorted(unknown))}" + ) + query_smiles = parameters.get("query_smiles") + query_id = parameters.get("query_id") + if isinstance(query_id, str) and not isinstance(query_smiles, str): + from rdkit import Chem + + from scimesh.chemistry.dataset import find_molecule_by_id + + record = find_molecule_by_id(input_path, query_id) + query_smiles = Chem.MolToSmiles(record.molecule, canonical=True) + if not isinstance(query_smiles, str) or not query_smiles.strip(): + raise ValueError("query_smiles is required for a distributed shard") + molecule = parse_smiles(query_smiles) + if molecule is None: + raise ValueError("query_smiles is invalid") + top_k = _positive_int(parameters.get("top_k", 20), "top_k") + threads = _nonnegative_int(parameters.get("threads", 0), "threads") + threshold = None + if "threshold" in parameters: + threshold = _unit_interval(parameters["threshold"], "threshold") + direction = parameters.get("threshold_direction", "greater") + if direction not in {"greater", "less"}: + raise ValueError("threshold_direction must be 'greater' or 'less'") + assert isinstance(direction, str) + + result = search_similar_parallel( + input_path, + MoleculeRecord("query", query_smiles, molecule), + top_k=top_k, + threads=threads, + threshold=threshold, + threshold_direction=direction, + ) + write_search_partial(output_path, result.matches) + return { + "scanned_rows": result.stats.scanned, + "valid_molecules": result.stats.valid, + "invalid_smiles": result.stats.invalid, + "matches_emitted": len(result.matches), + } + + +def _positive_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + return value + + +def _nonnegative_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer") + return value + + +def _unit_interval(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number between 0 and 1") + return float(value) + + +# The shared shard writer is re-exported so the workload definition can reuse +# the deterministic partitioning without importing search internals. +__all__ = [ + "search_similar_parallel", + "run_search_shard_parallel", + "write_search_shards", + "run_search_shard", +] diff --git a/scimesh/workloads/search_parallel/definition.py b/scimesh/workloads/search_parallel/definition.py new file mode 100644 index 0000000..31399cb --- /dev/null +++ b/scimesh/workloads/search_parallel/definition.py @@ -0,0 +1,144 @@ +"""SDK-built ``similarity-search-parallel`` workload definition and handlers. + +A subclass of ``SimilaritySearchSDKWorkload``: identical contract (plan-time +query resolution, deterministic sharding, top-k reduction, byte-identical +partials), but each shard's fingerprinting and scoring runs across a thread +pool (``threads`` parameter, default = CPU count). +""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from scimesh.sdk.batch import MapReduceWorkload +from scimesh.sdk.identity import WorkloadId +from scimesh.sdk.plans import JobRequest +from scimesh.sdk.registry import WorkloadDefinition +from scimesh.sdk.ui import UIElement + +from ..environment import current_environment_digest, current_scimesh_package_digest +from ..search.definition import SimilaritySearchSDKWorkload, _parameters_schema +from .core import run_search_shard_parallel + +MAP_ENTRY_POINT = "scimesh.workloads.search_parallel.definition:map_search_parallel@v1" + +# The parallel variant adds only the thread-count parameter on top of the +# search contract; everything else (schemas, entry points of the reduce stage, +# partitioning) is inherited. +_MAP_PARAMETERS = ( + "query_id", + "query_smiles", + "top_k", + "threshold", + "threshold_direction", + "progress_every", + "threads", +) + + +def _parallel_parameters_schema() -> dict[str, Any]: + schema = dict(_parameters_schema()) + properties = dict(schema["properties"]) + properties["threads"] = { + "type": "integer", + "minimum": 1, + "description": "Threads used to fingerprint and score one shard (default: CPU count).", + } + schema["properties"] = properties + return schema + + +class SimilaritySearchParallelSDKWorkload(SimilaritySearchSDKWorkload): + """Exact top-k Tanimoto search with a per-shard thread pool.""" + + workload_id = WorkloadId("similarity-search-parallel", "1.0.0") + description = ( + "Exact top-k Tanimoto molecular similarity search over deterministic " + "TSV shards with a bounded merge; each shard is fingerprinted and " + "scored across a thread pool. Output is byte-identical to " + "similarity-search." + ) + parameters_schema = _parallel_parameters_schema() + map_parameter_names = _MAP_PARAMETERS + map_entry_point = MAP_ENTRY_POINT + ui_elements = SimilaritySearchSDKWorkload.ui_elements + ( + UIElement( + "threads", + "number", + "Threads per shard", + help="Threads used to fingerprint and score one shard (default: CPU count).", + placeholder="auto", + order=6, + ), + ) + + def domain_validate(self, parameters: Mapping[str, Any]) -> None: + # The search base rejects unknown parameters; threads is our addition, + # so it is validated here and stripped before delegating. + rest = dict(parameters) + threads = rest.pop("threads", None) + if threads is not None and ( + isinstance(threads, bool) or not isinstance(threads, int) or threads < 1 + ): + raise ValueError("threads must be a positive integer") + super().domain_validate(rest) + + def resolved_parameters_for_plan( + self, + job, + input_path, + resolved, + ): + # threads is a map-stage-only knob; strip it from the plan-level + # resolved parameters so the reduce stage projection stays clean. + resolved = super().resolved_parameters_for_plan(job, input_path, resolved) + stripped = dict(resolved) + stripped.pop("threads", None) + return stripped + + def resolved_parameters(self, request: JobRequest) -> dict[str, Any]: + resolved = super().resolved_parameters(request) + if "threads" in request.parameters: + threads = request.parameters["threads"] + if isinstance(threads, bool) or not isinstance(threads, int) or threads < 1: + raise ValueError("threads must be a positive integer") + resolved["threads"] = threads + return resolved + + def compute_shard( + self, + inputs: Mapping[str, Path], + parameters: Mapping[str, Any], + output_path: Path, + ) -> Mapping[str, int | float]: + return run_search_shard_parallel(inputs["input"], parameters, output_path) + + +def similarity_search_parallel_sdk_definition( + *, + shard_rows: int = 10_000, + package_digest: str | None = None, + environment_digest: str | None = None, +) -> SimilaritySearchParallelSDKWorkload: + """Build the SDK-built parallel similarity-search definition for tests.""" + return SimilaritySearchParallelSDKWorkload( + shard_rows=shard_rows, + package_digest=package_digest or current_scimesh_package_digest(), + environment_digest=environment_digest or current_environment_digest(), + ) + + +def workload_definition() -> WorkloadDefinition: + """Installed entry-point factory for the SDK-built parallel search.""" + return similarity_search_parallel_sdk_definition().definition() + + +def map_search_parallel( + input_path: Path, + parameters: Mapping[str, object], + output_path: Path, +) -> dict[str, int]: + """Digest-pinned map entry point for the parallel search shard.""" + return run_search_shard_parallel(input_path, parameters, output_path) diff --git a/tests/test_cli_workload.py b/tests/test_cli_workload.py index 44abba2..1914363 100644 --- a/tests/test_cli_workload.py +++ b/tests/test_cli_workload.py @@ -244,7 +244,13 @@ def test_workload_cli_exports_the_library_as_json(tmp_path: Path) -> None: assert payload["schema_version"] == 2 names = [item["name"] for item in payload["workloads"]] assert names == sorted( - ["descriptor-batch", "molwt-filter", "similarity-graph", "similarity-search"] + [ + "descriptor-batch", + "molwt-filter", + "similarity-graph", + "similarity-search", + "similarity-search-parallel", + ] ) for item in payload["workloads"]: assert item["version"] == "1.0.0" diff --git a/tests/test_sdk_compatibility.py b/tests/test_sdk_compatibility.py index 3f9b2cb..c7bd1fa 100644 --- a/tests/test_sdk_compatibility.py +++ b/tests/test_sdk_compatibility.py @@ -60,7 +60,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) == 4 + assert len(descriptions) == 5 description = next( item for item in descriptions if item.workload.name == "similarity-search" ) diff --git a/tests/test_sdk_search_parallel.py b/tests/test_sdk_search_parallel.py new file mode 100644 index 0000000..8e6fda0 --- /dev/null +++ b/tests/test_sdk_search_parallel.py @@ -0,0 +1,185 @@ +"""Tests for the SDK-built similarity-search-parallel workload.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pytest + +from scimesh.sdk import ( + ArtifactCollection, + DeterminismProfile, + JobRequest, + LocalArtifactStore, + LocalCoreBatchExecutor, + LocalPlanningContext, + StageKind, +) +from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime +from scimesh.workloads.search.core import run_search_shard, write_search_shards +from scimesh.workloads.search_parallel import ( + run_search_shard_parallel, + search_similar_parallel, +) +from scimesh.workloads.similarity_search import ( + find_molecule_by_id, + search_similar, + write_search_results, +) + + +def _write_dataset(path: Path, molecules: list[tuple[str, str]]) -> None: + path.write_text( + "chembl_id\tcanonical_smiles\n" + + "".join(f"{mid}\t{smiles}\n" for mid, smiles in molecules), + encoding="utf-8", + ) + + +def _tie_dataset(path: Path) -> None: + # Deliberate similarity ties: propanol isomers and duplicated rows, so the + # parallel merge must reproduce the sequential row-order preference. + _write_dataset( + path, + [ + ("QUERY", "CCO"), + ("A1", "CCCO"), + ("A2", "C(CC)O"), + ("B", "CCN"), + ("C1", "CCC"), + ("C2", "CCC"), + ("D", "CCCC"), + ], + ) + + +def test_parallel_matches_sequential_byte_exactly(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _tie_dataset(dataset) + query = find_molecule_by_id(dataset, "QUERY") + + reference = search_similar(dataset, query, top_k=5, progress_every=0) + reference_path = tmp_path / "reference.csv" + write_search_results(reference_path, reference.matches) + + for threads in (1, 2, 4): + parallel = search_similar_parallel(dataset, query, top_k=5, threads=threads) + parallel_path = tmp_path / f"parallel-{threads}.csv" + write_search_results(parallel_path, parallel.matches) + assert parallel_path.read_bytes() == reference_path.read_bytes(), ( + f"threads={threads} diverged from the reference" + ) + + +def test_parallel_shard_matches_sequential_shard(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _tie_dataset(dataset) + shard_dir = tmp_path / "shards" + shard_dir.mkdir() + shards = write_search_shards(dataset, shard_dir, shard_rows=2) + parameters = {"query_smiles": "CCO", "top_k": 3, "threads": 4} + + sequential_out = tmp_path / "seq.tsv" + parallel_out = tmp_path / "par.tsv" + sequential_parameters = dict(parameters) + sequential_parameters.pop("threads") + run_search_shard(shards[0], sequential_parameters, sequential_out) + run_search_shard_parallel(shards[0], parameters, parallel_out) + assert parallel_out.read_bytes() == sequential_out.read_bytes() + + with parallel_out.open(encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert rows[0]["rank"] == "1" + assert rows[0]["similarity"].startswith("0.5") # CCO vs CCCO + assert len(rows) <= 3 + + +def test_parallel_rejects_bad_parameters(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _tie_dataset(dataset) + output = tmp_path / "out.tsv" + + with pytest.raises(ValueError, match="threads must be a non-negative integer"): + run_search_shard_parallel( + dataset, {"query_smiles": "CCO", "threads": -1}, output + ) + with pytest.raises(ValueError, match="unsupported"): + run_search_shard_parallel(dataset, {"query_smiles": "CCO", "nope": 1}, output) + with pytest.raises(ValueError, match="query_smiles is invalid"): + run_search_shard_parallel(dataset, {"query_smiles": "СС"}, output) + + +def _registered_parallel_search(shard_rows: int = 2): + registry = default_sdk_registry(shard_rows=shard_rows) + runtime = default_sdk_runtime() + description = next( + item + for item in registry.descriptions() + if item.workload.name == "similarity-search-parallel" + ) + definition, negotiated = registry.require( + description.workload.name, + description.workload.version, + description.package_digest, + runtime=runtime, + ) + return registry, runtime, description, definition, negotiated + + +def test_parallel_manifest_is_registered_and_negotiable() -> None: + _, runtime, description, definition, negotiated = _registered_parallel_search() + manifest = definition.manifest + + assert description.enabled is True + assert manifest.workload.name == "similarity-search-parallel" + 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 "threads" in manifest.parameters_schema["properties"] + assert negotiated is not None + assert runtime is not None + + +def test_parallel_executor_matches_reference(tmp_path: Path) -> None: + dataset = tmp_path / "molecules.tsv" + _tie_dataset(dataset) + registry, runtime, description, definition, _ = _registered_parallel_search() + artifact_store = LocalArtifactStore(tmp_path / "artifacts") + input_port = definition.manifest.inputs["input"] + dataset_artifact = artifact_store.import_file( + dataset, + declaration=input_port.schema, + ) + request = JobRequest( + workload=definition.manifest.workload, + parameters={"query_id": "QUERY", "top_k": 3, "threads": 2}, + inputs={"input": ArtifactCollection.single(dataset_artifact)}, + ) + + result = LocalCoreBatchExecutor( + registry, + runtime, + artifact_store, + tmp_path / "sdk-work", + ).execute(request, description.package_digest) + result_artifact = result.outputs["result"].items[0].artifact + + reference_path = tmp_path / "reference.csv" + query = find_molecule_by_id(dataset, "QUERY") + reference = search_similar(dataset, query, top_k=3, progress_every=0) + write_search_results(reference_path, reference.matches) + + assert ( + artifact_store.materialize(result_artifact).read_bytes() + == reference_path.read_bytes() + ) + assert result.task_key == "reduce/final"