324 lines
12 KiB
Python
324 lines
12 KiB
Python
"""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.
|
|
|
|
The runner is workload-generic: it loads definitions by name (from an
|
|
explicit mapping, built-in defaults, or installed-package discovery through
|
|
an administrator allowlist) and executes any workload whose map stage has a
|
|
single ``input`` port and a single ``partial`` output. Anything else fails
|
|
closed with a clear message, so adding a workload never requires touching
|
|
worker code.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import platform
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Mapping, Protocol
|
|
from uuid import NAMESPACE_URL, uuid5
|
|
|
|
from scimesh.sdk.artifacts import (
|
|
ArtifactCollection,
|
|
OutputManifest,
|
|
Provenance,
|
|
)
|
|
from scimesh.sdk.conformance import (
|
|
CancellationFlag,
|
|
LocalArtifactStore,
|
|
LocalTaskContext,
|
|
ScopedArtifactSink,
|
|
)
|
|
from scimesh.sdk._validation import canonical_json
|
|
from scimesh.sdk.identity import SDK_API_VERSION
|
|
from scimesh.sdk.manifest import TrustMode
|
|
from scimesh.sdk.plans import TaskSpec
|
|
from scimesh.sdk.registry import WorkloadDefinition, WorkloadRegistry
|
|
from scimesh.sdk.resources import ResourceAllocation, ResourceInventory, ResourcePool
|
|
from scimesh.sdk.runtime import (
|
|
NegotiatedWorkload,
|
|
RuntimeCapabilities,
|
|
negotiate_manifest,
|
|
)
|
|
from scimesh.sdk.workflow import StageKind
|
|
|
|
from .models import ClaimedTask, ProducedArtifact, RunResult
|
|
|
|
|
|
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: ...
|
|
|
|
|
|
def _definitions_from_environment() -> dict[str, WorkloadDefinition]:
|
|
"""Load definitions from ``SCIMESH_WORKLOAD_ALLOWLIST`` or the built-ins.
|
|
|
|
The environment-driven discovery mirrors the former worker configuration;
|
|
the Go agent passes the allowlist through unchanged. Without an allowlist
|
|
the worker serves every enabled built-in SDK workload of the installed
|
|
package — the same library the coordinator embeds as its catalog.
|
|
"""
|
|
import os
|
|
|
|
from scimesh.sdk.registry import workload_allowlist_from_json
|
|
|
|
allowlist = workload_allowlist_from_json(os.getenv("SCIMESH_WORKLOAD_ALLOWLIST"))
|
|
if allowlist:
|
|
registry = WorkloadRegistry()
|
|
registry.discover_installed(allowlist)
|
|
definitions: dict[str, WorkloadDefinition] = {}
|
|
for description in registry.descriptions():
|
|
definition, _ = registry.require(
|
|
description.workload.name,
|
|
description.workload.version,
|
|
description.package_digest,
|
|
)
|
|
definitions[description.workload.name] = definition
|
|
if not definitions:
|
|
raise ValueError("workload_allowlist discovered no workloads")
|
|
return definitions
|
|
from scimesh.workloads.library import default_sdk_registry
|
|
|
|
registry = default_sdk_registry()
|
|
return {
|
|
description.workload.name: registry.require(
|
|
description.workload.name,
|
|
description.workload.version,
|
|
description.package_digest,
|
|
)[0]
|
|
for description in registry.descriptions()
|
|
if description.enabled
|
|
}
|
|
|
|
|
|
def _inventory_for(
|
|
definitions: Mapping[str, WorkloadDefinition],
|
|
*,
|
|
cpu_cores: int,
|
|
memory_mb: int,
|
|
) -> ResourceInventory:
|
|
return ResourceInventory(
|
|
cpu_cores=cpu_cores,
|
|
memory_mb=memory_mb,
|
|
scratch_mb=memory_mb,
|
|
architecture=platform.machine().lower() or "unknown",
|
|
environment_digests=tuple(
|
|
dict.fromkeys(
|
|
definition.manifest.environment.digest
|
|
for definition in definitions.values()
|
|
)
|
|
),
|
|
)
|
|
|
|
|
|
def _runtime_for(
|
|
definitions: Mapping[str, WorkloadDefinition], inventory: ResourceInventory
|
|
) -> RuntimeCapabilities:
|
|
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=tuple(sorted(definitions)),
|
|
inventory=inventory,
|
|
)
|
|
|
|
|
|
class SciMeshRunner:
|
|
"""Execute claimed coordinator tasks through SDK-built workloads.
|
|
|
|
Definitions come from ``SCIMESH_WORKLOAD_ALLOWLIST`` when set, otherwise
|
|
the built-in ``similarity-search``; callers may supply an explicit mapping
|
|
for tests.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
definitions: Mapping[str, WorkloadDefinition] | None = None,
|
|
*,
|
|
inventory: ResourceInventory | None = None,
|
|
runtime: RuntimeCapabilities | None = None,
|
|
) -> None:
|
|
if definitions is None:
|
|
definitions = _definitions_from_environment()
|
|
self._definitions = dict(definitions)
|
|
self._inventory = inventory or _inventory_for(
|
|
self._definitions,
|
|
cpu_cores=1,
|
|
memory_mb=1024,
|
|
)
|
|
self._runtime = runtime or _runtime_for(self._definitions, self._inventory)
|
|
self._pool = ResourcePool(self._runtime.inventory, max_concurrency=1)
|
|
|
|
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
|
|
task_dir = task_dir.resolve()
|
|
workload = task.workload.replace("_", "-")
|
|
definition = self._definitions.get(workload)
|
|
if definition is None:
|
|
raise ValueError(f"unsupported workload: {task.workload}")
|
|
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
|
|
)
|
|
if set(map_stage.inputs) != {"input"} or len(map_stage.outputs) != 1:
|
|
raise ValueError(
|
|
f"workload {workload} is not executable through the v1 single-input contract"
|
|
)
|
|
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")
|
|
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=task.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)
|
|
|
|
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 _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,
|
|
)
|
|
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)
|