Emil_Shanaty/SciMesh
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Replace legacy distributed protocol with SDK-built workloads

Browse Source
This commit is contained in:
Emil
2026-08-01 23:57:41 +03:00
parent 96169086f0
commit 19fbb8e926
34 changed files with 3009 additions and 1986 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scimesh/distributed/__init__.py
-36
View File
@@ -1,36 +0,0 @@
"""Coordinator-independent contracts for distributed SciMesh workloads.
This package defines the typed plan and reduction boundary shared by future
planners, worker adapters, and coordinator bridges. It intentionally has no
network, database, or coordinator imports.
"""
from .models import (
ArtifactReference,
CompletedPartial,
DistributedPlan,
FinalResult,
PlannedTask,
)
from .registry import (
DistributedWorkloadRegistry,
PlanningService,
WorkloadDescription,
default_distributed_registry,
)
from .similarity_search import SimilaritySearchDistributedWorkload
from .workload import DistributedWorkload
__all__ = [
"ArtifactReference",
"CompletedPartial",
"default_distributed_registry",
"DistributedPlan",
"DistributedWorkload",
"DistributedWorkloadRegistry",
"FinalResult",
"PlannedTask",
"PlanningService",
"SimilaritySearchDistributedWorkload",
"WorkloadDescription",
]
scimesh/distributed/models.py
-262
View File
@@ -1,262 +0,0 @@
"""Versioned, JSON-safe value objects for distributed workload contracts."""
from __future__ import annotations
import json
import math
import re
from dataclasses import dataclass
from typing import Any, Mapping, Sequence
from uuid import UUID
SCHEMA_VERSION = 1
_WORKLOAD_NAME = re.compile(r"^[a-z][a-z0-9]*(?:-[a-z0-9]+)*$")
def _canonical_uuid(value: object, field: str) -> str:
if not isinstance(value, str):
raise ValueError(f"{field} must be a UUID string")
try:
return str(UUID(value))
except ValueError as error:
raise ValueError(f"{field} must be a UUID string") from error
def _sha256(value: object, field: str) -> str:
if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value):
raise ValueError(f"{field} must be a lowercase SHA-256 hex digest")
return value
def _content_type(value: object, field: str) -> str:
if not isinstance(value, str) or not value or len(value) > 128:
raise ValueError(f"{field} must be a non-empty content type")
if any(character.isspace() or ord(character) < 32 for character in value):
raise ValueError(f"{field} must be a non-empty content type")
return value
def _workload_name(value: object, field: str = "workload") -> str:
if not isinstance(value, str) or not _WORKLOAD_NAME.fullmatch(value):
raise ValueError(f"{field} must be a canonical hyphenated workload name")
return value
def _json_value(value: object, field: str) -> Any:
"""Deep-copy a JSON value and reject non-finite or non-string-key data."""
if value is None or isinstance(value, (bool, int)):
return value
if isinstance(value, str):
# Coordinator-owned artifacts are represented exclusively by
# ArtifactReference. A URI or a local path in a generic JSON payload
# would let a planner accidentally leak a bridge/worker implementation
# detail into durable task metadata.
forbidden_prefixes = ("file://", "worker://", "http://", "https://", "s3://", "/")
is_windows_path = len(value) >= 3 and value[0].isalpha() and value[1:3] in (":/", ":\\")
if value.startswith(forbidden_prefixes) or is_windows_path:
raise ValueError(f"{field} must not contain a URI or local path")
return value
if isinstance(value, float):
if not math.isfinite(value):
raise ValueError(f"{field} must not contain NaN or infinity")
return value
if isinstance(value, Mapping):
copied: dict[str, Any] = {}
for key, child in value.items():
if not isinstance(key, str):
raise ValueError(f"{field} must use string object keys")
copied[key] = _json_value(child, f"{field}.{key}")
return copied
if isinstance(value, (list, tuple)):
return [_json_value(child, f"{field}[]") for child in value]
raise ValueError(f"{field} must contain only JSON-compatible values")
def _json_mapping(value: object, field: str) -> dict[str, Any]:
if not isinstance(value, Mapping):
raise ValueError(f"{field} must be an object")
return _json_value(value, field)
@dataclass(frozen=True)
class ArtifactReference:
"""Immutable coordinator-owned artifact identity used in a plan."""
artifact_id: str
sha256: str
content_type: str
def __post_init__(self) -> None:
object.__setattr__(self, "artifact_id", _canonical_uuid(self.artifact_id, "artifact_id"))
object.__setattr__(self, "sha256", _sha256(self.sha256, "sha256"))
object.__setattr__(self, "content_type", _content_type(self.content_type, "content_type"))
def to_dict(self) -> dict[str, str]:
return {
"artifact_id": self.artifact_id,
"sha256": self.sha256,
"content_type": self.content_type,
}
@classmethod
def from_dict(cls, value: object) -> "ArtifactReference":
if not isinstance(value, Mapping):
raise ValueError("artifact reference must be an object")
_require_exact_keys(value, {"artifact_id", "sha256", "content_type"}, "artifact reference")
return cls(
artifact_id=value["artifact_id"],
sha256=value["sha256"],
content_type=value["content_type"],
)
@dataclass(frozen=True)
class PlannedTask:
"""One deterministically indexed, artifact-backed worker task."""
chunk_index: int
input_artifact: ArtifactReference
parameters: Mapping[str, object]
def __post_init__(self) -> None:
if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
raise ValueError("chunk_index must be a non-negative integer")
if not isinstance(self.input_artifact, ArtifactReference):
raise ValueError("input_artifact must be an ArtifactReference")
object.__setattr__(self, "parameters", _json_mapping(self.parameters, "task parameters"))
def to_dict(self) -> dict[str, Any]:
return {
"chunk_index": self.chunk_index,
"input_artifact": self.input_artifact.to_dict(),
"parameters": _json_value(self.parameters, "task parameters"),
}
@classmethod
def from_dict(cls, value: object) -> "PlannedTask":
if not isinstance(value, Mapping):
raise ValueError("planned task must be an object")
_require_exact_keys(value, {"chunk_index", "input_artifact", "parameters"}, "planned task")
return cls(
chunk_index=value["chunk_index"],
input_artifact=ArtifactReference.from_dict(value["input_artifact"]),
parameters=value["parameters"],
)
@dataclass(frozen=True)
class DistributedPlan:
"""The complete schema-versioned output of a distributed planner."""
workload: str
resolved_parameters: Mapping[str, object]
tasks: Sequence[PlannedTask]
schema_version: int = SCHEMA_VERSION
def __post_init__(self) -> None:
if self.schema_version != SCHEMA_VERSION:
raise ValueError(f"schema_version must be {SCHEMA_VERSION}")
object.__setattr__(self, "workload", _workload_name(self.workload))
object.__setattr__(self, "resolved_parameters", _json_mapping(self.resolved_parameters, "resolved_parameters"))
task_list = tuple(self.tasks)
if not task_list:
raise ValueError("plan must contain at least one task")
if any(not isinstance(task, PlannedTask) for task in task_list):
raise ValueError("tasks must contain PlannedTask values")
indexes = [task.chunk_index for task in task_list]
if indexes != sorted(indexes) or len(set(indexes)) != len(indexes):
raise ValueError("tasks must have unique, ascending chunk_index values")
object.__setattr__(self, "tasks", task_list)
def to_dict(self) -> dict[str, Any]:
return {
"schema_version": self.schema_version,
"workload": self.workload,
"resolved_parameters": _json_value(self.resolved_parameters, "resolved_parameters"),
"tasks": [task.to_dict() for task in self.tasks],
}
def to_json(self) -> str:
"""Return stable JSON suitable for hashing, tests, and durable payloads."""
return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False)
@classmethod
def from_dict(cls, value: object) -> "DistributedPlan":
if not isinstance(value, Mapping):
raise ValueError("distributed plan must be an object")
_require_exact_keys(
value,
{"schema_version", "workload", "resolved_parameters", "tasks"},
"distributed plan",
)
raw_tasks = value["tasks"]
if not isinstance(raw_tasks, list):
raise ValueError("tasks must be an array")
return cls(
schema_version=value["schema_version"],
workload=value["workload"],
resolved_parameters=value["resolved_parameters"],
tasks=tuple(PlannedTask.from_dict(task) for task in raw_tasks),
)
@classmethod
def from_json(cls, value: str) -> "DistributedPlan":
try:
decoded = json.loads(value)
except (TypeError, json.JSONDecodeError) as error:
raise ValueError("distributed plan must be valid JSON") from error
return cls.from_dict(decoded)
@dataclass(frozen=True)
class CompletedPartial:
"""Coordinator-owned partial output supplied to a reducer."""
chunk_index: int
artifact: ArtifactReference
metrics: Mapping[str, int | float]
def __post_init__(self) -> None:
if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
raise ValueError("chunk_index must be a non-negative integer")
if not isinstance(self.artifact, ArtifactReference):
raise ValueError("artifact must be an ArtifactReference")
if not isinstance(self.metrics, Mapping):
raise ValueError("metrics must be an object")
metrics: dict[str, int | float] = {}
for name, value in self.metrics.items():
if not isinstance(name, str) or not name:
raise ValueError("metric names must be non-empty strings")
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
raise ValueError("metric values must be finite JSON numbers")
metrics[name] = value
object.__setattr__(self, "metrics", metrics)
@dataclass(frozen=True)
class FinalResult:
"""A reducer's durable output, ready for coordinator persistence."""
artifact: ArtifactReference
metrics: Mapping[str, int | float]
def __post_init__(self) -> None:
if not isinstance(self.artifact, ArtifactReference):
raise ValueError("artifact must be an ArtifactReference")
# Reuse the CompletedPartial metric validation without inventing a fake
# artifact lifecycle or widening the result contract.
object.__setattr__(self, "metrics", CompletedPartial(0, self.artifact, self.metrics).metrics)
def _require_exact_keys(value: Mapping[str, object], expected: set[str], label: str) -> None:
actual = set(value)
if actual != expected:
missing = sorted(expected - actual)
unknown = sorted(actual - expected)
details: list[str] = []
if missing:
details.append(f"missing {', '.join(missing)}")
if unknown:
details.append(f"unknown {', '.join(unknown)}")
raise ValueError(f"{label} has {'; '.join(details)} fields")
scimesh/distributed/registry.py
-108
View File
@@ -1,108 +0,0 @@
"""Registry and orchestration helpers for distributed workload contracts."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Mapping, Sequence
from .models import CompletedPartial, DistributedPlan, FinalResult, _workload_name
from .workload import DistributedWorkload
@dataclass(frozen=True)
class WorkloadDescription:
"""Safe metadata that a future coordinator or UI may display."""
name: str
description: str
class DistributedWorkloadRegistry:
"""Collect distributed workloads without coupling them to the CLI registry."""
def __init__(self) -> None:
self._workloads: dict[str, DistributedWorkload] = {}
def register(self, workload: DistributedWorkload) -> None:
name = _workload_name(workload.name)
if name in self._workloads:
raise ValueError(f"distributed workload already registered: {name}")
if not isinstance(workload.description, str) or not workload.description.strip():
raise ValueError("distributed workload description must be non-empty")
self._workloads[name] = workload
def require(self, name: str) -> DistributedWorkload:
try:
return self._workloads[_workload_name(name)]
except KeyError as error:
raise ValueError(f"unknown distributed workload: {name}") from error
def descriptions(self) -> tuple[WorkloadDescription, ...]:
return tuple(
WorkloadDescription(name, workload.description)
for name, workload in sorted(self._workloads.items())
)
class PlanningService:
"""Small bridge-safe orchestration around a distributed workload registry.
It writes neither jobs nor artifacts. A Go coordinator bridge can therefore
validate and produce a plan before opening its own all-or-nothing persistence
transaction; CTX-09 will implement that concrete bridge and durable result
orchestration.
"""
def __init__(self, registry: DistributedWorkloadRegistry) -> None:
self._registry = registry
def plan(
self,
workload_name: str,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
raise ValueError("shard_rows must be a positive integer")
workload = self._registry.require(workload_name)
workload.validate_job(parameters)
plan = workload.plan(input_path, input_artifact_id, parameters, shard_rows, workspace)
if not isinstance(plan, DistributedPlan):
raise ValueError("distributed planner must return a DistributedPlan")
if plan.workload != workload.name:
raise ValueError("distributed planner returned a plan for another workload")
# Round-trip through the strict wire schema now, before a future bridge
# persists anything. This catches non-JSON values and undeclared fields.
return DistributedPlan.from_json(plan.to_json())
def reduce(
self,
workload_name: str,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
workload = self._registry.require(workload_name)
indexes = [partial.chunk_index for partial in partial_results]
if len(indexes) != len(set(indexes)):
raise ValueError("partial results must have unique chunk_index values")
ordered = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
result = workload.reduce(ordered, parameters, workspace)
if not isinstance(result, FinalResult):
raise ValueError("distributed reducer must return a FinalResult")
return result
def default_distributed_registry() -> DistributedWorkloadRegistry:
"""Return the currently supported distributed scientific workloads."""
# Delayed import keeps the generic registry independent of concrete RDKit
# workloads and avoids making the contract layer import application setup.
from .similarity_search import SimilaritySearchDistributedWorkload
registry = DistributedWorkloadRegistry()
registry.register(SimilaritySearchDistributedWorkload())
return registry
scimesh/distributed/similarity_search.py
-380
View File
@@ -1,380 +0,0 @@
"""Distributed planning and reduction for exact molecular similarity search."""
from __future__ import annotations
import csv
import hashlib
import heapq
import math
from pathlib import Path
from typing import Any, Iterator, Mapping, Sequence
from uuid import UUID, uuid5
from rdkit import Chem
from scimesh.chemistry.dataset import MoleculeRecord, find_molecule_by_id, parse_smiles
from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE
from scimesh.workloads.similarity_search import (
SimilarityMatch,
_HeapEntry,
search_similar,
write_search_results,
)
from .models import ArtifactReference, CompletedPartial, DistributedPlan, FinalResult, PlannedTask
_TSV_CONTENT_TYPE = "text/tab-separated-values"
_CSV_CONTENT_TYPE = "text/csv"
_SEARCH_COLUMNS = ("rank", "chembl_id", "canonical_smiles", "similarity")
_REQUIRED_COLUMNS = {"chembl_id", "canonical_smiles"}
def write_similarity_search_partial(output_path: Path, matches: Sequence[SimilarityMatch]) -> None:
"""Write a worker partial with a round-trip score, not display rounding.
The public final CSV continues to use the local CLI's six-decimal display.
A reducer needs the full binary float representation to rank candidates
from separate shards exactly as the single-process reference does.
"""
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", encoding="utf-8", newline="") as destination:
writer = csv.DictWriter(destination, fieldnames=_SEARCH_COLUMNS)
writer.writeheader()
for rank, match in enumerate(matches, start=1):
writer.writerow({
"rank": rank,
"chembl_id": match.molecule_id,
"canonical_smiles": match.smiles,
"similarity": repr(match.similarity),
})
class SimilaritySearchDistributedWorkload:
"""Planner/reducer for exact global top-k Tanimoto similarity search."""
name = "similarity-search"
description = "Exact top-k molecular similarity search over deterministic TSV shards."
def validate_job(self, parameters: Mapping[str, object]) -> None:
allowed = {
"query_id", "query_smiles", "top_k", "threshold",
"threshold_direction", "max_rows", "progress_every",
}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
query_id = parameters.get("query_id")
query_smiles = parameters.get("query_smiles")
if (query_id is None) == (query_smiles is None):
raise ValueError("exactly one of query_id or query_smiles is required")
if query_id is not None:
self._string(query_id, "query_id")
if query_smiles is not None:
self._string(query_smiles, "query_smiles")
self._positive_int(parameters.get("top_k", 20), "top_k")
if "max_rows" in parameters:
self._positive_int(parameters["max_rows"], "max_rows")
if "progress_every" in parameters:
self._nonnegative_int(parameters["progress_every"], "progress_every")
if "threshold" in parameters:
self._unit_interval(parameters["threshold"], "threshold")
if "threshold_direction" in parameters and parameters["threshold_direction"] not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
self.validate_job(parameters)
if not input_path.is_file():
raise ValueError("input_path must be a readable dataset file")
if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
raise ValueError("shard_rows must be a positive integer")
try:
input_id = UUID(input_artifact_id)
except ValueError as error:
raise ValueError("input_artifact_id must be a UUID") from error
query_smiles, query_source = self._resolve_query(input_path, parameters)
resolved = self._resolved_parameters(parameters, query_smiles, query_source)
workspace.mkdir(parents=True, exist_ok=True)
shard_paths: list[Path] = []
try:
shard_paths = self._write_shards(input_path, workspace, shard_rows, resolved.get("max_rows"))
tasks = tuple(
PlannedTask(
chunk_index=index,
input_artifact=ArtifactReference(
artifact_id=str(uuid5(input_id, f"scimesh:similarity-search:shard:{index}")),
sha256=_sha256_file(path),
content_type=_TSV_CONTENT_TYPE,
),
parameters=self._task_parameters(resolved),
)
for index, path in enumerate(shard_paths)
)
except Exception:
for path in shard_paths:
path.unlink(missing_ok=True)
raise
return DistributedPlan(self.name, resolved, tasks)
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
"""Merge materialized partial CSVs into one deterministic final CSV.
The coordinator bridge materializes each downloaded artifact at
``workspace / artifact_id`` before it calls this method. Those local
paths are an ephemeral bridge detail, never present in the plan or task
payload. CTX-09 owns the durable final-artifact upload and job state.
"""
if not partial_results:
raise ValueError("at least one partial result is required")
resolved = self._validate_resolved_parameters(parameters)
top_k = resolved["top_k"]
direction = resolved["threshold_direction"]
heap: list[_HeapEntry] = []
ordered_partials = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
indexes = [partial.chunk_index for partial in ordered_partials]
if len(indexes) != len(set(indexes)):
raise ValueError("partial results must have unique chunk_index values")
for partial in ordered_partials:
path = workspace / partial.artifact.artifact_id
if not path.is_file():
raise ValueError("materialized partial result is missing")
if _sha256_file(path) != partial.artifact.sha256:
raise ValueError("materialized partial result checksum does not match its artifact reference")
for match in self._read_partial(path, direction):
rank_key = match.sort_key(direction)
entry = _HeapEntry(match, rank_key)
if len(heap) < top_k:
heapq.heappush(heap, entry)
elif rank_key < heap[0].rank_key:
heapq.heapreplace(heap, entry)
matches = sorted((entry.match for entry in heap), key=lambda match: match.sort_key(direction))
output = workspace / "result.csv"
write_search_results(output, matches)
final_id = uuid5(
UUID(ordered_partials[0].artifact.artifact_id),
"scimesh:similarity-search:final:" + ",".join(
partial.artifact.artifact_id for partial in ordered_partials
),
)
return FinalResult(
ArtifactReference(str(final_id), _sha256_file(output), _CSV_CONTENT_TYPE),
{"matches_emitted": len(matches), "partial_count": len(ordered_partials)},
)
def _resolve_query(
self, input_path: Path, parameters: Mapping[str, object]
) -> tuple[str, dict[str, str]]:
query_id = parameters.get("query_id")
if isinstance(query_id, str):
record = find_molecule_by_id(input_path, query_id)
return Chem.MolToSmiles(record.molecule, canonical=True), {"kind": "chembl_id", "value": query_id}
supplied = parameters["query_smiles"]
assert isinstance(supplied, str) # checked by validate_job
molecule = parse_smiles(supplied)
if molecule is None:
raise ValueError("query_smiles is invalid")
return Chem.MolToSmiles(molecule, canonical=True), {"kind": "smiles", "value": supplied}
def _resolved_parameters(
self, parameters: Mapping[str, object], query_smiles: str, query_source: Mapping[str, str]
) -> dict[str, object]:
resolved: dict[str, object] = {
"query_smiles": query_smiles,
"query_source": dict(query_source),
"top_k": self._positive_int(parameters.get("top_k", 20), "top_k"),
"threshold_direction": parameters.get("threshold_direction", "greater"),
"fingerprint": {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE},
}
if "threshold" in parameters:
resolved["threshold"] = self._unit_interval(parameters["threshold"], "threshold")
if "max_rows" in parameters:
resolved["max_rows"] = self._positive_int(parameters["max_rows"], "max_rows")
if "progress_every" in parameters:
resolved["progress_every"] = self._nonnegative_int(parameters["progress_every"], "progress_every")
return resolved
def _validate_resolved_parameters(self, parameters: Mapping[str, object]) -> dict[str, object]:
query_smiles = self._string(parameters.get("query_smiles"), "query_smiles")
if parse_smiles(query_smiles) is None:
raise ValueError("query_smiles is invalid")
resolved = self._resolved_parameters(
parameters,
Chem.MolToSmiles(parse_smiles(query_smiles), canonical=True),
{"kind": "resolved", "value": query_smiles},
)
# A reducer receives immutable plan metadata, whose query source and
# fixed fingerprint are observational context rather than worker input.
if "fingerprint" in parameters:
fingerprint = parameters["fingerprint"]
if fingerprint != {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE}:
raise ValueError("resolved fingerprint does not match SciMesh defaults")
return resolved
@staticmethod
def _task_parameters(resolved: Mapping[str, object]) -> dict[str, object]:
# max_rows is applied before sharding. Passing it to each task would
# silently scan N rows per shard instead of the requested global prefix.
return {
key: value for key, value in resolved.items()
if key in {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
}
def _write_shards(
self, input_path: Path, workspace: Path, shard_rows: int, max_rows: object
) -> list[Path]:
limit = int(max_rows) if isinstance(max_rows, int) else None
paths: list[Path] = []
current: Path | None = None
destination = None
rows_in_shard = 0
seen_rows = 0
try:
with input_path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source, delimiter="\t")
fieldnames = reader.fieldnames or []
if not _REQUIRED_COLUMNS.issubset(set(fieldnames)):
missing = sorted(_REQUIRED_COLUMNS - set(fieldnames))
raise ValueError(f"dataset is missing required columns: {', '.join(missing)}")
for row in reader:
if limit is not None and seen_rows >= limit:
break
if destination is None or rows_in_shard == shard_rows:
if destination is not None:
destination.close()
current = workspace / f"shard-{len(paths)}.tsv"
destination = current.open("w", encoding="utf-8", newline="")
writer = csv.DictWriter(destination, fieldnames=fieldnames, delimiter="\t", lineterminator="\n")
writer.writeheader()
paths.append(current)
rows_in_shard = 0
writer.writerow(row)
rows_in_shard += 1
seen_rows += 1
finally:
if destination is not None:
destination.close()
if not paths:
raise ValueError("dataset has no data rows")
return paths
@staticmethod
def _read_partial(path: Path, direction: object) -> Iterator[SimilarityMatch]:
if not path.is_file():
raise ValueError("materialized partial result is missing")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
previous_key: tuple[float, str, str] | None = None
with path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source)
if tuple(reader.fieldnames or ()) != _SEARCH_COLUMNS:
raise ValueError("partial result has an invalid CSV header")
for expected_rank, row in enumerate(reader, start=1):
if set(row) != set(_SEARCH_COLUMNS) or row["rank"] != str(expected_rank):
raise ValueError("partial result has an invalid rank")
try:
similarity = float(row["similarity"])
except (TypeError, ValueError) as error:
raise ValueError("partial result has an invalid similarity") from error
if not math.isfinite(similarity) or not 0 <= similarity <= 1:
raise ValueError("partial result has an invalid similarity")
match = SimilarityMatch(similarity, row["chembl_id"], row["canonical_smiles"])
key = match.sort_key(direction)
if previous_key is not None and key < previous_key:
raise ValueError("partial result is not sorted deterministically")
previous_key = key
yield match
@staticmethod
def _string(value: object, name: str) -> str:
if not isinstance(value, str) or not value.strip() or len(value) > 200:
raise ValueError(f"{name} must be a non-empty string")
return value
@staticmethod
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
@staticmethod
def _nonnegative_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
@staticmethod
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or not 0 <= value <= 1:
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
def run_similarity_search_shard(
input_path: Path, parameters: Mapping[str, object], output_path: Path
) -> dict[str, int]:
"""Run one planned shard using the local reference implementation.
This is the worker adapter used by CTX-08. It deliberately accepts only
resolved ``query_smiles``: resolving an identifier independently in each
shard would make the distributed search scientifically invalid.
"""
allowed = {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
query_smiles = parameters.get("query_smiles")
if not isinstance(query_smiles, str) or not query_smiles.strip():
raise ValueError("query_smiles is required for a distributed shard")
molecule = parse_smiles(query_smiles)
if molecule is None:
raise ValueError("query_smiles is invalid")
top_k = SimilaritySearchDistributedWorkload._positive_int(parameters.get("top_k", 20), "top_k")
threshold = None
if "threshold" in parameters:
threshold = SimilaritySearchDistributedWorkload._unit_interval(parameters["threshold"], "threshold")
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
progress_every = 0
if "progress_every" in parameters:
progress_every = SimilaritySearchDistributedWorkload._nonnegative_int(
parameters["progress_every"], "progress_every"
)
result = search_similar(
input_path,
MoleculeRecord("query", query_smiles, molecule),
top_k=top_k,
progress_every=progress_every,
threshold=threshold,
threshold_direction=direction,
)
write_similarity_search_partial(output_path, result.matches)
return {
"scanned_rows": result.stats.scanned,
"valid_molecules": result.stats.valid,
"invalid_smiles": result.stats.invalid,
"matches_emitted": len(result.matches),
}
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for block in iter(lambda: source.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
scimesh/distributed/workload.py
-40
View File
@@ -1,40 +0,0 @@
"""Protocol implemented by coordinator-independent distributed workloads."""
from __future__ import annotations
from pathlib import Path
from typing import Mapping, Protocol, Sequence
from .models import CompletedPartial, DistributedPlan, FinalResult
class DistributedWorkload(Protocol):
"""Validate, plan, and reduce one explicit scientific workload.
``input_path`` and ``workspace`` are bridge-provided temporary local paths.
They must never be included in returned plans or persisted task payloads.
"""
name: str
description: str
def validate_job(self, parameters: Mapping[str, object]) -> None:
"""Reject invalid public parameters before the bridge writes metadata."""
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
"""Build a JSON-safe plan containing only coordinator artifact references."""
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
"""Reduce coordinator-owned partial artifacts in ascending chunk order."""
scimesh/sdk/__init__.py
-12
View File
@@ -19,13 +19,6 @@ from .artifacts import (
PortSpec,
Provenance,
)
from .builtins import (
current_environment_digest,
current_scimesh_package_digest,
default_sdk_registry,
default_sdk_runtime,
similarity_search_sdk_adapter,
)
from .conformance import (
CancellationFlag,
LocalArtifactStore,
@@ -215,11 +208,6 @@ __all__ = [
"WorkloadManifest",
"WorkloadRegistry",
"assert_manifest_round_trip",
"current_environment_digest",
"current_scimesh_package_digest",
"default_sdk_registry",
"default_sdk_runtime",
"installed_distribution_digest",
"negotiate_manifest",
"similarity_search_sdk_adapter",
]
scimesh/sdk/builtins.py
-150
View File
@@ -1,150 +0,0 @@
"""SDK definitions for existing SciMesh workloads and local core runtime."""
from __future__ import annotations
import hashlib
import os
import platform
import sys
from rdkit import rdBase
from scimesh.distributed.similarity_search import (
SimilaritySearchDistributedWorkload,
run_similarity_search_shard,
)
from .artifacts import ArtifactSchema, PortSpec
from .compat import LegacyDistributedWorkloadAdapter
from .identity import ComponentRef, SDK_API_VERSION, SchemaRef
from .integrity import installed_distribution_digest
from .registry import WorkloadRegistry
from .resources import ResourceInventory
from .runtime import RuntimeCapabilities
def current_scimesh_package_digest() -> str:
"""Hash installed SciMesh Python sources for the built-in trusted adapter.
This is a local immutable-code pin, not a package signature or container
attestation. Consequently the built-in compatibility manifest is trusted
only; an administrator must supply signed image metadata before enabling an
untrusted quorum policy.
"""
# Source/editable installs are allowed only for this explicit local
# development helper. Registry discovery keeps the secure default.
return installed_distribution_digest("scimesh", allow_editable=True)
def current_environment_digest() -> str:
payload = "\n".join(
(
current_scimesh_package_digest(),
f"python={sys.implementation.name}-{platform.python_version()}",
f"rdkit={rdBase.rdkitVersion}",
f"platform={sys.platform}-{platform.machine().lower()}",
)
)
return "sha256:" + hashlib.sha256(payload.encode("utf-8")).hexdigest()
def similarity_search_sdk_adapter(
*, shard_rows: int = 10_000
) -> LegacyDistributedWorkloadAdapter:
dataset_schema = ArtifactSchema(
SchemaRef("molecule-table", 1),
"text/tab-separated-values",
"utf-8",
max_bytes=10 * 1024 * 1024 * 1024,
validator=ComponentRef("delimited-table", 1),
validator_configuration={
"required_columns": ["canonical_smiles", "chembl_id"],
},
max_records=100_000_000,
canonicalizer="scimesh-tsv-v1",
)
partial_schema = ArtifactSchema(
SchemaRef("similarity-search-partial", 1),
"text/csv",
"utf-8",
max_bytes=1024 * 1024 * 1024,
validator=ComponentRef("delimited-table", 1),
validator_configuration={
"columns": ["rank", "chembl_id", "canonical_smiles", "similarity"],
},
max_records=100_000,
canonicalizer="scimesh-search-partial-v1",
)
result_schema = ArtifactSchema(
SchemaRef("similarity-search-result", 1),
"text/csv",
"utf-8",
max_bytes=1024 * 1024 * 1024,
validator=ComponentRef("delimited-table", 1),
validator_configuration={
"columns": ["rank", "chembl_id", "canonical_smiles", "similarity"],
},
max_records=100_000,
canonicalizer="scimesh-search-result-v1",
)
parameters_schema = {
"type": "object",
"additionalProperties": False,
"properties": {
"query_id": {"type": "string", "minLength": 1, "maxLength": 200},
"query_smiles": {"type": "string", "minLength": 1, "maxLength": 200},
"top_k": {"type": "integer", "minimum": 1},
"threshold": {"type": "number", "minimum": 0, "maximum": 1},
"threshold_direction": {"enum": ["greater", "less"]},
"max_rows": {"type": "integer", "minimum": 1},
"progress_every": {"type": "integer", "minimum": 0},
},
"oneOf": [
{"required": ["query_id"], "not": {"required": ["query_smiles"]}},
{"required": ["query_smiles"], "not": {"required": ["query_id"]}},
],
}
return LegacyDistributedWorkloadAdapter(
SimilaritySearchDistributedWorkload(),
run_similarity_search_shard,
version="1.0.0",
package_digest=current_scimesh_package_digest(),
environment_digest=current_environment_digest(),
parameters_schema=parameters_schema,
input_port=PortSpec(dataset_schema),
partial_port=PortSpec(partial_schema),
output_port=PortSpec(result_schema),
resolved_parameter_names=("query_source", "fingerprint"),
shard_rows=shard_rows,
)
def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry:
registry = WorkloadRegistry()
registry.register(
similarity_search_sdk_adapter(shard_rows=shard_rows).definition(), enabled=True
)
return registry
def similarity_search_workload_definition():
"""Installed entry-point factory for the default shard-size definition."""
return similarity_search_sdk_adapter().definition()
def default_sdk_runtime() -> RuntimeCapabilities:
architecture = platform.machine().lower() or "unknown"
return RuntimeCapabilities(
sdk_api_version=SDK_API_VERSION,
protocol_version="1.0.0",
profiles=("core-batch-v1",),
features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"},
workload_capabilities=("similarity-search", "descriptor-batch"),
inventory=ResourceInventory(
cpu_cores=max(os.cpu_count() or 1, 1),
memory_mb=4096,
scratch_mb=4096,
architecture=architecture,
environment_digests=(current_environment_digest(),),
),
)
scimesh/sdk/compat/__init__.py
-5
View File
@@ -1,5 +0,0 @@
"""Adapters for versioned pre-SDK SciMesh workload contracts."""
from .distributed_v1 import LegacyDistributedWorkloadAdapter
__all__ = ["LegacyDistributedWorkloadAdapter"]
scimesh/sdk/compat/distributed_v1.py
-369
View File
@@ -1,369 +0,0 @@
"""Compatibility adapter for the CTX-07 ``DistributedWorkload`` protocol."""
from __future__ import annotations
import hashlib
import shutil
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence
from scimesh.distributed.models import (
ArtifactReference as LegacyArtifactReference,
CompletedPartial,
FinalResult,
)
from scimesh.distributed.workload import DistributedWorkload
from ..artifacts import (
ArtifactCollection,
ArtifactItem,
ArtifactRef,
Cardinality,
CollectionKind,
OutputManifest,
PortSpec,
)
from ..execution import CheckpointPolicy, ExecutionProfile, NetworkPolicy, RetryPolicy
from ..identity import ComponentRef, SchemaRef, VersionRange, WorkloadId
from ..manifest import (
DeterminismProfile,
EnvironmentSpec,
PackageSpec,
TrustMode,
VerifierSpec,
WorkloadLimits,
WorkloadManifest,
)
from ..plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan
from ..protocols import PlanningContext, ReduceContext, TaskContext
from ..registry import WorkloadDefinition
from ..resources import ResourceRequirements
from ..verification import ExactArtifactVerifier
from ..workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec
ShardRunner = Callable[[Path, Mapping[str, object], Path], Mapping[str, int | float]]
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for block in iter(lambda: source.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
class LegacyDistributedWorkloadAdapter:
"""Expose a legacy map/reduce workload through the SDK core-batch profile.
The adapter preserves the old wire schema. Local files are materialized and
sealed only through bridge-owned contexts, and no path is included in a
``TaskSpec`` or ``WorkflowPlan``.
"""
MAP_ENTRY_POINT = "scimesh.sdk.compat.distributed_v1:run_legacy@v1"
REDUCE_ENTRY_POINT = "scimesh.sdk.compat.distributed_v1:reduce_legacy@v1"
def __init__(
self,
workload: DistributedWorkload,
shard_runner: ShardRunner,
*,
version: str,
package_digest: str,
environment_digest: str,
parameters_schema: Mapping[str, Any],
input_port: PortSpec,
partial_port: PortSpec,
output_port: PortSpec,
resolved_parameter_names: Sequence[str] = (),
shard_rows: int = 10_000,
resources: ResourceRequirements | None = None,
execution: ExecutionProfile | None = None,
limits: WorkloadLimits | None = None,
) -> None:
if not isinstance(workload.name, str) or not isinstance(workload.description, str):
raise ValueError("legacy workload must expose name and description")
if not callable(shard_runner):
raise ValueError("shard_runner must be callable")
if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
raise ValueError("shard_rows must be a positive integer")
self.workload = workload
self.shard_runner = shard_runner
self.shard_rows = shard_rows
self.input_port = input_port
self.partial_port = partial_port
self.output_port = output_port
resources = resources or ResourceRequirements(
profile="legacy-cpu-v1",
cpu_cores=1,
memory_mb=1024,
scratch_mb=1024,
max_duration_seconds=3600,
)
execution = execution or ExecutionProfile(
profile="legacy-python-process-v1",
network=NetworkPolicy.TRUSTED,
timeout_seconds=3600,
checkpoint=CheckpointPolicy(),
)
limits = limits or WorkloadLimits(
max_input_bytes=input_port.schema.max_bytes,
max_tasks=10_000,
max_output_bytes=output_port.schema.max_bytes,
)
parameter_names = tuple(sorted(parameters_schema.get("properties", {})))
reduce_parameter_names = tuple(sorted(set(parameter_names).union(resolved_parameter_names)))
map_stage = StageSpec(
stage_id="map",
kind=StageKind.MAP,
entry_point=self.MAP_ENTRY_POINT,
needs=(),
inputs={"input": input_port},
outputs={"partial": partial_port},
parameter_names=parameter_names,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=("trusted",),
max_fan_out=limits.max_tasks,
cacheable=True,
)
reduce_input = PortSpec(
schema=partial_port.schema,
cardinality=Cardinality.MANY,
collection=CollectionKind.KEYED,
)
reduce_stage = StageSpec(
stage_id="reduce",
kind=StageKind.REDUCE,
entry_point=self.REDUCE_ENTRY_POINT,
needs=("map",),
inputs={"partials": reduce_input},
outputs={"result": output_port},
parameter_names=reduce_parameter_names,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=("trusted",),
cacheable=True,
)
workflow = WorkflowSpec(
workflow_id="map-reduce-v1",
inputs={"input": input_port},
stages=(map_stage, reduce_stage),
edges=(
ArtifactEdge(PortRef("input"), PortRef("input", "map")),
ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")),
),
outputs={"result": PortRef("result", "reduce")},
max_tasks=limits.max_tasks,
max_output_bytes=limits.max_output_bytes,
)
self.manifest = WorkloadManifest(
sdk_api=VersionRange(">=1.0,<2.0"),
protocol=VersionRange(">=1,<2"),
workload=WorkloadId(workload.name, version),
description=workload.description,
package=PackageSpec("scimesh", package_digest),
environment=EnvironmentSpec("python-process", environment_digest, {"adapter": "distributed-v1"}),
parameters_schema=parameters_schema,
workflow=workflow,
inputs={"input": input_port},
outputs={"result": output_port},
determinism=DeterminismProfile.BYTE_EXACT,
trust_modes=(TrustMode.TRUSTED,),
verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}),
limits=limits,
capabilities=(workload.name,),
conformance_profiles=("core-batch-v1",),
)
self._exact_verifier = ExactArtifactVerifier()
def definition(self) -> WorkloadDefinition:
return WorkloadDefinition(
manifest=self.manifest,
planner=self,
runners={self.MAP_ENTRY_POINT: self},
reducers={self.REDUCE_ENTRY_POINT: self},
verifiers={self._exact_verifier.identity.canonical: self._exact_verifier},
)
def validate(self, request: JobRequest) -> ValidatedJob:
if request.workload != self.manifest.workload:
raise ValueError("legacy adapter received a request for another workload")
self.workload.validate_job(request.parameters)
return ValidatedJob(request, request.parameters)
def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan:
if not isinstance(job, ValidatedJob):
raise ValueError("job must be a ValidatedJob")
collection = job.request.inputs.get("input")
if collection is None:
raise ValueError("legacy adapter requires the input port")
self.input_port.validate_collection(collection, "job input")
input_artifact = collection.items[0].artifact
input_path = context.catalog.materialize(input_artifact)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
legacy = self.workload.plan(
input_path,
input_artifact.artifact_id,
job.request.parameters,
self.shard_rows,
workspace,
)
if legacy.workload != self.workload.name:
raise ValueError("legacy planner returned a plan for another workload")
tasks: list[TaskSpec] = []
negotiated = context.negotiated
map_stage = self.manifest.workflow.stages[0]
used_paths: set[Path] = set()
for planned in legacy.tasks:
path = self._find_planned_file(workspace, planned.input_artifact.sha256, used_paths)
sealed = context.sink.seal(
path,
declaration=self.input_port.schema,
)
if sealed.sha256 != planned.input_artifact.sha256:
raise ValueError("artifact sink returned a checksum that differs from the legacy plan")
tasks.append(
TaskSpec(
workload=self.manifest.workload,
package_digest=self.manifest.package.digest,
manifest_digest=self.manifest.digest,
trust_mode=job.request.trust_mode,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=map_stage.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
task_key=f"map/{planned.chunk_index:08d}",
stage_id="map",
parameters=planned.parameters,
inputs={"input": ArtifactCollection.single(sealed)},
expected_outputs={"partial": self.partial_port},
resources=map_stage.resources,
execution=map_stage.execution,
)
)
return WorkflowPlan(
workload=self.manifest.workload,
package_digest=self.manifest.package.digest,
manifest_digest=self.manifest.digest,
trust_mode=job.request.trust_mode,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=self.manifest.verifier.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
workflow_id=self.manifest.workflow.workflow_id,
resolved_parameters=legacy.resolved_parameters,
tasks=tuple(tasks),
)
@staticmethod
def _find_planned_file(workspace: Path, expected_sha256: str, used: set[Path]) -> Path:
for candidate in sorted(workspace.rglob("*")):
if candidate in used or not candidate.is_file() or candidate.is_symlink():
continue
if _sha256_file(candidate) == expected_sha256:
used.add(candidate)
return candidate
raise ValueError("legacy planner did not materialize its planned artifact")
def run(self, context: TaskContext) -> OutputManifest:
context.cancellation.raise_if_cancelled()
collection = context.task.inputs.get("input")
if collection is None:
raise ValueError("legacy map task requires one input collection")
self.input_port.validate_collection(collection, "legacy map input")
source = context.catalog.materialize(collection.items[0].artifact)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
input_path = workspace / "input"
output_path = workspace / "result"
if source.resolve() != input_path.resolve():
shutil.copyfile(source, input_path)
metrics = self.shard_runner(input_path, context.task.parameters, output_path)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
output_path,
declaration=self.partial_port.schema,
)
return OutputManifest(
context.task.task_key,
{"partial": ArtifactCollection.single(sealed)},
metrics,
context.provenance,
).validate_against(context.task.expected_outputs, max_output_bytes=self.manifest.limits.max_output_bytes)
def reduce(self, context: ReduceContext) -> OutputManifest:
context.cancellation.raise_if_cancelled()
collection = context.accepted_inputs.get("partials")
if collection is None or collection.kind is not CollectionKind.KEYED or not collection.items:
raise ValueError("legacy reducer requires a non-empty keyed partial collection")
self.manifest.workflow.stages[1].inputs["partials"].validate_collection(
collection,
"legacy reducer partials",
)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
partials: list[CompletedPartial] = []
indexed_items: list[tuple[int, ArtifactItem]] = []
for item in collection.items:
key = item.key or ""
prefix = "map."
raw_index = key[len(prefix):] if key.startswith(prefix) else ""
if len(raw_index) != 8 or not raw_index.isdigit():
raise ValueError("legacy partial key must use map.<eight-digit-index>")
indexed_items.append((int(raw_index), item))
indices = [index for index, _ in indexed_items]
expected_keys = context.task.expected_input_keys.get("partials")
if expected_keys is None or {item.key for item in collection.items} != set(expected_keys):
raise ValueError("legacy partial keys do not match the coordinator expected set")
if sorted(indices) != list(range(len(indexed_items))):
raise ValueError("legacy partial keys must be complete and contiguous")
for index, item in sorted(indexed_items):
artifact = item.artifact
source = context.catalog.materialize(artifact)
target = workspace / artifact.artifact_id
if source.resolve() != target.resolve():
shutil.copyfile(source, target)
if _sha256_file(target) != artifact.sha256:
raise ValueError("materialized partial checksum does not match")
partials.append(
CompletedPartial(
index,
LegacyArtifactReference(
artifact.artifact_id,
artifact.sha256,
artifact.media_type,
),
{},
)
)
result = self.workload.reduce(partials, context.task.parameters, workspace)
if not isinstance(result, FinalResult):
raise ValueError("legacy reducer must return a FinalResult")
path = self._find_planned_file(workspace, result.artifact.sha256, set())
sealed = context.sink.seal(
path,
declaration=self.output_port.schema,
)
if sealed.sha256 != result.artifact.sha256:
raise ValueError("artifact sink returned a checksum that differs from the legacy result")
return OutputManifest(
context.task.task_key,
{"result": ArtifactCollection.single(sealed)},
result.metrics,
context.provenance,
).validate_against(context.task.expected_outputs, max_output_bytes=self.manifest.limits.max_output_bytes)
scimesh/worker/runners.py
+243 -86
View File
@@ -1,110 +1,267 @@
"""Local workload adapters. They receive no arbitrary commands from the network."""
"""SDK-based local workload execution for claimed coordinator tasks.
The runner is a v1-wire bridge: the coordinator still claims flat tasks and
the worker still uploads one partial CSV, but execution goes through the
SDK-built workload's own Runner handler with a real ``TaskSpec``,
provenance, resource reservation, and a content-addressed local store. No
legacy distributed-protocol code is involved.
"""
from __future__ import annotations
import hashlib
from datetime import datetime, timezone
from pathlib import Path
import subprocess
import sys
from typing import Protocol
from typing import Mapping, Protocol
from uuid import NAMESPACE_URL, uuid5
from scimesh.distributed.similarity_search import run_similarity_search_shard
from scimesh.sdk.artifacts import (
ArtifactCollection,
ArtifactRef,
OutputManifest,
Provenance,
)
from scimesh.sdk.conformance import (
CancellationFlag,
LocalArtifactStore,
LocalTaskContext,
ScopedArtifactSink,
)
from scimesh.sdk._validation import canonical_json
from scimesh.sdk.manifest import TrustMode
from scimesh.sdk.plans import TaskSpec
from scimesh.sdk.registry import WorkloadDefinition
from scimesh.sdk.resources import ResourceAllocation, ResourcePool
from scimesh.sdk.runtime import (
NegotiatedWorkload,
RuntimeCapabilities,
negotiate_manifest,
)
from scimesh.sdk.workflow import StageKind
from scimesh.workloads.library import default_sdk_runtime
from scimesh.workloads.search import similarity_search_sdk_definition
from .models import ClaimedTask, ProducedArtifact, RunResult
#: Parameters the worker may hand to a map task. ``max_rows`` is a plan-time
#: option applied before sharding and is intentionally rejected here.
_RUNNER_PARAMETERS = frozenset(
{"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
)
def _utc_now() -> str:
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
class Runner(Protocol):
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: ...
class SciMeshRunner:
"""Allowlisted adapter from coordinator workloads to the local SciMesh CLI."""
"""Execute claimed coordinator tasks through the SDK-built workloads."""
def __init__(
self,
definitions: Mapping[str, WorkloadDefinition] | None = None,
runtime: RuntimeCapabilities | None = None,
) -> None:
self._definitions = dict(definitions or {})
if "similarity-search" not in self._definitions:
self._definitions["similarity-search"] = (
similarity_search_sdk_definition().definition()
)
self._runtime = runtime or default_sdk_runtime()
self._pool = ResourcePool(self._runtime.inventory, max_concurrency=1)
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
# The subprocess changes cwd to task_dir. Absolute paths keep a caller
# supplied relative work directory from being resolved twice.
task_dir = task_dir.resolve()
input_path = task_dir / "input"
output_path = task_dir / "result.csv"
# The coordinator contract historically used underscores while the
# public SciMesh CLI uses hyphens. Accept both spellings at this narrow
# boundary so an API job cannot turn into an opaque worker failure.
workload = task.workload.replace("_", "-")
command = [sys.executable, "-m", "scimesh.cli", workload, str(input_path)]
params = task.parameters
if workload == "similarity-search":
self._reject_unknown(params, {"query_id", "query_smiles", "top_k", "threshold", "threshold_direction", "max_rows", "progress_every"})
query_id, query_smiles = params.get("query_id"), params.get("query_smiles")
if (query_id is None) == (query_smiles is None):
raise ValueError("exactly one of query_id or query_smiles is required")
if query_smiles is not None and "max_rows" not in params:
metrics = run_similarity_search_shard(input_path, params, output_path)
return RunResult((ProducedArtifact(output_path, "text/csv"),), metrics)
# Legacy URI jobs may still use query_id or an explicitly task-local
# max_rows value. CTX-08 plans never create those payloads; retain
# CLI execution only for backwards compatibility at this boundary.
top_k = self._positive_int(params, "top_k", default=20)
command += ["--query-id", self._string(params, "query_id")] if query_id is not None else ["--query-smiles", self._string(params, "query_smiles")]
command += ["--top-k", str(top_k)]
self._append_common_options(command, params)
elif workload == "similarity-graph":
self._reject_unknown(params, {"threshold", "threshold_direction", "block_size", "max_rows", "progress_every"})
threshold = self._number(params, "threshold")
command += ["--threshold", str(threshold)]
self._append_common_options(command, params, include_threshold=False)
if "block_size" in params:
command += ["--block-size", str(self._positive_int(params, "block_size", default=1_000))]
else:
definition = self._definitions.get(workload)
if definition is None:
raise ValueError(f"unsupported workload: {task.workload}")
command += ["--output", str(output_path)]
subprocess.run(command, check=True, cwd=task_dir) # explicit list: never shell=True
if not output_path.is_file():
raise RuntimeError("SciMesh CLI did not create its result")
processed_rows = max(sum(1 for _ in output_path.open(encoding="utf-8")) - 1, 0)
return RunResult((ProducedArtifact(output_path, "text/csv"),), {"processed_rows": processed_rows})
def _append_common_options(self, command: list[str], params: dict[str, object], *, include_threshold: bool = True) -> None:
if include_threshold and "threshold" in params:
command += ["--threshold", str(self._number(params, "threshold"))]
if "threshold_direction" in params:
direction = params["threshold_direction"]
if direction not in ("greater", "less"):
raise ValueError("threshold_direction must be 'greater' or 'less'")
command += ["--threshold-direction", str(direction)]
if "max_rows" in params:
command += ["--max-rows", str(self._positive_int(params, "max_rows", default=1))]
if "progress_every" in params:
command += ["--progress-every", str(self._nonnegative_int(params, "progress_every"))]
manifest = definition.manifest
negotiated = negotiate_manifest(manifest, self._runtime)
map_stage = next(
stage for stage in manifest.workflow.stages if stage.kind is StageKind.MAP
)
assert map_stage.verifier is not None
input_path = task_dir / "input"
if not input_path.is_file():
raise ValueError("claimed task input is missing")
parameters = self._resolve_parameters(task, input_path)
store = LocalArtifactStore(task_dir / "sdk-store")
input_ref = store.import_file(
input_path,
declaration=manifest.inputs["input"].schema,
)
spec = TaskSpec(
workload=manifest.workload,
package_digest=manifest.package.digest,
manifest_digest=manifest.digest,
trust_mode=TrustMode.TRUSTED,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=manifest.manifest_schema_version,
workflow_schema_version=manifest.workflow.schema_version,
environment_digest=manifest.environment.digest,
verifier=map_stage.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
task_key="map/00000000",
stage_id=map_stage.stage_id,
parameters=parameters,
inputs={"input": ArtifactCollection.single(input_ref)},
expected_outputs=map_stage.outputs,
resources=map_stage.resources,
execution=map_stage.execution,
).validate_stage(map_stage)
allocation = self._pool.reserve(task.task_id, spec.resources)
try:
provenance = self._provenance(
definition, negotiated, spec, task, allocation
)
context = LocalTaskContext(
spec,
store,
store,
task_dir,
CancellationFlag(),
provenance,
spec.inputs,
transaction=None,
)
output = definition.runners[map_stage.entry_point].run(context)
self._validate_output(
output,
spec,
context,
store,
provenance,
max_output_bytes=manifest.limits.max_output_bytes,
)
partial_ref = output.outputs["partial"].items[0].artifact
partial_path = store.materialize(partial_ref)
return RunResult(
(ProducedArtifact(partial_path, "text/csv"),),
dict(output.metrics),
)
finally:
self._pool.release(allocation.allocation_id)
@staticmethod
def _reject_unknown(params: dict[str, object], allowed: set[str]) -> None:
unknown = set(params) - allowed
def _resolve_parameters(task: ClaimedTask, input_path: Path) -> dict[str, object]:
"""Resolve ``query_id`` once per task and reject plan-time options."""
parameters = dict(task.parameters)
query_id = parameters.get("query_id")
query_smiles = parameters.get("query_smiles")
if isinstance(query_id, str) and not isinstance(query_smiles, str):
from scimesh.chemistry.dataset import find_molecule_by_id
from rdkit import Chem
record = find_molecule_by_id(input_path, query_id)
parameters["query_smiles"] = Chem.MolToSmiles(
record.molecule, canonical=True
)
del parameters["query_id"]
unknown = set(parameters) - _RUNNER_PARAMETERS
if unknown:
raise ValueError(f"unsupported parameters: {', '.join(sorted(unknown))}")
raise ValueError(
"unsupported runner parameters: " + ", ".join(sorted(unknown))
)
return parameters
def _provenance(
self,
definition: WorkloadDefinition,
negotiated: NegotiatedWorkload,
spec: TaskSpec,
task: ClaimedTask,
allocation: ResourceAllocation,
) -> Provenance:
manifest = definition.manifest
started_at = _utc_now()
return Provenance(
workload=manifest.workload,
sdk_api_version=spec.sdk_api_version,
protocol_version=spec.protocol_version,
manifest_schema_version=spec.manifest_schema_version,
workflow_schema_version=spec.workflow_schema_version,
verifier=spec.verifier,
artifact_schemas=tuple(
sorted(
{
item.artifact.schema
for collection in spec.inputs.values()
for item in collection.items
}.union(port.schema.ref for port in spec.expected_outputs.values()),
key=lambda value: value.canonical,
)
),
package_digest=spec.package_digest,
manifest_digest=spec.manifest_digest,
environment_digest=spec.environment_digest,
worker_runtime={"kind": "worker-agent-v1"},
allocated_resource_ids=(allocation.allocation_id,),
parameters_digest=hashlib.sha256(
canonical_json(spec.parameters).encode("utf-8")
).hexdigest(),
input_collection_digest=spec.inputs["input"].digest,
execution_contract_digest=spec.digest,
selected_features=spec.selected_features,
optional_fallbacks=spec.optional_fallbacks,
job_id=str(uuid5(NAMESPACE_URL, f"scimesh:job:{task.task_id}")),
task_id=str(uuid5(NAMESPACE_URL, f"scimesh:task:{task.task_id}")),
started_at=started_at,
finished_at=started_at,
trust_mode=TrustMode.TRUSTED.value,
)
@staticmethod
def _string(params: dict[str, object], name: str) -> str:
value = params.get(name)
if not isinstance(value, str) or not value.strip() or len(value) > 200:
raise ValueError(f"{name} must be a non-empty string")
return value
@staticmethod
def _positive_int(params: dict[str, object], name: str, default: int) -> int:
value = params.get(name, default)
if isinstance(value, bool) or not isinstance(value, int) or value < 1 or value > 100_000:
raise ValueError(f"{name} must be a positive integer")
return value
@staticmethod
def _nonnegative_int(params: dict[str, object], name: str) -> int:
value = params[name]
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
@staticmethod
def _number(params: dict[str, object], name: str) -> float:
value = params.get(name)
if isinstance(value, bool) or not isinstance(value, (int, float)) or not 0 <= value <= 1:
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
def _validate_output(
output: object,
spec: TaskSpec,
context: LocalTaskContext,
store: LocalArtifactStore,
provenance: Provenance,
*,
max_output_bytes: int,
) -> None:
if not isinstance(output, OutputManifest):
raise ValueError("SDK workload must return an OutputManifest")
if output.task_key != spec.task_key:
raise ValueError(
"SDK workload output task_key does not match its trusted task"
)
if output.provenance != provenance:
raise ValueError(
"SDK workload output provenance does not match its context"
)
output.validate_against(
spec.expected_outputs,
max_output_bytes=max_output_bytes,
)
if output.provenance != provenance:
raise ValueError(
"SDK workload output provenance does not match its context"
)
output.validate_against(
spec.expected_outputs,
max_output_bytes=spec.resources.max_duration_seconds, # replaced below
)
sink = context.sink
if not isinstance(sink, ScopedArtifactSink):
raise ValueError("SDK execution requires a scoped artifact sink")
declared = {
item.artifact.artifact_id: item.artifact
for collection in output.outputs.values()
for item in collection.items
}
issued = {artifact.artifact_id: artifact for artifact in sink.sealed_references}
if issued != declared:
raise ValueError(
"SDK workload outputs must declare exactly the artifacts sealed by its attempt"
)
for collection in output.outputs.values():
for item in collection.items:
store.require(item.artifact)
scimesh/sdk/descriptors/__init__.py → scimesh/workloads/descriptors/__init__.py
+1 -1
View File
@@ -1,4 +1,4 @@
"""SDK-native ``descriptor-batch`` reference workload.
"""SDK-built ``descriptor-batch`` reference workload.
See ``core.py`` for the pinned scientific contract and ``definition.py`` for
the manifest-backed planner/runner/reducer handlers.
scimesh/sdk/descriptors/core.py → scimesh/workloads/descriptors/core.py
View File
scimesh/sdk/descriptors/definition.py → scimesh/workloads/descriptors/definition.py
+15 -15
View File
@@ -1,4 +1,4 @@
"""SDK-native ``descriptor-batch`` workload definition and handlers.
"""SDK-built ``descriptor-batch`` workload definition and handlers.
This module is the first non-adapter reference workload built directly on the
``core-batch-v1`` profile: an explicit immutable manifest, a static map/reduce
@@ -15,7 +15,7 @@ import shutil
from pathlib import Path
from typing import Any, Mapping, Sequence
from ..artifacts import (
from ...sdk.artifacts import (
ArtifactCollection,
ArtifactItem,
ArtifactRef,
@@ -25,15 +25,15 @@ from ..artifacts import (
OutputManifest,
PortSpec,
)
from ..builtins import current_environment_digest, current_scimesh_package_digest
from ..execution import (
from ..environment import current_environment_digest, current_scimesh_package_digest
from ...sdk.execution import (
CheckpointPolicy,
ExecutionProfile,
NetworkPolicy,
RetryPolicy,
)
from ..identity import ComponentRef, SchemaRef, VersionRange, WorkloadId
from ..manifest import (
from ...sdk.identity import ComponentRef, SchemaRef, VersionRange, WorkloadId
from ...sdk.manifest import (
DeterminismProfile,
EnvironmentSpec,
PackageSpec,
@@ -42,12 +42,12 @@ from ..manifest import (
WorkloadLimits,
WorkloadManifest,
)
from ..plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan
from ..protocols import PlanningContext, ReduceContext, TaskContext
from ..registry import WorkloadDefinition
from ..resources import ResourceRequirements
from ..verification import ExactArtifactVerifier
from ..workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec
from ...sdk.plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan
from ...sdk.protocols import PlanningContext, ReduceContext, TaskContext
from ...sdk.registry import WorkloadDefinition
from ...sdk.resources import ResourceRequirements
from ...sdk.verification import ExactArtifactVerifier
from ...sdk.workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec
from .core import (
DESCRIPTOR_COLUMNS,
compute_descriptor_batch,
@@ -56,8 +56,8 @@ from .core import (
write_descriptor_shards,
)
MAP_ENTRY_POINT = "scimesh.sdk.descriptors.definition:map_descriptors@v1"
REDUCE_ENTRY_POINT = "scimesh.sdk.descriptors.definition:reduce_descriptors@v1"
MAP_ENTRY_POINT = "scimesh.workloads.descriptors.definition:map_descriptors@v1"
REDUCE_ENTRY_POINT = "scimesh.workloads.descriptors.definition:reduce_descriptors@v1"
_DESCRIPTOR_PARAMETERS = ("skip_invalid",)
@@ -118,7 +118,7 @@ class DescriptorBatchWorkload:
"""Manifest-backed planner, runner, and reducer for descriptor-batch.
The class follows the legacy adapter's structural pattern (one object
registered under each stage entry point) while remaining fully SDK-native:
registered under each stage entry point) while remaining fully SDK-built:
sharding is explicit and deterministic, every artifact is sealed through
the bridge-owned sink, and no filesystem path ever enters a plan or task.
"""
scimesh/workloads/environment.py
+42
View File
@@ -0,0 +1,42 @@
"""Local development digest helpers for built-in SDK workload definitions.
This module is a leaf: it must not import other SDK workload modules, so the
workload definition packages (``search``, ``graph``, ``descriptors``) and the
``builtins`` registry wiring can import it without creating cycles.
"""
from __future__ import annotations
import hashlib
import os
import platform
import sys
from rdkit import rdBase
from scimesh.sdk.integrity import installed_distribution_digest
def current_scimesh_package_digest() -> str:
"""Hash installed SciMesh Python sources for the built-in trusted adapters.
This is a local immutable-code pin, not a package signature or container
attestation. Consequently the built-in manifests are trusted only; an
administrator must supply signed image metadata before enabling an
untrusted quorum policy.
"""
# Source/editable installs are allowed only for this explicit local
# development helper. Registry discovery keeps the secure default.
return installed_distribution_digest("scimesh", allow_editable=True)
def current_environment_digest() -> str:
payload = "\n".join(
(
current_scimesh_package_digest(),
f"python={sys.implementation.name}-{platform.python_version()}",
f"rdkit={rdBase.rdkitVersion}",
f"platform={sys.platform}-{platform.machine().lower()}",
)
)
return "sha256:" + hashlib.sha256(payload.encode("utf-8")).hexdigest()
scimesh/workloads/graph/__init__.py
+40
View File
@@ -0,0 +1,40 @@
"""SDK-built ``similarity-graph`` workload.
A user workload script built on the SciMesh Workload SDK. See ``core.py`` for
the block/pair scientific core and ``definition.py`` for the manifest-backed
planner/runner/reducer handlers.
"""
from .core import (
block_pair_from_key,
check_pair_coverage,
compute_block_edges,
merge_edge_partials,
parse_molecule_blocks,
read_block_rows,
write_block_tsv,
write_edge_csv,
)
from .definition import (
MAP_ENTRY_POINT,
REDUCE_ENTRY_POINT,
SimilarityGraphSDKWorkload,
similarity_graph_sdk_definition,
workload_definition,
)
__all__ = [
"MAP_ENTRY_POINT",
"REDUCE_ENTRY_POINT",
"SimilarityGraphSDKWorkload",
"block_pair_from_key",
"check_pair_coverage",
"compute_block_edges",
"merge_edge_partials",
"parse_molecule_blocks",
"read_block_rows",
"similarity_graph_sdk_definition",
"workload_definition",
"write_block_tsv",
"write_edge_csv",
]
scimesh/workloads/graph/core.py
+249
View File
@@ -0,0 +1,249 @@
"""Scientific core for the SDK-built ``similarity-graph`` workload.
Molecules are parsed once into deterministic row-ordered blocks; every block
pair ``(i, j)`` with ``i <= j`` becomes one map task (diagonal tasks compare
pairs ``a < b`` inside a block, off-diagonal tasks compare every molecule
across two blocks). The reducer enforces the CTX-10 pair-coverage invariant:
the union of task pair sets must equal all unordered molecule pairs exactly
once, and the merged edge set must contain no duplicate unordered pair.
"""
from __future__ import annotations
import csv
from pathlib import Path
from typing import Iterable, Sequence
from rdkit import Chem, DataStructs
from scimesh.chemistry.dataset import iter_valid_molecules
from scimesh.chemistry.fingerprints import fingerprint
EDGE_COLUMNS = ("source_id", "target_id", "similarity")
MoleculeBlock = list[tuple[str, str]] # (chembl_id, smiles), row-ordered
def parse_molecule_blocks(
input_path: Path,
block_size: int,
max_rows: int | None = None,
) -> tuple[list[MoleculeBlock], dict[str, int]]:
"""Parse valid molecules into deterministic row-ordered blocks.
Mirrors the local reference's strictness: an empty or duplicate
``chembl_id`` fails the run, because the edge identity is the molecule id.
Invalid SMILES rows are skipped and counted.
"""
if (
isinstance(block_size, bool)
or not isinstance(block_size, int)
or block_size < 1
):
raise ValueError("block_size must be a positive integer")
from scimesh.chemistry.dataset import DatasetStats
stats = DatasetStats()
blocks: list[MoleculeBlock] = []
current: MoleculeBlock = []
seen_ids: set[str] = set()
for record in iter_valid_molecules(input_path, stats, max_rows=max_rows):
if not record.molecule_id:
raise ValueError("dataset contains an empty chembl_id")
if record.molecule_id in seen_ids:
raise ValueError(
f"dataset contains a duplicate chembl_id: {record.molecule_id}"
)
seen_ids.add(record.molecule_id)
current.append((record.molecule_id, record.smiles))
if len(current) == block_size:
blocks.append(current)
current = []
if current:
blocks.append(current)
if not blocks:
raise ValueError("dataset has no valid molecules")
return blocks, {
"rows_scanned": stats.scanned,
"valid_molecules": stats.valid,
"invalid_smiles": stats.invalid,
"block_count": len(blocks),
}
def write_block_tsv(rows: MoleculeBlock, path: Path) -> None:
"""Write one molecule block as a header TSV with the input column names."""
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8", newline="") as destination:
writer = csv.DictWriter(
destination,
fieldnames=["chembl_id", "canonical_smiles"],
delimiter="\t",
lineterminator="\n",
)
writer.writeheader()
for molecule_id, smiles in rows:
writer.writerow({"chembl_id": molecule_id, "canonical_smiles": smiles})
def read_block_rows(path: Path) -> MoleculeBlock:
rows: MoleculeBlock = []
with path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source, delimiter="\t")
for row in reader:
molecule_id = row.get("chembl_id", "")
smiles = row.get("canonical_smiles", "")
if not molecule_id or not smiles:
raise ValueError("block artifact contains an invalid row")
rows.append((molecule_id, smiles))
return rows
def compute_block_edges(
left: MoleculeBlock,
right: MoleculeBlock,
threshold: float,
threshold_direction: str,
) -> list[tuple[str, str, float]]:
"""Compare one planned block pair and emit only thresholded edges."""
if not 0.0 <= threshold <= 1.0:
raise ValueError("threshold must be between 0 and 1")
if threshold_direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
left_fingerprints = [
(molecule_id, fingerprint(Chem.MolFromSmiles(smiles)))
for molecule_id, smiles in left
]
right_fingerprints = [
(molecule_id, fingerprint(Chem.MolFromSmiles(smiles)))
for molecule_id, smiles in right
]
diagonal = left is right or left == right
edges: list[tuple[str, str, float]] = []
for left_index, (left_id, left_fp) in enumerate(left_fingerprints):
right_start = left_index + 1 if diagonal else 0
for right_index in range(right_start, len(right_fingerprints)):
right_id, right_fp = right_fingerprints[right_index]
similarity = DataStructs.TanimotoSimilarity(left_fp, right_fp)
matches_threshold = (
similarity >= threshold
if threshold_direction == "greater"
else similarity <= threshold
)
if matches_threshold:
edges.append((left_id, right_id, similarity))
return edges
def write_edge_csv(output_path: Path, edges: Iterable[tuple[str, str, float]]) -> None:
"""Write an edge table CSV with six-decimal similarity values.
Uses the CSV module's default ``\\r\\n`` line terminator so the bytes match
the local ``write_graph_edges`` reference exactly.
"""
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", encoding="utf-8", newline="") as destination:
writer = csv.DictWriter(destination, fieldnames=list(EDGE_COLUMNS))
writer.writeheader()
for source_id, target_id, similarity in edges:
writer.writerow(
{
"source_id": source_id,
"target_id": target_id,
"similarity": f"{similarity:.6f}",
}
)
def read_edge_csv(path: Path) -> list[tuple[str, str, float]]:
"""Read a materialized edge CSV with strict row validation."""
if not path.is_file():
raise ValueError("materialized edge partial is missing")
edges: list[tuple[str, str, float]] = []
with path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source)
if tuple(reader.fieldnames or ()) != EDGE_COLUMNS:
raise ValueError("edge partial has an invalid CSV header")
for row in reader:
if set(row) != set(EDGE_COLUMNS):
raise ValueError("edge partial has an invalid row")
source_id = row["source_id"]
target_id = row["target_id"]
if not source_id or not target_id:
raise ValueError("edge partial contains an empty molecule id")
try:
similarity = float(row["similarity"])
except (TypeError, ValueError) as error:
raise ValueError("edge partial has an invalid similarity") from error
if not 0 <= similarity <= 1:
raise ValueError("edge partial has an invalid similarity")
edges.append((source_id, target_id, similarity))
return edges
def block_pair_from_key(key: str) -> tuple[int, int]:
"""Parse ``map.<i>x<j>`` partial keys into block indices."""
prefix = "map."
if not key.startswith(prefix):
raise ValueError("graph partial key must use map.<left>x<right>")
raw = key[len(prefix) :]
left_raw, separator, right_raw = raw.partition("x")
if not separator or not left_raw.isdigit() or not right_raw.isdigit():
raise ValueError("graph partial key must use map.<left>x<right>")
return int(left_raw), int(right_raw)
def check_pair_coverage(pairs: Sequence[tuple[int, int]]) -> None:
"""Enforce the pair-coverage invariant: every block pair exactly once.
``pairs`` must contain every ``(i, j)`` with ``0 <= i <= j < n`` exactly
once, where ``n`` is derived from the largest referenced block index.
"""
if not pairs:
raise ValueError("graph partial keys cover no block pairs")
unique = set(pairs)
if len(unique) != len(pairs):
raise ValueError("graph partial keys contain a duplicate block pair")
if any(left > right or left < 0 or right < 0 for left, right in unique):
raise ValueError("graph partial keys reference an invalid block pair")
n = max(right for _, right in unique) + 1
expected = {(left, right) for left in range(n) for right in range(left, n)}
missing = sorted(expected - unique)
if missing:
raise ValueError(
"graph partial keys do not cover the full block pair set: "
+ ", ".join(f"{left}x{right}" for left, right in missing)
)
unexpected = sorted(unique - expected)
if unexpected:
raise ValueError(
"graph partial keys cover pairs outside the block pair set: "
+ ", ".join(f"{left}x{right}" for left, right in unexpected)
)
def merge_edge_partials(
partial_paths: Sequence[Path],
output_path: Path,
) -> dict[str, int]:
"""Merge edge partials with duplicate detection and deterministic sort.
The merged edge list is sorted by ``(source_id, target_id, -similarity)``,
matching the local brute-force reference exactly.
"""
if not partial_paths:
raise ValueError("graph reducer requires at least one edge partial")
edges: list[tuple[str, str, float]] = []
seen_pairs: set[tuple[str, str]] = set()
for path in partial_paths:
for source_id, target_id, similarity in read_edge_csv(path):
unordered = (min(source_id, target_id), max(source_id, target_id))
if unordered in seen_pairs:
raise ValueError(
f"graph partials contain a duplicate unordered pair: {unordered[0]}, {unordered[1]}"
)
seen_pairs.add(unordered)
edges.append((source_id, target_id, similarity))
edges.sort(key=lambda edge: (edge[0], edge[1], -edge[2]))
write_edge_csv(output_path, edges)
return {"partial_count": len(partial_paths), "edges_emitted": len(edges)}
scimesh/workloads/graph/definition.py
+500
View File
@@ -0,0 +1,500 @@
"""SDK-built ``similarity-graph`` workload definition and handlers.
Built directly on the ``core-batch-v1`` profile: molecules are parsed once
into deterministic row-ordered blocks, every block pair ``(i, j)`` with
``i <= j`` becomes one map task, and the reducer enforces the CTX-10
pair-coverage invariant (every unordered molecule pair compared exactly once)
before emitting a deterministically sorted edge list that is byte-identical
to the local brute-force reference.
"""
from __future__ import annotations
import hashlib
import shutil
from pathlib import Path
from typing import Any, Mapping
from ...sdk.artifacts import (
ArtifactCollection,
ArtifactItem,
ArtifactRef,
ArtifactSchema,
Cardinality,
CollectionKind,
OutputManifest,
PortSpec,
)
from ...sdk.execution import (
CheckpointPolicy,
ExecutionProfile,
NetworkPolicy,
RetryPolicy,
)
from ...sdk.identity import ComponentRef, SchemaRef, VersionRange, WorkloadId
from ...sdk.manifest import (
DeterminismProfile,
EnvironmentSpec,
PackageSpec,
TrustMode,
VerifierSpec,
WorkloadLimits,
WorkloadManifest,
)
from ...sdk.plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan
from ...sdk.protocols import PlanningContext, ReduceContext, TaskContext
from ...sdk.registry import WorkloadDefinition
from ...sdk.resources import ResourceRequirements
from ...sdk.verification import ExactArtifactVerifier
from ...sdk.workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec
from ..environment import current_environment_digest, current_scimesh_package_digest
from .core import (
block_pair_from_key,
check_pair_coverage,
compute_block_edges,
merge_edge_partials,
parse_molecule_blocks,
read_block_rows,
write_block_tsv,
write_edge_csv,
)
MAP_ENTRY_POINT = "scimesh.workloads.graph.definition:map_graph@v1"
REDUCE_ENTRY_POINT = "scimesh.workloads.graph.definition:reduce_graph@v1"
_MAP_PARAMETERS = ("left_block", "right_block", "threshold", "threshold_direction")
_REDUCE_PARAMETERS = ("threshold", "threshold_direction", "block_size", "max_rows")
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for block in iter(lambda: source.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def _parameters_schema() -> dict[str, Any]:
return {
"type": "object",
"additionalProperties": False,
"properties": {
"threshold": {"type": "number", "minimum": 0, "maximum": 1},
"threshold_direction": {"enum": ["greater", "less"]},
"block_size": {"type": "integer", "minimum": 1},
"max_rows": {"type": "integer", "minimum": 1},
},
"required": ["threshold"],
}
def _molecule_schema() -> ArtifactSchema:
return ArtifactSchema(
SchemaRef("molecule-table", 1),
"text/tab-separated-values",
"utf-8",
max_bytes=10 * 1024 * 1024 * 1024,
validator=ComponentRef("delimited-table", 1),
validator_configuration={
"required_columns": ["canonical_smiles", "chembl_id"],
},
max_records=100_000_000,
canonicalizer="scimesh-tsv-v1",
)
def _edge_schema() -> ArtifactSchema:
return ArtifactSchema(
SchemaRef("similarity-edge-table", 1),
"text/csv",
"utf-8",
max_bytes=100 * 1024 * 1024 * 1024,
validator=ComponentRef("delimited-table", 1),
validator_configuration={
"columns": ["source_id", "target_id", "similarity"],
},
max_records=1_000_000_000,
canonicalizer="similarity-edge-table-v1",
)
class SimilarityGraphSDKWorkload:
"""Manifest-backed planner, runner, and reducer for similarity-graph."""
def __init__(
self,
*,
package_digest: str,
environment_digest: str,
) -> None:
self.entry_point = MAP_ENTRY_POINT
self.input_port = PortSpec(_molecule_schema())
self.block_port = PortSpec(_molecule_schema())
self.partial_port = PortSpec(_edge_schema())
self.output_port = PortSpec(_edge_schema())
resources = ResourceRequirements(
profile="graph-cpu-v1",
cpu_cores=1,
memory_mb=1024,
scratch_mb=1024,
max_duration_seconds=3600,
)
execution = ExecutionProfile(
profile="graph-python-process-v1",
network=NetworkPolicy.TRUSTED,
timeout_seconds=3600,
checkpoint=CheckpointPolicy(),
)
limits = WorkloadLimits(
max_input_bytes=self.input_port.schema.max_bytes,
max_tasks=10_000,
max_output_bytes=self.output_port.schema.max_bytes,
)
trust_modes = ("trusted", "untrusted_quorum")
map_stage = StageSpec(
stage_id="map",
kind=StageKind.MAP,
entry_point=MAP_ENTRY_POINT,
needs=(),
inputs={"left": self.block_port, "right": self.block_port},
outputs={"partial": self.partial_port},
parameter_names=_MAP_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=limits.max_tasks,
cacheable=True,
)
reduce_input = PortSpec(
schema=self.partial_port.schema,
cardinality=Cardinality.MANY,
collection=CollectionKind.KEYED,
)
reduce_stage = StageSpec(
stage_id="reduce",
kind=StageKind.REDUCE,
entry_point=REDUCE_ENTRY_POINT,
needs=("map",),
inputs={"partials": reduce_input},
outputs={"result": self.output_port},
parameter_names=_REDUCE_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=1,
cacheable=True,
)
workflow = WorkflowSpec(
workflow_id="graph-block-pairs-v1",
inputs={"input": self.input_port},
stages=(map_stage, reduce_stage),
edges=(
ArtifactEdge(PortRef("input"), PortRef("left", "map")),
ArtifactEdge(PortRef("input"), PortRef("right", "map")),
ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")),
),
outputs={"result": PortRef("result", "reduce")},
max_tasks=limits.max_tasks,
max_output_bytes=limits.max_output_bytes,
)
self.manifest = WorkloadManifest(
sdk_api=VersionRange(">=1.0,<2.0"),
protocol=VersionRange(">=1,<2"),
workload=WorkloadId("similarity-graph", "1.0.0"),
description=(
"Exact sparse Tanimoto similarity graph over deterministic "
"block pairs with a duplicate-safe, coverage-checked merge."
),
package=PackageSpec("scimesh", package_digest),
environment=EnvironmentSpec(
"python-process",
environment_digest,
{"adapter": "sdk-native"},
),
parameters_schema=_parameters_schema(),
workflow=workflow,
inputs={"input": self.input_port},
outputs={"result": self.output_port},
determinism=DeterminismProfile.BYTE_EXACT,
trust_modes=(TrustMode.TRUSTED, TrustMode.UNTRUSTED_QUORUM),
verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}),
limits=limits,
capabilities=("similarity-graph",),
conformance_profiles=("core-batch-v1",),
)
self._exact_verifier = ExactArtifactVerifier()
def definition(self) -> WorkloadDefinition:
return WorkloadDefinition(
manifest=self.manifest,
planner=self,
runners={MAP_ENTRY_POINT: self},
reducers={REDUCE_ENTRY_POINT: self},
verifiers={self._exact_verifier.identity.canonical: self._exact_verifier},
)
@staticmethod
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
@staticmethod
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
def validate(self, request: JobRequest) -> ValidatedJob:
if request.workload != self.manifest.workload:
raise ValueError("similarity-graph received a request for another workload")
parameters = request.parameters
unknown = set(parameters) - {
"threshold",
"threshold_direction",
"block_size",
"max_rows",
}
if unknown:
raise ValueError(
"unsupported similarity-graph parameters: " + ", ".join(sorted(unknown))
)
threshold = parameters.get("threshold")
if threshold is None:
raise ValueError("threshold is required")
self._unit_interval(threshold, "threshold")
if "threshold_direction" in parameters and parameters[
"threshold_direction"
] not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
if "block_size" in parameters:
self._positive_int(parameters["block_size"], "block_size")
if "max_rows" in parameters:
self._positive_int(parameters["max_rows"], "max_rows")
return ValidatedJob(request, request.parameters)
def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan:
if not isinstance(job, ValidatedJob):
raise ValueError("job must be a ValidatedJob")
collection = job.request.inputs.get("input")
if collection is None:
raise ValueError("similarity-graph requires the input port")
self.input_port.validate_collection(collection, "job input")
input_artifact = collection.items[0].artifact
input_path = context.catalog.materialize(input_artifact)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
parameters = job.resolved_parameters
threshold = self._unit_interval(parameters.get("threshold"), "threshold")
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
block_size = int(parameters.get("block_size", 1_000))
max_rows = parameters.get("max_rows")
blocks, stats = parse_molecule_blocks(
input_path,
block_size,
int(max_rows) if isinstance(max_rows, int) else None,
)
task_parameters = {
"threshold": threshold,
"threshold_direction": direction,
}
negotiated = context.negotiated
map_stage = self.manifest.workflow.stages[0]
assert map_stage.verifier is not None
block_refs: list[ArtifactRef] = []
for index, block in enumerate(blocks):
path = workspace / f"block-{index:04d}.tsv"
write_block_tsv(block, path)
block_refs.append(
context.sink.seal(
path,
declaration=self.block_port.schema,
)
)
tasks: list[TaskSpec] = []
for left in range(len(blocks)):
for right in range(left, len(blocks)):
tasks.append(
TaskSpec(
workload=self.manifest.workload,
package_digest=self.manifest.package.digest,
manifest_digest=self.manifest.digest,
trust_mode=job.request.trust_mode,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=map_stage.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
task_key=f"map/{left:04d}x{right:04d}",
stage_id="map",
parameters={
**task_parameters,
"left_block": left,
"right_block": right,
},
inputs={
"left": ArtifactCollection.single(block_refs[left]),
"right": ArtifactCollection.single(block_refs[right]),
},
expected_outputs={"partial": self.partial_port},
resources=map_stage.resources,
execution=map_stage.execution,
)
)
return WorkflowPlan(
workload=self.manifest.workload,
package_digest=self.manifest.package.digest,
manifest_digest=self.manifest.digest,
trust_mode=job.request.trust_mode,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=self.manifest.verifier.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
workflow_id=self.manifest.workflow.workflow_id,
resolved_parameters=dict(parameters),
tasks=tuple(tasks),
)
def run(self, context: TaskContext) -> OutputManifest:
context.cancellation.raise_if_cancelled()
parameters = context.task.parameters
left_block = parameters.get("left_block")
right_block = parameters.get("right_block")
if (
isinstance(left_block, bool)
or not isinstance(left_block, int)
or isinstance(right_block, bool)
or not isinstance(right_block, int)
):
raise ValueError("graph map task requires block indices")
diagonal = left_block == right_block
left_collection = context.task.inputs.get("left")
right_collection = context.task.inputs.get("right")
if left_collection is None or right_collection is None:
raise ValueError("graph map task requires left and right block inputs")
self.block_port.validate_collection(left_collection, "graph map left input")
self.block_port.validate_collection(right_collection, "graph map right input")
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
left_path = context.catalog.materialize(left_collection.items[0].artifact)
right_path = context.catalog.materialize(right_collection.items[0].artifact)
left_rows = read_block_rows(left_path)
right_rows = (
left_rows
if diagonal and left_path.resolve() == right_path.resolve()
else read_block_rows(right_path)
)
threshold = self._unit_interval(parameters.get("threshold"), "threshold")
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
checked_pairs = (
len(left_rows) * (len(left_rows) - 1) // 2
if diagonal
else len(left_rows) * len(right_rows)
)
edges = compute_block_edges(
left_rows,
right_rows,
threshold,
direction,
)
output_path = workspace / "result.csv"
write_edge_csv(output_path, edges)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
output_path,
declaration=self.partial_port.schema,
)
return OutputManifest(
context.task.task_key,
{"partial": ArtifactCollection.single(sealed)},
{"checked_pairs": checked_pairs, "edges_emitted": len(edges)},
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def reduce(self, context: ReduceContext) -> OutputManifest:
context.cancellation.raise_if_cancelled()
collection = context.accepted_inputs.get("partials")
if (
collection is None
or collection.kind is not CollectionKind.KEYED
or not collection.items
):
raise ValueError(
"graph reducer requires a non-empty keyed partial collection"
)
self.manifest.workflow.stages[1].inputs["partials"].validate_collection(
collection,
"graph reducer partials",
)
pairs = [block_pair_from_key(item.key or "") for item in collection.items]
check_pair_coverage(pairs)
expected_keys = context.task.expected_input_keys.get("partials")
if expected_keys is None or {item.key for item in collection.items} != set(
expected_keys
):
raise ValueError(
"graph partial keys do not match the coordinator expected set"
)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
partial_paths: list[Path] = []
for item in sorted(collection.items, key=lambda value: value.key or ""):
artifact: ArtifactRef = item.artifact
source = context.catalog.materialize(artifact)
target = workspace / artifact.artifact_id
if source.resolve() != target.resolve():
shutil.copyfile(source, target)
if _sha256_file(target) != artifact.sha256:
raise ValueError("materialized partial checksum does not match")
partial_paths.append(target)
result_path = workspace / "result.csv"
metrics = merge_edge_partials(partial_paths, result_path)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
result_path,
declaration=self.output_port.schema,
)
return OutputManifest(
context.task.task_key,
{"result": ArtifactCollection.single(sealed)},
metrics,
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def similarity_graph_sdk_definition(
*,
package_digest: str | None = None,
environment_digest: str | None = None,
) -> SimilarityGraphSDKWorkload:
"""Build the SDK-based similarity-graph definition for tests."""
return SimilarityGraphSDKWorkload(
package_digest=package_digest or current_scimesh_package_digest(),
environment_digest=environment_digest or current_environment_digest(),
)
def workload_definition() -> WorkloadDefinition:
"""Installed entry-point factory for the SDK-based similarity-graph."""
return similarity_graph_sdk_definition().definition()
scimesh/workloads/library.py
+68
View File
@@ -0,0 +1,68 @@
"""Built-in workload library: default registry and runtime.
This module is workload code, not SDK framework code. It composes the
installed SciMesh workload packages (``search``, ``graph``, ``descriptors``)
with the SDK registry and runtime. External workload libraries can follow the
same pattern with their own packages and entry points.
"""
from __future__ import annotations
import os
import platform
from scimesh.sdk.identity import SDK_API_VERSION
from scimesh.sdk.registry import WorkloadRegistry
from scimesh.sdk.resources import ResourceInventory
from scimesh.sdk.runtime import RuntimeCapabilities
from .descriptors import descriptor_batch_sdk_definition
from .environment import current_environment_digest
from .graph import similarity_graph_sdk_definition
from .search import similarity_search_sdk_definition
__all__ = [
"default_sdk_registry",
"default_sdk_runtime",
]
def default_sdk_registry(*, shard_rows: int = 10_000) -> WorkloadRegistry:
"""Registry of every built-in SDK-built workload, all enabled."""
registry = WorkloadRegistry()
registry.register(
similarity_search_sdk_definition(shard_rows=shard_rows).definition(),
enabled=True,
)
registry.register(
similarity_graph_sdk_definition().definition(),
enabled=True,
)
registry.register(
descriptor_batch_sdk_definition(shard_rows=shard_rows).definition(),
enabled=True,
)
return registry
def default_sdk_runtime() -> RuntimeCapabilities:
"""Runtime advertising the built-in workloads' capabilities and inventory."""
architecture = platform.machine().lower() or "unknown"
return RuntimeCapabilities(
sdk_api_version=SDK_API_VERSION,
protocol_version="1.0.0",
profiles=("core-batch-v1",),
features={"artifact-collections": "1.0.0", "exact-verifier": "1.0.0"},
workload_capabilities=(
"similarity-search",
"similarity-graph",
"descriptor-batch",
),
inventory=ResourceInventory(
cpu_cores=max(os.cpu_count() or 1, 1),
memory_mb=4096,
scratch_mb=4096,
architecture=architecture,
environment_digests=(current_environment_digest(),),
),
)
scimesh/workloads/search/__init__.py
+31
View File
@@ -0,0 +1,31 @@
"""SDK-built ``similarity-search`` reference workload.
See ``core.py`` for the sharding/merge scientific core and ``definition.py``
for the manifest-backed planner/runner/reducer handlers.
"""
from .core import (
merge_search_partials,
run_search_shard,
write_search_partial,
write_search_shards,
)
from .definition import (
MAP_ENTRY_POINT,
REDUCE_ENTRY_POINT,
SimilaritySearchSDKWorkload,
similarity_search_sdk_definition,
workload_definition,
)
__all__ = [
"MAP_ENTRY_POINT",
"REDUCE_ENTRY_POINT",
"SimilaritySearchSDKWorkload",
"merge_search_partials",
"run_search_shard",
"similarity_search_sdk_definition",
"workload_definition",
"write_search_partial",
"write_search_shards",
]
scimesh/workloads/search/core.py
+261
View File
@@ -0,0 +1,261 @@
"""Scientific core for the SDK-built ``similarity-search`` workload.
Reuses the local reference implementation (``search_similar``) and the shared
CTX-08 partial format (full-precision ``repr`` scores) so that shard outputs
and the merged final CSV are byte-identical to the legacy distributed path and
to the single-process reference.
"""
from __future__ import annotations
import csv
import heapq
from pathlib import Path
from typing import Any, Iterator, Mapping, Sequence
from scimesh.chemistry.dataset import MoleculeRecord, parse_smiles
from scimesh.workloads.similarity_search import (
SimilarityMatch,
_HeapEntry,
search_similar,
write_search_results,
)
SEARCH_COLUMNS = ("rank", "chembl_id", "canonical_smiles", "similarity")
REQUIRED_COLUMNS = {"chembl_id", "canonical_smiles"}
def write_search_partial(
output_path: Path,
matches: Sequence[SimilarityMatch],
) -> None:
"""Write a worker partial with a round-trip score, not display rounding.
The public final CSV continues to use the local CLI's six-decimal display.
A reducer needs the full binary float representation to rank candidates
from separate shards exactly as the single-process reference does.
"""
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", encoding="utf-8", newline="") as destination:
writer = csv.DictWriter(destination, fieldnames=list(SEARCH_COLUMNS))
writer.writeheader()
for rank, match in enumerate(matches, start=1):
writer.writerow(
{
"rank": rank,
"chembl_id": match.molecule_id,
"canonical_smiles": match.smiles,
"similarity": repr(match.similarity),
}
)
def run_search_shard(
input_path: Path,
parameters: Mapping[str, object],
output_path: Path,
) -> dict[str, int]:
"""Run one planned shard with the local reference implementation.
This is the worker entry used by the SDK-built runner. It deliberately
accepts only resolved ``query_smiles``: resolving an identifier
independently in each shard would make the distributed search
scientifically invalid, so identifier resolution happens once in the
planner (or at the worker bridge for v1-wire tasks that still carry
``query_id``).
"""
allowed = {
"query_smiles",
"top_k",
"threshold",
"threshold_direction",
"progress_every",
}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(
f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}"
)
query_smiles = parameters.get("query_smiles")
if not isinstance(query_smiles, str) or not query_smiles.strip():
raise ValueError("query_smiles is required for a distributed shard")
molecule = parse_smiles(query_smiles)
if molecule is None:
raise ValueError("query_smiles is invalid")
top_k = _positive_int(parameters.get("top_k", 20), "top_k")
threshold = None
if "threshold" in parameters:
threshold = _unit_interval(parameters["threshold"], "threshold")
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
assert isinstance(direction, str)
progress_every = 0
if "progress_every" in parameters:
progress_every = _nonnegative_int(
parameters["progress_every"], "progress_every"
)
result = search_similar(
input_path,
MoleculeRecord("query", query_smiles, molecule),
top_k=top_k,
progress_every=progress_every,
threshold=threshold,
threshold_direction=direction,
)
write_search_partial(output_path, result.matches)
return {
"scanned_rows": result.stats.scanned,
"valid_molecules": result.stats.valid,
"invalid_smiles": result.stats.invalid,
"matches_emitted": len(result.matches),
}
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
def _nonnegative_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
def write_search_shards(
input_path: Path,
workspace: Path,
shard_rows: int,
max_rows: int | None = None,
) -> list[Path]:
"""Split the input TSV into deterministic row-bounded shards with headers."""
if (
isinstance(shard_rows, bool)
or not isinstance(shard_rows, int)
or shard_rows < 1
):
raise ValueError("shard_rows must be a positive integer")
if max_rows is not None and (
isinstance(max_rows, bool) or not isinstance(max_rows, int) or max_rows < 1
):
raise ValueError("max_rows must be a positive integer")
paths: list[Path] = []
current: Path | None = None
destination = None
writer = None
rows_in_shard = 0
seen_rows = 0
try:
with input_path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source, delimiter="\t")
fieldnames = tuple(reader.fieldnames or ())
if not REQUIRED_COLUMNS.issubset(set(fieldnames)):
missing = sorted(REQUIRED_COLUMNS - set(fieldnames))
raise ValueError(
f"dataset is missing required columns: {', '.join(missing)}"
)
for row in reader:
if max_rows is not None and seen_rows >= max_rows:
break
if destination is None or rows_in_shard == shard_rows:
if destination is not None:
destination.close()
current = workspace / f"shard-{len(paths)}.tsv"
destination = current.open("w", encoding="utf-8", newline="")
writer = csv.DictWriter(
destination,
fieldnames=list(fieldnames),
delimiter="\t",
lineterminator="\n",
)
writer.writeheader()
paths.append(current)
rows_in_shard = 0
assert writer is not None
writer.writerow(row)
rows_in_shard += 1
seen_rows += 1
finally:
if destination is not None:
destination.close()
if not paths:
raise ValueError("dataset has no data rows")
return paths
def iter_search_partial(
path: Path,
threshold_direction: str,
) -> Iterator[SimilarityMatch]:
"""Yield strictly ordered partial matches with full-precision scores."""
if not path.is_file():
raise ValueError("materialized partial result is missing")
if threshold_direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
previous_key: tuple[float, str, str] | None = None
with path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source)
if tuple(reader.fieldnames or ()) != SEARCH_COLUMNS:
raise ValueError("partial result has an invalid CSV header")
for expected_rank, row in enumerate(reader, start=1):
if set(row) != set(SEARCH_COLUMNS) or row["rank"] != str(expected_rank):
raise ValueError("partial result has an invalid rank")
try:
similarity = float(row["similarity"])
except (TypeError, ValueError) as error:
raise ValueError("partial result has an invalid similarity") from error
if not 0 <= similarity <= 1:
raise ValueError("partial result has an invalid similarity")
match = SimilarityMatch(
similarity, row["chembl_id"], row["canonical_smiles"]
)
key = match.sort_key(threshold_direction)
if previous_key is not None and key < previous_key:
raise ValueError("partial result is not sorted deterministically")
previous_key = key
yield match
def merge_search_partials(
partial_paths: Sequence[Path],
parameters: Mapping[str, Any],
output_path: Path,
) -> dict[str, int]:
"""Merge sorted shard partials into one deterministic final top-k CSV.
Mirrors the CTX-08/CTX-09 reducer: a bounded heap with the local
tie-breaker, so the merged file equals the single-process reference
byte-for-byte for the same input and options.
"""
if not partial_paths:
raise ValueError("at least one partial result is required")
raw_top_k = parameters.get("top_k", 20)
if isinstance(raw_top_k, bool) or not isinstance(raw_top_k, int) or raw_top_k < 1:
raise ValueError("top_k must be a positive integer")
top_k = raw_top_k
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
heap: list[_HeapEntry] = []
for path in partial_paths:
for match in iter_search_partial(path, direction):
rank_key = match.sort_key(direction)
entry = _HeapEntry(match, rank_key)
if len(heap) < top_k:
heapq.heappush(heap, entry)
elif rank_key < heap[0].rank_key:
heapq.heapreplace(heap, entry)
matches = sorted(
(entry.match for entry in heap),
key=lambda match: match.sort_key(direction),
)
write_search_results(output_path, matches)
return {"matches_emitted": len(matches), "partial_count": len(partial_paths)}
scimesh/workloads/search/definition.py
+565
View File
@@ -0,0 +1,565 @@
"""SDK-built ``similarity-search`` workload definition and handlers.
A direct ``core-batch-v1`` definition (not the legacy adapter): the planner
resolves the query once, shards deterministically, each map task computes the
local top-k with the reference implementation, and the reducer merges the
sorted partials with the same bounded heap and tie-breakers as the local CLI.
The manifest declares ``byte_exact`` with the exact-artifact verifier and both
``trusted`` and ``untrusted_quorum`` trust modes.
"""
from __future__ import annotations
import hashlib
import shutil
from pathlib import Path
from typing import Any, Mapping
from rdkit import Chem
from scimesh.chemistry.dataset import find_molecule_by_id, parse_smiles
from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE
from ...sdk.artifacts import (
ArtifactCollection,
ArtifactItem,
ArtifactRef,
ArtifactSchema,
Cardinality,
CollectionKind,
OutputManifest,
PortSpec,
)
from ..environment import current_environment_digest, current_scimesh_package_digest
from ...sdk.execution import (
CheckpointPolicy,
ExecutionProfile,
NetworkPolicy,
RetryPolicy,
)
from ...sdk.identity import ComponentRef, SchemaRef, VersionRange, WorkloadId
from ...sdk.manifest import (
DeterminismProfile,
EnvironmentSpec,
PackageSpec,
TrustMode,
VerifierSpec,
WorkloadLimits,
WorkloadManifest,
)
from ...sdk.plans import JobRequest, TaskSpec, ValidatedJob, WorkflowPlan
from ...sdk.protocols import PlanningContext, ReduceContext, TaskContext
from ...sdk.registry import WorkloadDefinition
from ...sdk.resources import ResourceRequirements
from ...sdk.verification import ExactArtifactVerifier
from ...sdk.workflow import ArtifactEdge, PortRef, StageKind, StageSpec, WorkflowSpec
from .core import merge_search_partials, run_search_shard, write_search_shards
MAP_ENTRY_POINT = "scimesh.workloads.search.definition:map_search@v1"
REDUCE_ENTRY_POINT = "scimesh.workloads.search.definition:reduce_search@v1"
_MAP_PARAMETERS = (
"query_smiles",
"top_k",
"threshold",
"threshold_direction",
"progress_every",
)
_REDUCE_PARAMETERS = _MAP_PARAMETERS + ("query_source", "fingerprint")
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for block in iter(lambda: source.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def _parameters_schema() -> dict[str, Any]:
return {
"type": "object",
"additionalProperties": False,
"properties": {
"query_id": {"type": "string", "minLength": 1, "maxLength": 200},
"query_smiles": {"type": "string", "minLength": 1, "maxLength": 200},
"top_k": {"type": "integer", "minimum": 1},
"threshold": {"type": "number", "minimum": 0, "maximum": 1},
"threshold_direction": {"enum": ["greater", "less"]},
"max_rows": {"type": "integer", "minimum": 1},
"progress_every": {"type": "integer", "minimum": 0},
},
"oneOf": [
{"required": ["query_id"], "not": {"required": ["query_smiles"]}},
{"required": ["query_smiles"], "not": {"required": ["query_id"]}},
],
}
def _dataset_schema() -> ArtifactSchema:
return ArtifactSchema(
SchemaRef("molecule-table", 1),
"text/tab-separated-values",
"utf-8",
max_bytes=10 * 1024 * 1024 * 1024,
validator=ComponentRef("delimited-table", 1),
validator_configuration={
"required_columns": ["canonical_smiles", "chembl_id"],
},
max_records=100_000_000,
canonicalizer="scimesh-tsv-v1",
)
def _search_table_schema(ref: SchemaRef, canonicalizer: str) -> ArtifactSchema:
return ArtifactSchema(
ref,
"text/csv",
"utf-8",
max_bytes=1024 * 1024 * 1024,
validator=ComponentRef("delimited-table", 1),
validator_configuration={
"columns": ["rank", "chembl_id", "canonical_smiles", "similarity"],
},
max_records=100_000,
canonicalizer=canonicalizer,
)
class SimilaritySearchSDKWorkload:
"""Manifest-backed planner, runner, and reducer for similarity-search."""
def __init__(
self,
*,
shard_rows: int,
package_digest: str,
environment_digest: str,
) -> None:
if (
isinstance(shard_rows, bool)
or not isinstance(shard_rows, int)
or shard_rows < 1
):
raise ValueError("shard_rows must be a positive integer")
self.entry_point = MAP_ENTRY_POINT
self.shard_rows = shard_rows
self.input_port = PortSpec(_dataset_schema())
self.partial_port = PortSpec(
_search_table_schema(
SchemaRef("similarity-search-partial", 1), "scimesh-search-partial-v1"
)
)
self.output_port = PortSpec(
_search_table_schema(
SchemaRef("similarity-search-result", 1), "scimesh-search-result-v1"
)
)
resources = ResourceRequirements(
profile="search-cpu-v1",
cpu_cores=1,
memory_mb=1024,
scratch_mb=1024,
max_duration_seconds=3600,
)
execution = ExecutionProfile(
profile="search-python-process-v1",
network=NetworkPolicy.TRUSTED,
timeout_seconds=3600,
checkpoint=CheckpointPolicy(),
)
limits = WorkloadLimits(
max_input_bytes=self.input_port.schema.max_bytes,
max_tasks=10_000,
max_output_bytes=self.output_port.schema.max_bytes,
)
trust_modes = ("trusted", "untrusted_quorum")
map_stage = StageSpec(
stage_id="map",
kind=StageKind.MAP,
entry_point=MAP_ENTRY_POINT,
needs=(),
inputs={"input": self.input_port},
outputs={"partial": self.partial_port},
parameter_names=_MAP_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=limits.max_tasks,
cacheable=True,
)
reduce_input = PortSpec(
schema=self.partial_port.schema,
cardinality=Cardinality.MANY,
collection=CollectionKind.KEYED,
)
reduce_stage = StageSpec(
stage_id="reduce",
kind=StageKind.REDUCE,
entry_point=REDUCE_ENTRY_POINT,
needs=("map",),
inputs={"partials": reduce_input},
outputs={"result": self.output_port},
parameter_names=_REDUCE_PARAMETERS,
resources=resources,
execution=execution,
retry=RetryPolicy(),
verifier=ComponentRef("exact-artifact", 1),
trust_modes=trust_modes,
max_fan_out=1,
cacheable=True,
)
workflow = WorkflowSpec(
workflow_id="search-map-reduce-v1",
inputs={"input": self.input_port},
stages=(map_stage, reduce_stage),
edges=(
ArtifactEdge(PortRef("input"), PortRef("input", "map")),
ArtifactEdge(PortRef("partial", "map"), PortRef("partials", "reduce")),
),
outputs={"result": PortRef("result", "reduce")},
max_tasks=limits.max_tasks,
max_output_bytes=limits.max_output_bytes,
)
self.manifest = WorkloadManifest(
sdk_api=VersionRange(">=1.0,<2.0"),
protocol=VersionRange(">=1,<2"),
workload=WorkloadId("similarity-search", "1.0.0"),
description=(
"Exact top-k Tanimoto molecular similarity search over "
"deterministic TSV shards with a bounded merge."
),
package=PackageSpec("scimesh", package_digest),
environment=EnvironmentSpec(
"python-process",
environment_digest,
{"adapter": "sdk-native"},
),
parameters_schema=_parameters_schema(),
workflow=workflow,
inputs={"input": self.input_port},
outputs={"result": self.output_port},
determinism=DeterminismProfile.BYTE_EXACT,
trust_modes=(TrustMode.TRUSTED, TrustMode.UNTRUSTED_QUORUM),
verifier=VerifierSpec(ComponentRef("exact-artifact", 1), {}),
limits=limits,
capabilities=("similarity-search",),
conformance_profiles=("core-batch-v1",),
)
self._exact_verifier = ExactArtifactVerifier()
def definition(self) -> WorkloadDefinition:
return WorkloadDefinition(
manifest=self.manifest,
planner=self,
runners={MAP_ENTRY_POINT: self},
reducers={REDUCE_ENTRY_POINT: self},
verifiers={self._exact_verifier.identity.canonical: self._exact_verifier},
)
@staticmethod
def _string(value: object, name: str) -> str:
if not isinstance(value, str) or not value.strip() or len(value) > 200:
raise ValueError(f"{name} must be a non-empty string")
return value
@staticmethod
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
@staticmethod
def _nonnegative_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
@staticmethod
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
def validate(self, request: JobRequest) -> ValidatedJob:
if request.workload != self.manifest.workload:
raise ValueError(
"similarity-search received a request for another workload"
)
parameters = request.parameters
unknown = set(parameters) - {
"query_id",
"query_smiles",
"top_k",
"threshold",
"threshold_direction",
"max_rows",
"progress_every",
}
if unknown:
raise ValueError(
"unsupported similarity-search parameters: "
+ ", ".join(sorted(unknown))
)
query_id = parameters.get("query_id")
query_smiles = parameters.get("query_smiles")
if (query_id is None) == (query_smiles is None):
raise ValueError("exactly one of query_id or query_smiles is required")
if query_id is not None:
self._string(query_id, "query_id")
if query_smiles is not None:
self._string(query_smiles, "query_smiles")
self._positive_int(parameters.get("top_k", 20), "top_k")
if "max_rows" in parameters:
self._positive_int(parameters["max_rows"], "max_rows")
if "progress_every" in parameters:
self._nonnegative_int(parameters["progress_every"], "progress_every")
if "threshold" in parameters:
self._unit_interval(parameters["threshold"], "threshold")
if "threshold_direction" in parameters and parameters[
"threshold_direction"
] not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
return ValidatedJob(request, self._resolved_parameters(request))
def _resolved_parameters(self, request: JobRequest) -> dict[str, object]:
parameters = request.parameters
query_id = parameters.get("query_id")
if isinstance(query_id, str):
query_source: dict[str, str] = {"kind": "chembl_id", "value": query_id}
else:
query_source = {
"kind": "smiles",
"value": self._string(parameters.get("query_smiles"), "query_smiles"),
}
resolved: dict[str, object] = {
"query_source": query_source,
"top_k": self._positive_int(parameters.get("top_k", 20), "top_k"),
"threshold_direction": parameters.get("threshold_direction", "greater"),
"fingerprint": {
"algorithm": "morgan",
"radius": FP_RADIUS,
"fp_size": FP_SIZE,
},
}
if "threshold" in parameters:
resolved["threshold"] = self._unit_interval(
parameters["threshold"], "threshold"
)
if "max_rows" in parameters:
resolved["max_rows"] = self._positive_int(
parameters["max_rows"], "max_rows"
)
if "progress_every" in parameters:
resolved["progress_every"] = self._nonnegative_int(
parameters["progress_every"], "progress_every"
)
return resolved
def plan(self, job: ValidatedJob, context: PlanningContext) -> WorkflowPlan:
if not isinstance(job, ValidatedJob):
raise ValueError("job must be a ValidatedJob")
collection = job.request.inputs.get("input")
if collection is None:
raise ValueError("similarity-search requires the input port")
self.input_port.validate_collection(collection, "job input")
input_artifact = collection.items[0].artifact
input_path = context.catalog.materialize(input_artifact)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
resolved = dict(job.resolved_parameters)
query_smiles = self._resolve_query(input_path, job.request.parameters)
resolved["query_smiles"] = query_smiles
max_rows = resolved.get("max_rows")
shard_paths = write_search_shards(
input_path,
workspace,
self.shard_rows,
int(max_rows) if isinstance(max_rows, int) else None,
)
task_parameters = {
key: value for key, value in resolved.items() if key in set(_MAP_PARAMETERS)
}
negotiated = context.negotiated
map_stage = self.manifest.workflow.stages[0]
assert map_stage.verifier is not None
tasks: list[TaskSpec] = []
for index, path in enumerate(shard_paths):
sealed = context.sink.seal(
path,
declaration=self.input_port.schema,
)
tasks.append(
TaskSpec(
workload=self.manifest.workload,
package_digest=self.manifest.package.digest,
manifest_digest=self.manifest.digest,
trust_mode=job.request.trust_mode,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=map_stage.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
task_key=f"map/{index:08d}",
stage_id="map",
parameters=task_parameters,
inputs={"input": ArtifactCollection.single(sealed)},
expected_outputs={"partial": self.partial_port},
resources=map_stage.resources,
execution=map_stage.execution,
)
)
return WorkflowPlan(
workload=self.manifest.workload,
package_digest=self.manifest.package.digest,
manifest_digest=self.manifest.digest,
trust_mode=job.request.trust_mode,
sdk_api_version=negotiated.sdk_api_version,
protocol_version=negotiated.protocol_version,
manifest_schema_version=self.manifest.manifest_schema_version,
workflow_schema_version=self.manifest.workflow.schema_version,
environment_digest=self.manifest.environment.digest,
verifier=self.manifest.verifier.verifier,
selected_features=negotiated.selected_features,
optional_fallbacks=negotiated.optional_fallbacks,
workflow_id=self.manifest.workflow.workflow_id,
resolved_parameters=resolved,
tasks=tuple(tasks),
)
@staticmethod
def _resolve_query(input_path: Path, parameters: Mapping[str, object]) -> str:
query_id = parameters.get("query_id")
if isinstance(query_id, str):
record = find_molecule_by_id(input_path, query_id)
return Chem.MolToSmiles(record.molecule, canonical=True)
supplied = parameters["query_smiles"]
assert isinstance(supplied, str)
molecule = parse_smiles(supplied)
if molecule is None:
raise ValueError("query_smiles is invalid")
return Chem.MolToSmiles(molecule, canonical=True)
def run(self, context: TaskContext) -> OutputManifest:
context.cancellation.raise_if_cancelled()
collection = context.task.inputs.get("input")
if collection is None:
raise ValueError("search map task requires one input collection")
self.input_port.validate_collection(collection, "search map input")
source = context.catalog.materialize(collection.items[0].artifact)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
input_path = workspace / "input"
output_path = workspace / "result.csv"
if source.resolve() != input_path.resolve():
shutil.copyfile(source, input_path)
metrics = run_search_shard(
input_path,
context.task.parameters,
output_path,
)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
output_path,
declaration=self.partial_port.schema,
)
return OutputManifest(
context.task.task_key,
{"partial": ArtifactCollection.single(sealed)},
metrics,
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def reduce(self, context: ReduceContext) -> OutputManifest:
context.cancellation.raise_if_cancelled()
collection = context.accepted_inputs.get("partials")
if (
collection is None
or collection.kind is not CollectionKind.KEYED
or not collection.items
):
raise ValueError(
"search reducer requires a non-empty keyed partial collection"
)
self.manifest.workflow.stages[1].inputs["partials"].validate_collection(
collection,
"search reducer partials",
)
workspace = context.workspace
workspace.mkdir(parents=True, exist_ok=True)
indexed_items: list[tuple[int, ArtifactItem]] = []
for item in collection.items:
key = item.key or ""
prefix = "map."
raw_index = key[len(prefix) :] if key.startswith(prefix) else ""
if len(raw_index) != 8 or not raw_index.isdigit():
raise ValueError("search partial key must use map.<eight-digit-index>")
indexed_items.append((int(raw_index), item))
expected_keys = context.task.expected_input_keys.get("partials")
if expected_keys is None or {item.key for item in collection.items} != set(
expected_keys
):
raise ValueError(
"search partial keys do not match the coordinator expected set"
)
if sorted(index for index, _ in indexed_items) != list(
range(len(indexed_items))
):
raise ValueError("search partial keys must be complete and contiguous")
partial_paths: list[Path] = []
for index, item in sorted(indexed_items):
artifact: ArtifactRef = item.artifact
source = context.catalog.materialize(artifact)
target = workspace / artifact.artifact_id
if source.resolve() != target.resolve():
shutil.copyfile(source, target)
if _sha256_file(target) != artifact.sha256:
raise ValueError("materialized partial checksum does not match")
partial_paths.append(target)
result_path = workspace / "result.csv"
metrics = merge_search_partials(
partial_paths,
context.task.parameters,
result_path,
)
context.cancellation.raise_if_cancelled()
sealed = context.sink.seal(
result_path,
declaration=self.output_port.schema,
)
return OutputManifest(
context.task.task_key,
{"result": ArtifactCollection.single(sealed)},
metrics,
context.provenance,
).validate_against(
context.task.expected_outputs,
max_output_bytes=self.manifest.limits.max_output_bytes,
)
def similarity_search_sdk_definition(
*,
shard_rows: int = 10_000,
package_digest: str | None = None,
environment_digest: str | None = None,
) -> SimilaritySearchSDKWorkload:
"""Build the SDK-built similarity-search definition for tests."""
return SimilaritySearchSDKWorkload(
shard_rows=shard_rows,
package_digest=package_digest or current_scimesh_package_digest(),
environment_digest=environment_digest or current_environment_digest(),
)
def workload_definition() -> WorkloadDefinition:
"""Installed entry-point factory for the SDK-built similarity-search."""
return similarity_search_sdk_definition().definition()