diff --git a/amber/LICENSE-binary-python b/amber/LICENSE-binary-python index 976c2f702d7..e785d1bdf25 100644 --- a/amber/LICENSE-binary-python +++ b/amber/LICENSE-binary-python @@ -290,8 +290,8 @@ Dependencies under the BSD 3-Clause License Python packages: - cached-property==2.0.1 - - click==8.4.2 - cloudpickle==3.1.2 + - click==8.4.2 - contourpy==1.3.3 - cycler==0.12.1 - fsspec==2026.6.0 diff --git a/amber/requirements.txt b/amber/requirements.txt index 25c30339eaa..2e59fd56218 100644 --- a/amber/requirements.txt +++ b/amber/requirements.txt @@ -25,6 +25,7 @@ overrides==7.7.0 typing_extensions==4.14.1 bidict==0.22.0 cached_property==2.0.1 +cloudpickle==3.1.2 psutil==7.2.2 tzlocal==2.1 # Not imported directly: s3fs (with aiobotocore) backs pyiceberg's diff --git a/amber/src/main/python/pytexera/workflow/__init__.py b/amber/src/main/python/pytexera/workflow/__init__.py new file mode 100644 index 00000000000..10e4e67edd7 --- /dev/null +++ b/amber/src/main/python/pytexera/workflow/__init__.py @@ -0,0 +1,23 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Minimal keyed workflow runtime used by generated Python UDFs.""" + +from pytexera.workflow.operators import TupleOperator +from pytexera.workflow.runtime import Heap, InputPort, Runtime + +__all__ = ["Heap", "InputPort", "Runtime", "TupleOperator"] diff --git a/amber/src/main/python/pytexera/workflow/codec.py b/amber/src/main/python/pytexera/workflow/codec.py new file mode 100644 index 00000000000..1dcc86951b5 --- /dev/null +++ b/amber/src/main/python/pytexera/workflow/codec.py @@ -0,0 +1,130 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Cloudpickle transport for explicit workflow boundary payloads.""" + +from __future__ import annotations + +import pickle +from dataclasses import dataclass + +import cloudpickle + + +@dataclass(frozen=True, order=True) +class BoundaryPayload: + """One boundary contract, its path-present fields, and encoded values.""" + + boundary_id: str + fields: tuple[str, ...] + present: tuple[str, ...] + payload: bytes + + def __post_init__(self) -> None: + if not self.boundary_id: + raise ValueError("boundary ID must be nonempty") + if self.fields != tuple(sorted(set(self.fields))) or any( + not field.isidentifier() for field in self.fields + ): + raise ValueError("boundary fields must be canonical Python names") + if self.present != tuple( + field for field in self.fields if field in frozenset(self.present) + ): + raise ValueError("present fields must be a canonical contract subset") + if not isinstance(self.payload, bytes): + raise TypeError("boundary payload must be bytes") + + +@dataclass(frozen=True) +class WorkflowEnvelope: + """All selected boundary payloads for one independent execution key.""" + + execution_key: str + boundaries: tuple[BoundaryPayload, ...] + + def __post_init__(self) -> None: + if not self.execution_key: + raise ValueError("execution key must be nonempty") + if not isinstance(self.boundaries, tuple) or any( + not isinstance(row, BoundaryPayload) for row in self.boundaries + ): + raise TypeError("workflow boundaries must be a typed tuple") + ids = tuple(row.boundary_id for row in self.boundaries) + if ids != tuple(sorted(set(ids))): + raise ValueError("workflow boundaries must be canonical and unique") + + +def encode_boundary( + boundary_id: str, + fields: tuple[str, ...], + values: tuple[object, ...], + *, + present: tuple[str, ...] | None = None, +) -> BoundaryPayload: + """Encode values present on this path under one selected field contract.""" + + present = fields if present is None else present + if len(present) != len(values): + raise ValueError("present boundary fields and values must have equal length") + payload = cloudpickle.dumps(values, protocol=pickle.HIGHEST_PROTOCOL) + return BoundaryPayload(boundary_id, fields, present, payload) + + +def decode_boundary( + boundary: BoundaryPayload, + fields: tuple[str, ...], +) -> tuple[object, ...]: + """Decode a payload only under its exact selected field contract.""" + + if boundary.fields != fields: + raise ValueError("boundary fields do not match the requested contract") + values = cloudpickle.loads(boundary.payload) + if not isinstance(values, tuple) or len(values) != len(boundary.present): + raise ValueError("decoded boundary payload has an invalid shape") + return values + + +def dumps_envelope(envelope: WorkflowEnvelope) -> bytes: + """Validate and encode a workflow envelope.""" + + if not isinstance(envelope, WorkflowEnvelope): + raise TypeError("envelope codec requires WorkflowEnvelope") + return cloudpickle.dumps(envelope, protocol=pickle.HIGHEST_PROTOCOL) + + +def loads_envelope(payload: bytes) -> WorkflowEnvelope: + """Decode and type-check one workflow envelope.""" + + envelope = cloudpickle.loads(payload) + if not isinstance(envelope, WorkflowEnvelope): + raise TypeError("decoded payload is not WorkflowEnvelope") + return envelope + + +def merge_envelopes( + left: WorkflowEnvelope, + right: WorkflowEnvelope, +) -> WorkflowEnvelope: + """Merge independent fan-in payloads for the same execution key.""" + + if left.execution_key != right.execution_key: + raise ValueError("cannot merge envelopes with different execution keys") + boundaries = (*left.boundaries, *right.boundaries) + return WorkflowEnvelope( + left.execution_key, + tuple(sorted(boundaries, key=lambda row: row.boundary_id)), + ) diff --git a/amber/src/main/python/pytexera/workflow/operators.py b/amber/src/main/python/pytexera/workflow/operators.py new file mode 100644 index 00000000000..50775494f15 --- /dev/null +++ b/amber/src/main/python/pytexera/workflow/operators.py @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Generic keyed fan-in adapter between Amber tuples and workflow Runtime.""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass + +from pytexera.udf.udf_operator import UDFOperatorV2 +from pytexera.workflow.codec import ( + WorkflowEnvelope, + dumps_envelope, + loads_envelope, + merge_envelopes, +) +from pytexera.workflow.runtime import Runtime + +ENVELOPE_FIELD = "workflow_envelope" +EXECUTION_KEY_FIELD = "execution_key" + + +@dataclass(frozen=True) +class _PendingExecution: + envelope: WorkflowEnvelope + arrived_ports: frozenset[int] + + +class TupleOperator(UDFOperatorV2): + """Execute once after all configured boundaries for one key have arrived.""" + + runtime: Runtime + + def __init__(self) -> None: + super().__init__() + self._pending: dict[str, _PendingExecution] = {} + self._finished_ports: set[int] = set() + + def process_tuple(self, tuple_, port: int) -> Iterator[dict[str, object]]: + """Join one keyed port arrival and execute when all ports are present.""" + + if not self.runtime.input_ports: + key = _field(tuple_, EXECUTION_KEY_FIELD, "default") + yield _output(self.runtime.execute(WorkflowEnvelope(str(key), ()))) + return + incoming = self._incoming_envelope(tuple_, port) + pending = self._merge_pending(incoming, port) + if len(pending.arrived_ports) != len(self.runtime.input_ports): + return + del self._pending[incoming.execution_key] + yield _output(self.runtime.execute(pending.envelope)) + + def _incoming_envelope(self, tuple_, port: int) -> WorkflowEnvelope: + """Validate and retain only boundaries owned by one physical input port.""" + + _validate_port(port, len(self.runtime.input_ports)) + if port in self._finished_ports: + raise RuntimeError(f"input arrived after port {port} finished") + raw = _field(tuple_, ENVELOPE_FIELD) + envelope = loads_envelope(raw) + relevant = tuple( + row + for row in envelope.boundaries + if row.boundary_id in self.runtime.incoming + ) + expected = frozenset(self.runtime.input_ports[port].boundaries) + found = frozenset(row.boundary_id for row in relevant) + if found != expected: + raise ValueError( + f"port {port} boundaries differ from its Runtime contract: " + f"expected {tuple(sorted(expected))!r}, found {tuple(sorted(found))!r}" + ) + return WorkflowEnvelope(envelope.execution_key, relevant) + + def _merge_pending( + self, + incoming: WorkflowEnvelope, + port: int, + ) -> _PendingExecution: + """Merge one validated arrival into its execution-key accumulator.""" + + execution_key = incoming.execution_key + previous = self._pending.get(execution_key) + if previous is not None and port in previous.arrived_ports: + raise RuntimeError( + f"execution {execution_key!r} arrived twice on port {port}" + ) + merged = ( + incoming + if previous is None + else merge_envelopes(previous.envelope, incoming) + ) + arrived = ( + frozenset((port,)) if previous is None else previous.arrived_ports | {port} + ) + pending = _PendingExecution(merged, arrived) + self._pending[execution_key] = pending + return pending + + def on_finish(self, port: int) -> Iterator[None]: + """Close one input port and reject incomplete keyed executions.""" + + if not self.runtime.input_ports: + return + yield None + _validate_port(port, len(self.runtime.input_ports)) + if port in self._finished_ports: + raise RuntimeError(f"input port {port} finished twice") + self._finished_ports.add(port) + if len(self._finished_ports) == len(self.runtime.input_ports) and self._pending: + keys = tuple(sorted(self._pending)) + raise RuntimeError(f"incomplete workflow executions at finish: {keys!r}") + return + yield None + + +def _validate_port(port: int, count: int) -> None: + if not isinstance(port, int) or isinstance(port, bool) or not 0 <= port < count: + raise ValueError(f"input port must be an integer in [0, {count})") + + +def _field(tuple_, name: str, default=...) -> object: + try: + return tuple_[name] + except (KeyError, IndexError): + if default is ...: + raise ValueError(f"input tuple is missing field {name!r}") from None + return default + + +def _output(envelope: WorkflowEnvelope) -> dict[str, object]: + return {ENVELOPE_FIELD: dumps_envelope(envelope)} diff --git a/amber/src/main/python/pytexera/workflow/runtime.py b/amber/src/main/python/pytexera/workflow/runtime.py new file mode 100644 index 00000000000..484745056b4 --- /dev/null +++ b/amber/src/main/python/pytexera/workflow/runtime.py @@ -0,0 +1,286 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Stateless driver execution over local heaps and explicit boundaries.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass + +from pytexera.workflow.codec import ( + BoundaryPayload, + WorkflowEnvelope, + decode_boundary, + encode_boundary, +) + +Driver = Callable[["Heap"], "Heap"] + + +@dataclass(frozen=True) +class InputPort: + """Boundary IDs delivered together by one physical input port.""" + + boundaries: tuple[str, ...] + + def __post_init__(self) -> None: + object.__setattr__( + self, + "boundaries", + _boundary_ids(self.boundaries, "input-port"), + ) + if not self.boundaries: + raise ValueError("input port must own at least one boundary") + + +class Heap: + """One process-local namespace whose user values live behind string keys.""" + + __slots__ = ( + "_incoming", + "_outgoing", + "_values", + "_field_owners", + ) + + def __init__(self, envelope: WorkflowEnvelope) -> None: + object.__setattr__( + self, + "_incoming", + {row.boundary_id: row for row in envelope.boundaries}, + ) + object.__setattr__(self, "_outgoing", {}) + object.__setattr__(self, "_values", {}) + object.__setattr__(self, "_field_owners", {}) + + def __getattr__(self, name: str) -> object: + """Read one non-reserved source binding through attribute syntax.""" + + return _read_value(self._values, name) + + def __setattr__(self, name: str, value: object) -> None: + """Write one non-reserved source binding through attribute syntax.""" + + _write_value(self._values, name, value) + + def __delattr__(self, name: str) -> None: + """Delete one non-reserved source binding with Python name semantics.""" + + _delete_value(self._values, name) + + def __getitem__(self, name: str) -> object: + """Read a source binding whose name collides with the Heap API.""" + + return _read_value(self._values, name) + + def __setitem__(self, name: str, value: object) -> None: + """Write a source binding whose name collides with the Heap API.""" + + _write_value(self._values, name, value) + + def __delitem__(self, name: str) -> None: + """Delete a source binding whose name collides with the Heap API.""" + + _delete_value(self._values, name) + + +class Runtime: + """Configured boundary contract and exactly one generated driver.""" + + def __init__( + self, + *, + input_ports: tuple[InputPort, ...] = (), + outgoing: tuple[str, ...] = (), + ) -> None: + if not isinstance(input_ports, tuple) or any( + not isinstance(port, InputPort) for port in input_ports + ): + raise TypeError("input_ports must contain InputPort values") + incoming = tuple( + boundary for port in input_ports for boundary in port.boundaries + ) + if len(incoming) != len(set(incoming)): + raise ValueError( + "each incoming boundary must belong to exactly one input port" + ) + self.input_ports = input_ports + self.incoming = incoming + self.outgoing = _boundary_ids(outgoing, "outgoing") + self._driver: Driver | None = None + + def driver(self, function: Driver, /) -> Driver: + """Register the generated driver and return it as a decorator.""" + + if self._driver is not None: + raise RuntimeError("workflow Runtime already has a driver") + self._driver = function + return function + + def import_boundary( + self, + heap: Heap, + boundary_id: str, + fields: tuple[str, ...], + ) -> None: + """Restore exactly the values selected for one incoming boundary.""" + + if boundary_id not in self.incoming: + raise ValueError(f"boundary {boundary_id!r} is not incoming") + try: + boundary = heap._incoming[boundary_id] + except KeyError as error: + raise RuntimeError( + f"incoming boundary {boundary_id!r} is missing" + ) from error + values = decode_boundary(boundary, fields) + _claim_fields(heap, boundary_id, fields) + for field, value in zip( + boundary.present, + values, + strict=True, + ): + _write_value(heap._values, field, value) + + def export_boundary( + self, + heap: Heap, + boundary_id: str, + fields: tuple[str, ...], + namespace: Mapping[str, object] | None = None, + ) -> None: + """Encode exactly the values selected for one outgoing boundary.""" + + if boundary_id not in self.outgoing: + raise ValueError(f"boundary {boundary_id!r} is not outgoing") + if boundary_id in heap._outgoing: + raise RuntimeError(f"outgoing boundary {boundary_id!r} was exported twice") + _merge_namespace(heap, fields, namespace) + present = tuple( + field for field in fields if _contains_value(heap._values, field) + ) + values = tuple(_read_value(heap._values, field) for field in present) + heap._outgoing[boundary_id] = encode_boundary( + boundary_id, + fields, + values, + present=present, + ) + + def execute(self, envelope: WorkflowEnvelope) -> WorkflowEnvelope: + """Run the driver locally and return only newly exported boundaries.""" + + if self._driver is None: + raise RuntimeError("workflow Runtime has no driver") + available = frozenset(row.boundary_id for row in envelope.boundaries) + missing = set(self.incoming) - available + if missing: + raise RuntimeError( + f"incoming boundaries are missing: {tuple(sorted(missing))}" + ) + heap = Heap(envelope) + returned = self._driver(heap) + if returned is not heap: + raise RuntimeError("workflow driver must return its input Heap") + produced = frozenset(heap._outgoing) + if produced != frozenset(self.outgoing): + raise RuntimeError( + "outgoing boundaries differ from the Runtime contract: " + f"expected {self.outgoing!r}, found {tuple(sorted(produced))!r}" + ) + rows: tuple[BoundaryPayload, ...] = tuple( + heap._outgoing[boundary_id] for boundary_id in self.outgoing + ) + return WorkflowEnvelope(envelope.execution_key, rows) + + +def _boundary_ids(values: tuple[str, ...], label: str) -> tuple[str, ...]: + if not isinstance(values, tuple) or values != tuple(sorted(set(values))): + raise ValueError(f"{label} boundary IDs must be canonical and unique") + if any(not value for value in values): + raise ValueError(f"{label} boundary IDs must be nonempty") + return values + + +def _validate_field_name(name: str) -> None: + """Reject ambiguous or non-string keys at the Heap API boundary.""" + + if not isinstance(name, str): + raise TypeError("heap field name must be a string") + + +def _read_value(values: dict[str, object], name: str, /) -> object: + """Read one source binding with Python's missing-name semantics.""" + + _validate_field_name(name) + try: + return values[name] + except KeyError as error: + raise NameError(f"name {name!r} is not defined") from error + + +def _write_value(values: dict[str, object], name: str, value: object, /) -> None: + """Write one exact source binding into private Heap storage.""" + + _validate_field_name(name) + values[name] = value + + +def _contains_value(values: dict[str, object], name: str, /) -> bool: + """Return whether one exact source binding exists.""" + + _validate_field_name(name) + return name in values + + +def _claim_fields(heap: Heap, boundary_id: str, fields: tuple[str, ...], /) -> None: + """Atomically assign imported fields to one incoming boundary.""" + + for field in fields: + owner = heap._field_owners.get(field) + if owner is not None and owner != boundary_id: + raise RuntimeError( + f"heap field {field!r} is already owned by boundary {owner!r}" + ) + for field in fields: + heap._field_owners[field] = boundary_id + + +def _delete_value(values: dict[str, object], name: str, /) -> None: + """Delete one exact source binding or reproduce Python's NameError.""" + + _validate_field_name(name) + try: + del values[name] + except KeyError as error: + raise NameError(f"name {name!r} is not defined") from error + + +def _merge_namespace( + heap: Heap, + fields: tuple[str, ...], + namespace: Mapping[str, object] | None, + /, +) -> None: + """Publish selected ordinary locals without generated per-field branches.""" + + if namespace is None: + return + for field in fields: + if field in namespace: + _write_value(heap._values, field, namespace[field]) diff --git a/amber/src/test/python/pytexera/workflow/test_codec.py b/amber/src/test/python/pytexera/workflow/test_codec.py new file mode 100644 index 00000000000..bee2dd64c0d --- /dev/null +++ b/amber/src/test/python/pytexera/workflow/test_codec.py @@ -0,0 +1,82 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pytest +from pytexera.workflow.codec import ( + BoundaryPayload, + WorkflowEnvelope, + decode_boundary, + dumps_envelope, + encode_boundary, + loads_envelope, + merge_envelopes, +) + + +def test_boundary_cloudpickle_preserves_aliases_and_cycles() -> None: + shared = [] + shared.append(shared) + + boundary = encode_boundary("edge", ("left", "right"), (shared, shared)) + values = decode_boundary(boundary, ("left", "right")) + + assert values[0] is values[1] + assert values[0][0] is values[0] + + +def test_envelope_round_trip_contains_only_explicit_boundaries() -> None: + boundary = encode_boundary("edge", ("value",), ({"large": [1, 2, 3]},)) + envelope = WorkflowEnvelope("run-1", (boundary,)) + + decoded = loads_envelope(dumps_envelope(envelope)) + + assert decoded == envelope + assert decoded.boundaries[0].fields == ("value",) + + +def test_decode_rejects_field_contract_mismatch() -> None: + boundary = encode_boundary("edge", ("value",), (1,)) + + with pytest.raises(ValueError, match="fields"): + decode_boundary(boundary, ("other",)) + + +def test_boundary_cloudpickle_preserves_an_absent_selected_field() -> None: + """The wire contract and the values present on one path remain distinct.""" + + boundary = encode_boundary( + "edge", + ("left", "right"), + (41,), + present=("left",), + ) + + assert boundary.fields == ("left", "right") + assert boundary.present == ("left",) + assert decode_boundary(boundary, ("left", "right")) == (41,) + + +def test_envelope_rejects_duplicate_boundaries_and_cross_key_merge() -> None: + boundary = BoundaryPayload("edge", ("value",), ("value",), b"payload") + with pytest.raises(ValueError, match="canonical"): + WorkflowEnvelope("run", (boundary, boundary)) + + with pytest.raises(ValueError, match="execution key"): + merge_envelopes( + WorkflowEnvelope("left", ()), + WorkflowEnvelope("right", ()), + ) diff --git a/amber/src/test/python/pytexera/workflow/test_operators.py b/amber/src/test/python/pytexera/workflow/test_operators.py new file mode 100644 index 00000000000..728a7294b91 --- /dev/null +++ b/amber/src/test/python/pytexera/workflow/test_operators.py @@ -0,0 +1,171 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pytest +from pytexera.workflow.codec import ( + WorkflowEnvelope, + decode_boundary, + dumps_envelope, + encode_boundary, + loads_envelope, +) +from pytexera.workflow.operators import ENVELOPE_FIELD, TupleOperator +from pytexera.workflow.runtime import Heap, InputPort, Runtime + + +def test_workflow_tuple_field_is_runtime_neutral() -> None: + """The tuple field name is independent of a workflow generator.""" + + assert ENVELOPE_FIELD == "workflow_envelope" + + +def test_tuple_operator_assembles_fan_in_by_execution_key() -> None: + runtime = Runtime( + input_ports=(InputPort(("left",)), InputPort(("right",))), + outgoing=("result",), + ) + + @runtime.driver + def driver(heap: Heap) -> Heap: + runtime.import_boundary(heap, "left", ("left_value",)) + runtime.import_boundary(heap, "right", ("right_value",)) + heap.total = heap.left_value + heap.right_value + runtime.export_boundary(heap, "result", ("total",)) + return heap + + class Operator(TupleOperator): + pass + + Operator.runtime = runtime + operator = Operator() + left = WorkflowEnvelope( + "run", + (encode_boundary("left", ("left_value",), (20,)),), + ) + right = WorkflowEnvelope( + "run", + (encode_boundary("right", ("right_value",), (22,)),), + ) + + assert list(operator.process_tuple({ENVELOPE_FIELD: dumps_envelope(left)}, 0)) == [] + outputs = list(operator.process_tuple({ENVELOPE_FIELD: dumps_envelope(right)}, 1)) + + assert len(outputs) == 1 + envelope = loads_envelope(outputs[0][ENVELOPE_FIELD]) + assert decode_boundary(envelope.boundaries[0], ("total",)) == (42,) + + +def test_tuple_operator_allows_one_input_to_finish_before_another() -> None: + runtime = Runtime( + input_ports=(InputPort(("left",)), InputPort(("right",))), + outgoing=("result",), + ) + + @runtime.driver + def driver(heap: Heap) -> Heap: + runtime.import_boundary(heap, "left", ("left_value",)) + runtime.import_boundary(heap, "right", ("right_value",)) + heap.total = heap.left_value + heap.right_value + runtime.export_boundary(heap, "result", ("total",)) + return heap + + class Operator(TupleOperator): + pass + + Operator.runtime = runtime + operator = Operator() + left = WorkflowEnvelope( + "run", + (encode_boundary("left", ("left_value",), (20,)),), + ) + right = WorkflowEnvelope( + "run", + (encode_boundary("right", ("right_value",), (22,)),), + ) + + assert list(operator.process_tuple({ENVELOPE_FIELD: dumps_envelope(left)}, 0)) == [] + assert list(operator.on_finish(0)) == [] + outputs = list(operator.process_tuple({ENVELOPE_FIELD: dumps_envelope(right)}, 1)) + assert len(outputs) == 1 + assert list(operator.on_finish(1)) == [] + + +def test_tuple_operator_fails_after_all_ports_finish_with_incomplete_key() -> None: + runtime = Runtime( + input_ports=(InputPort(("left",)), InputPort(("right",))), + ) + + @runtime.driver + def driver(heap: Heap) -> Heap: + return heap + + class Operator(TupleOperator): + pass + + Operator.runtime = runtime + operator = Operator() + left = WorkflowEnvelope( + "run", + (encode_boundary("left", ("left_value",), (20,)),), + ) + + assert list(operator.process_tuple({ENVELOPE_FIELD: dumps_envelope(left)}, 0)) == [] + assert list(operator.on_finish(0)) == [] + with pytest.raises(RuntimeError, match="incomplete"): + list(operator.on_finish(1)) + + +def test_tuple_operator_rejects_boundaries_arriving_on_the_wrong_port() -> None: + runtime = Runtime( + input_ports=(InputPort(("left",)), InputPort(("right",))), + ) + + @runtime.driver + def driver(heap: Heap) -> Heap: + return heap + + class Operator(TupleOperator): + pass + + Operator.runtime = runtime + right = WorkflowEnvelope( + "run", + (encode_boundary("right", ("right_value",), (22,)),), + ) + + with pytest.raises(ValueError, match="port 0"): + list(Operator().process_tuple({ENVELOPE_FIELD: dumps_envelope(right)}, 0)) + + +def test_entry_tuple_operator_executes_without_incoming_boundaries() -> None: + runtime = Runtime(outgoing=("result",)) + + @runtime.driver + def driver(heap: Heap) -> Heap: + heap.value = 7 + runtime.export_boundary(heap, "result", ("value",)) + return heap + + class Operator(TupleOperator): + pass + + Operator.runtime = runtime + output = list(Operator().process_tuple({"execution_key": "run"}, 0)) + + assert len(output) == 1 + envelope = loads_envelope(output[0][ENVELOPE_FIELD]) + assert envelope.execution_key == "run" diff --git a/amber/src/test/python/pytexera/workflow/test_runtime.py b/amber/src/test/python/pytexera/workflow/test_runtime.py new file mode 100644 index 00000000000..9e882130229 --- /dev/null +++ b/amber/src/test/python/pytexera/workflow/test_runtime.py @@ -0,0 +1,206 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pickle + +import pytest +from pytexera.workflow.codec import WorkflowEnvelope, decode_boundary, encode_boundary +from pytexera.workflow.runtime import Heap, InputPort, Runtime + + +def test_runtime_imports_and_exports_only_declared_fields() -> None: + runtime = Runtime(input_ports=(InputPort(("input",)),), outgoing=("output",)) + + @runtime.driver + def driver(heap: Heap) -> Heap: + runtime.import_boundary(heap, "input", ("value",)) + heap.result = heap.value + 1 + heap.unselected = "must not cross" + runtime.export_boundary(heap, "output", ("result",)) + return heap + + inbound = WorkflowEnvelope( + "run", + (encode_boundary("input", ("value",), (41,)),), + ) + + outbound = runtime.execute(inbound) + + assert tuple(row.boundary_id for row in outbound.boundaries) == ("output",) + assert decode_boundary(outbound.boundaries[0], ("result",)) == (42,) + + +def test_heap_exposes_ordinary_fields_as_a_python_namespace() -> None: + heap = Heap(WorkflowEnvelope("run", ())) + heap.summary = 41 + + assert heap.summary == 41 + heap.summary = 42 + assert heap.summary == 42 + del heap.summary + with pytest.raises(NameError, match="summary"): + _ = heap.summary + + +def test_export_collects_local_values_without_generated_presence_branches() -> None: + runtime = Runtime(outgoing=("edge",)) + + @runtime.driver + def driver(heap: Heap) -> Heap: + present = 42 + runtime.export_boundary( + heap, + "edge", + ("missing", "present"), + locals(), + ) + return heap + + boundary = runtime.execute(WorkflowEnvelope("run", ())).boundaries[0] + + assert boundary.present == ("present",) + assert decode_boundary(boundary, ("missing", "present")) == (42,) + + +def test_runtime_round_trips_only_fields_present_on_this_path() -> None: + """An absent selected field stays absent instead of becoming a failure.""" + + producer = Runtime(outgoing=("edge",)) + + @producer.driver + def produce(heap: Heap) -> Heap: + heap.left = 41 + producer.export_boundary(heap, "edge", ("left", "right")) + return heap + + exported = producer.execute(WorkflowEnvelope("run", ())) + boundary = exported.boundaries[0] + assert boundary.present == ("left",) + + consumer = Runtime(input_ports=(InputPort(("edge",)),)) + + @consumer.driver + def consume(heap: Heap) -> Heap: + consumer.import_boundary(heap, "edge", ("left", "right")) + heap.result = heap.left + 1 + with pytest.raises(NameError, match="right"): + _ = heap.right + return heap + + consumed = consumer.execute(exported) + + assert consumed.boundaries == () + + +def test_runtime_rejects_two_boundaries_claiming_the_same_field() -> None: + """Fan-in field ownership is explicit and independent of import order.""" + + runtime = Runtime( + input_ports=(InputPort(("left",)), InputPort(("right",))), + ) + + @runtime.driver + def driver(heap: Heap) -> Heap: + runtime.import_boundary(heap, "left", ("value",)) + runtime.import_boundary(heap, "right", ("value",)) + return heap + + envelope = WorkflowEnvelope( + "run", + ( + encode_boundary("left", ("value",), (1,)), + encode_boundary("right", ("value",), (2,)), + ), + ) + + with pytest.raises(RuntimeError, match="already owned"): + runtime.execute(envelope) + + +def test_failed_multi_field_claim_does_not_reserve_earlier_fields() -> None: + """A conflicting fan-in claim validates every field before committing any.""" + + runtime = Runtime(input_ports=(InputPort(("left", "right", "third")),)) + heap = Heap( + WorkflowEnvelope( + "run", + ( + encode_boundary("left", ("x",), (1,)), + encode_boundary("right", ("x", "y"), (2, 3)), + encode_boundary("third", ("y",), (4,)), + ), + ) + ) + runtime.import_boundary(heap, "left", ("x",)) + + with pytest.raises(RuntimeError, match="already owned"): + runtime.import_boundary(heap, "right", ("x", "y")) + + runtime.import_boundary(heap, "third", ("y",)) + assert heap.y == 4 + + +def test_failed_decode_does_not_claim_boundary_fields() -> None: + """Invalid bytes cannot mutate fan-in ownership before decode succeeds.""" + + malformed = encode_boundary("bad", ("x",), (1,)) + object.__setattr__(malformed, "payload", b"not-a-cloudpickle-payload") + recovery = encode_boundary("recovery", ("x",), (2,)) + heap = Heap(WorkflowEnvelope("run", (malformed, recovery))) + runtime = Runtime(input_ports=(InputPort(("bad", "recovery")),)) + + with pytest.raises(pickle.UnpicklingError): + runtime.import_boundary(heap, "bad", ("x",)) + + runtime.import_boundary(heap, "recovery", ("x",)) + assert heap.x == 2 + + +def test_runtime_fails_closed_on_missing_export() -> None: + runtime = Runtime(outgoing=("required",)) + + @runtime.driver + def driver(heap: Heap) -> Heap: + return heap + + with pytest.raises(RuntimeError, match="outgoing"): + runtime.execute(WorkflowEnvelope("run", ())) + + +def test_heap_rejects_missing_source_binding() -> None: + heap = Heap(WorkflowEnvelope("run", ())) + + with pytest.raises(NameError, match="missing"): + _ = heap.missing + + +def test_heap_keyed_access_is_reserved_only_for_private_source_names() -> None: + heap = Heap(WorkflowEnvelope("run", ())) + heap["_value"] = 1 + + assert heap["_value"] == 1 + del heap["_value"] + + with pytest.raises(NameError, match="_value"): + _ = heap["_value"] + + +def test_runtime_rejects_boundary_ownership_by_multiple_input_ports() -> None: + with pytest.raises(ValueError, match="exactly one input port"): + Runtime( + input_ports=(InputPort(("shared",)), InputPort(("shared",))), + )