109 lines
4.3 KiB
Python
109 lines
4.3 KiB
Python
"""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
|