"""Tests for the SDK-built similarity-search workload.""" from __future__ import annotations import csv from pathlib import Path import pytest from scimesh.sdk import ( ArtifactCollection, DeterminismProfile, JobRequest, LocalArtifactStore, LocalCoreBatchExecutor, LocalPlanningContext, StageKind, assert_manifest_round_trip, ) from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime from scimesh.workloads.similarity_search import ( find_molecule_by_id, search_similar, write_search_results, ) def _write_tiny_dataset(path: Path) -> None: path.write_text( "chembl_id\tcanonical_smiles\textra\n" "QUERY\tCCO\tquery\n" "ALCOHOL\tCCCO\talcohol\n" "ALKANE\tCCCC\talkane\n" "BROKEN\tnot-a-smiles\tinvalid\n" "DUPLICATE\tCCO\tduplicate\n" "AMINE\tCCN\tamine\n", encoding="utf-8", ) def _registered_similarity_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" ) definition, negotiated = registry.require( description.workload.name, description.workload.version, description.package_digest, runtime=runtime, ) return registry, runtime, description, definition, negotiated def _request_for( dataset: Path, artifact_store: LocalArtifactStore, definition, ) -> JobRequest: input_port = definition.manifest.inputs["input"] dataset_artifact = artifact_store.import_file( dataset, declaration=input_port.schema, ) return JobRequest( workload=definition.manifest.workload, parameters={"query_id": "QUERY", "top_k": 3, "progress_every": 0}, inputs={"input": ArtifactCollection.single(dataset_artifact)}, ) def test_similarity_search_manifest_is_registered_and_negotiable() -> None: _, runtime, description, definition, negotiated = _registered_similarity_search() manifest = definition.manifest assert description.enabled is True assert manifest.workload.name == "similarity-search" 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 set(definition.runners) == {manifest.workflow.stages[0].entry_point} assert set(definition.reducers) == {manifest.workflow.stages[1].entry_point} assert negotiated is not None assert negotiated.manifest == manifest assert_manifest_round_trip(manifest) assert runtime is not None def test_local_sdk_executor_matches_similarity_search_reference(tmp_path: Path) -> None: dataset = tmp_path / "molecules.tsv" _write_tiny_dataset(dataset) registry, runtime, description, definition, _ = _registered_similarity_search() artifact_store = LocalArtifactStore(tmp_path / "artifacts") request = _request_for(dataset, artifact_store, definition) 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" assert dict(result.metrics) == {"matches_emitted": 3, "partial_count": 3} def test_similarity_search_planning_is_deterministic_ordered_and_path_free( tmp_path: Path, ) -> None: dataset = tmp_path / "molecules.tsv" _write_tiny_dataset(dataset) registry, runtime, description, definition, _ = _registered_similarity_search() artifact_store = LocalArtifactStore(tmp_path / "artifacts") request = _request_for(dataset, artifact_store, definition) input_artifact = request.inputs["input"].items[0].artifact first = registry.plan( request, description.package_digest, runtime, LocalPlanningContext( artifact_store, artifact_store, tmp_path / "first-plan", allowed_artifacts=(input_artifact,), ), ) second = registry.plan( request, description.package_digest, runtime, LocalPlanningContext( artifact_store, artifact_store, tmp_path / "second-plan", allowed_artifacts=(input_artifact,), ), ) assert first.to_json() == second.to_json() assert first.digest == second.digest assert first.package_digest == definition.manifest.package.digest assert first.manifest_digest == definition.manifest.digest assert [task.task_key for task in first.tasks] == [ "map/00000000", "map/00000001", "map/00000002", ] assert all(task.stage_id == "map" for task in first.tasks) assert all("query_id" not in task.parameters for task in first.tasks) assert all(task.parameters["query_smiles"] == "CCO" for task in first.tasks) assert all(task.parameters["top_k"] == 3 for task in first.tasks) assert all(task.package_digest == first.package_digest for task in first.tasks) assert all(task.manifest_digest == first.manifest_digest for task in first.tasks) assert first.resolved_parameters["query_source"] == { "kind": "chembl_id", "value": "QUERY", } shard_ids: list[list[str]] = [] for task in first.tasks: artifact = task.inputs["input"].items[0].artifact with artifact_store.materialize(artifact).open( encoding="utf-8", newline="" ) as source: shard_ids.append( [row["chembl_id"] for row in csv.DictReader(source, delimiter="\t")] ) assert shard_ids == [ ["QUERY", "ALCOHOL"], ["ALKANE", "BROKEN"], ["DUPLICATE", "AMINE"], ] wire_payload = first.to_json() assert str(tmp_path) not in wire_payload assert "file://" not in wire_payload assert "worker://" not in wire_payload assert "workspace" not in wire_payload def test_similarity_search_rejects_ambiguous_or_mistyped_parameters( tmp_path: Path, ) -> None: dataset = tmp_path / "molecules.tsv" _write_tiny_dataset(dataset) registry, runtime, _, definition, _ = _registered_similarity_search() artifact_store = LocalArtifactStore(tmp_path / "artifacts") input_artifact = artifact_store.import_file( dataset, declaration=definition.manifest.inputs["input"].schema, ) base = JobRequest( workload=definition.manifest.workload, parameters={"query_id": "QUERY", "top_k": 3}, inputs={"input": ArtifactCollection.single(input_artifact)}, ) for bad_parameters, message in ( ({"query_id": "QUERY", "query_smiles": "CCO"}, "oneOf did not match"), ({"query_id": "QUERY", "top_k": 0}, "violates minimum"), ): request = JobRequest( workload=base.workload, parameters=bad_parameters, inputs=base.inputs, ) with pytest.raises(ValueError, match=message): registry.plan( request, definition.manifest.package.digest, runtime, LocalPlanningContext( artifact_store, artifact_store, tmp_path / "bad-plan", allowed_artifacts=(input_artifact,), ), ) def test_similarity_search_workload_definition_is_discoverable() -> None: from scimesh.sdk import WorkloadDefinition from scimesh.workloads.search import workload_definition definition = workload_definition() assert isinstance(definition, WorkloadDefinition) assert definition.manifest.workload.name == "similarity-search" assert definition.manifest.workload.version == "1.0.0"