295 lines
9.2 KiB
Python
295 lines
9.2 KiB
Python
"""Tests for the SDK-built similarity-graph 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.graph import (
|
|
check_pair_coverage,
|
|
merge_edge_partials,
|
|
)
|
|
from scimesh.workloads.library import default_sdk_registry, default_sdk_runtime
|
|
from scimesh.workloads.similarity_graph import (
|
|
build_similarity_graph,
|
|
write_graph_edges,
|
|
)
|
|
|
|
|
|
def _write_tiny_dataset(path: Path) -> None:
|
|
path.write_text(
|
|
"chembl_id\tcanonical_smiles\n"
|
|
"A\tCCO\n"
|
|
"B\tCCCC\n"
|
|
"C\tCCN\n"
|
|
"D\tCCCCCC\n"
|
|
"E\tnot-a-smiles\n"
|
|
"F\tCCOCC\n"
|
|
"G\tc1ccccc1\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def _registered_similarity_graph():
|
|
registry = default_sdk_registry()
|
|
runtime = default_sdk_runtime()
|
|
description = next(
|
|
item
|
|
for item in registry.descriptions()
|
|
if item.workload.name == "similarity-graph"
|
|
)
|
|
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,
|
|
*,
|
|
threshold: float = 0.3,
|
|
threshold_direction: str = "greater",
|
|
block_size: int = 2,
|
|
) -> 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={
|
|
"threshold": threshold,
|
|
"threshold_direction": threshold_direction,
|
|
"block_size": block_size,
|
|
},
|
|
inputs={"input": ArtifactCollection.single(dataset_artifact)},
|
|
)
|
|
|
|
|
|
def test_similarity_graph_manifest_is_registered_and_negotiable() -> None:
|
|
_, runtime, description, definition, negotiated = _registered_similarity_graph()
|
|
manifest = definition.manifest
|
|
|
|
assert description.enabled is True
|
|
assert manifest.workload.name == "similarity-graph"
|
|
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_manifest_round_trip(manifest)
|
|
assert runtime is not None
|
|
|
|
|
|
@pytest.mark.parametrize("threshold_direction", ("greater", "less"))
|
|
def test_local_sdk_executor_matches_similarity_graph_reference(
|
|
tmp_path: Path,
|
|
threshold_direction: str,
|
|
) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, description, definition, _ = _registered_similarity_graph()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(
|
|
dataset,
|
|
artifact_store,
|
|
definition,
|
|
threshold_direction=threshold_direction,
|
|
)
|
|
|
|
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"
|
|
reference = build_similarity_graph(
|
|
dataset,
|
|
threshold=0.3,
|
|
block_size=1_000,
|
|
threshold_direction=threshold_direction,
|
|
)
|
|
write_graph_edges(reference_path, reference.edges)
|
|
|
|
assert (
|
|
artifact_store.materialize(result_artifact).read_bytes()
|
|
== reference_path.read_bytes()
|
|
)
|
|
assert result.task_key == "reduce/final"
|
|
assert result.metrics["partial_count"] == 6
|
|
assert result.metrics["edges_emitted"] == len(reference.edges)
|
|
|
|
|
|
def test_similarity_graph_result_is_invariant_to_block_size(tmp_path: Path) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, _, definition, _ = _registered_similarity_graph()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
|
|
outputs = []
|
|
for block_size in (2, 3):
|
|
request = _request_for(
|
|
dataset,
|
|
artifact_store,
|
|
definition,
|
|
block_size=block_size,
|
|
)
|
|
result = LocalCoreBatchExecutor(
|
|
registry,
|
|
runtime,
|
|
artifact_store,
|
|
tmp_path / f"sdk-work-{block_size}",
|
|
).execute(request, definition.manifest.package.digest)
|
|
artifact = result.outputs["result"].items[0].artifact
|
|
outputs.append(artifact_store.materialize(artifact).read_bytes())
|
|
assert result.metrics["partial_count"] == {2: 6, 3: 3}[block_size]
|
|
|
|
assert outputs[0] == outputs[1]
|
|
|
|
|
|
def test_similarity_graph_planning_covers_each_block_pair_once(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, description, definition, _ = _registered_similarity_graph()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(dataset, artifact_store, definition)
|
|
input_artifact = request.inputs["input"].items[0].artifact
|
|
|
|
plan = registry.plan(
|
|
request,
|
|
description.package_digest,
|
|
runtime,
|
|
LocalPlanningContext(
|
|
artifact_store,
|
|
artifact_store,
|
|
tmp_path / "plan",
|
|
allowed_artifacts=(input_artifact,),
|
|
),
|
|
)
|
|
|
|
assert [task.task_key for task in plan.tasks] == [
|
|
"map/0000x0000",
|
|
"map/0000x0001",
|
|
"map/0000x0002",
|
|
"map/0001x0001",
|
|
"map/0001x0002",
|
|
"map/0002x0002",
|
|
]
|
|
for task in plan.tasks:
|
|
assert set(task.inputs) == {"left", "right"}
|
|
assert task.inputs["left"].items[0].artifact is not None
|
|
pairs = {
|
|
(int(left), int(right))
|
|
for task in plan.tasks
|
|
for left, right in (task.task_key[len("map/") :].split("x"),)
|
|
}
|
|
check_pair_coverage(tuple(sorted(pairs)))
|
|
|
|
wire_payload = plan.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_graph_rejects_duplicate_molecule_ids(tmp_path: Path) -> None:
|
|
dataset = tmp_path / "duplicates.tsv"
|
|
dataset.write_text(
|
|
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nA\tCCN\n",
|
|
encoding="utf-8",
|
|
)
|
|
registry, runtime, _, definition, _ = _registered_similarity_graph()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(dataset, artifact_store, definition)
|
|
|
|
with pytest.raises(ValueError, match="duplicate chembl_id"):
|
|
LocalCoreBatchExecutor(
|
|
registry,
|
|
runtime,
|
|
artifact_store,
|
|
tmp_path / "work",
|
|
).execute(request, definition.manifest.package.digest)
|
|
|
|
|
|
def test_similarity_graph_reducer_rejects_duplicate_unordered_pairs(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
first = tmp_path / "first.csv"
|
|
first.write_text(
|
|
"source_id,target_id,similarity\nA,B,0.500000\nC,D,0.100000\n",
|
|
encoding="utf-8",
|
|
)
|
|
second = tmp_path / "second.csv"
|
|
second.write_text(
|
|
"source_id,target_id,similarity\nB,A,0.500000\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="duplicate unordered pair"):
|
|
merge_edge_partials((first, second), tmp_path / "result.csv")
|
|
|
|
|
|
def test_similarity_graph_pair_coverage_rejects_missing_block_pair() -> None:
|
|
with pytest.raises(ValueError, match="do not cover the full block pair set"):
|
|
check_pair_coverage(((0, 0), (0, 1)))
|
|
|
|
with pytest.raises(ValueError, match="duplicate block pair"):
|
|
check_pair_coverage(((0, 0), (0, 0), (0, 1), (1, 1)))
|
|
|
|
|
|
def test_similarity_graph_merge_is_deterministically_sorted(tmp_path: Path) -> None:
|
|
first = tmp_path / "first.csv"
|
|
first.write_text(
|
|
"source_id,target_id,similarity\nC,A,0.200000\nB,C,0.400000\n",
|
|
encoding="utf-8",
|
|
)
|
|
second = tmp_path / "second.csv"
|
|
second.write_text(
|
|
"source_id,target_id,similarity\nA,B,0.900000\n",
|
|
encoding="utf-8",
|
|
)
|
|
result_path = tmp_path / "result.csv"
|
|
metrics = merge_edge_partials((first, second), result_path)
|
|
|
|
assert metrics == {"partial_count": 2, "edges_emitted": 3}
|
|
rows = list(csv.DictReader(result_path.open(encoding="utf-8", newline="")))
|
|
assert [
|
|
(row["source_id"], row["target_id"], row["similarity"]) for row in rows
|
|
] == [
|
|
("A", "B", "0.900000"),
|
|
("B", "C", "0.400000"),
|
|
("C", "A", "0.200000"),
|
|
]
|