Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5b2d5b5f6e |
@@ -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"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user