test(stream): validate P07.01 combined with P07.04 (OpenDCAI#550) and merged OpenDCAI#551

Author: rikkurea-boop

.github/workflows/p07-01-validation.yml

@@ -0,0 +1,46 @@
+name: P07.01 transactional resume validation
+
+on:
+  push:
+    branches:
+      - validation/p07-01-with-pr550
+  workflow_dispatch:
+
+jobs:
+  contract:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - os: ubuntu-latest
+            python: "3.10"
+          - os: ubuntu-latest
+            python: "3.12"
+          - os: windows-latest
+            python: "3.10"
+          - os: windows-latest
+            python: "3.12"
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+      - uses: astral-sh/setup-uv@v7
+      - name: Install dependencies
+        env:
+          UV_SYSTEM_PYTHON: "1"
+        run: |
+          uv pip install --no-cache -r requirements.txt
+          uv pip install -e .[test]
+      - name: P07.01 focused tests
+        run: |
+          python -m pytest -q test/cpu_only/test_stream_transactional_resume.py
+      - name: CPU/API suite
+        env:
+          PYTHONDONTWRITEBYTECODE: "1"
+        run: |
+          python -m pytest -m "cpu or api"
+      - name: Syntax compile
+        run: |
+          python -m compileall -q dataflow/pipeline/Pipeline.py dataflow/utils/storage.py dataflow/utils/stream_transaction.py test/cpu_only/test_stream_transactional_resume.py

dataflow/pipeline/Pipeline.py

@@ -625,7 +625,64 @@ class StreamBatchedPipelineABC(BatchedPipelineABC):
     def __init__(self):
         super().__init__()
 
-    def _compiled_forward(self, resume_step: int=0, batch_size: int|None=None, resume_from_last: bool=True):
+    def _compiled_forward(self, resume_step: int=0, batch_size: int|None=None, resume_from_last: bool=True, transactional_resume: bool=False):
+        if transactional_resume:
+            if not resume_from_last or resume_step != 0:
+                raise ValueError("Transactional streaming requires resume_from_last=True and resume_step=0")
+            return self._transactional_forward(batch_size)
+        return self._legacy_stream_forward(resume_step, batch_size, resume_from_last)
+
+    def _transactional_forward(self, batch_size):
+        from dataflow.utils.stream_transaction import StreamTransactionCoordinator
+        from dataflow.utils.stream_transaction import TransactionError
+
+        if not self.compiled:
+            raise RuntimeError("Pipeline is not compiled yet")
+        nodes = [node for node in self.op_nodes_list if node.op_obj is not None]
+        with StreamTransactionCoordinator(nodes, batch_size) as coordinator:
+            while coordinator.manifest["stage"] < len(nodes):
+                stage = coordinator.manifest["stage"]
+                node = nodes[stage]
+                storage = node.storage
+                storage.batch_size = batch_size
+                batch = coordinator.manifest["batch"]
+                count = coordinator.input_rows(stage)
+                total_batches = (count + batch_size - 1) // batch_size
+                if batch > total_batches:
+                    raise TransactionError("Committed cursor exceeds stage input")
+                if batch == total_batches:
+                    coordinator.transition(stage)
+                    continue
+                storage.batch_step = batch
+                stream = storage.iter_chunks()
+                try:
+                    for _ in range(batch):
+                        if next(stream, None) is None:
+                            raise TransactionError("Input stream shorter than committed cursor")
+                    chunk = next(stream, None)
+                finally:
+                    if "chunk" not in locals() and hasattr(stream, "close"):
+                        stream.close()
+                if chunk is None:
+                    raise TransactionError("Input stream shorter than manifest row count")
+                storage._current_streaming_chunk = chunk
+                coordinator.begin(stage, batch)
+                try:
+                    node.op_obj.run(storage=storage, **node.kwargs)
+                    coordinator.commit(stage, batch)
+                except BaseException:
+                    coordinator.abort()
+                    raise
+                finally:
+                    storage._current_streaming_chunk = None
+                    if hasattr(stream, "close"):
+                        stream.close()
+                    if "chunk" in locals():
+                        del chunk
+                storage.batch_step = coordinator.manifest["batch"]
+            return None
+
+    def _legacy_stream_forward(self, resume_step: int=0, batch_size: int|None=None, resume_from_last: bool=True):
         """
             resume_step (int): resume inference from this step
             batch_size (int|None): if set, run the pipeline in batch mode with this batch size

dataflow/utils/storage.py

@@ -1013,6 +1013,12 @@ def read(self, output_type: Literal["dataframe", "dict"]="dataframe") -> Any:
         return self._convert_output(dataframe, output_type)
 
     def write(self, data: Any) -> Any:
+        transaction = getattr(self, "_stream_transaction", None)
+        if transaction is not None:
+            return transaction.write(self, data)
+        return self._legacy_batched_write(data)
+
+    def _legacy_batched_write(self, data: Any) -> Any:
         """
         Write data to current file managed by storage.
         data: Any, the data to write, it should be a dataframe, List[dict], etc.