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/codec.py b/amber/src/main/python/pytexera/workflow/codec.py new file mode 100644 index 00000000000..ca0ab81c8c2 --- /dev/null +++ b/amber/src/main/python/pytexera/workflow/codec.py @@ -0,0 +1,158 @@ +# 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) +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 isinstance(self.boundary_id, str): + raise TypeError("boundary ID must be a string") + if not self.boundary_id: + raise ValueError("boundary ID must be nonempty") + if not isinstance(self.fields, tuple) or any( + not isinstance(field, str) for field in self.fields + ): + raise TypeError("boundary fields must be a tuple of strings") + 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 not isinstance(self.present, tuple) or any( + not isinstance(field, str) for field in self.present + ): + raise TypeError("present fields must be a tuple of strings") + present_fields = frozenset(self.present) + if self.present != tuple( + field for field in self.fields if field in present_fields + ): + 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 isinstance(self.execution_key, str): + raise TypeError("execution key must be a string") + 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") + for row in self.boundaries: + row.__post_init__() + 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.""" + + if not isinstance(values, tuple): + raise TypeError("boundary values must be a tuple") + 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 under the exact contract; pickle/loading exceptions propagate unchanged. + + Empty payloads raise EOFError; invalid opcodes raise pickle.UnpicklingError. + Other loading failures also propagate. Only load trusted workflow payloads. + """ + + boundary.__post_init__() + 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: + """Encode the inert envelope after validating its complete shape.""" + + if not isinstance(envelope, WorkflowEnvelope): + raise TypeError("envelope codec requires WorkflowEnvelope") + envelope.__post_init__() + return cloudpickle.dumps(envelope, protocol=pickle.HIGHEST_PROTOCOL) + + +def loads_envelope(payload: bytes) -> WorkflowEnvelope: + """Decode trusted bytes and revalidate envelope and nested boundary metadata. + + Pickle/loading exceptions propagate unchanged, as in decode_boundary. + Post-load validation is not a security boundary for untrusted pickle data. + """ + + envelope = cloudpickle.loads(payload) + if not isinstance(envelope, WorkflowEnvelope): + raise TypeError("decoded payload is not WorkflowEnvelope") + envelope.__post_init__() + 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/test/python/pytexera/workflow/test_codec.py b/amber/src/test/python/pytexera/workflow/test_codec.py new file mode 100644 index 00000000000..21a0e576e50 --- /dev/null +++ b/amber/src/test/python/pytexera/workflow/test_codec.py @@ -0,0 +1,196 @@ +# 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 cloudpickle +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,) + + +@pytest.mark.parametrize("size", [0, 1, 128]) +@pytest.mark.parametrize("step", [1, 2]) +def test_present_validation_scans_present_fields_linearly(size, step): + visits = 0 + + class CountedPresent(tuple): + def __iter__(self): + nonlocal visits + for field in super().__iter__(): + visits += 1 + yield field + + fields = tuple(f"field{index:03d}" for index in range(size)) + present = CountedPresent(fields[::step]) + + boundary = BoundaryPayload("edge", fields, present, b"payload") + + assert boundary.present == present + assert visits <= 2 * len(present) + + +@pytest.mark.parametrize("present", [("雪", "a"), ("a", "a"), ("missing",)]) +def test_present_validation_rejects_noncanonical_subsets(present): + with pytest.raises(ValueError, match="canonical contract subset"): + BoundaryPayload("edge", ("a", "雪"), present, b"payload") + + +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", ()), + ) + + +@pytest.mark.parametrize("identity", [None, 1, True, b"edge", []]) +def test_identities_require_strings(identity): + with pytest.raises(TypeError, match="string"): + BoundaryPayload(identity, (), (), b"") + with pytest.raises(TypeError, match="string"): + WorkflowEnvelope(identity, ()) + + +@pytest.mark.parametrize("identity", ["", "\u96ea"]) +def test_empty_and_unicode_identities(identity): + if not identity: + with pytest.raises(ValueError, match="nonempty"): + WorkflowEnvelope(identity, ()) + with pytest.raises(ValueError, match="nonempty"): + BoundaryPayload(identity, (), (), b"") + else: + envelope = WorkflowEnvelope(identity, (encode_boundary(identity, (), ()),)) + assert loads_envelope(dumps_envelope(envelope)) == envelope + + +@pytest.mark.parametrize( + ("target", "field", "value", "error"), + [ + ("envelope", "execution_key", 42, TypeError), + ("envelope", "execution_key", "", ValueError), + ("envelope", "boundaries", [], TypeError), + ("boundary", "boundary_id", 42, TypeError), + ("boundary", "fields", ["value"], TypeError), + ("boundary", "fields", (1,), TypeError), + ("boundary", "present", ["value"], TypeError), + ("boundary", "present", ("other",), ValueError), + ("boundary", "payload", "bytes", TypeError), + ], +) +def test_envelope_codec_revalidates_unpickled_metadata(target, field, value, error): + boundary = encode_boundary("edge", ("value",), (42,)) + envelope = WorkflowEnvelope("run", (boundary,)) + object.__setattr__(envelope if target == "envelope" else boundary, field, value) + # Simulate an existing serialized instance: pickle bypasses constructors. + payload = cloudpickle.dumps(envelope) + with pytest.raises(error): + loads_envelope(payload) + with pytest.raises(error): + dumps_envelope(envelope) + + +def test_merge_rejects_overlapping_ids_and_orders_disjoint_ids(): + first = encode_boundary("z", ("value",), (1,)) + second = encode_boundary("z", ("value",), (2,)) + with pytest.raises(ValueError, match="unique"): + merge_envelopes( + WorkflowEnvelope("run", (first,)), WorkflowEnvelope("run", (second,)) + ) + other = encode_boundary("a", ("value",), (3,)) + merged = merge_envelopes( + WorkflowEnvelope("run", (first,)), WorkflowEnvelope("run", (other,)) + ) + assert merged.boundaries == (other, first) + assert decode_boundary(merged.boundaries[0], ("value",)) == (3,) + assert decode_boundary(merged.boundaries[1], ("value",)) == (1,) + with pytest.raises(TypeError): + _ = first < other + + +@pytest.mark.parametrize("values", [[42], {"value": 42}, "x", None]) +def test_encoder_rejects_non_tuple_values_before_serialization(values, monkeypatch): + def unexpected_serialization(*args, **kwargs): + pytest.fail("invalid values reached serialization") + + monkeypatch.setattr(cloudpickle, "dumps", unexpected_serialization) + with pytest.raises(TypeError, match="tuple"): + encode_boundary("edge", ("value",), values) + + +@pytest.mark.parametrize( + ("payload", "error"), [(b"", EOFError), (b"!", pickle.UnpicklingError)] +) +def test_decoders_preserve_pickle_exceptions(payload, error): + with pytest.raises(error): + loads_envelope(payload) + with pytest.raises(error): + decode_boundary(BoundaryPayload("edge", (), (), payload), ())