Add MapReduceWorkload scaffold and generic workload execution
This commit is contained in:
Emil
@@ -0,0 +1,229 @@
|
||||
"""Tests for the generic ``scimesh workload`` CLI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from scimesh.cli import main
|
||||
|
||||
|
||||
def test_workload_cli_lists_sdk_workloads(capsys: object) -> None:
|
||||
assert main(["workload", "list"]) == 0
|
||||
output = capsys.readouterr().out
|
||||
assert "descriptor-batch" in output
|
||||
assert "similarity-graph" in output
|
||||
assert "similarity-search" in output
|
||||
assert "enabled" in output
|
||||
|
||||
|
||||
def test_workload_cli_runs_descriptor_batch(tmp_path: Path, capsys: object) -> None:
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
dataset.write_text(
|
||||
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nC\tCCN\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
output = tmp_path / "descriptors.csv"
|
||||
|
||||
code = main(
|
||||
[
|
||||
"workload",
|
||||
"run",
|
||||
"descriptor-batch",
|
||||
"--input",
|
||||
str(dataset),
|
||||
"--params",
|
||||
'{"skip_invalid": true}',
|
||||
"--shard-rows",
|
||||
"2",
|
||||
"-o",
|
||||
str(output),
|
||||
]
|
||||
)
|
||||
|
||||
assert code == 0
|
||||
lines = output.read_text(encoding="utf-8").splitlines()
|
||||
assert lines[0].startswith("chembl_id,canonical_smiles,ExactMolWt")
|
||||
assert len(lines) == 4
|
||||
assert "rows_emitted" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_workload_cli_runs_similarity_search(tmp_path: Path) -> None:
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
dataset.write_text(
|
||||
"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
output = tmp_path / "search.csv"
|
||||
|
||||
code = main(
|
||||
[
|
||||
"workload",
|
||||
"run",
|
||||
"similarity-search",
|
||||
"--input",
|
||||
str(dataset),
|
||||
"--params",
|
||||
'{"query_smiles": "CCO", "top_k": 5, "progress_every": 0}',
|
||||
"--shard-rows",
|
||||
"2",
|
||||
"-o",
|
||||
str(output),
|
||||
]
|
||||
)
|
||||
|
||||
assert code == 0
|
||||
lines = output.read_text(encoding="utf-8").splitlines()
|
||||
assert lines[0] == "rank,chembl_id,canonical_smiles,similarity"
|
||||
assert len(lines) == 2
|
||||
|
||||
|
||||
def test_workload_cli_rejects_unknown_or_missing_workload(tmp_path: Path) -> None:
|
||||
import pytest
|
||||
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
dataset.write_text("chembl_id\tcanonical_smiles\nA\tCCO\n", encoding="utf-8")
|
||||
assert main(["workload", "run", "no-such-workload", "--input", str(dataset)]) == 1
|
||||
with pytest.raises(SystemExit):
|
||||
main(["workload", "run", "descriptor-batch"])
|
||||
|
||||
|
||||
def test_workload_cli_rejects_invalid_params_json(tmp_path: Path) -> None:
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
dataset.write_text("chembl_id\tcanonical_smiles\nA\tCCO\n", encoding="utf-8")
|
||||
assert (
|
||||
main(
|
||||
[
|
||||
"workload",
|
||||
"run",
|
||||
"descriptor-batch",
|
||||
"--input",
|
||||
str(dataset),
|
||||
"--params",
|
||||
"{broken",
|
||||
]
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
def test_workload_cli_runs_an_allowlisted_custom_workload(
|
||||
tmp_path: Path, monkeypatch: object, capsys: object
|
||||
) -> None:
|
||||
import csv
|
||||
|
||||
from scimesh.sdk import (
|
||||
ArtifactSchema,
|
||||
ComponentRef,
|
||||
MapReduceWorkload,
|
||||
PortSpec,
|
||||
SchemaRef,
|
||||
WorkloadId,
|
||||
)
|
||||
from scimesh.sdk.registry import WorkloadRegistry
|
||||
from scimesh.workloads.environment import (
|
||||
current_environment_digest,
|
||||
current_scimesh_package_digest,
|
||||
)
|
||||
|
||||
class CountRowsWorkload(MapReduceWorkload):
|
||||
workload_id = WorkloadId("count-rows", "1.0.0")
|
||||
description = "Count TSV data rows per shard."
|
||||
parameters_schema = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {},
|
||||
}
|
||||
input_port = PortSpec(
|
||||
ArtifactSchema(
|
||||
SchemaRef("molecule-table", 1),
|
||||
"text/tab-separated-values",
|
||||
"utf-8",
|
||||
10**9,
|
||||
ComponentRef("delimited-table", 1),
|
||||
validator_configuration={
|
||||
"required_columns": ["canonical_smiles", "chembl_id"]
|
||||
},
|
||||
max_records=10**8,
|
||||
)
|
||||
)
|
||||
partial_port = output_port = PortSpec(
|
||||
ArtifactSchema(
|
||||
SchemaRef("count-table", 1),
|
||||
"text/csv",
|
||||
"utf-8",
|
||||
10**9,
|
||||
ComponentRef("delimited-table", 1),
|
||||
validator_configuration={"columns": ["id", "rows"]},
|
||||
max_records=10**8,
|
||||
)
|
||||
)
|
||||
|
||||
def partition_input(self, input_path, parameters, workspace):
|
||||
paths = []
|
||||
with input_path.open(encoding="utf-8", newline="") as source:
|
||||
for index, row in enumerate(csv.DictReader(source, delimiter="\t")):
|
||||
path = workspace / f"shard-{index}.tsv"
|
||||
path.write_text(
|
||||
"chembl_id\tcanonical_smiles\n"
|
||||
+ row["chembl_id"]
|
||||
+ "\t"
|
||||
+ row["canonical_smiles"]
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
def compute_shard(self, inputs, parameters, output_path):
|
||||
lines = inputs["input"].read_text(encoding="utf-8").splitlines()
|
||||
rows = max(len(lines) - 1, 0)
|
||||
output_path.write_text(
|
||||
"id,rows\nshard," + str(rows) + "\n", encoding="utf-8"
|
||||
)
|
||||
return {"rows": rows}
|
||||
|
||||
def reduce_partials(self, partial_paths, parameters, output_path):
|
||||
with output_path.open("w", encoding="utf-8") as destination:
|
||||
destination.write("id,rows\n")
|
||||
total = 0
|
||||
for partial in partial_paths:
|
||||
rows = partial.read_text(encoding="utf-8").splitlines()[1:]
|
||||
destination.write("".join(row + "\n" for row in rows))
|
||||
total += len(rows)
|
||||
return {"rows_total": total, "partial_count": len(partial_paths)}
|
||||
|
||||
definition = CountRowsWorkload(
|
||||
package_digest=current_scimesh_package_digest(),
|
||||
environment_digest=current_environment_digest(),
|
||||
).definition()
|
||||
|
||||
def fake_discover(self: WorkloadRegistry, allowlist) -> None:
|
||||
self.register(definition, enabled=True)
|
||||
|
||||
monkeypatch.setattr(WorkloadRegistry, "discover_installed", fake_discover)
|
||||
monkeypatch.setenv(
|
||||
"SCIMESH_WORKLOAD_ALLOWLIST",
|
||||
'[{"distribution": "scimesh", "name": "count-rows", "version": "1.0.0", '
|
||||
'"digest": "' + current_scimesh_package_digest() + '"}]',
|
||||
)
|
||||
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
dataset.write_text(
|
||||
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n", encoding="utf-8"
|
||||
)
|
||||
output = tmp_path / "counts.csv"
|
||||
assert (
|
||||
main(
|
||||
[
|
||||
"workload",
|
||||
"run",
|
||||
"count-rows",
|
||||
"--input",
|
||||
str(dataset),
|
||||
"-o",
|
||||
str(output),
|
||||
]
|
||||
)
|
||||
== 0
|
||||
)
|
||||
assert output.read_text(encoding="utf-8") == "id,rows\nshard,1\nshard,1\n"
|
||||
@@ -0,0 +1,305 @@
|
||||
"""Tests for the MapReduceWorkload authoring scaffold."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from scimesh.sdk import (
|
||||
ArtifactCollection,
|
||||
ArtifactSchema,
|
||||
ComponentRef,
|
||||
DeterminismProfile,
|
||||
JobRequest,
|
||||
LocalArtifactStore,
|
||||
LocalCoreBatchExecutor,
|
||||
LocalPlanningContext,
|
||||
MapReduceWorkload,
|
||||
PortSpec,
|
||||
SchemaRef,
|
||||
StageKind,
|
||||
TrustMode,
|
||||
WorkloadId,
|
||||
WorkloadRegistry,
|
||||
assert_manifest_round_trip,
|
||||
)
|
||||
from scimesh.workloads.library import default_sdk_runtime
|
||||
from scimesh.workloads.environment import (
|
||||
current_environment_digest,
|
||||
current_scimesh_package_digest,
|
||||
)
|
||||
|
||||
|
||||
def _molecule_port() -> PortSpec:
|
||||
return PortSpec(
|
||||
ArtifactSchema(
|
||||
SchemaRef("molecule-table", 1),
|
||||
"text/tab-separated-values",
|
||||
"utf-8",
|
||||
10**9,
|
||||
ComponentRef("delimited-table", 1),
|
||||
validator_configuration={
|
||||
"required_columns": ["canonical_smiles", "chembl_id"]
|
||||
},
|
||||
max_records=10**8,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _count_port() -> PortSpec:
|
||||
return PortSpec(
|
||||
ArtifactSchema(
|
||||
SchemaRef("count-table", 1),
|
||||
"text/csv",
|
||||
"utf-8",
|
||||
10**9,
|
||||
ComponentRef("delimited-table", 1),
|
||||
validator_configuration={"columns": ["id", "rows"]},
|
||||
max_records=10**8,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class CountRowsWorkload(MapReduceWorkload):
|
||||
"""A minimal author-written workload: three scientific hooks only."""
|
||||
|
||||
workload_id = WorkloadId("count-rows", "1.0.0")
|
||||
description = "Count TSV data rows per shard and concatenate the counts."
|
||||
parameters_schema = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {"prefix": {"type": "string", "minLength": 1, "maxLength": 50}},
|
||||
}
|
||||
input_port = _molecule_port()
|
||||
partial_port = _count_port()
|
||||
output_port = _count_port()
|
||||
map_parameter_names = ("prefix",)
|
||||
reduce_parameter_names = ("prefix",)
|
||||
|
||||
def partition_input(self, input_path, parameters, workspace):
|
||||
paths = []
|
||||
with input_path.open(encoding="utf-8", newline="") as source:
|
||||
for index, row in enumerate(csv.DictReader(source, delimiter="\t")):
|
||||
path = workspace / f"shard-{index}.tsv"
|
||||
path.write_text(
|
||||
"chembl_id\tcanonical_smiles\n"
|
||||
+ row["chembl_id"]
|
||||
+ "\t"
|
||||
+ row["canonical_smiles"]
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
def compute_shard(self, inputs, parameters, output_path):
|
||||
lines = inputs["input"].read_text(encoding="utf-8").splitlines()
|
||||
rows = max(len(lines) - 1, 0)
|
||||
output_path.write_text(
|
||||
"id,rows\n" + parameters.get("prefix", "shard") + "," + str(rows) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return {"rows": rows} # type: ignore[return-value]
|
||||
|
||||
def reduce_partials(self, partial_paths, parameters, output_path):
|
||||
total = 0
|
||||
with output_path.open("w", encoding="utf-8") as destination:
|
||||
destination.write("id,rows\n")
|
||||
for partial in partial_paths:
|
||||
for index, line in enumerate(
|
||||
partial.read_text(encoding="utf-8").splitlines()
|
||||
):
|
||||
if index == 0:
|
||||
continue
|
||||
destination.write(line + "\n")
|
||||
total += int(line.split(",")[1])
|
||||
return {"rows_total": total, "partial_count": len(partial_paths)} # type: ignore[return-value]
|
||||
|
||||
|
||||
def _registered_count_rows():
|
||||
workload = CountRowsWorkload(
|
||||
package_digest=current_scimesh_package_digest(),
|
||||
environment_digest=current_environment_digest(),
|
||||
)
|
||||
registry = WorkloadRegistry()
|
||||
registry.register(workload.definition(), enabled=True)
|
||||
runtime = replace(
|
||||
default_sdk_runtime(),
|
||||
workload_capabilities=(
|
||||
*default_sdk_runtime().workload_capabilities,
|
||||
"count-rows",
|
||||
),
|
||||
)
|
||||
return workload, registry, runtime
|
||||
|
||||
|
||||
def _write_dataset(path: Path) -> None:
|
||||
path.write_text(
|
||||
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nC\tCCN\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def _request(workload, store, dataset) -> JobRequest:
|
||||
artifact = store.import_file(
|
||||
dataset,
|
||||
declaration=workload.manifest.inputs["input"].schema,
|
||||
)
|
||||
return JobRequest(
|
||||
workload=workload.manifest.workload,
|
||||
parameters={"prefix": "x"},
|
||||
inputs={"input": ArtifactCollection.single(artifact)},
|
||||
)
|
||||
|
||||
|
||||
def test_map_reduce_scaffold_assembles_the_manifest_and_runs(tmp_path: Path) -> None:
|
||||
workload, registry, runtime = _registered_count_rows()
|
||||
manifest = workload.manifest
|
||||
|
||||
assert manifest.workload.name == "count-rows"
|
||||
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(workload.definition().runners) == {
|
||||
manifest.workflow.stages[0].entry_point
|
||||
}
|
||||
assert set(workload.definition().reducers) == {
|
||||
manifest.workflow.stages[1].entry_point
|
||||
}
|
||||
assert_manifest_round_trip(manifest)
|
||||
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
_write_dataset(dataset)
|
||||
store = LocalArtifactStore(tmp_path / "artifacts")
|
||||
result = LocalCoreBatchExecutor(
|
||||
registry,
|
||||
runtime,
|
||||
store,
|
||||
tmp_path / "work",
|
||||
).execute(_request(workload, store, dataset), workload.manifest.package.digest)
|
||||
|
||||
assert result.task_key == "reduce/final"
|
||||
assert dict(result.metrics) == {"rows_total": 3, "partial_count": 3}
|
||||
text = store.materialize(result.outputs["result"].items[0].artifact).read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
assert text == "id,rows\nx,1\nx,1\nx,1\n"
|
||||
|
||||
|
||||
def test_map_reduce_scaffold_derives_pinned_plans_and_parameters(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
workload, registry, runtime = _registered_count_rows()
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
_write_dataset(dataset)
|
||||
store = LocalArtifactStore(tmp_path / "artifacts")
|
||||
request = _request(workload, store, dataset)
|
||||
input_artifact = request.inputs["input"].items[0].artifact
|
||||
|
||||
plan = registry.plan(
|
||||
request,
|
||||
workload.manifest.package.digest,
|
||||
runtime,
|
||||
LocalPlanningContext(
|
||||
store,
|
||||
store,
|
||||
tmp_path / "plan",
|
||||
allowed_artifacts=(input_artifact,),
|
||||
),
|
||||
)
|
||||
assert [task.task_key for task in plan.tasks] == [
|
||||
"map/00000000",
|
||||
"map/00000001",
|
||||
"map/00000002",
|
||||
]
|
||||
assert all(task.parameters == {"prefix": "x"} for task in plan.tasks)
|
||||
assert all(task.package_digest == plan.package_digest for task in plan.tasks)
|
||||
assert all(task.manifest_digest == plan.manifest_digest for task in plan.tasks)
|
||||
assert all(task.trust_mode is TrustMode.TRUSTED for task in plan.tasks)
|
||||
|
||||
|
||||
def test_map_reduce_scaffold_requires_scientific_hooks() -> None:
|
||||
class MissingHooksWorkload(MapReduceWorkload):
|
||||
workload_id = WorkloadId("missing-hooks", "1.0.0")
|
||||
description = "A workload that forgets its scientific hooks."
|
||||
parameters_schema = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {},
|
||||
}
|
||||
input_port = _molecule_port()
|
||||
partial_port = _count_port()
|
||||
output_port = _count_port()
|
||||
|
||||
workload = MissingHooksWorkload(
|
||||
package_digest=current_scimesh_package_digest(),
|
||||
environment_digest=current_environment_digest(),
|
||||
)
|
||||
with pytest.raises(NotImplementedError, match="partition_input"):
|
||||
workload.partition_input(Path("input"), {}, Path("workspace"))
|
||||
with pytest.raises(NotImplementedError, match="compute_shard"):
|
||||
workload.compute_shard({}, {}, Path("output"))
|
||||
with pytest.raises(NotImplementedError, match="reduce_partials"):
|
||||
workload.reduce_partials([], {}, Path("output"))
|
||||
|
||||
|
||||
def test_map_reduce_scaffold_default_partial_keys_are_contiguous() -> None:
|
||||
workload = CountRowsWorkload(
|
||||
package_digest=current_scimesh_package_digest(),
|
||||
environment_digest=current_environment_digest(),
|
||||
)
|
||||
assert workload.parse_partial_key("map.00000000") == 0
|
||||
assert workload.parse_partial_key("map.00000002") == 2
|
||||
with pytest.raises(ValueError, match="eight-digit-index"):
|
||||
workload.parse_partial_key("map.0")
|
||||
workload.validate_partial_keys((0, 1, 2))
|
||||
with pytest.raises(ValueError, match="complete and contiguous"):
|
||||
workload.validate_partial_keys((0, 2))
|
||||
with pytest.raises(ValueError, match="complete and contiguous"):
|
||||
workload.validate_partial_keys((0, 0, 1))
|
||||
|
||||
|
||||
def test_map_reduce_scaffold_rejects_domain_invalid_parameters(tmp_path: Path) -> None:
|
||||
workload, registry, runtime = _registered_count_rows()
|
||||
dataset = tmp_path / "molecules.tsv"
|
||||
_write_dataset(dataset)
|
||||
store = LocalArtifactStore(tmp_path / "artifacts")
|
||||
artifact = store.import_file(
|
||||
dataset,
|
||||
declaration=workload.manifest.inputs["input"].schema,
|
||||
)
|
||||
|
||||
class StrictCountRows(CountRowsWorkload):
|
||||
def domain_validate(self, parameters):
|
||||
if "prefix" not in parameters:
|
||||
raise ValueError("prefix is required")
|
||||
|
||||
strict = StrictCountRows(
|
||||
package_digest=current_scimesh_package_digest(),
|
||||
environment_digest=current_environment_digest(),
|
||||
)
|
||||
registry2 = WorkloadRegistry()
|
||||
registry2.register(strict.definition(), enabled=True)
|
||||
request = JobRequest(
|
||||
workload=strict.manifest.workload,
|
||||
parameters={},
|
||||
inputs={"input": ArtifactCollection.single(artifact)},
|
||||
)
|
||||
with pytest.raises(ValueError, match="prefix is required"):
|
||||
registry2.plan(
|
||||
request,
|
||||
strict.manifest.package.digest,
|
||||
runtime,
|
||||
LocalPlanningContext(store, store, tmp_path / "plan"),
|
||||
)
|
||||
+111
-3
@@ -10,7 +10,7 @@ from urllib.request import Request
|
||||
|
||||
import pytest
|
||||
|
||||
from scimesh.worker.config import WorkerConfig
|
||||
from scimesh.worker.config import WorkerConfig, _workload_allowlist
|
||||
from scimesh.worker import cli as worker_cli
|
||||
from scimesh.worker.cli import build_parser
|
||||
from scimesh.worker.coordinator import CoordinatorTransientError
|
||||
@@ -579,7 +579,7 @@ def test_runner_executes_search_through_the_sdk_and_rejects_graph(
|
||||
)
|
||||
|
||||
|
||||
def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_options(
|
||||
def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_parameters(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
task_dir = tmp_path / "search"
|
||||
@@ -607,7 +607,7 @@ def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_options(
|
||||
InputArtifact("https://example/input", "a" * 64),
|
||||
{"query_smiles": "CCO", "max_rows": 1},
|
||||
)
|
||||
with pytest.raises(ValueError, match="unsupported runner parameters"):
|
||||
with pytest.raises(ValueError, match="outside the stage projection"):
|
||||
SciMeshRunner().run(with_max_rows, task_dir)
|
||||
|
||||
|
||||
@@ -732,3 +732,111 @@ def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None:
|
||||
worker._register_worker()
|
||||
assert worker.worker_id == "11111111-1111-4111-8111-111111111111"
|
||||
assert worker.config.heartbeat_interval == 15
|
||||
|
||||
|
||||
def test_runner_executes_an_arbitrary_sdk_workload(tmp_path: Path) -> None:
|
||||
from scimesh.workloads.descriptors import descriptor_batch_sdk_definition
|
||||
|
||||
content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n"
|
||||
task = ClaimedTask(
|
||||
"task-1",
|
||||
1,
|
||||
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
|
||||
"descriptor-batch",
|
||||
InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()),
|
||||
{"skip_invalid": True},
|
||||
)
|
||||
task_dir = tmp_path / "task-1" / "1"
|
||||
task_dir.mkdir(parents=True)
|
||||
(task_dir / "input").write_bytes(content)
|
||||
runner = SciMeshRunner(
|
||||
definitions={
|
||||
"descriptor-batch": descriptor_batch_sdk_definition().definition()
|
||||
}
|
||||
)
|
||||
|
||||
result = runner.run(task, task_dir)
|
||||
header = result.artifacts[0].path.read_text(encoding="utf-8").splitlines()[0]
|
||||
assert header.startswith("chembl_id,canonical_smiles,ExactMolWt")
|
||||
assert result.metrics["rows_emitted"] == 2
|
||||
|
||||
|
||||
def test_runner_rejects_workloads_outside_the_v1_single_input_contract(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
from scimesh.workloads.graph import similarity_graph_sdk_definition
|
||||
|
||||
content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n"
|
||||
task = ClaimedTask(
|
||||
"graph-task",
|
||||
1,
|
||||
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
|
||||
"similarity-graph",
|
||||
InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()),
|
||||
{"threshold": 0.5},
|
||||
)
|
||||
task_dir = tmp_path / "graph"
|
||||
task_dir.mkdir(parents=True)
|
||||
(task_dir / "input").write_bytes(content)
|
||||
runner = SciMeshRunner(
|
||||
definitions={
|
||||
"similarity-graph": similarity_graph_sdk_definition().definition()
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="v1 single-input contract"):
|
||||
runner.run(task, task_dir)
|
||||
|
||||
|
||||
def test_runner_for_worker_discovers_allowlisted_installed_workloads(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
from scimesh.sdk.registry import WorkloadRegistry
|
||||
from scimesh.workloads.search import similarity_search_sdk_definition
|
||||
|
||||
definition = similarity_search_sdk_definition().definition()
|
||||
|
||||
def fake_discover(self: WorkloadRegistry, allowlist) -> None:
|
||||
assert len(allowlist) == 1
|
||||
self.register(definition, enabled=True)
|
||||
|
||||
monkeypatch.setattr(WorkloadRegistry, "discover_installed", fake_discover)
|
||||
allowlist = _workload_allowlist(
|
||||
'[{"distribution": "scimesh", "name": "similarity-search", '
|
||||
'"version": "1.0.0", "digest": "sha256:' + "a" * 64 + '"}]'
|
||||
)
|
||||
config = WorkerConfig(
|
||||
"https://example.test", "worker-1", tmp_path / "work",
|
||||
capabilities=("similarity-search",),
|
||||
workload_allowlist=allowlist,
|
||||
)
|
||||
runner = SciMeshRunner.for_worker(config)
|
||||
assert set(runner._definitions) == {"similarity-search"}
|
||||
|
||||
|
||||
def test_worker_config_parses_capabilities_and_workload_allowlist(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
monkeypatch.setenv("SCIMESH_CAPABILITIES", "similarity-search,descriptor-batch")
|
||||
config = WorkerConfig.from_environment(
|
||||
{"coordinator_url": "https://example.test", "work_dir": tmp_path}
|
||||
)
|
||||
assert config.capabilities == ("similarity-search", "descriptor-batch")
|
||||
assert config.workload_allowlist == ()
|
||||
|
||||
monkeypatch.setenv(
|
||||
"SCIMESH_WORKLOAD_ALLOWLIST",
|
||||
'[{"distribution": "scimesh", "name": "descriptor-batch", '
|
||||
'"version": "1.0.0", "digest": "sha256:' + "b" * 64 + '"}]',
|
||||
)
|
||||
config = WorkerConfig.from_environment(
|
||||
{"coordinator_url": "https://example.test", "work_dir": tmp_path}
|
||||
)
|
||||
assert len(config.workload_allowlist) == 1
|
||||
assert config.workload_allowlist[0].workload.name == "descriptor-batch"
|
||||
|
||||
monkeypatch.setenv("SCIMESH_WORKLOAD_ALLOWLIST", "not-json")
|
||||
with pytest.raises(ValueError, match="valid JSON"):
|
||||
WorkerConfig.from_environment(
|
||||
{"coordinator_url": "https://example.test", "work_dir": tmp_path}
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user