Add distributed similarity search

scimesh/distributed/__init__.py

@@ -12,17 +12,25 @@ from .models import (
     FinalResult,
     PlannedTask,
 )
-from .registry import DistributedWorkloadRegistry, PlanningService, WorkloadDescription
+from .registry import (
+    DistributedWorkloadRegistry,
+    PlanningService,
+    WorkloadDescription,
+    default_distributed_registry,
+)
+from .similarity_search import SimilaritySearchDistributedWorkload
 from .workload import DistributedWorkload
 
 __all__ = [
     "ArtifactReference",
     "CompletedPartial",
+    "default_distributed_registry",
     "DistributedPlan",
     "DistributedWorkload",
     "DistributedWorkloadRegistry",
     "FinalResult",
     "PlannedTask",
     "PlanningService",
+    "SimilaritySearchDistributedWorkload",
     "WorkloadDescription",
 ]

scimesh/distributed/registry.py

@@ -50,7 +50,8 @@ class PlanningService:
 
     It writes neither jobs nor artifacts.  A Go coordinator bridge can therefore
     validate and produce a plan before opening its own all-or-nothing persistence
-    transaction; CTX-08/09 will implement that concrete bridge and reducers.
+    transaction; CTX-09 will implement that concrete bridge and durable result
+    orchestration.
     """
 
     def __init__(self, registry: DistributedWorkloadRegistry) -> None:
@@ -94,3 +95,14 @@ class PlanningService:
         if not isinstance(result, FinalResult):
             raise ValueError("distributed reducer must return a FinalResult")
         return result
+
+
+def default_distributed_registry() -> DistributedWorkloadRegistry:
+    """Return the currently supported distributed scientific workloads."""
+    # Delayed import keeps the generic registry independent of concrete RDKit
+    # workloads and avoids making the contract layer import application setup.
+    from .similarity_search import SimilaritySearchDistributedWorkload
+
+    registry = DistributedWorkloadRegistry()
+    registry.register(SimilaritySearchDistributedWorkload())
+    return registry

scimesh/distributed/similarity_search.py

@@ -0,0 +1,380 @@
+"""Distributed planning and reduction for exact molecular similarity search."""
+
+from __future__ import annotations
+
+import csv
+import hashlib
+import heapq
+import math
+from pathlib import Path
+from typing import Any, Iterator, Mapping, Sequence
+from uuid import UUID, uuid5
+
+from rdkit import Chem
+
+from scimesh.chemistry.dataset import MoleculeRecord, find_molecule_by_id, parse_smiles
+from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE
+from scimesh.workloads.similarity_search import (
+    SimilarityMatch,
+    _HeapEntry,
+    search_similar,
+    write_search_results,
+)
+
+from .models import ArtifactReference, CompletedPartial, DistributedPlan, FinalResult, PlannedTask
+
+
+_TSV_CONTENT_TYPE = "text/tab-separated-values"
+_CSV_CONTENT_TYPE = "text/csv"
+_SEARCH_COLUMNS = ("rank", "chembl_id", "canonical_smiles", "similarity")
+_REQUIRED_COLUMNS = {"chembl_id", "canonical_smiles"}
+
+
+def write_similarity_search_partial(output_path: Path, matches: Sequence[SimilarityMatch]) -> None:
+    """Write a worker partial with a round-trip score, not display rounding.
+
+    The public final CSV continues to use the local CLI's six-decimal display.
+    A reducer needs the full binary float representation to rank candidates
+    from separate shards exactly as the single-process reference does.
+    """
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w", encoding="utf-8", newline="") as destination:
+        writer = csv.DictWriter(destination, fieldnames=_SEARCH_COLUMNS)
+        writer.writeheader()
+        for rank, match in enumerate(matches, start=1):
+            writer.writerow({
+                "rank": rank,
+                "chembl_id": match.molecule_id,
+                "canonical_smiles": match.smiles,
+                "similarity": repr(match.similarity),
+            })
+
+
+class SimilaritySearchDistributedWorkload:
+    """Planner/reducer for exact global top-k Tanimoto similarity search."""
+
+    name = "similarity-search"
+    description = "Exact top-k molecular similarity search over deterministic TSV shards."
+
+    def validate_job(self, parameters: Mapping[str, object]) -> None:
+        allowed = {
+            "query_id", "query_smiles", "top_k", "threshold",
+            "threshold_direction", "max_rows", "progress_every",
+        }
+        unknown = set(parameters) - allowed
+        if unknown:
+            raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
+        query_id = parameters.get("query_id")
+        query_smiles = parameters.get("query_smiles")
+        if (query_id is None) == (query_smiles is None):
+            raise ValueError("exactly one of query_id or query_smiles is required")
+        if query_id is not None:
+            self._string(query_id, "query_id")
+        if query_smiles is not None:
+            self._string(query_smiles, "query_smiles")
+        self._positive_int(parameters.get("top_k", 20), "top_k")
+        if "max_rows" in parameters:
+            self._positive_int(parameters["max_rows"], "max_rows")
+        if "progress_every" in parameters:
+            self._nonnegative_int(parameters["progress_every"], "progress_every")
+        if "threshold" in parameters:
+            self._unit_interval(parameters["threshold"], "threshold")
+        if "threshold_direction" in parameters and parameters["threshold_direction"] not in {"greater", "less"}:
+            raise ValueError("threshold_direction must be 'greater' or 'less'")
+
+    def plan(
+        self,
+        input_path: Path,
+        input_artifact_id: str,
+        parameters: Mapping[str, object],
+        shard_rows: int,
+        workspace: Path,
+    ) -> DistributedPlan:
+        self.validate_job(parameters)
+        if not input_path.is_file():
+            raise ValueError("input_path must be a readable dataset file")
+        if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
+            raise ValueError("shard_rows must be a positive integer")
+        try:
+            input_id = UUID(input_artifact_id)
+        except ValueError as error:
+            raise ValueError("input_artifact_id must be a UUID") from error
+
+        query_smiles, query_source = self._resolve_query(input_path, parameters)
+        resolved = self._resolved_parameters(parameters, query_smiles, query_source)
+        workspace.mkdir(parents=True, exist_ok=True)
+        shard_paths: list[Path] = []
+        try:
+            shard_paths = self._write_shards(input_path, workspace, shard_rows, resolved.get("max_rows"))
+            tasks = tuple(
+                PlannedTask(
+                    chunk_index=index,
+                    input_artifact=ArtifactReference(
+                        artifact_id=str(uuid5(input_id, f"scimesh:similarity-search:shard:{index}")),
+                        sha256=_sha256_file(path),
+                        content_type=_TSV_CONTENT_TYPE,
+                    ),
+                    parameters=self._task_parameters(resolved),
+                )
+                for index, path in enumerate(shard_paths)
+            )
+        except Exception:
+            for path in shard_paths:
+                path.unlink(missing_ok=True)
+            raise
+        return DistributedPlan(self.name, resolved, tasks)
+
+    def reduce(
+        self,
+        partial_results: Sequence[CompletedPartial],
+        parameters: Mapping[str, object],
+        workspace: Path,
+    ) -> FinalResult:
+        """Merge materialized partial CSVs into one deterministic final CSV.
+
+        The coordinator bridge materializes each downloaded artifact at
+        ``workspace / artifact_id`` before it calls this method. Those local
+        paths are an ephemeral bridge detail, never present in the plan or task
+        payload. CTX-09 owns the durable final-artifact upload and job state.
+        """
+        if not partial_results:
+            raise ValueError("at least one partial result is required")
+        resolved = self._validate_resolved_parameters(parameters)
+        top_k = resolved["top_k"]
+        direction = resolved["threshold_direction"]
+        heap: list[_HeapEntry] = []
+
+        ordered_partials = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
+        indexes = [partial.chunk_index for partial in ordered_partials]
+        if len(indexes) != len(set(indexes)):
+            raise ValueError("partial results must have unique chunk_index values")
+        for partial in ordered_partials:
+            path = workspace / partial.artifact.artifact_id
+            if not path.is_file():
+                raise ValueError("materialized partial result is missing")
+            if _sha256_file(path) != partial.artifact.sha256:
+                raise ValueError("materialized partial result checksum does not match its artifact reference")
+            for match in self._read_partial(path, direction):
+                rank_key = match.sort_key(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 = sorted((entry.match for entry in heap), key=lambda match: match.sort_key(direction))
+        output = workspace / "result.csv"
+        write_search_results(output, matches)
+        final_id = uuid5(
+            UUID(ordered_partials[0].artifact.artifact_id),
+            "scimesh:similarity-search:final:" + ",".join(
+                partial.artifact.artifact_id for partial in ordered_partials
+            ),
+        )
+        return FinalResult(
+            ArtifactReference(str(final_id), _sha256_file(output), _CSV_CONTENT_TYPE),
+            {"matches_emitted": len(matches), "partial_count": len(ordered_partials)},
+        )
+
+    def _resolve_query(
+        self, input_path: Path, parameters: Mapping[str, object]
+    ) -> tuple[str, dict[str, str]]:
+        query_id = parameters.get("query_id")
+        if isinstance(query_id, str):
+            record = find_molecule_by_id(input_path, query_id)
+            return Chem.MolToSmiles(record.molecule, canonical=True), {"kind": "chembl_id", "value": query_id}
+        supplied = parameters["query_smiles"]
+        assert isinstance(supplied, str)  # checked by validate_job
+        molecule = parse_smiles(supplied)
+        if molecule is None:
+            raise ValueError("query_smiles is invalid")
+        return Chem.MolToSmiles(molecule, canonical=True), {"kind": "smiles", "value": supplied}
+
+    def _resolved_parameters(
+        self, parameters: Mapping[str, object], query_smiles: str, query_source: Mapping[str, str]
+    ) -> dict[str, object]:
+        resolved: dict[str, object] = {
+            "query_smiles": query_smiles,
+            "query_source": dict(query_source),
+            "top_k": self._positive_int(parameters.get("top_k", 20), "top_k"),
+            "threshold_direction": parameters.get("threshold_direction", "greater"),
+            "fingerprint": {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE},
+        }
+        if "threshold" in parameters:
+            resolved["threshold"] = self._unit_interval(parameters["threshold"], "threshold")
+        if "max_rows" in parameters:
+            resolved["max_rows"] = self._positive_int(parameters["max_rows"], "max_rows")
+        if "progress_every" in parameters:
+            resolved["progress_every"] = self._nonnegative_int(parameters["progress_every"], "progress_every")
+        return resolved
+
+    def _validate_resolved_parameters(self, parameters: Mapping[str, object]) -> dict[str, object]:
+        query_smiles = self._string(parameters.get("query_smiles"), "query_smiles")
+        if parse_smiles(query_smiles) is None:
+            raise ValueError("query_smiles is invalid")
+        resolved = self._resolved_parameters(
+            parameters,
+            Chem.MolToSmiles(parse_smiles(query_smiles), canonical=True),
+            {"kind": "resolved", "value": query_smiles},
+        )
+        # A reducer receives immutable plan metadata, whose query source and
+        # fixed fingerprint are observational context rather than worker input.
+        if "fingerprint" in parameters:
+            fingerprint = parameters["fingerprint"]
+            if fingerprint != {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE}:
+                raise ValueError("resolved fingerprint does not match SciMesh defaults")
+        return resolved
+
+    @staticmethod
+    def _task_parameters(resolved: Mapping[str, object]) -> dict[str, object]:
+        # max_rows is applied before sharding. Passing it to each task would
+        # silently scan N rows per shard instead of the requested global prefix.
+        return {
+            key: value for key, value in resolved.items()
+            if key in {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
+        }
+
+    def _write_shards(
+        self, input_path: Path, workspace: Path, shard_rows: int, max_rows: object
+    ) -> list[Path]:
+        limit = int(max_rows) if isinstance(max_rows, int) else None
+        paths: list[Path] = []
+        current: Path | None = None
+        destination = None
+        rows_in_shard = 0
+        seen_rows = 0
+        try:
+            with input_path.open("r", encoding="utf-8", newline="") as source:
+                reader = csv.DictReader(source, delimiter="\t")
+                fieldnames = reader.fieldnames or []
+                if not _REQUIRED_COLUMNS.issubset(set(fieldnames)):
+                    missing = sorted(_REQUIRED_COLUMNS - set(fieldnames))
+                    raise ValueError(f"dataset is missing required columns: {', '.join(missing)}")
+                for row in reader:
+                    if limit is not None and seen_rows >= limit:
+                        break
+                    if destination is None or rows_in_shard == 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=fieldnames, delimiter="\t", lineterminator="\n")
+                        writer.writeheader()
+                        paths.append(current)
+                        rows_in_shard = 0
+                    writer.writerow(row)
+                    rows_in_shard += 1
+                    seen_rows += 1
+        finally:
+            if destination is not None:
+                destination.close()
+        if not paths:
+            raise ValueError("dataset has no data rows")
+        return paths
+
+    @staticmethod
+    def _read_partial(path: Path, direction: object) -> Iterator[SimilarityMatch]:
+        if not path.is_file():
+            raise ValueError("materialized partial result is missing")
+        if direction not in {"greater", "less"}:
+            raise ValueError("threshold_direction must be 'greater' or 'less'")
+        previous_key: tuple[float, str, str] | None = None
+        with path.open("r", encoding="utf-8", newline="") as source:
+            reader = csv.DictReader(source)
+            if tuple(reader.fieldnames or ()) != _SEARCH_COLUMNS:
+                raise ValueError("partial result has an invalid CSV header")
+            for expected_rank, row in enumerate(reader, start=1):
+                if set(row) != set(_SEARCH_COLUMNS) or row["rank"] != str(expected_rank):
+                    raise ValueError("partial result has an invalid rank")
+                try:
+                    similarity = float(row["similarity"])
+                except (TypeError, ValueError) as error:
+                    raise ValueError("partial result has an invalid similarity") from error
+                if not math.isfinite(similarity) or not 0 <= similarity <= 1:
+                    raise ValueError("partial result has an invalid similarity")
+                match = SimilarityMatch(similarity, row["chembl_id"], row["canonical_smiles"])
+                key = match.sort_key(direction)
+                if previous_key is not None and key < previous_key:
+                    raise ValueError("partial result is not sorted deterministically")
+                previous_key = key
+                yield match
+
+    @staticmethod
+    def _string(value: object, name: str) -> str:
+        if not isinstance(value, str) or not value.strip() or len(value) > 200:
+            raise ValueError(f"{name} must be a non-empty string")
+        return value
+
+    @staticmethod
+    def _positive_int(value: object, name: str) -> int:
+        if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+            raise ValueError(f"{name} must be a positive integer")
+        return value
+
+    @staticmethod
+    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
+
+    @staticmethod
+    def _unit_interval(value: object, name: str) -> float:
+        if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or not 0 <= value <= 1:
+            raise ValueError(f"{name} must be a number between 0 and 1")
+        return float(value)
+
+
+def run_similarity_search_shard(
+    input_path: Path, parameters: Mapping[str, object], output_path: Path
+) -> dict[str, int]:
+    """Run one planned shard using the local reference implementation.
+
+    This is the worker adapter used by CTX-08. It deliberately accepts only
+    resolved ``query_smiles``: resolving an identifier independently in each
+    shard would make the distributed search scientifically invalid.
+    """
+    allowed = {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
+    unknown = set(parameters) - allowed
+    if unknown:
+        raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
+    query_smiles = parameters.get("query_smiles")
+    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 = SimilaritySearchDistributedWorkload._positive_int(parameters.get("top_k", 20), "top_k")
+    threshold = None
+    if "threshold" in parameters:
+        threshold = SimilaritySearchDistributedWorkload._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'")
+    progress_every = 0
+    if "progress_every" in parameters:
+        progress_every = SimilaritySearchDistributedWorkload._nonnegative_int(
+            parameters["progress_every"], "progress_every"
+        )
+    result = search_similar(
+        input_path,
+        MoleculeRecord("query", query_smiles, molecule),
+        top_k=top_k,
+        progress_every=progress_every,
+        threshold=threshold,
+        threshold_direction=direction,
+    )
+    write_similarity_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 _sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as source:
+        for block in iter(lambda: source.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()