refactor: support RayOrch 0.1.0

dataflow/rayorch/accelerated_op.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any, Generic, Optional, Protocol, Type, ParamSpec
+import math
+from typing import Any, Generic, Mapping, Optional, Protocol, Type, ParamSpec
 
 import pandas as pd
 
@@ -36,21 +37,58 @@ def run(
 class _OpRunner:
     """Actor-side worker: each replica holds an independent operator instance.
 
-    Receives a chunk of records (``list[dict]``), wraps it in
-    :class:`InMemoryStorage`, delegates to the DataFlow operator's ``run``,
-    and returns the result as ``list[dict]``.
+    Receives logical record groups, runs each group through
+    :class:`InMemoryStorage`, and returns one output group per input group.
     """
 
     def __init__(self, op_cls: type, op_init_args: tuple, op_init_kwargs: dict):
         self.op = op_cls(*op_init_args, **op_init_kwargs)
 
-    def run(self, records: list[dict], run_params: dict) -> list[dict]:
-        if not records:
-            return []
-        df = pd.DataFrame(records)
-        storage = InMemoryStorage(df)
-        self.op.run(storage, *run_params.get("args", ()), **run_params.get("kwargs", {}))
-        return storage.result.to_dict("records")
+    def run(
+        self,
+        record_groups: list[list[dict]],
+        run_params: list[dict],
+    ) -> list[list[dict]]:
+        outputs = []
+        for records, params in zip(record_groups, run_params, strict=True):
+            storage = InMemoryStorage(pd.DataFrame(records))
+            self.op.run(storage, *params["args"], **params["kwargs"])
+            outputs.append(storage.result.to_dict("records"))
+        return outputs
+
+
+def _build_operator_pipeline(
+    op_cls: type,
+    op_init_args: tuple,
+    op_init_kwargs: dict,
+    *,
+    replicas: int,
+    num_gpus_per_replica: float,
+    runtime_env: Mapping[str, Any] | None,
+):
+    """Build the RayOrch pipeline only when the optional dependency is used."""
+
+    from rayorch import Pipeline, RayModule
+
+    class OperatorPipeline(Pipeline):
+        def __init__(self) -> None:
+            options: dict[str, Any] = {
+                "replicas": replicas,
+                "batch_size": 1,
+                "num_gpus": num_gpus_per_replica,
+            }
+            if runtime_env is not None:
+                options["runtime_env"] = dict(runtime_env)
+            self.operator = (
+                RayModule(_OpRunner)
+                .pre_init(op_cls, op_init_args, op_init_kwargs)
+                .ray_options(**options)
+            )
+
+        def forward(self, records, run_params):
+            return self.operator(records, run_params)
+
+    return OperatorPipeline()
 
 
 class RayAcceleratedOperator(OperatorABC, Generic[_INITP, _RUNP]):
@@ -81,7 +119,13 @@ class RayAcceleratedOperator(OperatorABC, Generic[_INITP, _RUNP]):
         Fractional GPU allocation per replica (e.g. ``0.25`` to share one
         GPU across four replicas).
     env:
-        Optional RayOrch ``EnvRegistry`` key for a custom ``runtime_env``.
+        Backward-compatible shorthand for ``runtime_env={"conda": env}``.
+    batch_size:
+        Maximum rows passed to one actor call. By default each input is
+        divided evenly across replicas, matching the previous contiguous-shard
+        behavior.
+    runtime_env:
+        Optional Ray ``runtime_env`` forwarded to every actor.
 
     Example
     -------
@@ -106,15 +150,32 @@ def __init__(
         replicas: int = 1,
         num_gpus_per_replica: float = 0.0,
         env: Optional[str] = None,
+        batch_size: int | None = None,
+        runtime_env: Mapping[str, Any] | None = None,
     ):
         super().__init__()
+        if type(replicas) is not int or replicas <= 0:
+            raise ValueError("replicas must be a positive integer")
+        if batch_size is not None and (
+            type(batch_size) is not int or batch_size <= 0
+        ):
+            raise ValueError("batch_size must be a positive integer or None")
+        if num_gpus_per_replica < 0:
+            raise ValueError("num_gpus_per_replica cannot be negative")
         self._op_cls = op_cls
         self._op_init_args: tuple = ()
         self._op_init_kwargs: dict = {}
         self._replicas = replicas
+        self._batch_size = batch_size
         self._num_gpus_per_replica = num_gpus_per_replica
-        self._env = env
-        self._module = None  # created lazily
+        if env is not None and runtime_env is not None:
+            raise ValueError("pass env or runtime_env, not both")
+        self._runtime_env = (
+            {"conda": env}
+            if env is not None
+            else None if runtime_env is None else dict(runtime_env)
+        )
+        self._executor = None  # Created lazily and reused across run() calls.
 
         # PipelineABC.compile() compatibility:
         # compile() → AutoOP uses inspect.signature(operator.run) to bind()
@@ -140,22 +201,21 @@ def op_cls_init(
         return self
 
     def _ensure_initialized(self) -> None:
-        if self._module is not None:
+        if self._executor is not None:
             return
-        from rayorch import Dispatch, RayModule
+        from rayorch import Executor
 
-        self._module = RayModule(
-            _OpRunner,
+        pipeline = _build_operator_pipeline(
+            self._op_cls,
+            self._op_init_args,
+            self._op_init_kwargs,
             replicas=self._replicas,
             num_gpus_per_replica=self._num_gpus_per_replica,
-            dispatch_mode=Dispatch.SHARD_CONTIGUOUS,
-            env=self._env,
-        )
-        self._module.pre_init(
-            op_cls=self._op_cls,
-            op_init_args=self._op_init_args,
-            op_init_kwargs=self._op_init_kwargs,
+            runtime_env=self._runtime_env,
         )
+        # DataFlow may invoke one operator repeatedly, so keep the Executor
+        # alive instead of using Pipeline.run(), which is intentionally one-shot.
+        self._executor = Executor(pipeline)
 
     # --- inner signature propagation ---
 
@@ -214,27 +274,43 @@ def _run_impl(
         *args: _RUNP.args,
         **kwargs: _RUNP.kwargs,
     ) -> None:
-        self._ensure_initialized()
         df = storage.read("dataframe")
+        if df.empty:
+            storage.write(df.copy())
+            return
         records: list[dict] = df.to_dict("records")
+        chunk_size = self._batch_size or max(
+            1,
+            math.ceil(len(records) / self._replicas),
+        )
+        record_groups = [
+            records[start : start + chunk_size]
+            for start in range(0, len(records), chunk_size)
+        ]
+        self._ensure_initialized()
         run_params: dict = {"args": args, "kwargs": kwargs}
-        result_records = self._module(records, run_params)
-        storage.write(pd.DataFrame(result_records))
+        result = self._executor.run(
+            record_groups,
+            [run_params] * len(record_groups),
+        )
+        output_records = [
+            record
+            for group in result.outputs
+            for record in group
+        ]
+        storage.write(pd.DataFrame(output_records))
 
     # --- lifecycle helpers ---
 
     def shutdown(self) -> None:
         """Terminate all Ray actors held by this operator."""
-        if self._module is None:
+        if self._executor is None:
             return
-        import ray
-
-        for actor in self._module.actors:
-            ray.kill(actor)
-        self._module = None
+        self._executor.close()
+        self._executor = None
 
     def __repr__(self) -> str:
-        state = "initialized" if self._module is not None else "lazy"
+        state = "initialized" if self._executor is not None else "lazy"
         return (
             f"RayAcceleratedOperator({self._op_cls.__name__}, "
             f"replicas={self._replicas}, state={state})"

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=42", "wheel"]
+requires = ["setuptools>=61", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -43,16 +43,17 @@ dataflow = "dataflow.cli:app"
 
 [tool.setuptools]
 include-package-data = true
-packages = ["dataflow"]  # 显式指定主包
 
+[tool.setuptools.packages.find]
+include = ["dataflow*"]
 
 [tool.setuptools.dynamic]
 version = {attr = "dataflow.version.__version__"}
 dependencies = {file = "requirements.txt"}
 
 
 [project.optional-dependencies]
-test = ["flask", "setuptools<=81.0.0", "rayorch"] # https://setuptools.pypa.io/en/latest/pkg_resources.html & https://github.com/dgtlmoon/changedetection.io/pull/2424
+test = ["flask", "setuptools<=81.0.0", "rayorch==0.1.0; python_version >= '3.11'"] # https://setuptools.pypa.io/en/latest/pkg_resources.html & https://github.com/dgtlmoon/changedetection.io/pull/2424
 vllm =["vllm>=0.7.0,<=0.9.2", "numpy<2.0.0"]
 vllm07 = ["vllm<0.8", "numpy<2.0.0"]
 vllm08 = ["vllm<0.9"]
@@ -94,4 +95,4 @@ rag = ["lightrag-hku", "asyncio"]
 pdf2vqa=["mineru[vlm]>=2.5.0,<2.7.0"]
 
 flash-mineru = ["flash-mineru"]
-ray = ["rayorch"] 
+ray = ["rayorch==0.1.0"]

test/rayorch/README-zh.md

@@ -56,9 +56,12 @@ Actor 在首次 `run()` 时**懒加载**创建，pipeline compile 阶段不会
 pip install -e .
 
 # 或仅额外安装 RayOrch
-pip install rayorch==0.0.1
+pip install rayorch==0.1.0
 ```
 
+RayOrch 加速功能要求 Python 3.11 或更高版本；未安装可选 `ray` extra 时，
+DataFlow 的其他功能仍可正常导入和使用。
+
 ## 测试文件
 
 | 文件 | 说明 |

test/rayorch/README.md

@@ -59,9 +59,12 @@ compilation does not trigger model loading.
 pip install -e .
 
 # Or install RayOrch separately
-pip install rayorch==0.0.1
+pip install rayorch==0.1.0
 ```
 
+RayOrch acceleration requires Python 3.11 or newer. Other DataFlow features
+remain importable without installing the optional `ray` extra.
+
 ## Test Files
 
 | File | Description |

test/rayorch/test_compile_cpu.py

@@ -10,11 +10,18 @@
 from __future__ import annotations
 
 import os
+import sys
 
 import numpy as np
 import pandas as pd
 import pytest
 
+if sys.version_info < (3, 11):
+    pytest.skip(
+        "RayOrch 0.1.0 requires Python 3.11 or newer",
+        allow_module_level=True,
+    )
+
 from dataflow.pipeline.Pipeline import (
     BatchedPipelineABC,
     PipelineABC,
@@ -533,7 +540,7 @@ def test_auto_shutdown_pipeline_abc(ray_env, test_data):
     assert len(df) == N_ROWS
 
     for node in _get_ray_op_nodes(pipe):
-        assert node.op_obj._module is None, (
+        assert node.op_obj._executor is None, (
             f"{node.op_name} was not auto-shutdown after compiled forward"
         )
 
@@ -556,7 +563,7 @@ def test_auto_shutdown_batched(ray_env, test_data):
     assert len(df) == N_ROWS
 
     for node in _get_ray_op_nodes(pipe):
-        assert node.op_obj._module is None, (
+        assert node.op_obj._executor is None, (
             f"{node.op_name} was not auto-shutdown after compiled forward"
         )
 
@@ -579,6 +586,6 @@ def test_auto_shutdown_stream_batched(ray_env, test_data):
     assert len(df) == N_ROWS
 
     for node in _get_ray_op_nodes(pipe):
-        assert node.op_obj._module is None, (
+        assert node.op_obj._executor is None, (
             f"{node.op_name} was not auto-shutdown after compiled forward"
         )