DistributedTaskIR API

DistributedTaskIR owns logical meshes, typed sharding, collective/P2P/local task variants, distributed dependencies, and source lineage. Physical devices and target implementation choices are forbidden here.

Logical distributed program over a virtual device mesh.

NonEmptyText module-attribute

NonEmptyText: TypeAlias = Annotated[str, NON_EMPTY]

NonNegativeInt module-attribute

NonNegativeInt: TypeAlias = Annotated[int, NON_NEGATIVE]

PositiveInt module-attribute

PositiveInt: TypeAlias = Annotated[int, POSITIVE]

Scalar module-attribute

Scalar: TypeAlias = int | FiniteFloat | Symbol | ScalarExprVariant

EMPTY_SEMANTIC module-attribute

EMPTY_SEMANTIC = EmptySemantic()

MeshAxes module-attribute

MeshAxes: TypeAlias = Annotated[tuple[MeshAxis, ...], NON_EMPTY]

CollectiveParticipants module-attribute

CollectiveParticipants: TypeAlias = Annotated[tuple[int, ...], NON_EMPTY, UNIQUE_ITEMS, NON_NEGATIVE_ITEMS]

CollectiveSpecVariant module-attribute

CollectiveSpecVariant: TypeAlias = AllReduce | ReduceScatter | AllGather | AllToAll | Broadcast

TaskBodyVariant module-attribute

TaskBodyVariant: TypeAlias = LocalCompute | Collective | PointToPoint | Reshard | Control

_DISTRIBUTED_RESERVED module-attribute

_DISTRIBUTED_RESERVED = frozenset({'model_spec', 'workload_spec', 'mapping_spec', 'inference_mapping_spec', 'inference_phase', 'batch_size', 'query_tokens', 'context_tokens', 'datatype', 'invocation', 'block_memory', 'scope', 'physical_device', 'device_id', 'route', 'queue', 'stream', 'kernel', 'implementation_id', 'start', 'start_time', 'end', 'end_time', 'duration', 'latency', 'bandwidth'})

__all__ module-attribute

__all__ = ['AllGather', 'AllReduce', 'AllToAll', 'Broadcast', 'Collective', 'CollectiveKind', 'CollectiveSpec', 'CollectiveSpecVariant', 'Control', 'DistributedTask', 'DistributedTaskIR', 'DistributedValue', 'LocalCompute', 'LogicalMesh', 'MeshAxis', 'PeerTransfer', 'PointToPoint', 'ReductionKind', 'ReduceScatter', 'Reshard', 'ShardingSpec', 'TaskBody', 'TaskBodyVariant', 'collective_kind', 'make_collective_spec']

ValueConstraint

Bases: Enum

Small closed vocabulary of reusable structural refinements.

VariantSpec dataclass

VariantSpec(family: ADTSpec, local_tag: str, wire_tag: str, constructor: type[Any])

Manifest entry for one explicitly named ADT constructor.

FrozenDict

FrozenDict(source: Mapping[str, V] | None = None, *, items: Iterable[tuple[str, V]] | None = None)

Bases: Mapping[str, V], Generic[V]

A compact, hashable mapping with recursively frozen values.

Source code in src/blueprinting/schema/frozen.py

def __init__(
    self,
    source: Mapping[str, V] | None = None,
    *,
    items: Iterable[tuple[str, V]] | None = None,
) -> None:
    if source is not None and items is not None:
        raise TypeError("provide either source or items, not both")
    raw_items = source.items() if source is not None else (items or ())
    copied: dict[str, V] = {}
    for key, value in raw_items:
        if not isinstance(key, str):
            raise TypeError("FrozenDict keys must be strings")
        if key in copied:
            raise ValueError(f"duplicate FrozenDict key: {key!r}")
        copied[key] = cast(V, freeze(value))
    self._items: tuple[tuple[str, V], ...] = tuple(sorted(copied.items(), key=lambda pair: pair[0]))
    self._hash: int | None = None

__reduce__

__reduce__() -> tuple[type[FrozenDict[V]], tuple[dict[str, V]]]

Use the public constructor for process and UI cache round-trips.

Source code in src/blueprinting/schema/frozen.py

def __reduce__(self) -> tuple[type[FrozenDict[V]], tuple[dict[str, V]]]:
    """Use the public constructor for process and UI cache round-trips."""

    return FrozenDict, (dict(self._items),)

DiagnosticBag

DiagnosticBag()

Bases: DiagnosticBag

Compatibility builder that publishes VerificationReport.

Source code in src/blueprinting/schema/diagnostics.py

def __init__(self) -> None:
    self._items: list[Diagnostic] = []

VerificationReport dataclass

VerificationReport(diagnostics: tuple[Diagnostic, ...] = ())

Bases: DiagnosticSet

Backward-compatible name for the domain-free immutable diagnostics.

ScalarExpr

Bases: _ExpressionOperators

Closed family of exact scalar expression constructors.

Symbol

Bases: _ExpressionOperators

Lineage

Typed provenance from source entities to one lowering product.

NodeId dataclass

NodeId(value: str)

Bases: StableId

ValueId dataclass

ValueId(value: str)

Bases: StableId

DistributedTaskSemantic

Bases: SemanticPayload

Dialect semantics attached to a DistributedTaskIR task.

ProgramSemantic

Bases: SemanticPayload

Dialect semantics shared by one canonical program snapshot.

CanonicalIRMixin

Behavior shared by immutable canonical IR roots.

to_json

to_json() -> str

Serialize through a self-describing, digest-checked envelope.

Source code in src/blueprinting/synthesizer/stages/common.py

def to_json(self) -> str:
    """Serialize through a self-describing, digest-checked envelope."""

    self.require_valid()
    snapshot = IRSnapshot(
        schema_name=self.header.schema_name,
        schema_version=self.header.schema_version,
        producer_version=self.header.producer_version,
        feature_set=self.header.feature_set,
        content_digest=self.digest,
        payload=self,
    )
    return canonical_dumps(snapshot)

from_json classmethod

from_json(payload: str) -> Checked[IR]

Decode an exact-schema snapshot without exception control flow.

Source code in src/blueprinting/synthesizer/stages/common.py

@classmethod
def from_json(cls: type[IR], payload: str) -> Checked[IR]:
    """Decode an exact-schema snapshot without exception control flow."""

    try:
        return Ok(cls._decode_json(payload))
    except IRVerificationError as error:
        return Err(DiagnosticSet(tuple(error.diagnostics)))
    except SerializationError as error:
        return Err(
            DiagnosticSet.of(
                Diagnostic("serialization.snapshot", str(error), ("snapshot",)),
            )
        )

require_from_json classmethod

require_from_json(payload: str) -> IR

Explicit exception adapter for trusted internal/replay boundaries.

Source code in src/blueprinting/synthesizer/stages/common.py

@classmethod
def require_from_json(cls: type[IR], payload: str) -> IR:
    """Explicit exception adapter for trusted internal/replay boundaries."""

    return cls.from_json(payload).or_raise(
        lambda diagnostics: SerializationError("; ".join(item.render() for item in diagnostics.errors))
    )

load_migrated classmethod

load_migrated(payload: str, *, registry: Any = None) -> Checked[IR]

Load through an explicit registered migration path.

Source code in src/blueprinting/synthesizer/stages/common.py

@classmethod
def load_migrated(cls: type[IR], payload: str, *, registry: Any = None) -> Checked[IR]:
    """Load through an explicit registered migration path."""

    if registry is None:
        from ..schema_migration import DEFAULT_SCHEMA_MIGRATIONS

        registry = DEFAULT_SCHEMA_MIGRATIONS
    try:
        result = registry.migrate_json(
            payload,
            schema_name=cls.SCHEMA_NAME,
            target_version=cls.SCHEMA_VERSION,
        )
    except SerializationError as error:
        return Err(DiagnosticSet.of(Diagnostic("serialization.migration", str(error), ("snapshot",))))
    return cls.from_json(result.payload)

verify

verify() -> Checked[IR]

Return the immutable snapshot or all expected verifier failures.

Source code in src/blueprinting/synthesizer/stages/common.py

def verify(self: IR) -> Checked[IR]:
    """Return the immutable snapshot or all expected verifier failures."""

    return checked(self, self.diagnostics())

Effect

IRHeader

Version and provenance header embedded in every canonical IR.

OperationName

Structured operation identity; dialect is never inferred from a string.

SchemaVersion

Semantic version of one serialized IR schema.

TensorType

Target-neutral logical tensor type.

ValueRole

Bases: Enum

Semantic ownership role of a model-level SSA value.

MeshAxis

One named dimension of the logical, target-neutral device mesh.

LogicalMesh

Cartesian logical-rank space used by sharding and collectives.

ShardingSpec

Mapping from tensor dimensions to logical mesh axes.

CollectiveKind

Bases: Enum

Logical collective semantics independent of a communication library.

ReductionKind

Bases: Enum

Associative reduction operation required by a logical collective.

CollectiveSpec

Closed collective semantics without conditional reduction/root fields.

AllReduce

Bases: CollectiveSpec

ReduceScatter

Bases: CollectiveSpec

AllGather

Bases: CollectiveSpec

AllToAll

Bases: CollectiveSpec

Broadcast

Bases: CollectiveSpec

PeerTransfer

Logical point-to-point transfer between two virtual ranks.

TaskBody

Closed family of mutually exclusive distributed task semantics.

LocalCompute

Bases: TaskBody

Execute a target-neutral operation independently on the logical ranks.

Collective

Bases: TaskBody

Task body carrying a well-formed logical collective specification.

PointToPoint

Bases: TaskBody

Task body carrying one logical peer transfer.

Reshard

Bases: TaskBody

Change logical ownership/sharding through explicit dataflow values.

Control

Bases: TaskBody

Represent a dependency-only logical coordination task.

DistributedValue

Model value specialized with logical ownership and sharding.

DistributedTask

Common graph envelope around one typed distributed task body.

body_tag property

body_tag: str

Stable short constructor identity derived from the ADT manifest.

DistributedTaskIR

Bases: CanonicalIRMixin

Logical task graph whose ranks are virtual, never physical devices.

adt

adt(*, wire: str) -> Callable[[type[T]], type[T]]

Declare the shared semantic wire namespace for a closed sum type.

Source code in src/blueprinting/schema/deriving.py

@dataclass_transform(frozen_default=True)
def adt(*, wire: str) -> Callable[[type[T]], type[T]]:
    """Declare the shared semantic wire namespace for a closed sum type."""

    if not isinstance(wire, str) or _WIRE_RE.fullmatch(wire) is None:
        raise TypeError("ADT wire namespace must be a dotted lowercase identity")

    def decorate(cls: type[T]) -> type[T]:
        if "__dataclass_fields__" in cls.__dict__:
            raise TypeError("@adt derives its own frozen dataclass; do not combine it with @dataclass")
        cls = dataclass(frozen=True, slots=True)(cls)
        if cls in _ADT_SPECS:
            raise RuntimeError(f"ADT family {cls.__name__} is already registered")
        if any(item.wire == wire for item in _ADT_SPECS.values()):
            raise RuntimeError(f"ADT wire namespace {wire!r} is already registered")
        spec = ADTSpec(wire, cls)
        _ADT_SPECS[cls] = spec
        setattr(cls, "__adt_spec__", spec)  # noqa: B010

        root = cls

        def adt_new(constructor: type[Any], *_args: Any, **_kwargs: Any) -> Any:
            if constructor is root:
                raise TypeError(f"ADT family {root.__name__} is abstract; instantiate one of its sealed variants")
            return object.__new__(constructor)

        setattr(cls, "__new__", staticmethod(adt_new))  # noqa: B010
        return cls

    return decorate

enum

enum(tag: str) -> Callable[[type[T]], type[T]]

Register one closed enumeration through the public authoring surface.

Source code in src/blueprinting/schema/deriving.py

def enum(tag: str) -> Callable[[type[T]], type[T]]:
    """Register one closed enumeration through the public authoring surface."""

    if not isinstance(tag, str) or _WIRE_RE.fullmatch(tag) is None:
        raise TypeError("canonical enum tag must be a dotted lowercase wire identity")
    return enum_type(tag)

record

record(tag: str, *, order: bool = False) -> Callable[[type[T]], type[T]]

Derive an immutable canonical record from an annotated class declaration.

Source code in src/blueprinting/schema/deriving.py

@dataclass_transform(frozen_default=True)
def record(
    tag: str,
    *,
    order: bool = False,
) -> Callable[[type[T]], type[T]]:
    """Derive an immutable canonical record from an annotated class declaration."""

    if not isinstance(tag, str) or _WIRE_RE.fullmatch(tag) is None:
        raise TypeError("canonical record tag must be a dotted lowercase wire identity")

    def decorate(cls: type[T]) -> type[T]:
        if "__dataclass_fields__" in cls.__dict__:
            raise TypeError("@record derives its own frozen dataclass; do not combine it with @dataclass")
        _install_structural_post_init(cls)
        cls = dataclass(frozen=True, slots=True, order=order)(cls)
        return record_type(tag)(cls)

    return decorate

seal_adt

seal_adt(family: type[Any], variants: Any) -> Any

Declare the explicit runtime closure corresponding to a static Union alias.

Source code in src/blueprinting/schema/deriving.py

def seal_adt(family: type[Any], variants: Any) -> Any:
    """Declare the explicit runtime closure corresponding to a static Union alias."""

    if family not in _ADT_SPECS:
        raise TypeError(f"{getattr(family, '__name__', family)!r} is not a declared ADT family")
    members = _union_members(variants)
    if any(_VARIANT_SPECS.get(item) is None for item in members):
        raise TypeError("an ADT closure may contain only registered variants")
    if any(_VARIANT_SPECS[item].family.root is not family for item in members):
        raise TypeError("an ADT closure cannot mix constructors from different families")
    if len(set(members)) != len(members):
        raise ValueError("an ADT closure cannot repeat a constructor")
    registered = tuple(item.constructor for item in adt_manifest(family))
    if set(members) != set(registered):
        missing = tuple(item.__name__ for item in registered if item not in members)
        unknown = tuple(item.__name__ for item in members if item not in registered)
        detail = []
        if missing:
            detail.append(f"missing {', '.join(missing)}")
        if unknown:
            detail.append(f"unknown {', '.join(unknown)}")
        raise ValueError(f"ADT closure must contain exactly its registered variants: {'; '.join(detail)}")
    previous = _ADT_CLOSURES.get(family)
    if previous is not None and previous != registered:
        raise RuntimeError(f"ADT family {family.__name__} is already sealed with another closure")
    _ADT_CLOSURES[family] = registered
    return variants

variant

variant(local_tag: str) -> Callable[[type[T]], type[T]]

Derive and register one explicitly named constructor of an ADT family.

Source code in src/blueprinting/schema/deriving.py

@dataclass_transform(frozen_default=True)
def variant(local_tag: str) -> Callable[[type[T]], type[T]]:
    """Derive and register one explicitly named constructor of an ADT family."""

    if not isinstance(local_tag, str) or _LOCAL_TAG_RE.fullmatch(local_tag) is None:
        raise TypeError("variant tag must be a short lowercase kebab-case identity")

    def decorate(cls: type[T]) -> type[T]:
        family = _family_of(cls)
        if family.root in _ADT_CLOSURES:
            raise RuntimeError(f"ADT family {family.root.__name__} is sealed and cannot accept late variants")
        if any(item.family == family and item.local_tag == local_tag for item in _VARIANT_SPECS.values()):
            raise RuntimeError(f"ADT family {family.root.__name__} already declares variant {local_tag!r}")
        constructor = record(family.tag(local_tag))(cls)
        spec = VariantSpec(family, local_tag, family.tag(local_tag), constructor)
        _VARIANT_SPECS[constructor] = spec
        setattr(constructor, "__variant_spec__", spec)  # noqa: B010
        return constructor

    return decorate

is_content_digest

is_content_digest(value: str) -> bool

Source code in src/blueprinting/synthesizer/stages/common.py

def is_content_digest(value: str) -> bool:
    return isinstance(value, str) and _DIGEST_RE.fullmatch(value) is not None

is_known_target_dialect

is_known_target_dialect(operation: OperationName) -> bool

Recognize built-in target dialects forbidden before target binding.

Source code in src/blueprinting/synthesizer/stages/common.py

def is_known_target_dialect(operation: OperationName) -> bool:
    """Recognize built-in target dialects forbidden before target binding."""

    return operation.dialect.lower() in _KNOWN_TARGET_DIALECTS

make_header

make_header(schema_name: str, schema_version: SchemaVersion, *, parent_digests: Iterable[str] = (), features: Iterable[str] = (), producer_version: str = '0.0.0') -> IRHeader

Source code in src/blueprinting/synthesizer/stages/common.py

def make_header(
    schema_name: str,
    schema_version: SchemaVersion,
    *,
    parent_digests: Iterable[str] = (),
    features: Iterable[str] = (),
    producer_version: str = "0.0.0",
) -> IRHeader:
    return IRHeader(
        schema_name=schema_name,
        schema_version=schema_version,
        producer_version=producer_version,
        feature_set=REQUIRED_IR_FEATURES | frozenset(features),
        parent_digests=tuple(parent_digests),
    )

reject_reserved_attributes

reject_reserved_attributes(bag: DiagnosticBag, attributes: FrozenDict, reserved: frozenset[str], *path: str) -> None

Reject semantic fields smuggled through extension dictionaries.

Source code in src/blueprinting/synthesizer/stages/common.py

def reject_reserved_attributes(
    bag: DiagnosticBag,
    attributes: FrozenDict,
    reserved: frozenset[str],
    *path: str,
) -> None:
    """Reject semantic fields smuggled through extension dictionaries."""

    def walk(value: Any, location: tuple[str, ...]) -> None:
        if isinstance(value, FrozenDict):
            for key, item in value.items():
                if key.lower() in reserved:
                    bag.error(
                        "attribute.reserved",
                        f"{key!r} is a semantic field and is illegal in this IR layer",
                        *(location + (key,)),
                    )
                walk(item, location + (key,))
        elif isinstance(value, (tuple, frozenset)):
            values = value if isinstance(value, tuple) else tuple(sorted(value, key=repr))
            for index, item in enumerate(values):
                walk(item, location + (str(index),))

    walk(attributes, tuple(path))

verify_known_references

verify_known_references(bag: DiagnosticBag, references: Iterable[StableId], known: Iterable[StableId], *path: str) -> None

Source code in src/blueprinting/synthesizer/stages/common.py

def verify_known_references(
    bag: DiagnosticBag,
    references: Iterable[StableId],
    known: Iterable[StableId],
    *path: str,
) -> None:
    known_set = set(known)
    for reference in references:
        if reference not in known_set:
            bag.error("reference.unknown", f"unknown reference {reference}", *path)

verify_nonnegative_scalar

verify_nonnegative_scalar(bag: DiagnosticBag, value: Scalar, *path: str) -> None

Source code in src/blueprinting/synthesizer/stages/common.py

def verify_nonnegative_scalar(bag: DiagnosticBag, value: Scalar, *path: str) -> None:
    if isinstance(value, bool):
        bag.error("scalar.invalid", "boolean is not a scalar quantity", *path)
    elif not (isinstance(value, (int, float, Symbol)) or is_adt_variant(value, ScalarExpr)):
        bag.error("scalar.invalid", f"unsupported scalar quantity {value!r}", *path)
    elif isinstance(value, float) and not math.isfinite(value):
        bag.error("scalar.nonfinite", "quantity must be finite", *path)
    elif isinstance(value, (int, float)) and value < 0:
        bag.error("scalar.negative", "quantity must not be negative", *path)

verify_ordered_dag

verify_ordered_dag(bag: DiagnosticBag, entities: Sequence[Entity], id_of: Callable[[Entity], StableId], dependencies_of: Callable[[Entity], Iterable[StableId]], path: str) -> None

Verify references and require canonical topological sequence order.

Source code in src/blueprinting/synthesizer/stages/common.py

def verify_ordered_dag(
    bag: DiagnosticBag,
    entities: Sequence[Entity],
    id_of: Callable[[Entity], StableId],
    dependencies_of: Callable[[Entity], Iterable[StableId]],
    path: str,
) -> None:
    """Verify references and require canonical topological sequence order."""

    identifiers = {id_of(entity) for entity in entities}
    seen = set()
    for index, entity in enumerate(entities):
        identifier = id_of(entity)
        dependencies = tuple(dependencies_of(entity))
        if len(set(dependencies)) != len(dependencies):
            bag.error("dag.duplicate_dependency", "dependencies must be unique", path, str(index), "dependencies")
        if identifier in dependencies:
            bag.error("dag.self_dependency", f"{identifier} depends on itself", path, str(index), "dependencies")
        for dependency in dependencies:
            if dependency not in identifiers:
                bag.error(
                    "reference.unknown",
                    f"dependency {dependency} is not declared",
                    path,
                    str(index),
                    "dependencies",
                )
            elif dependency not in seen:
                bag.error(
                    "dag.not_topological",
                    f"dependency {dependency} must precede {identifier}",
                    path,
                    str(index),
                    "dependencies",
                    hint="store canonical DAG nodes in topological order",
                )
        seen.add(identifier)

verify_unique_ids

verify_unique_ids(bag: DiagnosticBag, entities: Sequence[Entity], id_of: Callable[[Entity], StableId], path: str) -> None

Source code in src/blueprinting/synthesizer/stages/common.py

def verify_unique_ids(
    bag: DiagnosticBag,
    entities: Sequence[Entity],
    id_of: Callable[[Entity], StableId],
    path: str,
) -> None:
    seen = set()
    for index, entity in enumerate(entities):
        identifier = id_of(entity)
        if identifier in seen:
            bag.error("id.duplicate", f"duplicate identifier {identifier}", path, str(index), "id")
        seen.add(identifier)

collective_kind

collective_kind(spec: CollectiveSpecVariant) -> CollectiveKind

Source code in src/blueprinting/synthesizer/stages/distributed/ir.py

def collective_kind(spec: CollectiveSpecVariant) -> CollectiveKind:
    match spec:
        case AllReduce():
            return CollectiveKind.ALL_REDUCE
        case ReduceScatter():
            return CollectiveKind.REDUCE_SCATTER
        case AllGather():
            return CollectiveKind.ALL_GATHER
        case AllToAll():
            return CollectiveKind.ALL_TO_ALL
        case Broadcast():
            return CollectiveKind.BROADCAST
    assert_never(spec)

make_collective_spec

make_collective_spec(kind: CollectiveKind, participants: tuple[int, ...], message_bytes: Scalar, *, reduction: ReductionKind | None = None, root: int | None = None) -> CollectiveSpecVariant

Boundary adapter from enum-oriented inputs into the canonical ADT.

Source code in src/blueprinting/synthesizer/stages/distributed/ir.py

def make_collective_spec(
    kind: CollectiveKind,
    participants: tuple[int, ...],
    message_bytes: Scalar,
    *,
    reduction: ReductionKind | None = None,
    root: int | None = None,
) -> CollectiveSpecVariant:
    """Boundary adapter from enum-oriented inputs into the canonical ADT."""

    match kind:
        case CollectiveKind.ALL_REDUCE:
            if reduction is None or root is not None:
                raise ValueError("all-reduce requires reduction and does not accept root")
            return AllReduce(participants, message_bytes, reduction)
        case CollectiveKind.REDUCE_SCATTER:
            if reduction is None or root is not None:
                raise ValueError("reduce-scatter requires reduction and does not accept root")
            return ReduceScatter(participants, message_bytes, reduction)
        case CollectiveKind.ALL_GATHER:
            if reduction is not None or root is not None:
                raise ValueError("all-gather does not accept reduction or root")
            return AllGather(participants, message_bytes)
        case CollectiveKind.ALL_TO_ALL:
            if reduction is not None or root is not None:
                raise ValueError("all-to-all does not accept reduction or root")
            return AllToAll(participants, message_bytes)
        case CollectiveKind.BROADCAST:
            if reduction is not None or root is None:
                raise ValueError("broadcast requires root and does not accept reduction")
            return Broadcast(participants, message_bytes, root)
    assert_never(kind)