474 lines
16 KiB
Python
474 lines
16 KiB
Python
"""Tests for the SDK-native descriptor-batch reference workload."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from scimesh.sdk import (
|
|
AllowedPackage,
|
|
ArtifactCollection,
|
|
CandidateOutput,
|
|
CandidateOutputs,
|
|
DeterminismProfile,
|
|
ExactArtifactVerifier,
|
|
JobRequest,
|
|
LocalArtifactStore,
|
|
LocalCoreBatchExecutor,
|
|
LocalPlanningContext,
|
|
TrustMode,
|
|
VerificationBinding,
|
|
VerificationStatus,
|
|
VerifyContext,
|
|
WorkloadRegistry,
|
|
assert_manifest_round_trip,
|
|
)
|
|
from scimesh.workloads.descriptors import (
|
|
DESCRIPTOR_COLUMNS,
|
|
descriptor_batch_sdk_definition,
|
|
compute_descriptor_batch,
|
|
)
|
|
from scimesh.workloads.library import default_sdk_runtime
|
|
|
|
|
|
def _write_tiny_dataset(path: Path) -> None:
|
|
path.write_text(
|
|
"chembl_id\tcanonical_smiles\textra\n"
|
|
"ALCOHOL\tCCO\talcohol\n"
|
|
"ALKANE\tCCCC\talkane\n"
|
|
"AMINE\tCCN\tamine\n"
|
|
"BROKEN\tnot-a-smiles\tinvalid\n"
|
|
"HEXANE\tCCCCCC\thexane\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def _registered_descriptor_batch(shard_rows: int = 2):
|
|
workload = descriptor_batch_sdk_definition(shard_rows=shard_rows)
|
|
registry = WorkloadRegistry()
|
|
registry.register(workload.definition(), enabled=True)
|
|
runtime = default_sdk_runtime()
|
|
definition, negotiated = registry.require(
|
|
workload.manifest.workload.name,
|
|
workload.manifest.workload.version,
|
|
workload.manifest.package.digest,
|
|
runtime=runtime,
|
|
)
|
|
return registry, runtime, workload, definition, negotiated
|
|
|
|
|
|
def _request_for(
|
|
dataset: Path,
|
|
artifact_store: LocalArtifactStore,
|
|
workload,
|
|
*,
|
|
skip_invalid: bool = True,
|
|
) -> JobRequest:
|
|
input_port = workload.manifest.inputs["input"]
|
|
dataset_artifact = artifact_store.import_file(
|
|
dataset,
|
|
declaration=input_port.schema,
|
|
)
|
|
return JobRequest(
|
|
workload=workload.manifest.workload,
|
|
parameters={"skip_invalid": skip_invalid},
|
|
inputs={"input": ArtifactCollection.single(dataset_artifact)},
|
|
)
|
|
|
|
|
|
def test_descriptor_batch_manifest_is_registered_and_negotiable() -> None:
|
|
_, runtime, workload, definition, negotiated = _registered_descriptor_batch()
|
|
manifest = definition.manifest
|
|
|
|
assert manifest.workload.name == "descriptor-batch"
|
|
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 manifest.conformance_profiles == ("core-batch-v1",)
|
|
assert manifest.capabilities == ("descriptor-batch",)
|
|
assert [stage.kind.value for stage in manifest.workflow.stages] == ["map", "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_descriptor_batch_reference(tmp_path: Path) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, workload, definition, _ = _registered_descriptor_batch()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(dataset, artifact_store, workload)
|
|
|
|
result = LocalCoreBatchExecutor(
|
|
registry,
|
|
runtime,
|
|
artifact_store,
|
|
tmp_path / "sdk-work",
|
|
).execute(request, definition.manifest.package.digest)
|
|
result_artifact = result.outputs["result"].items[0].artifact
|
|
|
|
reference_path = tmp_path / "reference.csv"
|
|
reference_metrics = compute_descriptor_batch(
|
|
dataset, reference_path, skip_invalid=True
|
|
)
|
|
|
|
assert result.task_key == "reduce/final"
|
|
assert dict(result.metrics) == {
|
|
"partial_count": 3,
|
|
"rows_emitted": reference_metrics["rows_emitted"],
|
|
}
|
|
assert (
|
|
artifact_store.materialize(result_artifact).read_bytes()
|
|
== reference_path.read_bytes()
|
|
)
|
|
|
|
with artifact_store.materialize(result_artifact).open(
|
|
encoding="utf-8", newline=""
|
|
) as source:
|
|
rows = list(csv.reader(source))
|
|
assert rows[0] == list(DESCRIPTOR_COLUMNS)
|
|
assert [row[0] for row in rows[1:]] == ["ALCOHOL", "ALKANE", "AMINE", "HEXANE"]
|
|
assert len(rows[1:]) == reference_metrics["rows_emitted"]
|
|
assert any(len(row) == len(DESCRIPTOR_COLUMNS) for row in rows[1:])
|
|
|
|
|
|
def test_descriptor_batch_planning_is_deterministic_ordered_and_path_free(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, workload, definition, _ = _registered_descriptor_batch()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(dataset, artifact_store, workload)
|
|
input_artifact = request.inputs["input"].items[0].artifact
|
|
|
|
first = registry.plan(
|
|
request,
|
|
definition.manifest.package.digest,
|
|
runtime,
|
|
LocalPlanningContext(
|
|
artifact_store,
|
|
artifact_store,
|
|
tmp_path / "first-plan",
|
|
allowed_artifacts=(input_artifact,),
|
|
),
|
|
)
|
|
second = registry.plan(
|
|
request,
|
|
definition.manifest.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(task.parameters == {"skip_invalid": True} 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)
|
|
|
|
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 == [
|
|
["ALCOHOL", "ALKANE"],
|
|
["AMINE", "BROKEN"],
|
|
["HEXANE"],
|
|
]
|
|
|
|
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_descriptor_batch_skip_invalid_policy_is_explicit(tmp_path: Path) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, workload, definition, _ = _registered_descriptor_batch()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(dataset, artifact_store, workload, skip_invalid=False)
|
|
|
|
with pytest.raises(ValueError, match="invalid canonical_smiles"):
|
|
LocalCoreBatchExecutor(
|
|
registry, runtime, artifact_store, tmp_path / "work"
|
|
).execute(
|
|
request,
|
|
definition.manifest.package.digest,
|
|
)
|
|
|
|
|
|
def test_descriptor_batch_rejects_unknown_or_mistyped_parameters(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, workload, definition, _ = _registered_descriptor_batch()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
base = _request_for(dataset, artifact_store, workload)
|
|
|
|
for bad_parameters, message in (
|
|
({"skip_invalid": True, "bogus": 1}, "unknown field"),
|
|
({"skip_invalid": "yes"}, "type mismatch"),
|
|
):
|
|
request = replace(base, parameters=bad_parameters)
|
|
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=(base.inputs["input"].items[0].artifact,),
|
|
),
|
|
)
|
|
|
|
|
|
def _binding_from(provenance, trust_mode: TrustMode) -> VerificationBinding:
|
|
return VerificationBinding(
|
|
workload=provenance.workload,
|
|
task_key="reduce/final",
|
|
package_digest=provenance.package_digest,
|
|
manifest_digest=provenance.manifest_digest,
|
|
environment_digest=provenance.environment_digest,
|
|
parameters_digest=provenance.parameters_digest,
|
|
input_collection_digest=provenance.input_collection_digest,
|
|
execution_contract_digest=provenance.execution_contract_digest,
|
|
selected_features=provenance.selected_features,
|
|
optional_fallbacks=provenance.optional_fallbacks,
|
|
job_id=provenance.job_id,
|
|
task_id=provenance.task_id,
|
|
verifier=provenance.verifier,
|
|
sdk_api_version=provenance.sdk_api_version,
|
|
protocol_version=provenance.protocol_version,
|
|
manifest_schema_version=provenance.manifest_schema_version,
|
|
workflow_schema_version=provenance.workflow_schema_version,
|
|
artifact_schemas=provenance.artifact_schemas,
|
|
trust_mode=trust_mode,
|
|
)
|
|
|
|
|
|
def _candidate_for(
|
|
manifest,
|
|
candidate_id: str,
|
|
owner_id: str,
|
|
authentication_key: bytes,
|
|
) -> CandidateOutput:
|
|
return CandidateOutput.from_coordinator_record(
|
|
candidate_id,
|
|
owner_id,
|
|
manifest,
|
|
authentication_key,
|
|
)
|
|
|
|
|
|
def test_descriptor_batch_accepts_two_owner_quorum_on_identical_outputs(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, workload, definition, _ = _registered_descriptor_batch()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(dataset, artifact_store, workload)
|
|
|
|
final = LocalCoreBatchExecutor(
|
|
registry,
|
|
runtime,
|
|
artifact_store,
|
|
tmp_path / "sdk-work",
|
|
).execute(request, definition.manifest.package.digest)
|
|
|
|
provenance = replace(
|
|
final.provenance,
|
|
trust_mode="untrusted_quorum",
|
|
worker_runtime={"kind": "worker-one"},
|
|
)
|
|
first = replace(final, provenance=provenance)
|
|
second = replace(
|
|
final,
|
|
provenance=replace(
|
|
provenance,
|
|
worker_runtime={"kind": "worker-two"},
|
|
),
|
|
)
|
|
binding = _binding_from(provenance, TrustMode.UNTRUSTED_QUORUM)
|
|
assert binding.matches(first)
|
|
assert binding.matches(second)
|
|
|
|
key = b"coordinator-authentication-key-32-bytes"
|
|
decision = ExactArtifactVerifier().verify(
|
|
VerifyContext(
|
|
expected_outputs=definition.manifest.outputs,
|
|
max_output_bytes=definition.manifest.limits.max_output_bytes,
|
|
minimum_matches=2,
|
|
binding=binding,
|
|
trust_mode=TrustMode.UNTRUSTED_QUORUM,
|
|
),
|
|
CandidateOutputs(
|
|
(
|
|
_candidate_for(first, "candidate-one", "owner-one", key),
|
|
_candidate_for(second, "candidate-two", "owner-two", key),
|
|
)
|
|
),
|
|
)
|
|
|
|
assert decision.status is VerificationStatus.ACCEPTED
|
|
assert decision.reason_code == "quorum-match"
|
|
assert decision.accepted_digest == first.digest
|
|
assert decision.evidence["matched"] == 2
|
|
assert decision.evidence["distinct_digests"] == 1
|
|
|
|
|
|
def test_descriptor_batch_quorum_rejects_conflicting_worker_outputs(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
dataset = tmp_path / "molecules.tsv"
|
|
_write_tiny_dataset(dataset)
|
|
registry, runtime, workload, definition, _ = _registered_descriptor_batch()
|
|
artifact_store = LocalArtifactStore(tmp_path / "artifacts")
|
|
request = _request_for(dataset, artifact_store, workload)
|
|
|
|
final = LocalCoreBatchExecutor(
|
|
registry,
|
|
runtime,
|
|
artifact_store,
|
|
tmp_path / "sdk-work",
|
|
).execute(request, definition.manifest.package.digest)
|
|
provenance = replace(
|
|
final.provenance,
|
|
trust_mode="untrusted_quorum",
|
|
worker_runtime={"kind": "worker-one"},
|
|
)
|
|
good = replace(final, provenance=provenance)
|
|
forged_output = replace(
|
|
good.outputs["result"],
|
|
items=(
|
|
replace(
|
|
good.outputs["result"].items[0],
|
|
artifact=replace(
|
|
good.outputs["result"].items[0].artifact,
|
|
sha256="a" * 64,
|
|
),
|
|
),
|
|
),
|
|
)
|
|
bad = replace(good, outputs={"result": forged_output})
|
|
assert bad.digest != good.digest
|
|
binding = _binding_from(provenance, TrustMode.UNTRUSTED_QUORUM)
|
|
assert binding.matches(bad)
|
|
|
|
key = b"coordinator-authentication-key-32-bytes"
|
|
decision = ExactArtifactVerifier().verify(
|
|
VerifyContext(
|
|
expected_outputs=definition.manifest.outputs,
|
|
max_output_bytes=definition.manifest.limits.max_output_bytes,
|
|
minimum_matches=2,
|
|
binding=binding,
|
|
trust_mode=TrustMode.UNTRUSTED_QUORUM,
|
|
),
|
|
CandidateOutputs(
|
|
(
|
|
_candidate_for(good, "candidate-good-one", "owner-one", key),
|
|
_candidate_for(good, "candidate-good-two", "owner-two", key),
|
|
_candidate_for(bad, "candidate-bad-one", "owner-three", key),
|
|
_candidate_for(bad, "candidate-bad-two", "owner-four", key),
|
|
)
|
|
),
|
|
)
|
|
|
|
assert decision.status is VerificationStatus.REJECTED
|
|
assert decision.reason_code == "conflicting-quorums"
|
|
assert decision.evidence["largest_group"] == 2
|
|
assert decision.evidence["distinct_digests"] == 2
|
|
assert decision.accepted_digest is None
|
|
|
|
|
|
def test_descriptor_batch_discovery_imports_an_allowlisted_installed_entry_point(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from importlib import metadata
|
|
|
|
from scimesh.sdk.registry import WorkloadRegistry as RegistryClass
|
|
|
|
definition = descriptor_batch_sdk_definition().definition()
|
|
loaded: list[str] = []
|
|
|
|
class EntryPoint:
|
|
name = "descriptor-batch@1.0.0"
|
|
dist = metadata.distribution("scimesh")
|
|
value = "scimesh.workloads.descriptors:workload_definition"
|
|
|
|
@property
|
|
def module(self) -> str:
|
|
return self.value.partition(":")[0]
|
|
|
|
def load(self):
|
|
loaded.append(self.name)
|
|
return lambda: definition
|
|
|
|
class EntryPoints:
|
|
def __init__(self, values: tuple) -> None:
|
|
self._values = values
|
|
|
|
def __iter__(self):
|
|
return iter(self._values)
|
|
|
|
def select(self, *, group: str):
|
|
assert group == RegistryClass.ENTRY_POINT_GROUP
|
|
return self
|
|
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.metadata.entry_points",
|
|
lambda: EntryPoints((EntryPoint(),)),
|
|
)
|
|
monkeypatch.setattr(
|
|
"scimesh.sdk.registry.installed_distribution_digest",
|
|
lambda _distribution: definition.manifest.package.digest,
|
|
)
|
|
registry = WorkloadRegistry()
|
|
registry.discover_installed(
|
|
(
|
|
AllowedPackage(
|
|
"scimesh",
|
|
definition.manifest.workload,
|
|
definition.manifest.package.digest,
|
|
),
|
|
)
|
|
)
|
|
|
|
assert loaded == ["descriptor-batch@1.0.0"]
|
|
description = registry.descriptions()[0]
|
|
assert description.workload.name == "descriptor-batch"
|
|
assert description.enabled
|