Add distributed similarity search

tests/test_distributed_similarity_search.py

@@ -0,0 +1,184 @@
+"""Scientific reference tests for the CTX-08 distributed search workload."""
+
+from __future__ import annotations
+
+import csv
+import hashlib
+from pathlib import Path
+from uuid import NAMESPACE_URL, uuid5
+
+import pytest
+
+from scimesh.chemistry.dataset import find_molecule_by_id
+from scimesh.distributed import (
+    ArtifactReference,
+    CompletedPartial,
+    PlanningService,
+    default_distributed_registry,
+)
+from scimesh.distributed.registry import DistributedWorkloadRegistry
+from scimesh.distributed.similarity_search import (
+    SimilaritySearchDistributedWorkload,
+    run_similarity_search_shard,
+    write_similarity_search_partial,
+)
+from scimesh.workloads.similarity_search import search_similar, write_search_results
+
+
+def make_dataset(path: Path) -> None:
+    path.write_text(
+        "chembl_id\tcanonical_smiles\textra\n"
+        "CHEMBL_QUERY\tCCO\tquery\n"
+        "CHEMBL_A\tCCCO\ta\n"
+        "CHEMBL_B\tCCCC\tb\n"
+        "CHEMBL_INVALID\tnot-a-smiles\tbad\n"
+        "CHEMBL_DUPLICATE\tCCO\tduplicate\n"
+        "CHEMBL_C\tCCN\tc\n",
+        encoding="utf-8",
+    )
+
+
+def planner() -> tuple[PlanningService, SimilaritySearchDistributedWorkload]:
+    workload = SimilaritySearchDistributedWorkload()
+    registry = DistributedWorkloadRegistry()
+    registry.register(workload)
+    return PlanningService(registry), workload
+
+
+def test_default_registry_exposes_only_supported_distributed_search() -> None:
+    assert [item.name for item in default_distributed_registry().descriptions()] == ["similarity-search"]
+
+
+def checksum(path: Path) -> str:
+    return hashlib.sha256(path.read_bytes()).hexdigest()
+
+
+def test_query_id_is_resolved_once_before_deterministic_shards(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    dataset = tmp_path / "chembl.tsv"
+    workspace = tmp_path / "workspace"
+    make_dataset(dataset)
+    service, _ = planner()
+    calls = 0
+    real_find = find_molecule_by_id
+
+    def count_find(path: Path, query_id: str) -> MoleculeRecord:
+        nonlocal calls
+        calls += 1
+        return real_find(path, query_id)
+
+    monkeypatch.setattr("scimesh.distributed.similarity_search.find_molecule_by_id", count_find)
+    input_id = str(uuid5(NAMESPACE_URL, "dataset"))
+    plan = service.plan(
+        "similarity-search", dataset, input_id,
+        {"query_id": "CHEMBL_QUERY", "top_k": 3, "max_rows": 5, "progress_every": 0},
+        2, workspace,
+    )
+
+    assert calls == 1
+    assert plan.resolved_parameters["query_smiles"] == "CCO"
+    assert plan.resolved_parameters["query_source"] == {"kind": "chembl_id", "value": "CHEMBL_QUERY"}
+    assert [task.chunk_index for task in plan.tasks] == [0, 1, 2]
+    assert all("query_id" not in task.parameters for task in plan.tasks)
+    assert all("max_rows" not in task.parameters for task in plan.tasks)
+    assert all(task.parameters["query_smiles"] == "CCO" for task in plan.tasks)
+    assert [
+        sum(1 for _ in path.open(encoding="utf-8")) - 1
+        for path in sorted(workspace.glob("shard-*.tsv"))
+    ] == [2, 2, 1]
+
+
+def test_distributed_reduction_matches_single_process_reference(tmp_path: Path) -> None:
+    dataset = tmp_path / "chembl.tsv"
+    workspace = tmp_path / "workspace"
+    make_dataset(dataset)
+    service, workload = planner()
+    plan = service.plan(
+        "similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
+        {"query_smiles": "CCO", "top_k": 3, "threshold": 0.0}, 2, workspace,
+    )
+
+    partials: list[CompletedPartial] = []
+    # Worker two finishes the latter shards first. Worker one loses its first
+    # attempt for shard zero, then retries it last. The reducer must remain
+    # independent of both completion and retry order.
+    for task in reversed(plan.tasks):
+        shard = workspace / f"shard-{task.chunk_index}.tsv"
+        temporary_partial = workspace / f"worker-output-{task.chunk_index}.csv"
+        metrics = run_similarity_search_shard(shard, task.parameters, temporary_partial)
+        partial_id = str(uuid5(NAMESPACE_URL, f"partial:{task.chunk_index}"))
+        partials.append(
+            CompletedPartial(
+                task.chunk_index,
+                ArtifactReference(
+                    partial_id, checksum(temporary_partial), "text/csv",
+                ),
+                metrics,
+            )
+        )
+        # The reducer materializes result files under their own coordinator IDs,
+        # not shard input IDs. Keep this fixture faithful to that boundary.
+        temporary_partial.rename(workspace / partial_id)
+
+    final = workload.reduce(tuple(partials), plan.resolved_parameters, workspace)
+    reference = tmp_path / "reference.csv"
+    query_record = find_molecule_by_id(dataset, "CHEMBL_QUERY")
+    write_search_results(reference, search_similar(dataset, query_record, top_k=3, threshold=0.0).matches)
+
+    assert (workspace / "result.csv").read_bytes() == reference.read_bytes()
+    assert final.metrics == {"matches_emitted": 3, "partial_count": 3}
+    rows = list(csv.DictReader((workspace / "result.csv").open(encoding="utf-8")))
+    assert {row["chembl_id"] for row in rows}.isdisjoint({"CHEMBL_QUERY", "CHEMBL_DUPLICATE"})
+
+
+def test_reducer_rejects_unsorted_or_invalid_partial_csv(tmp_path: Path) -> None:
+    workspace = tmp_path / "workspace"
+    workspace.mkdir()
+    workload = SimilaritySearchDistributedWorkload()
+    artifact_id = str(uuid5(NAMESPACE_URL, "bad"))
+    partial_path = workspace / artifact_id
+    partial_path.write_text(
+        "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.1\n2,B,CCC,0.9\n",
+        encoding="utf-8",
+    )
+    artifact = ArtifactReference(artifact_id, checksum(partial_path), "text/csv")
+
+    with pytest.raises(ValueError, match="not sorted"):
+        workload.reduce(
+            (CompletedPartial(0, artifact, {"scanned_rows": 2}),),
+            {"query_smiles": "CCO", "top_k": 2, "threshold_direction": "greater", "fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}},
+            workspace,
+        )
+
+
+def test_partial_csv_preserves_exact_scores_for_global_ranking(tmp_path: Path) -> None:
+    partial = tmp_path / "partial.csv"
+    # Both values look identical in a six-decimal final CSV. The exact value
+    # must survive shard transport so the global reducer can still rank them.
+    from scimesh.workloads.similarity_search import SimilarityMatch
+
+    write_similarity_search_partial(
+        partial,
+        [SimilarityMatch(0.50000049, "A", "CC"), SimilarityMatch(0.50000048, "B", "CCC")],
+    )
+    values = list(csv.DictReader(partial.open(encoding="utf-8")))
+    assert values[0]["similarity"] == repr(0.50000049)
+    assert values[1]["similarity"] == repr(0.50000048)
+
+
+def test_planner_rejects_fingerprint_override_and_invalid_query(tmp_path: Path) -> None:
+    dataset = tmp_path / "chembl.tsv"
+    make_dataset(dataset)
+    service, _ = planner()
+
+    with pytest.raises(ValueError, match="unsupported similarity-search parameters"):
+        service.plan(
+            "similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
+            {"query_smiles": "CCO", "fingerprint": {"radius": 1}}, 2, tmp_path / "workspace",
+        )
+    with pytest.raises(ValueError, match="query_smiles is invalid"):
+        service.plan(
+            "similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
+            {"query_smiles": "invalid"}, 2, tmp_path / "workspace",
+        )

tests/test_worker_daemon.py

@@ -106,6 +106,90 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
     }
 
 
+def test_worker_executes_a_resolved_similarity_search_shard(tmp_path: Path) -> None:
+    content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\nINVALID\tnot-a-smiles\n"
+    task = make_task(content)
+    task = ClaimedTask(
+        task.task_id, task.attempt, task.lease_expires_at, task.workload, task.input,
+        {"query_smiles": "CCO", "top_k": 5, "progress_every": 0},
+    )
+    worker, coordinator, artifacts, _, _ = daemon(tmp_path, task, content)
+    worker.runner = SciMeshRunner()
+
+    assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
+    output = artifacts.uploaded[0][2].read_text(encoding="utf-8")
+    assert output.startswith("rank,chembl_id,canonical_smiles,similarity\n")
+    metrics = coordinator.submissions[0]["metrics"]
+    assert metrics["scanned_rows"] == 3
+    assert metrics["valid_molecules"] == 2
+    assert metrics["invalid_smiles"] == 1
+    assert metrics["matches_emitted"] == 1
+    assert isinstance(metrics["elapsed_seconds"], float)
+
+
+def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> None:
+    content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n"
+    first = make_task(content)
+    first = ClaimedTask(
+        "retry-task", 1, first.lease_expires_at, "similarity-search", first.input,
+        {"query_smiles": "CCO", "top_k": 5},
+    )
+    second = ClaimedTask(
+        "other-task", 1, first.lease_expires_at, "similarity-search", first.input,
+        {"query_smiles": "CCO", "top_k": 5},
+    )
+
+    class RetryCoordinator(FakeCoordinator):
+        def __init__(self) -> None:
+            super().__init__(None)
+            self.queue = [first, second]
+            self.claimants: list[str] = []
+
+        def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
+            self.claimants.append(worker_id)
+            return self.queue.pop(0) if self.queue else None
+
+        def fail(self, task: ClaimedTask, payload: dict) -> None:
+            self.failures.append(payload)
+            if task.task_id == "retry-task" and task.attempt == 1 and payload["retryable"]:
+                self.queue.append(
+                    ClaimedTask(
+                        task.task_id, 2, task.lease_expires_at, task.workload, task.input,
+                        task.parameters,
+                    )
+                )
+
+    class FailFirstAttempt:
+        def __init__(self) -> None:
+            self.calls = 0
+            self.delegate = SciMeshRunner()
+
+        def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
+            self.calls += 1
+            if self.calls == 1:
+                raise RuntimeError("simulated retryable shard failure")
+            return self.delegate.run(task, task_dir)
+
+    coordinator = RetryCoordinator()
+    artifacts = FakeArtifacts(content)
+    worker_a = WorkerDaemon(
+        WorkerConfig("https://example.test", "worker-a", tmp_path / "worker-a"),
+        coordinator, artifacts, FailFirstAttempt(),
+    )
+    worker_b = WorkerDaemon(
+        WorkerConfig("https://example.test", "worker-b", tmp_path / "worker-b"),
+        coordinator, artifacts, SciMeshRunner(),
+    )
+
+    assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=False)
+    assert worker_b.run_once() == RunOnceOutcome(claimed=True, completed=True)
+    assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=True)
+    assert coordinator.claimants == ["worker-a", "worker-b", "worker-a"]
+    assert len(coordinator.failures) == 1
+    assert coordinator.failures[0]["retryable"] is True
+    assert len(coordinator.submissions) == 2
+
+
 def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
     worker, _, _, runner, config = daemon(tmp_path, None, b"")
     assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False)
@@ -161,6 +245,7 @@ def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path)
             "attempt": 1,
             "error_code": "InterruptedError",
             "error_message": "worker interrupted by operator",
+            "retryable": True,
         }
     ]
 
@@ -234,6 +319,7 @@ def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None:
     assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
     assert runner.calls == 0
     assert coordinator.failures[0]["error_code"] == "ValueError"
+    assert coordinator.failures[0]["retryable"] is False
     assert not coordinator.submissions
 
 
@@ -356,30 +442,35 @@ def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypa
     runner = SciMeshRunner()
     graph = ClaimedTask("graph", 1, "2026-07-30T00:00:00Z", "similarity-graph", InputArtifact("https://example/input", "x"), {"threshold": 0.2, "threshold_direction": "less", "block_size": 42, "max_rows": 7, "progress_every": 0})
     search = ClaimedTask("search", 1, "2026-07-30T00:00:00Z", "similarity-search", InputArtifact("https://example/input", "x"), {"query_smiles": "CCO", "top_k": 3})
+    search_dir = tmp_path / "search"
+    search_dir.mkdir()
+    (search_dir / "input").write_text(
+        "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
+    )
     runner.run(graph, tmp_path / "graph")
-    runner.run(search, tmp_path / "search")
+    result = runner.run(search, search_dir)
     assert "--threshold-direction" in commands[0] and "less" in commands[0]
     assert "--block-size" in commands[0] and "42" in commands[0]
     assert "--max-rows" in commands[0] and "7" in commands[0]
-    assert "--query-smiles" in commands[1] and "CCO" in commands[1]
+    assert len(commands) == 1
+    assert result.metrics == {
+        "scanned_rows": 2, "valid_molecules": 2, "invalid_smiles": 0, "matches_emitted": 1,
+    }
 
 
 def test_runner_accepts_coordinator_workload_names(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
-    commands: list[list[str]] = []
-
-    def fake_run(command: list[str], **_: object) -> None:
-        commands.append(command)
-        output = Path(command[command.index("--output") + 1])
-        output.parent.mkdir(parents=True, exist_ok=True)
-        output.write_text("a,b\\n", encoding="utf-8")
-
-    monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run)
+    task_dir = tmp_path / "search"
+    task_dir.mkdir()
+    (task_dir / "input").write_text(
+        "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
+    )
     task = ClaimedTask(
         "search", 1, "2026-07-30T00:00:00Z", "similarity_search",
         InputArtifact("https://example/input", "a" * 64), {"query_smiles": "CCO"},
     )
-    SciMeshRunner().run(task, tmp_path / "search")
-    assert commands[0][3] == "similarity-search"
+    result = SciMeshRunner().run(task, task_dir)
+    assert result.metrics["matches_emitted"] == 1
+    assert (task_dir / "result.csv").is_file()
 
 
 def test_claimed_task_rejects_path_traversal_and_invalid_metadata() -> None:
@@ -457,21 +548,15 @@ def test_relative_work_dir_is_normalized_for_runner_subprocesses(
 
     task_dir = config.work_dir / "task" / "1"
     task_dir.mkdir(parents=True)
-    (task_dir / "input").write_text("fixture", encoding="utf-8")
-    command: list[str] = []
-
-    def fake_run(args: list[str], **_: object) -> None:
-        command.extend(args)
-        output = Path(args[args.index("--output") + 1])
-        output.write_text("id,score\n", encoding="utf-8")
-
-    monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run)
+    (task_dir / "input").write_text(
+        "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
+    )
     task = ClaimedTask(
         "task", 1, "2026-07-30T00:00:00Z", "similarity-search",
         InputArtifact("https://example.test/input", "a" * 64), {"query_smiles": "CCO"},
     )
     SciMeshRunner().run(task, task_dir)
-    assert command[4] == str(task_dir / "input")
+    assert (task_dir / "result.csv").is_file()
 
 
 def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None: