Compare commits

..
Author SHA1 Message Date
Emil 5b2d5b5f6e Add similarity-search-parallel: thread-pool scoring, byte-identical to similarity-search
coordinator / test (push) Waiting to run
python / test (push) Waiting to run
release / binaries (amd64, darwin) (push) Waiting to run
release / binaries (amd64, linux) (push) Waiting to run
release / binaries (amd64, windows) (push) Waiting to run
release / binaries (arm64, darwin) (push) Waiting to run
release / binaries (arm64, linux) (push) Waiting to run
release / binaries (arm64, windows) (push) Waiting to run
release / wheel (push) Waiting to run
release / release (push) Blocked by required conditions
release / image (push) Waiting to run
users / test (push) Waiting to run
2026-08-03 15:12:18 +03:00
9 changed files with 771 additions and 2 deletions
@@ -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"
}
]
}
+1
View File
@@ -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"
+6
View File
@@ -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",
@@ -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",
]
+199
View File
@@ -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",
]
@@ -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)
+7 -1
View File
@@ -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"
+1 -1
View File
@@ -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"
)
+185
View File
@@ -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"