1367 lines
54 KiB
Python
1367 lines
54 KiB
Python
"""Local core-batch execution and reusable SDK conformance checks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import shutil
|
|
import stat
|
|
import tempfile
|
|
from dataclasses import dataclass, replace
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from threading import Event, Lock
|
|
from typing import Callable, Mapping
|
|
from uuid import NAMESPACE_URL, uuid4, uuid5
|
|
|
|
from ._validation import canonical_json, require_positive_int
|
|
from .artifacts import (
|
|
ArtifactCollection,
|
|
ArtifactItem,
|
|
ArtifactRef,
|
|
ArtifactSchema,
|
|
CollectionKind,
|
|
OutputManifest,
|
|
Provenance,
|
|
)
|
|
from .identity import ComponentRef, SDK_API_VERSION
|
|
from .execution import NetworkPolicy, ProcessModel
|
|
from .manifest import TrustMode, WorkloadManifest
|
|
from .plans import JobRequest, TaskSpec
|
|
from .protocols import ArtifactCatalog, ArtifactSink
|
|
from .registry import WorkloadDefinition, WorkloadRegistry
|
|
from .resources import ResourceAllocation, ResourcePool
|
|
from .runtime import RuntimeCapabilities
|
|
from .verification import (
|
|
CandidateOutputs,
|
|
VerificationBinding,
|
|
VerificationDecision,
|
|
VerificationStatus,
|
|
VerifyContext,
|
|
)
|
|
from .workflow import StageKind, WorkflowFailurePolicy
|
|
|
|
|
|
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 LocalArtifactStore:
|
|
"""Credential-free content store for local SDK/conformance execution."""
|
|
|
|
def __init__(
|
|
self,
|
|
root: Path,
|
|
*,
|
|
inspectors: Mapping[
|
|
str,
|
|
tuple[
|
|
ComponentRef,
|
|
Callable[
|
|
[Path, Mapping[str, object]],
|
|
tuple[int | None, tuple[int, ...]],
|
|
],
|
|
],
|
|
]
|
|
| None = None,
|
|
) -> None:
|
|
self.root = root.resolve()
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
self._paths: dict[str, Path] = {}
|
|
self._references: dict[str, ArtifactRef] = {}
|
|
self._refcounts: dict[str, int] = {}
|
|
self._inspectors = dict(inspectors or {})
|
|
if any(
|
|
not isinstance(binding, tuple)
|
|
or len(binding) != 2
|
|
or not isinstance(binding[0], ComponentRef)
|
|
or not callable(binding[1])
|
|
for binding in self._inspectors.values()
|
|
):
|
|
raise ValueError("artifact inspectors must bind an identity and callable")
|
|
self._lock = Lock()
|
|
|
|
def seal(
|
|
self,
|
|
path: Path,
|
|
*,
|
|
declaration: ArtifactSchema,
|
|
records: int | None = None,
|
|
dimensions: tuple[int, ...] = (),
|
|
) -> ArtifactRef:
|
|
source_flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
|
|
try:
|
|
source_fd = os.open(path, source_flags)
|
|
except OSError as error:
|
|
raise ValueError(
|
|
"artifact sink could not open a regular non-symlink file"
|
|
) from error
|
|
try:
|
|
return self.seal_descriptor(
|
|
source_fd,
|
|
declaration=declaration,
|
|
records=records,
|
|
dimensions=dimensions,
|
|
)
|
|
finally:
|
|
os.close(source_fd)
|
|
|
|
def seal_descriptor(
|
|
self,
|
|
descriptor: int,
|
|
*,
|
|
declaration: ArtifactSchema,
|
|
records: int | None = None,
|
|
dimensions: tuple[int, ...] = (),
|
|
) -> ArtifactRef:
|
|
"""Copy and validate one already safely opened regular-file descriptor."""
|
|
if (
|
|
isinstance(descriptor, bool)
|
|
or not isinstance(descriptor, int)
|
|
or descriptor < 0
|
|
):
|
|
raise ValueError("artifact descriptor must be an open file descriptor")
|
|
if not isinstance(declaration, ArtifactSchema):
|
|
raise ValueError("artifact declaration must be an ArtifactSchema")
|
|
source_fd = os.dup(descriptor)
|
|
temporary_fd, temporary_name = tempfile.mkstemp(prefix=".seal-", dir=self.root)
|
|
temporary = Path(temporary_name)
|
|
digest_builder = hashlib.sha256()
|
|
try:
|
|
if not stat.S_ISREG(os.fstat(source_fd).st_mode):
|
|
raise ValueError("artifact sink accepts only regular files")
|
|
with (
|
|
os.fdopen(source_fd, "rb", closefd=True) as source_file,
|
|
os.fdopen(temporary_fd, "wb", closefd=True) as destination_file,
|
|
):
|
|
source_fd = -1
|
|
temporary_fd = -1
|
|
for block in iter(lambda: source_file.read(1024 * 1024), b""):
|
|
digest_builder.update(block)
|
|
destination_file.write(block)
|
|
destination_file.flush()
|
|
os.fsync(destination_file.fileno())
|
|
digest = digest_builder.hexdigest()
|
|
artifact_id = str(
|
|
uuid5(NAMESPACE_URL, f"scimesh:{declaration.ref.canonical}:{digest}")
|
|
)
|
|
destination = self.root / artifact_id
|
|
size_bytes = temporary.stat().st_size
|
|
if size_bytes > declaration.max_bytes:
|
|
raise ValueError("sealed artifact exceeds its declared byte limit")
|
|
measured_records, measured_dimensions = self._inspect_content(
|
|
temporary,
|
|
declaration,
|
|
)
|
|
if records is not None and records != measured_records:
|
|
raise ValueError(
|
|
"artifact record summary does not match inspected content"
|
|
)
|
|
if dimensions and dimensions != measured_dimensions:
|
|
raise ValueError(
|
|
"artifact dimension summary does not match inspected content"
|
|
)
|
|
if declaration.max_records is not None:
|
|
if measured_records is None:
|
|
raise ValueError(
|
|
"artifact validator did not produce a required record count"
|
|
)
|
|
if measured_records > declaration.max_records:
|
|
raise ValueError(
|
|
"sealed artifact exceeds its declared record limit"
|
|
)
|
|
if declaration.max_dimensions:
|
|
if len(measured_dimensions) != len(declaration.max_dimensions) or any(
|
|
actual > maximum
|
|
for actual, maximum in zip(
|
|
measured_dimensions,
|
|
declaration.max_dimensions,
|
|
)
|
|
):
|
|
raise ValueError(
|
|
"sealed artifact exceeds its declared dimension limits"
|
|
)
|
|
reference = ArtifactRef(
|
|
artifact_id,
|
|
digest,
|
|
declaration.ref,
|
|
declaration.media_type,
|
|
size_bytes,
|
|
records=measured_records,
|
|
dimensions=measured_dimensions,
|
|
)
|
|
with self._lock:
|
|
if destination.is_symlink():
|
|
raise ValueError(
|
|
"local artifact destination must not be a symbolic link"
|
|
)
|
|
if destination.exists():
|
|
if not destination.is_file() or _sha256_file(destination) != digest:
|
|
raise ValueError("local artifact identity collision")
|
|
temporary.unlink()
|
|
else:
|
|
os.replace(temporary, destination)
|
|
destination.chmod(0o444)
|
|
existing = self._references.get(artifact_id)
|
|
if existing is not None and existing != reference:
|
|
raise ValueError(
|
|
"local artifact identity was reused with different metadata"
|
|
)
|
|
self._paths[artifact_id] = destination
|
|
self._references[artifact_id] = reference
|
|
self._refcounts[artifact_id] = self._refcounts.get(artifact_id, 0) + 1
|
|
return reference
|
|
finally:
|
|
if source_fd >= 0:
|
|
os.close(source_fd)
|
|
if temporary_fd >= 0:
|
|
os.close(temporary_fd)
|
|
if temporary.exists():
|
|
temporary.unlink()
|
|
|
|
def _inspect_content(
|
|
self,
|
|
path: Path,
|
|
declaration: ArtifactSchema,
|
|
) -> tuple[int | None, tuple[int, ...]]:
|
|
validator = declaration.validator
|
|
configuration = declaration.validator_configuration
|
|
if validator == ComponentRef("delimited-table", 1):
|
|
return self._inspect_delimited(path, declaration)
|
|
if validator == ComponentRef("json-document", 1):
|
|
return self._inspect_json(path, declaration)
|
|
if validator == ComponentRef("opaque-bytes", 1):
|
|
if configuration:
|
|
raise ValueError("opaque-bytes validator does not accept configuration")
|
|
return None, ()
|
|
binding = self._inspectors.get(declaration.ref.canonical)
|
|
if binding is None or binding[0] != validator:
|
|
raise ValueError("artifact schema has no matching registered validator")
|
|
inspected = binding[1](path, configuration)
|
|
if (
|
|
not isinstance(inspected, tuple)
|
|
or len(inspected) != 2
|
|
or (
|
|
inspected[0] is not None
|
|
and (
|
|
isinstance(inspected[0], bool)
|
|
or not isinstance(inspected[0], int)
|
|
or inspected[0] < 0
|
|
)
|
|
)
|
|
or not isinstance(inspected[1], tuple)
|
|
or any(
|
|
isinstance(value, bool) or not isinstance(value, int) or value < 0
|
|
for value in inspected[1]
|
|
)
|
|
):
|
|
raise ValueError("artifact inspector returned an invalid summary")
|
|
return inspected
|
|
|
|
@staticmethod
|
|
def _inspect_delimited(
|
|
path: Path,
|
|
declaration: ArtifactSchema,
|
|
) -> tuple[int, tuple[int, ...]]:
|
|
import csv
|
|
|
|
if declaration.media_type not in {"text/csv", "text/tab-separated-values"}:
|
|
raise ValueError("delimited-table validator requires CSV or TSV media type")
|
|
if declaration.encoding != "utf-8":
|
|
raise ValueError("delimited-table@1 requires utf-8 encoding")
|
|
configuration = dict(declaration.validator_configuration)
|
|
unknown = set(configuration) - {"columns", "required_columns"}
|
|
if unknown:
|
|
raise ValueError(
|
|
"delimited-table validator configuration has unknown fields"
|
|
)
|
|
columns = configuration.get("columns")
|
|
required = configuration.get("required_columns", ())
|
|
if columns is not None and not isinstance(columns, (list, tuple)):
|
|
raise ValueError("delimited-table columns must be an array")
|
|
if not isinstance(required, (list, tuple)):
|
|
raise ValueError("delimited-table required_columns must be an array")
|
|
expected_columns = tuple(columns) if columns is not None else None
|
|
required_columns = tuple(required)
|
|
for values, field_name in (
|
|
(expected_columns or (), "columns"),
|
|
(required_columns, "required_columns"),
|
|
):
|
|
if any(not isinstance(value, str) or not value for value in values) or len(
|
|
values
|
|
) != len(set(values)):
|
|
raise ValueError(f"delimited-table {field_name} must be unique strings")
|
|
delimiter = (
|
|
"\t" if declaration.media_type == "text/tab-separated-values" else ","
|
|
)
|
|
try:
|
|
with path.open("r", encoding="utf-8", newline="") as source_file:
|
|
reader = csv.reader(source_file, delimiter=delimiter)
|
|
try:
|
|
header = tuple(next(reader))
|
|
except StopIteration as error:
|
|
raise ValueError(
|
|
"delimited-table artifact must contain a header"
|
|
) from error
|
|
if (
|
|
not header
|
|
or any(not value for value in header)
|
|
or len(header) != len(set(header))
|
|
):
|
|
raise ValueError("delimited-table artifact has an invalid header")
|
|
if expected_columns is not None and header != expected_columns:
|
|
raise ValueError(
|
|
"delimited-table artifact header does not match its schema"
|
|
)
|
|
if not set(required_columns).issubset(header):
|
|
raise ValueError(
|
|
"delimited-table artifact is missing required columns"
|
|
)
|
|
count = 0
|
|
for row in reader:
|
|
if len(row) != len(header):
|
|
raise ValueError(
|
|
"delimited-table artifact has an inconsistent row width"
|
|
)
|
|
count += 1
|
|
if (
|
|
declaration.max_records is not None
|
|
and count > declaration.max_records
|
|
):
|
|
raise ValueError(
|
|
"sealed artifact exceeds its declared record limit"
|
|
)
|
|
except (UnicodeError, csv.Error) as error:
|
|
raise ValueError(
|
|
"sealed tabular artifact is not valid bounded text"
|
|
) from error
|
|
return count, ()
|
|
|
|
@staticmethod
|
|
def _inspect_json(
|
|
path: Path,
|
|
declaration: ArtifactSchema,
|
|
) -> tuple[int | None, tuple[int, ...]]:
|
|
import json
|
|
|
|
if not (
|
|
declaration.media_type == "application/json"
|
|
or declaration.media_type.endswith("+json")
|
|
):
|
|
raise ValueError("json-document validator requires a JSON media type")
|
|
if declaration.encoding != "utf-8":
|
|
raise ValueError("json-document@1 requires utf-8 encoding")
|
|
configuration = dict(declaration.validator_configuration)
|
|
if set(configuration) - {"top_level"}:
|
|
raise ValueError("json-document validator configuration has unknown fields")
|
|
top_level = configuration.get("top_level", "any")
|
|
if top_level not in {"any", "array", "object"}:
|
|
raise ValueError("json-document top_level is unsupported")
|
|
try:
|
|
with path.open("r", encoding="utf-8") as source_file:
|
|
value = json.load(
|
|
source_file,
|
|
parse_constant=lambda _value: (_ for _ in ()).throw(
|
|
ValueError("non-finite JSON number")
|
|
),
|
|
)
|
|
except (
|
|
UnicodeError,
|
|
json.JSONDecodeError,
|
|
ValueError,
|
|
RecursionError,
|
|
) as error:
|
|
raise ValueError(
|
|
"sealed JSON artifact is not a valid bounded document"
|
|
) from error
|
|
if top_level == "array" and not isinstance(value, list):
|
|
raise ValueError("JSON artifact must contain a top-level array")
|
|
if top_level == "object" and not isinstance(value, dict):
|
|
raise ValueError("JSON artifact must contain a top-level object")
|
|
|
|
def dimensions(current: object, depth: int = 0) -> tuple[int, ...]:
|
|
if depth > 8 or not isinstance(current, list):
|
|
return ()
|
|
if not current:
|
|
return (0,)
|
|
children = tuple(dimensions(child, depth + 1) for child in current)
|
|
if len(set(children)) != 1:
|
|
raise ValueError("JSON array dimensions must be rectangular")
|
|
return (len(current),) + children[0]
|
|
|
|
measured_dimensions = dimensions(value) if declaration.max_dimensions else ()
|
|
measured_records = len(value) if isinstance(value, list) else 1
|
|
return measured_records, measured_dimensions
|
|
|
|
def release(self, artifact: ArtifactRef) -> bool:
|
|
"""Release one seal reference and remove an unreferenced local blob."""
|
|
if not isinstance(artifact, ArtifactRef):
|
|
raise ValueError("artifact must be an ArtifactRef")
|
|
with self._lock:
|
|
if self._references.get(artifact.artifact_id) != artifact:
|
|
return False
|
|
remaining = self._refcounts[artifact.artifact_id] - 1
|
|
if remaining > 0:
|
|
self._refcounts[artifact.artifact_id] = remaining
|
|
return True
|
|
path = self._paths.pop(artifact.artifact_id)
|
|
self._references.pop(artifact.artifact_id, None)
|
|
self._refcounts.pop(artifact.artifact_id, None)
|
|
path.chmod(0o600)
|
|
path.unlink()
|
|
return True
|
|
|
|
def import_file(
|
|
self,
|
|
path: Path,
|
|
*,
|
|
declaration: ArtifactSchema,
|
|
records: int | None = None,
|
|
dimensions: tuple[int, ...] = (),
|
|
) -> ArtifactRef:
|
|
return self.seal(
|
|
path,
|
|
declaration=declaration,
|
|
records=records,
|
|
dimensions=dimensions,
|
|
)
|
|
|
|
def materialize(self, artifact: ArtifactRef) -> Path:
|
|
self.require(artifact)
|
|
with self._lock:
|
|
return self._paths[artifact.artifact_id]
|
|
|
|
def require(self, artifact: ArtifactRef) -> None:
|
|
if not isinstance(artifact, ArtifactRef):
|
|
raise ValueError("artifact must be an ArtifactRef")
|
|
with self._lock:
|
|
try:
|
|
path = self._paths[artifact.artifact_id]
|
|
stored = self._references[artifact.artifact_id]
|
|
except KeyError as error:
|
|
raise ValueError(
|
|
"artifact is not present in the local store"
|
|
) from error
|
|
if stored != artifact:
|
|
raise ValueError(
|
|
"artifact metadata does not match the sealed local reference"
|
|
)
|
|
if (
|
|
path.is_symlink()
|
|
or not path.is_file()
|
|
or path.stat().st_size != artifact.size_bytes
|
|
or _sha256_file(path) != artifact.sha256
|
|
):
|
|
raise ValueError("local artifact checksum mismatch")
|
|
|
|
|
|
class LocalArtifactTransaction:
|
|
"""Track store references until one local Job is fully accepted."""
|
|
|
|
def __init__(self, store: LocalArtifactStore) -> None:
|
|
self._store = store
|
|
self._references: list[ArtifactRef] = []
|
|
self._closed = False
|
|
self._lock = Lock()
|
|
|
|
def track(self, artifact: ArtifactRef) -> None:
|
|
with self._lock:
|
|
if self._closed:
|
|
raise ValueError("artifact transaction is already closed")
|
|
self._references.append(artifact)
|
|
|
|
def commit(self) -> None:
|
|
with self._lock:
|
|
if self._closed:
|
|
raise ValueError("artifact transaction is already closed")
|
|
self._closed = True
|
|
self._references.clear()
|
|
|
|
def rollback(self) -> None:
|
|
with self._lock:
|
|
if self._closed:
|
|
return
|
|
references = tuple(reversed(self._references))
|
|
self._references.clear()
|
|
self._closed = True
|
|
for artifact in references:
|
|
self._store.release(artifact)
|
|
|
|
|
|
class ScopedArtifactSink:
|
|
"""Restrict a local planning/attempt sink to one workspace tree."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: LocalArtifactStore,
|
|
workspace: Path,
|
|
*,
|
|
max_artifacts: int = 100_000,
|
|
max_bytes: int = 1 << 50,
|
|
transaction: LocalArtifactTransaction | None = None,
|
|
) -> None:
|
|
self._store = store
|
|
self._workspace = workspace.resolve()
|
|
self._workspace.mkdir(parents=True, exist_ok=True)
|
|
self._max_artifacts = require_positive_int(max_artifacts, "sink.max_artifacts")
|
|
self._max_bytes = require_positive_int(max_bytes, "sink.max_bytes")
|
|
self._sealed: dict[str, ArtifactRef] = {}
|
|
self._sealed_bytes = 0
|
|
self._transaction = transaction
|
|
self._lock = Lock()
|
|
|
|
@property
|
|
def sealed_references(self) -> tuple[ArtifactRef, ...]:
|
|
with self._lock:
|
|
return tuple(self._sealed[key] for key in sorted(self._sealed))
|
|
|
|
def seal(
|
|
self,
|
|
path: Path,
|
|
*,
|
|
declaration: ArtifactSchema,
|
|
records: int | None = None,
|
|
dimensions: tuple[int, ...] = (),
|
|
) -> ArtifactRef:
|
|
candidate = path if path.is_absolute() else self._workspace / path
|
|
if not isinstance(declaration, ArtifactSchema):
|
|
raise ValueError("artifact declaration must be an ArtifactSchema")
|
|
lexical = Path(os.path.abspath(candidate))
|
|
try:
|
|
lexical_relative = lexical.relative_to(self._workspace)
|
|
except ValueError as error:
|
|
raise ValueError(
|
|
"attempt artifact must remain inside its workspace"
|
|
) from error
|
|
if not lexical_relative.parts:
|
|
raise ValueError("attempt artifact must name a file inside its workspace")
|
|
directory_flags = (
|
|
os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0)
|
|
)
|
|
opened_directories: list[int] = []
|
|
file_descriptor = -1
|
|
try:
|
|
current_fd = os.open(self._workspace, directory_flags)
|
|
opened_directories.append(current_fd)
|
|
for component in lexical_relative.parts[:-1]:
|
|
current_fd = os.open(
|
|
component,
|
|
directory_flags,
|
|
dir_fd=current_fd,
|
|
)
|
|
opened_directories.append(current_fd)
|
|
file_descriptor = os.open(
|
|
lexical_relative.parts[-1],
|
|
os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
|
|
dir_fd=current_fd,
|
|
)
|
|
candidate_size = os.fstat(file_descriptor).st_size
|
|
except OSError as error:
|
|
if file_descriptor >= 0:
|
|
os.close(file_descriptor)
|
|
for directory_fd in reversed(opened_directories):
|
|
os.close(directory_fd)
|
|
raise ValueError(
|
|
"attempt artifact path must contain only real workspace directories"
|
|
) from error
|
|
try:
|
|
with self._lock:
|
|
if len(self._sealed) >= self._max_artifacts:
|
|
raise ValueError("attempt artifact count exceeds its sink limit")
|
|
if self._sealed_bytes + candidate_size > self._max_bytes:
|
|
raise ValueError("attempt artifact bytes exceed their sink limit")
|
|
if candidate_size > declaration.max_bytes:
|
|
raise ValueError("attempt artifact exceeds its schema byte limit")
|
|
reference = self._store.seal_descriptor(
|
|
file_descriptor,
|
|
declaration=declaration,
|
|
records=records,
|
|
dimensions=dimensions,
|
|
)
|
|
existing = self._sealed.get(reference.artifact_id)
|
|
if existing is not None and existing != reference:
|
|
self._store.release(reference)
|
|
raise ValueError(
|
|
"attempt sealed conflicting metadata for one artifact"
|
|
)
|
|
if existing is None:
|
|
if self._sealed_bytes + reference.size_bytes > self._max_bytes:
|
|
self._store.release(reference)
|
|
raise ValueError(
|
|
"attempt artifact bytes exceed their sink limit"
|
|
)
|
|
self._sealed[reference.artifact_id] = reference
|
|
self._sealed_bytes += reference.size_bytes
|
|
if self._transaction is not None:
|
|
self._transaction.track(reference)
|
|
else:
|
|
# ``seal_descriptor`` acquired another store reference for
|
|
# identical content; one attempt owns only one reference.
|
|
self._store.release(reference)
|
|
return reference
|
|
finally:
|
|
if file_descriptor >= 0:
|
|
os.close(file_descriptor)
|
|
for directory_fd in reversed(opened_directories):
|
|
os.close(directory_fd)
|
|
|
|
|
|
class ScopedArtifactCatalog:
|
|
"""Materialize verified, read-only copies inside one attempt workspace."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: LocalArtifactStore,
|
|
workspace: Path,
|
|
allowed_artifacts: tuple[ArtifactRef, ...],
|
|
) -> None:
|
|
self.__store = store
|
|
allowed: dict[str, ArtifactRef] = {}
|
|
for artifact in allowed_artifacts:
|
|
if not isinstance(artifact, ArtifactRef):
|
|
raise ValueError("catalog allowlist must contain ArtifactRef values")
|
|
existing = allowed.get(artifact.artifact_id)
|
|
if existing is not None and existing != artifact:
|
|
raise ValueError(
|
|
"catalog allowlist contains conflicting artifact metadata"
|
|
)
|
|
allowed[artifact.artifact_id] = artifact
|
|
self.__allowed = allowed
|
|
resolved_workspace = workspace.resolve()
|
|
input_root = resolved_workspace / "inputs"
|
|
if input_root.is_symlink():
|
|
raise ValueError("attempt input directory must not be a symbolic link")
|
|
self.__input_root = input_root
|
|
self.__input_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
def materialize(self, artifact: ArtifactRef) -> Path:
|
|
if self.__allowed.get(artifact.artifact_id) != artifact:
|
|
raise ValueError("artifact is outside this context's input allowlist")
|
|
source = self.__store.materialize(artifact)
|
|
destination = self.__input_root / artifact.artifact_id
|
|
if destination.is_symlink():
|
|
raise ValueError("attempt input destination must not be a symbolic link")
|
|
if destination.exists():
|
|
if (
|
|
not destination.is_file()
|
|
or destination.stat().st_size != artifact.size_bytes
|
|
or _sha256_file(destination) != artifact.sha256
|
|
):
|
|
raise ValueError("existing attempt input does not match its artifact")
|
|
else:
|
|
temporary_fd, temporary_name = tempfile.mkstemp(
|
|
prefix=".input-",
|
|
dir=self.__input_root,
|
|
)
|
|
os.close(temporary_fd)
|
|
temporary = Path(temporary_name)
|
|
try:
|
|
shutil.copyfile(source, temporary)
|
|
if (
|
|
temporary.stat().st_size != artifact.size_bytes
|
|
or _sha256_file(temporary) != artifact.sha256
|
|
):
|
|
raise ValueError("copied attempt input does not match its artifact")
|
|
temporary.chmod(0o444)
|
|
os.replace(temporary, destination)
|
|
finally:
|
|
if temporary.exists():
|
|
temporary.unlink()
|
|
if (
|
|
destination.stat().st_size != artifact.size_bytes
|
|
or _sha256_file(destination) != artifact.sha256
|
|
):
|
|
raise ValueError("copied attempt input does not match its artifact")
|
|
destination.chmod(0o444)
|
|
return destination
|
|
|
|
|
|
class CancellationFlag:
|
|
"""A thread-safe cancellation signal for local conformance runs."""
|
|
|
|
def __init__(self) -> None:
|
|
self._event = Event()
|
|
|
|
def cancelled(self) -> bool:
|
|
return self._event.is_set()
|
|
|
|
def raise_if_cancelled(self) -> None:
|
|
if self.cancelled():
|
|
raise RuntimeError("task-cancelled")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class LocalPlanningContext:
|
|
"""Scoped catalog/sink/workspace handed to a planner by the local executor.
|
|
|
|
Wraps a ``LocalArtifactStore`` into attempt-scoped catalog and sink so a
|
|
planner can materialize job inputs and seal shards.
|
|
"""
|
|
|
|
catalog: ArtifactCatalog
|
|
sink: ArtifactSink
|
|
workspace: Path
|
|
allowed_artifacts: tuple[ArtifactRef, ...] = ()
|
|
max_artifacts: int = 100_000
|
|
max_bytes: int = 1 << 50
|
|
transaction: LocalArtifactTransaction | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
workspace = self.workspace.resolve()
|
|
object.__setattr__(self, "workspace", workspace)
|
|
allowed_artifacts = tuple(self.allowed_artifacts)
|
|
if any(not isinstance(value, ArtifactRef) for value in allowed_artifacts):
|
|
raise ValueError("allowed_artifacts must contain ArtifactRef values")
|
|
object.__setattr__(self, "allowed_artifacts", allowed_artifacts)
|
|
object.__setattr__(
|
|
self,
|
|
"max_artifacts",
|
|
require_positive_int(self.max_artifacts, "max_artifacts"),
|
|
)
|
|
object.__setattr__(
|
|
self, "max_bytes", require_positive_int(self.max_bytes, "max_bytes")
|
|
)
|
|
if isinstance(self.catalog, LocalArtifactStore):
|
|
object.__setattr__(
|
|
self,
|
|
"catalog",
|
|
ScopedArtifactCatalog(self.catalog, workspace, allowed_artifacts),
|
|
)
|
|
if isinstance(self.sink, LocalArtifactStore):
|
|
object.__setattr__(
|
|
self,
|
|
"sink",
|
|
ScopedArtifactSink(
|
|
self.sink,
|
|
workspace,
|
|
max_artifacts=self.max_artifacts,
|
|
max_bytes=self.max_bytes,
|
|
transaction=self.transaction,
|
|
),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class LocalTaskContext:
|
|
"""Scoped context handed to a map/reduce handler by the local executor.
|
|
|
|
Carries the digest-pinned ``TaskSpec``, attempt workspace, cancellation
|
|
flag, provenance, and the accepted inputs for reducers.
|
|
"""
|
|
|
|
task: TaskSpec
|
|
catalog: ArtifactCatalog
|
|
sink: ArtifactSink
|
|
workspace: Path
|
|
cancellation: CancellationFlag
|
|
provenance: Provenance
|
|
accepted_inputs: Mapping[str, ArtifactCollection]
|
|
max_artifacts: int = 100_000
|
|
max_bytes: int = 1 << 50
|
|
transaction: LocalArtifactTransaction | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
workspace = self.workspace.resolve()
|
|
object.__setattr__(self, "workspace", workspace)
|
|
object.__setattr__(
|
|
self,
|
|
"max_artifacts",
|
|
require_positive_int(self.max_artifacts, "max_artifacts"),
|
|
)
|
|
object.__setattr__(
|
|
self, "max_bytes", require_positive_int(self.max_bytes, "max_bytes")
|
|
)
|
|
if isinstance(self.catalog, LocalArtifactStore):
|
|
allowed_artifacts = tuple(
|
|
item.artifact
|
|
for collection in self.task.inputs.values()
|
|
for item in collection.items
|
|
)
|
|
object.__setattr__(
|
|
self,
|
|
"catalog",
|
|
ScopedArtifactCatalog(self.catalog, workspace, allowed_artifacts),
|
|
)
|
|
if isinstance(self.sink, LocalArtifactStore):
|
|
object.__setattr__(
|
|
self,
|
|
"sink",
|
|
ScopedArtifactSink(
|
|
self.sink,
|
|
workspace,
|
|
max_artifacts=self.max_artifacts,
|
|
max_bytes=self.max_bytes,
|
|
transaction=self.transaction,
|
|
),
|
|
)
|
|
|
|
|
|
def assert_manifest_round_trip(manifest: WorkloadManifest) -> None:
|
|
"""Assert strict canonical serialization and reconstructive equality."""
|
|
reconstructed = WorkloadManifest.from_json(manifest.to_json())
|
|
if reconstructed != manifest or reconstructed.to_json() != manifest.to_json():
|
|
raise AssertionError("manifest canonical round-trip changed its value")
|
|
|
|
|
|
def _input_digest(inputs: Mapping[str, ArtifactCollection]) -> str:
|
|
value = {name: collection.digest for name, collection in sorted(inputs.items())}
|
|
return hashlib.sha256(canonical_json(value).encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _utc_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def _provenance(
|
|
definition: WorkloadDefinition,
|
|
runtime: RuntimeCapabilities,
|
|
task: TaskSpec,
|
|
allocation: ResourceAllocation,
|
|
started_at: str,
|
|
job_id: str,
|
|
task_id: str,
|
|
) -> Provenance:
|
|
parameters_digest = hashlib.sha256(
|
|
canonical_json(task.parameters).encode("utf-8")
|
|
).hexdigest()
|
|
return Provenance(
|
|
workload=definition.manifest.workload,
|
|
sdk_api_version=task.sdk_api_version,
|
|
protocol_version=task.protocol_version,
|
|
manifest_schema_version=task.manifest_schema_version,
|
|
workflow_schema_version=task.workflow_schema_version,
|
|
verifier=task.verifier,
|
|
artifact_schemas=tuple(
|
|
sorted(
|
|
{
|
|
item.artifact.schema
|
|
for collection in task.inputs.values()
|
|
for item in collection.items
|
|
}.union(port.schema.ref for port in task.expected_outputs.values()),
|
|
key=lambda value: value.canonical,
|
|
)
|
|
),
|
|
package_digest=task.package_digest,
|
|
manifest_digest=task.manifest_digest,
|
|
environment_digest=task.environment_digest,
|
|
worker_runtime={"kind": "local-conformance", "sdk_api": SDK_API_VERSION},
|
|
allocated_resource_ids=(allocation.allocation_id,) + allocation.accelerator_ids,
|
|
parameters_digest=parameters_digest,
|
|
input_collection_digest=_input_digest(task.inputs),
|
|
execution_contract_digest=hashlib.sha256(
|
|
task.to_json().encode("utf-8")
|
|
).hexdigest(),
|
|
selected_features=task.selected_features,
|
|
optional_fallbacks=task.optional_fallbacks,
|
|
job_id=job_id,
|
|
task_id=task_id,
|
|
started_at=started_at,
|
|
finished_at=started_at,
|
|
trust_mode=task.trust_mode.value,
|
|
)
|
|
|
|
|
|
def _verification_binding(manifest: OutputManifest) -> VerificationBinding:
|
|
provenance = manifest.provenance
|
|
return VerificationBinding(
|
|
workload=provenance.workload,
|
|
task_key=manifest.task_key,
|
|
package_digest=provenance.package_digest,
|
|
manifest_digest=provenance.manifest_digest,
|
|
environment_digest=provenance.environment_digest,
|
|
parameters_digest=provenance.parameters_digest,
|
|
input_collection_digest=provenance.input_collection_digest,
|
|
execution_contract_digest=provenance.execution_contract_digest,
|
|
selected_features=provenance.selected_features,
|
|
optional_fallbacks=provenance.optional_fallbacks,
|
|
job_id=provenance.job_id,
|
|
task_id=provenance.task_id,
|
|
verifier=provenance.verifier,
|
|
sdk_api_version=provenance.sdk_api_version,
|
|
protocol_version=provenance.protocol_version,
|
|
manifest_schema_version=provenance.manifest_schema_version,
|
|
workflow_schema_version=provenance.workflow_schema_version,
|
|
artifact_schemas=provenance.artifact_schemas,
|
|
trust_mode=TrustMode(provenance.trust_mode),
|
|
)
|
|
|
|
|
|
class LocalCoreBatchExecutor:
|
|
"""Trusted in-process correctness runtime for the static map/reduce profile.
|
|
|
|
This executor deliberately accepts only profiles that declare trusted host
|
|
execution. It is useful for SDK conformance and scientific parity tests;
|
|
it is not a process, network, credential, lease, or timeout isolation
|
|
boundary.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
registry: WorkloadRegistry,
|
|
runtime: RuntimeCapabilities,
|
|
artifact_store: LocalArtifactStore,
|
|
work_root: Path,
|
|
) -> None:
|
|
self.registry = registry
|
|
self.runtime = runtime
|
|
self.artifact_store = artifact_store
|
|
self.work_root = work_root.resolve()
|
|
self.work_root.mkdir(parents=True, exist_ok=True)
|
|
self.resources = ResourcePool(runtime.inventory, max_concurrency=1)
|
|
|
|
@staticmethod
|
|
def _assert_supported_profile(
|
|
request: JobRequest, definition: WorkloadDefinition
|
|
) -> None:
|
|
if request.trust_mode is not TrustMode.TRUSTED:
|
|
raise ValueError(
|
|
"local conformance execution supports only trusted workloads"
|
|
)
|
|
workflow = definition.manifest.workflow
|
|
if workflow.failure_policy is not WorkflowFailurePolicy.FAIL_FAST:
|
|
raise ValueError(
|
|
"local conformance execution supports only fail-fast workflows"
|
|
)
|
|
for stage in workflow.stages:
|
|
if stage.kind not in {StageKind.MAP, StageKind.REDUCE}:
|
|
raise ValueError(
|
|
"local conformance execution does not implement advanced stages"
|
|
)
|
|
execution = stage.execution
|
|
if (
|
|
execution.process_model is not ProcessModel.SINGLE
|
|
or execution.max_processes != 1
|
|
or execution.threads_per_process != 1
|
|
or execution.native_threads != 1
|
|
or execution.nested_parallelism
|
|
):
|
|
raise ValueError(
|
|
"local conformance execution supports one non-nested host thread"
|
|
)
|
|
if execution.network is not NetworkPolicy.TRUSTED:
|
|
raise ValueError(
|
|
"local conformance execution cannot enforce a restricted network policy"
|
|
)
|
|
if (
|
|
execution.checkpoint.enabled
|
|
or execution.allowed_egress
|
|
or execution.secret_handles
|
|
or stage.gang is not None
|
|
or stage.resources.accelerator_count
|
|
):
|
|
raise ValueError(
|
|
"local conformance execution cannot enforce this stage profile"
|
|
)
|
|
if stage.retry.max_attempts != 1:
|
|
raise ValueError(
|
|
"local conformance execution does not implement retries"
|
|
)
|
|
reducers = tuple(
|
|
stage for stage in workflow.stages if stage.kind is StageKind.REDUCE
|
|
)
|
|
if len(reducers) == 1:
|
|
reducer = reducers[0]
|
|
if set(workflow.outputs) != set(reducer.outputs) or any(
|
|
external_name != reference.port
|
|
or reference.stage_id != reducer.stage_id
|
|
for external_name, reference in workflow.outputs.items()
|
|
):
|
|
raise ValueError(
|
|
"local conformance execution requires identity-mapped reducer outputs"
|
|
)
|
|
|
|
def _track_artifacts(
|
|
self,
|
|
collections: Mapping[str, ArtifactCollection],
|
|
*,
|
|
known: dict[str, ArtifactRef],
|
|
max_artifacts: int,
|
|
output_ids: set[str] | None = None,
|
|
) -> int:
|
|
added_output_bytes = 0
|
|
for collection in collections.values():
|
|
for item in collection.items:
|
|
artifact = item.artifact
|
|
self.artifact_store.require(artifact)
|
|
existing = known.get(artifact.artifact_id)
|
|
if existing is not None and existing != artifact:
|
|
raise ValueError("one artifact ID carries conflicting metadata")
|
|
known[artifact.artifact_id] = artifact
|
|
if len(known) > max_artifacts:
|
|
raise ValueError("job exceeds the manifest artifact limit")
|
|
if output_ids is not None and artifact.artifact_id not in output_ids:
|
|
output_ids.add(artifact.artifact_id)
|
|
added_output_bytes += artifact.size_bytes
|
|
return added_output_bytes
|
|
|
|
def _run_task(
|
|
self,
|
|
definition: WorkloadDefinition,
|
|
task: TaskSpec,
|
|
workspace: Path,
|
|
operation: Callable[[LocalTaskContext], OutputManifest],
|
|
*,
|
|
job_id: str,
|
|
transaction: LocalArtifactTransaction,
|
|
max_artifacts: int,
|
|
max_output_bytes: int,
|
|
) -> OutputManifest:
|
|
task_id = str(uuid4())
|
|
allocation = self.resources.reserve(task_id, task.resources)
|
|
try:
|
|
started_at = _utc_now()
|
|
stage = next(
|
|
stage
|
|
for stage in definition.manifest.workflow.stages
|
|
if stage.stage_id == task.stage_id
|
|
)
|
|
if stage.verifier is None:
|
|
raise ValueError("local task stage has no declared acceptance verifier")
|
|
if (
|
|
task.workload != definition.manifest.workload
|
|
or task.package_digest != definition.manifest.package.digest
|
|
or task.manifest_digest != definition.manifest.digest
|
|
or task.sdk_api_version != self.runtime.sdk_api_version
|
|
or task.protocol_version != self.runtime.protocol_version
|
|
or task.manifest_schema_version
|
|
!= definition.manifest.manifest_schema_version
|
|
or task.workflow_schema_version
|
|
!= definition.manifest.workflow.schema_version
|
|
or task.environment_digest != definition.manifest.environment.digest
|
|
or task.verifier != stage.verifier
|
|
):
|
|
raise ValueError(
|
|
"task resolved pins do not match the selected runtime and manifest"
|
|
)
|
|
provenance = _provenance(
|
|
definition,
|
|
self.runtime,
|
|
task,
|
|
allocation,
|
|
started_at,
|
|
job_id,
|
|
task_id,
|
|
)
|
|
context = LocalTaskContext(
|
|
task,
|
|
self.artifact_store,
|
|
self.artifact_store,
|
|
workspace,
|
|
CancellationFlag(),
|
|
provenance,
|
|
task.inputs,
|
|
max_artifacts,
|
|
max_output_bytes,
|
|
transaction,
|
|
)
|
|
manifest = operation(context)
|
|
if not isinstance(manifest, OutputManifest):
|
|
raise ValueError("workload handler must return an OutputManifest")
|
|
if manifest.task_key != task.task_key:
|
|
raise ValueError(
|
|
"handler output task_key does not match its trusted task"
|
|
)
|
|
if manifest.provenance != provenance:
|
|
raise ValueError(
|
|
"handler output provenance does not match its trusted context"
|
|
)
|
|
manifest.validate_against(
|
|
task.expected_outputs,
|
|
max_output_bytes=max_output_bytes,
|
|
)
|
|
if not isinstance(context.sink, ScopedArtifactSink):
|
|
raise ValueError("local execution requires a scoped artifact sink")
|
|
declared = {
|
|
item.artifact.artifact_id: item.artifact
|
|
for collection in manifest.outputs.values()
|
|
for item in collection.items
|
|
}
|
|
issued = {
|
|
artifact.artifact_id: artifact
|
|
for artifact in context.sink.sealed_references
|
|
}
|
|
if issued != declared:
|
|
raise ValueError(
|
|
"handler outputs must declare exactly the artifacts sealed by its attempt"
|
|
)
|
|
for collection in manifest.outputs.values():
|
|
for item in collection.items:
|
|
self.artifact_store.require(item.artifact)
|
|
finished = replace(provenance, finished_at=_utc_now())
|
|
completed = replace(manifest, provenance=finished)
|
|
self._verify_output(
|
|
definition,
|
|
stage.verifier,
|
|
completed,
|
|
task.expected_outputs,
|
|
max_output_bytes,
|
|
)
|
|
return completed
|
|
finally:
|
|
self.resources.release(allocation.allocation_id)
|
|
|
|
@staticmethod
|
|
def _verify_output(
|
|
definition: WorkloadDefinition,
|
|
verifier_ref: ComponentRef,
|
|
output: OutputManifest,
|
|
expected_outputs: Mapping[str, object],
|
|
max_output_bytes: int,
|
|
) -> None:
|
|
verifier = definition.verifiers[verifier_ref.canonical]
|
|
decision = verifier.verify(
|
|
VerifyContext(
|
|
expected_outputs, # type: ignore[arg-type]
|
|
max_output_bytes,
|
|
binding=_verification_binding(output),
|
|
trust_mode=TrustMode(output.provenance.trust_mode),
|
|
),
|
|
CandidateOutputs((output,)),
|
|
)
|
|
if not isinstance(decision, VerificationDecision):
|
|
raise ValueError("declared verifier must return a VerificationDecision")
|
|
if decision.verifier != verifier_ref:
|
|
raise ValueError(
|
|
"verification decision identity does not match the declared verifier"
|
|
)
|
|
if decision.status is not VerificationStatus.ACCEPTED:
|
|
raise ValueError("task output did not pass its declared verifier")
|
|
|
|
def execute(self, request: JobRequest, package_digest: str) -> OutputManifest:
|
|
run_root = self.work_root / f"run-{uuid4()}"
|
|
run_root.mkdir(parents=False, exist_ok=False)
|
|
transaction = LocalArtifactTransaction(self.artifact_store)
|
|
try:
|
|
result = self._execute_run(
|
|
request,
|
|
package_digest,
|
|
run_root,
|
|
transaction,
|
|
)
|
|
shutil.rmtree(run_root, ignore_errors=False)
|
|
except BaseException:
|
|
transaction.rollback()
|
|
if run_root.exists():
|
|
shutil.rmtree(run_root, ignore_errors=True)
|
|
raise
|
|
transaction.commit()
|
|
return result
|
|
|
|
def _execute_run(
|
|
self,
|
|
request: JobRequest,
|
|
package_digest: str,
|
|
run_root: Path,
|
|
transaction: LocalArtifactTransaction,
|
|
) -> OutputManifest:
|
|
definition, _ = self.registry.require(
|
|
request.workload.name,
|
|
request.workload.version,
|
|
package_digest,
|
|
runtime=self.runtime,
|
|
)
|
|
self._assert_supported_profile(request, definition)
|
|
workflow = definition.manifest.workflow
|
|
map_stages = [stage for stage in workflow.stages if stage.kind is StageKind.MAP]
|
|
reduce_stages = [
|
|
stage for stage in workflow.stages if stage.kind is StageKind.REDUCE
|
|
]
|
|
unsupported = [
|
|
stage
|
|
for stage in workflow.stages
|
|
if stage.kind not in {StageKind.MAP, StageKind.REDUCE}
|
|
]
|
|
if len(map_stages) != 1 or len(reduce_stages) != 1 or unsupported:
|
|
raise ValueError(
|
|
"local core-batch executor supports one static map stage and one reducer"
|
|
)
|
|
limits = definition.manifest.limits
|
|
output_limit = min(limits.max_output_bytes, workflow.max_output_bytes)
|
|
job_id = str(uuid4())
|
|
known_artifacts: dict[str, ArtifactRef] = {}
|
|
output_artifact_ids: set[str] = set()
|
|
output_bytes = 0
|
|
self._track_artifacts(
|
|
request.inputs,
|
|
known=known_artifacts,
|
|
max_artifacts=limits.max_artifacts,
|
|
)
|
|
planning = LocalPlanningContext(
|
|
self.artifact_store,
|
|
self.artifact_store,
|
|
run_root / "planning",
|
|
allowed_artifacts=tuple(
|
|
item.artifact
|
|
for collection in request.inputs.values()
|
|
for item in collection.items
|
|
),
|
|
max_artifacts=limits.max_artifacts,
|
|
max_bytes=limits.max_input_bytes,
|
|
transaction=transaction,
|
|
)
|
|
plan = self.registry.plan(request, package_digest, self.runtime, planning)
|
|
if len(plan.tasks) + 1 > min(workflow.max_tasks, limits.max_tasks):
|
|
raise ValueError("core map/reduce execution exceeds the total task limit")
|
|
if not isinstance(planning.sink, ScopedArtifactSink):
|
|
raise ValueError("local planning requires a scoped artifact sink")
|
|
planned_references = {
|
|
item.artifact.artifact_id: item.artifact
|
|
for task in plan.tasks
|
|
for collection in task.inputs.values()
|
|
for item in collection.items
|
|
}
|
|
for issued in planning.sink.sealed_references:
|
|
if planned_references.get(issued.artifact_id) != issued:
|
|
raise ValueError(
|
|
"planner sealed an artifact that is not referenced by its plan"
|
|
)
|
|
authorized_plan_inputs = {
|
|
item.artifact.artifact_id: item.artifact
|
|
for collection in request.inputs.values()
|
|
for item in collection.items
|
|
}
|
|
authorized_plan_inputs.update(
|
|
{
|
|
artifact.artifact_id: artifact
|
|
for artifact in planning.sink.sealed_references
|
|
}
|
|
)
|
|
for artifact_id, artifact in planned_references.items():
|
|
if authorized_plan_inputs.get(artifact_id) != artifact:
|
|
raise ValueError(
|
|
"workflow plan references an artifact outside job inputs and planning outputs"
|
|
)
|
|
map_stage = map_stages[0]
|
|
reducer_stage = reduce_stages[0]
|
|
if (
|
|
len(workflow.stages) != 2
|
|
or map_stage.needs
|
|
or reducer_stage.needs != (map_stage.stage_id,)
|
|
or any(
|
|
edge.source.stage_id is not None
|
|
for edge in workflow.edges
|
|
if edge.target.stage_id == map_stage.stage_id
|
|
)
|
|
or any(
|
|
edge.source.stage_id != map_stage.stage_id
|
|
for edge in workflow.edges
|
|
if edge.target.stage_id == reducer_stage.stage_id
|
|
)
|
|
or any(
|
|
reference.stage_id != reducer_stage.stage_id
|
|
for reference in workflow.outputs.values()
|
|
)
|
|
):
|
|
raise ValueError(
|
|
"local core-batch executor requires a canonical map-to-reduce DAG"
|
|
)
|
|
runner = definition.runners[map_stage.entry_point]
|
|
map_results: list[OutputManifest] = []
|
|
for task_index, task in enumerate(plan.tasks):
|
|
task.validate_stage(map_stage)
|
|
self._track_artifacts(
|
|
task.inputs,
|
|
known=known_artifacts,
|
|
max_artifacts=limits.max_artifacts,
|
|
)
|
|
manifest = self._run_task(
|
|
definition,
|
|
task,
|
|
run_root / "tasks" / f"map-{task_index:08d}",
|
|
runner.run,
|
|
job_id=job_id,
|
|
transaction=transaction,
|
|
max_artifacts=limits.max_artifacts,
|
|
max_output_bytes=output_limit - output_bytes,
|
|
)
|
|
output_bytes += self._track_artifacts(
|
|
manifest.outputs,
|
|
known=known_artifacts,
|
|
max_artifacts=limits.max_artifacts,
|
|
output_ids=output_artifact_ids,
|
|
)
|
|
if output_bytes > output_limit:
|
|
raise ValueError("job exceeds the cumulative output byte limit")
|
|
map_results.append(manifest)
|
|
if len(map_results) != len(plan.tasks):
|
|
raise ValueError(
|
|
"map execution did not produce exactly one accepted result per task"
|
|
)
|
|
if len(map_stage.outputs) != 1 or len(reducer_stage.inputs) != 1:
|
|
raise ValueError(
|
|
"core map/reduce adapter requires one map output and one reducer input"
|
|
)
|
|
map_port = next(iter(map_stage.outputs))
|
|
reducer_input_name = next(iter(reducer_stage.inputs))
|
|
partial_items: list[ArtifactItem] = []
|
|
for task, result in zip(plan.tasks, map_results):
|
|
collection = result.outputs[map_port]
|
|
if len(collection.items) != 1:
|
|
raise ValueError(
|
|
"core map stage must produce exactly one partial per planned task"
|
|
)
|
|
partial_items.append(
|
|
ArtifactItem(
|
|
collection.items[0].artifact,
|
|
key=task.task_key.replace("/", "."),
|
|
)
|
|
)
|
|
if len(partial_items) != len(plan.tasks):
|
|
raise ValueError(
|
|
"core map stage must produce exactly one partial per planned task"
|
|
)
|
|
accepted = ArtifactCollection(CollectionKind.KEYED, tuple(partial_items))
|
|
assert reducer_stage.verifier is not None
|
|
reducer_task = TaskSpec(
|
|
workload=plan.workload,
|
|
package_digest=plan.package_digest,
|
|
manifest_digest=plan.manifest_digest,
|
|
trust_mode=plan.trust_mode,
|
|
sdk_api_version=plan.sdk_api_version,
|
|
protocol_version=plan.protocol_version,
|
|
manifest_schema_version=plan.manifest_schema_version,
|
|
workflow_schema_version=plan.workflow_schema_version,
|
|
environment_digest=plan.environment_digest,
|
|
verifier=reducer_stage.verifier,
|
|
selected_features=plan.selected_features,
|
|
optional_fallbacks=plan.optional_fallbacks,
|
|
task_key="reduce/final",
|
|
stage_id=reducer_stage.stage_id,
|
|
parameters=plan.resolved_parameters,
|
|
inputs={reducer_input_name: accepted},
|
|
expected_outputs=reducer_stage.outputs,
|
|
resources=reducer_stage.resources,
|
|
execution=reducer_stage.execution,
|
|
expected_input_keys={
|
|
reducer_input_name: tuple(
|
|
item.key for item in accepted.items if item.key is not None
|
|
)
|
|
},
|
|
).validate_stage(reducer_stage)
|
|
final = self._run_task(
|
|
definition,
|
|
reducer_task,
|
|
run_root / "tasks" / "reduce-final",
|
|
definition.reducers[reducer_stage.entry_point].reduce,
|
|
job_id=job_id,
|
|
transaction=transaction,
|
|
max_artifacts=limits.max_artifacts,
|
|
max_output_bytes=output_limit - output_bytes,
|
|
)
|
|
output_bytes += self._track_artifacts(
|
|
final.outputs,
|
|
known=known_artifacts,
|
|
max_artifacts=limits.max_artifacts,
|
|
output_ids=output_artifact_ids,
|
|
)
|
|
if output_bytes > output_limit:
|
|
raise ValueError("job exceeds the cumulative output byte limit")
|
|
final.validate_against(
|
|
definition.manifest.outputs, max_output_bytes=output_limit
|
|
)
|
|
return final
|