Short answer: you don’t need a special deferral pattern. In Prefect 3.x the

prefect-ray

integration’s

RayTaskRunner

does not call

ray.init()

at import/decoration time. It only initializes Ray when the flow actually starts running (inside the task runner’s context manager). So you can safely write

@flow(task_runner=RayTaskRunner(...))

in modules that are imported by tests without secrets/infra; no connection to Ray is attempted until the flow is invoked. Most common patterns teams use: - Rely on the built‑in deferral (recommended) - Define your flow with `RayTaskRunner`; Ray is initialized in

__enter__

when the flow runs, not when the module is imported. - Example:

from prefect import flow, task
    from prefect_ray import RayTaskRunner

    @task
    def work(x): return x * 2

    @flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS", None), init_kwargs={"namespace": "my-ns"}))
    def my_flow():
        return work.submit(21).result()

- In tests that just import this module, nothing tries to connect to Ray. - Environment-based runner selection (for CI/local dev) - Choose a non-Ray runner in environments where you don’t want to touch Ray; still no Ray init at import time.

import os
    from prefect import flow
    from prefect.task_runners import ThreadPoolTaskRunner
    try:
        from prefect_ray import RayTaskRunner
    except Exception:
        RayTaskRunner = None

    def pick_runner():
        if os.getenv("DISABLE_RAY") or RayTaskRunner is None:
            return ThreadPoolTaskRunner()
        return RayTaskRunner(
            address=os.getenv("RAY_ADDRESS"), 
            init_kwargs={"namespace": os.getenv("RAY_NAMESPACE", "default")}
        )

    @flow(task_runner=pick_runner())
    def my_flow():
        ...

- Test without Ray entirely - For unit tests, many teams just use the default

ThreadPoolTaskRunner

(i.e., don’t specify a task runner) or mark Ray tests to run only when Ray is available. - Example:

import pytest
    ray = pytest.importorskip("ray")  # skip if ray not installed

    def test_with_ray():
        from prefect_ray import RayTaskRunner
        @flow(task_runner=RayTaskRunner())
        def f(): ...
        f()

Key notes - Importing

prefect_ray

and instantiating

RayTaskRunner()

is safe at import time; the integration lazily imports Ray and defers

ray.init()

until flow execution. - If you connect to a remote cluster (e.g.,

address="ray://..."

), the connection is also attempted only when the flow runs, so secrets/infra can be provided by your worker/runner environment instead of your test process. References - Prefect Ray integration: PrefectHQ/prefect-ray - Task runners concept: Prefect docs: Task runners

Short answer: read the address from an env var and let it be unset locally. - In your flow, construct the task runner from RAY_ADDRESS. If it’s set (remote job), it connects to the cluster; if it’s unset (local dev), Prefect will start a local ephemeral Ray runtime.

import os
from prefect import flow, task
from prefect_ray import RayTaskRunner

@task
def work(i: int) -> int:
    return i * 2

@flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS")))
def my_flow(n: int = 8):
    futures = [work.submit(i) for i in range(n)]
    return [f.result() for f in futures]

if __name__ == "__main__":
    my_flow()

Ways to provide RAY_ADDRESS only for remote jobs - Deployment (prefect.yaml) via job variables:

deployments:
  - name: my-ray-deployment
    entrypoint: my_flow.py:my_flow
    work_pool:
      name: my-work-pool
      job_variables:
        env:
          RAY_ADDRESS: "<ray://10.0.0.12:10001>"

- Programmatic deploy:

my_flow.deploy(
  name="my-ray-deployment",
  work_pool_name="my-work-pool",
  job_variables={"env": {"RAY_ADDRESS": "<ray://10.0.0.12:10001>"}},
)

- Work pool defaults (UI): set RAY_ADDRESS in the pool’s default job variables (env). Every job from that pool will get it automatically. Notes - You can pass additional Ray init options with

init_kwargs

on RayTaskRunner if needed (e.g., namespace, runtime_env). - If Ray is already initialized in the process (e.g., you call ray.init yourself), the task runner will reuse it. Docs - Ray integration overview: prefect-ray - RayTaskRunner API: RayTaskRunner API - Customize job variables (env): Customize job variables - Manage work pools: Manage work pools If you share which work pool/worker type you’re using (Docker, Kubernetes, etc.), I can show the exact spot to set the env var in that environment.

Absolutely — here’s a minimal, runnable example that shows RayTaskRunner reusing an already-initialized Ray instance. What it demonstrates - We call

ray.init(...)

before running the flow. - The

RayTaskRunner

checks

ray.is_initialized()

and, since it’s True, it does not call

ray.init()

again — it reuses the existing runtime. - With Prefect logging at DEBUG, you’ll see the “already initialized… using existing local instance” message from the task runner.

import os
import ray
from prefect import flow, task
from prefect_ray import RayTaskRunner

# Turn up Prefect logs so you can see the reuse message
os.environ["PREFECT_LOGGING_LEVEL"] = "DEBUG"

# Pre-initialize Ray in this process (simulating "you called ray.init yourself")
ray.init(num_cpus=2)
print("Driver before flow -> ray.is_initialized():", ray.is_initialized())

@task
def whoami():
    # Runs on Ray workers; just show Ray is active there too
    import ray as _ray
    return {
        "worker_initialized": _ray.is_initialized(),
        "job_id": _ray.get_runtime_context().job_id,
    }

@flow(task_runner=RayTaskRunner())  # no address passed; runner will reuse existing Ray
def my_flow():
    # This executes on the driver; Ray should already be initialized
    import ray as _ray
    print("Driver inside flow -> ray.is_initialized():", _ray.is_initialized())
    return whoami.submit().result()

if __name__ == "__main__":
    result = my_flow()
    print("Task result:", result)

    # After the flow, RayTaskRunner will exit its context; see note below about shutdown.

What you should see - A Prefect DEBUG log line like: “Local Ray instance is already initialized. Using existing local instance.” -

ray.is_initialized()

prints True both before the flow and inside the flow. - The task confirms Ray is initialized on workers and shows a job_id. Notes and gotchas - If you pre-initialize Ray and also pass an address to

RayTaskRunner(address=...)

, the runner still reuses the existing local Ray (it checks

ray.is_initialized()

first and returns early). - By default, when the flow finishes,

RayTaskRunner

will call

ray.shutdown()

on the driver during its teardown. If you want to keep your pre-initialized Ray runtime alive after the flow, run the flow in a separate process or reinitialize Ray afterward. Docs - RayTaskRunner API: RayTaskRunner API - Integration overview: prefect-ray

Short answer: Pre-initializing Ray and letting RayTaskRunner reuse it is a valid pattern, but it’s more common in notebooks/dev or when you need Ray features that the runner doesn’t surface. For production deployments, the more common pattern is to let the task runner manage lifecycle (ephemeral locally,

RAY_ADDRESS

for remote) via job variables. Common patterns - Most common in prod: - In code:

RayTaskRunner(address=os.getenv("RAY_ADDRESS"))

- In deployment/work pool: set

RAY_ADDRESS

via job variables/env - Result: local dev spins up an ephemeral Ray; remote jobs connect to your cluster - Common in notebooks/dev or special setups: - You call

ray.init(...)

(e.g.,

address="auto"

on K8s Ray cluster, or custom options) - RayTaskRunner detects

ray.is_initialized()

and reuses it Pros of pre-initializing Ray (ray.init yourself) - Full control over Ray session: - Use Ray features not directly exposed by the runner (e.g.,

address="auto"

autodiscovery in Kubernetes, detailed

runtime_env

, placement groups, custom resources). - Reuse across multiple flows in the same long-lived process (e.g., a notebook running several flows) without repeated init overhead. - Easy to attach to a specific already-running Ray session for ad-hoc experimentation. Cons of pre-initializing Ray - Lifecycle surprises: - RayTaskRunner calls

ray.shutdown()

on flow exit. This disconnects the driver; if the cluster was local/in-process it’s torn down, but for a remote cluster it just detaches (the cluster keeps running). If you want to keep the same session for subsequent flows in the same process, you’ll need to reinit or override shutdown behavior. - Portability: - Code that relies on a pre-initialized global session is less portable to workers/containers where each run starts in a fresh process. Deployments typically expect the runner to manage Ray lifecycle with an address or start ephemeral. - Hidden state: - If some other part of your process (or earlier code in a session) initialized Ray pointing to a different cluster, the runner will silently reuse that, which can be surprising. - Operational clarity: - Env-driven

RAY_ADDRESS

via job variables/work pools is easier to reason about and standardize across environments than embedding cluster setup in code. When I’d recommend pre-init - You need Ray autodiscovery (

ray.init(address="auto")

) or advanced init that the runner’s

init_kwargs

can’t replicate. - You’re in a notebook or a long-running process and want a single Ray session across many flow runs. - You explicitly want to manage the driver lifecycle yourself. When I’d avoid it - Standard CI/CD or containerized workers running deployments. Prefer

RayTaskRunner(address=os.getenv("RAY_ADDRESS"))

with

RAY_ADDRESS

set via work pool/deployment job variables. - Teams optimizing for reproducibility and minimizing hidden global state. Tip: If you do pre-init and want to avoid the runner calling

ray.shutdown()

on exit, you can subclass the runner to skip shutdown:

from prefect_ray.task_runners import RayTaskRunner as _RayTaskRunner

class NoShutdownRayTaskRunner(_RayTaskRunner):
    def __exit__(self, *exc_info):
        # Skip ray.shutdown(); just run Prefect’s teardown
        return super().__exit__(*exc_info)

Then use

NoShutdownRayTaskRunner()

in your flow. This keeps your pre-initialized Ray session alive after the flow finishes. Docs - Ray integration overview: prefect-ray - RayTaskRunner API: RayTaskRunner API - Customize job variables (env): Customize job variables If you share where you’re running (e.g., Docker, Kubernetes + Ray Operator, bare VMs), I can suggest the pattern that tends to work best there.

Great question. Assuming you’re using Prefect 3.x with RayTaskRunner(address="ray://...") to connect to a Ray cluster managed by the Ray Operator on Kubernetes, here’s how to approach each goal: 1) Cache pip package installs - Best practice: bake all dependencies into the Ray worker image used by your RayCluster. This guarantees cache hits (no runtime pip) and is the most reliable/fast path. - If you need dynamic deps, set a job-level Ray runtime_env so Ray installs once per driver job and reuses the cache across tasks/actors and subsequent jobs with the same env hash: - Pin versions and keep the runtime_env stable to maximize cache reuse. - Ray caches runtime_envs per node; if nodes/pods stay up, subsequent runs hit the cache. - Optional: mount a PVC and set PIP_CACHE_DIR to persist wheel cache across pod restarts. Example: job-level runtime_env via RayTaskRunner (preferred over per-task)

import os
from prefect import flow, task
from prefect_ray import RayTaskRunner
from prefect_ray.context import remote_options

@task
def compute(x: int) -> int:
    return x * 2

@flow(
    task_runner=RayTaskRunner(
        address=os.getenv("RAY_ADDRESS"),
        init_kwargs={
            "runtime_env": {
                "pip": [
                    "pandas==2.2.2",
                    "pyarrow==15.0.2",
                ],
                # Optional: keep a shared pip cache if you mount a PVC in worker pods
                # "env_vars": {"PIP_CACHE_DIR": "/opt/pip-cache"},
            }
        },
    )
)
def my_flow(n: int = 100):
    # Optional: defaults for Ray tasks (num_cpus, resources, etc.)
    with remote_options(num_cpus=0.5):
        futures = [compute.submit(i) for i in range(n)]
        return [f.result() for f in futures]

2) Keep Ray workers alive and reused across multiple calls - Configure the RayCluster autoscaler to maintain a warm baseline: - Set workerGroupSpecs[].minReplicas to the number of workers you want to keep hot. - Increase or disable the autoscaler’s idle timeout so it doesn’t scale down between runs. - Optional keepalive: if you can’t change autoscaler settings, keep a tiny “keepalive” Ray actor running so the autoscaler sees the cluster as active. Keepalive actor pattern (optional)

import ray
from prefect import task

@task
def ensure_keepalive():
    @ray.remote(num_cpus=0)
    class KeepAlive:
        def ping(self):
            return "ok"

    # Create a named actor once; future jobs can find it by name
    try:
        ray.get_actor("prefect-keepalive", namespace="prefect")
    except ValueError:
        KeepAlive.options(name="prefect-keepalive", namespace="prefect").remote()

Note: RayTaskRunner’s teardown calls ray.shutdown() on the driver only; it won’t terminate your remote cluster. Scaling down is controlled by the Ray autoscaler. 3) Reuse pods (don’t spin them down) - Ray Operator autoscaler: set workerGroupSpecs[].minReplicas > 0 and a generous idle timeout so pods stay up and get reused. If you must scale down, consider a small baseline (e.g., minReplicas=1–2). - Kubernetes cluster autoscaler: ensure the node pool backing Ray workers does not scale to zero if you want pods to persist; keep a non-zero min node count. - Image caching: use imagePullPolicy: IfNotPresent so nodes reuse cached layers; this speeds up pod restarts and reduces cold starts. Practical checklist - Dependencies - Prefer: bake into Ray worker image. - Else: set job-level runtime_env on RayTaskRunner with pinned versions; optionally mount PVC + PIP_CACHE_DIR. - Ray autoscaler (Operator) - workerGroupSpecs[].minReplicas >= desired warm baseline - Set/raise idle timeout so workers don’t scale down between runs - Kubernetes - Node group minSize > 0 (avoid scale-to-zero) - imagePullPolicy IfNotPresent for faster pod reuse - Optional: a named “keepalive” actor to keep the cluster non-idle if autoscaler tuning isn’t possible Relevant Prefect docs

- Ray integration overview: prefect-ray - RayTaskRunner API: RayTaskRunner API - Customize job variables/env: Customize job variables If you share your RayCluster (Operator) spec (workerGroupSpecs/autoscaler settings), I can point to the exact fields to tweak and suggest concrete values for your workload profile.

Great questions — here’s how to think about each part. Why use remote_options and list of submit calls vs just task.map? - You can use either. Under RayTaskRunner, both task.map(...) and a list of task.submit(...) calls end up submitting a bunch of Ray tasks. - Prefer task.map for readability and orchestration features (result lists, automatic waiting, cleaner logs). Internally, Prefect’s map just loops over submit anyway. - Use remote_options when you need to set Ray scheduling options (num_cpus/num_gpus/resources/runtime_env) for the tasks you are about to submit/map. It’s the supported way to pass Ray’s @remote options through Prefect. Example showing both patterns (equivalent scheduling under Ray)

from prefect import flow, task
from prefect_ray import RayTaskRunner
from prefect_ray.context import remote_options

@task
def compute(x): 
    return x * 2

@flow(task_runner=RayTaskRunner())
def my_flow(n: int = 10):
    # Pattern A: idiomatic
    with remote_options(num_cpus=0.5):
        results = compute.map(range(n)).wait()

    # Pattern B: explicit submits (same effect)
    with remote_options(num_cpus=0.5):
        futures = [compute.submit(i) for i in range(n)]
        results2 = [f.result() for f in futures]

    return results

Does prefect-ray expose a ray.map_batches? - No. prefect-ray does not wrap ray.data APIs like map_batches. - To use ray.data, call it inside a Prefect task. RayTaskRunner ensures the Ray connection exists; your ray.data pipeline then runs on the connected cluster. Two ways to integrate ray.data with RayTaskRunner 1) Keep the whole ray.data pipeline inside a single Prefect task - Best when you want Ray to handle partitioning/parallelism internally.

import os
from prefect import flow, task
from prefect_ray import RayTaskRunner

@task
def transform_dataset(input_uri: str, output_uri: str):
    import ray
    import ray.data as rd

    ds = rd.read_parquet(input_uri)

    def enrich(batch):
        # mutate a Pandas or Arrow batch
        batch["value_x2"] = batch["value"] * 2
        return batch

    # Run distributed transform on the Ray cluster
    out = ds.map_batches(enrich, batch_size=1000)
    # Materialize results — don’t return Ray objects from the task
    out.write_parquet(output_uri)
    return {"rows": out.count(), "written_to": output_uri}

@flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS")))
def pipeline(input_uri: str, output_uri: str):
    return transform_dataset.submit(input_uri, output_uri).result()

Tips: - Don’t return Ray Datasets from a Prefect task; write to storage or convert to small Python objects first (e.g., ds.take(), to_pandas() for small results). - If you need specific Ray resources for the task driver itself, wrap the submit/map with remote_options (affects the Ray tasks Prefect creates). 2) Orchestrate chunked work with Prefect map and let Ray run each chunk - Best when you want Prefect-level retries/observability per chunk and explicit fan-out/fan-in.

from typing import Sequence
from prefect import flow, task
from prefect_ray import RayTaskRunner
from prefect_ray.context import remote_options

@task
def process_chunk(files: Sequence[str], output_uri: str):
    import ray
    import ray.data as rd

    ds = rd.read_parquet(files)
    def fn(batch):
        batch["z"] = batch["x"] + batch["y"]
        return batch
    ds = ds.map_batches(fn, batch_size=2000)
    # Persist per-chunk output; return metadata only
    ds.write_parquet(output_uri)
    return {"files": len(files)}

@flow(task_runner=RayTaskRunner())
def process_all(chunks: list[list[str]], base_out: str):
    with remote_options(num_cpus=2):  # shape Ray resources per Prefect task
        futs = [
            process_chunk.submit(files, f"{base_out}/chunk_{i}")
            for i, files in enumerate(chunks)
        ]
    return [f.result() for f in futs]

Key takeaways

- Use task.map for parallel collections; use submit for individual calls or when you want explicit control. Both work with RayTaskRunner; map is usually cleaner. - Use remote_options to pass Ray @remote options to the tasks you’re about to submit/map. There isn’t a per-task submit kwarg for Ray options; the context manager is the supported approach. - prefect-ray does not provide ray.data helpers like map_batches. Wrap ray.data pipelines inside Prefect tasks and ensure you materialize results (don’t return Ray objects). Docs - Ray integration overview: prefect-ray - RayTaskRunner API: RayTaskRunner API

Totally—there isn’t a built-in map_batches wrapper in prefect-ray, but you can create a small helper that turns a batch transform function into a Prefect task that runs ray.data Dataset.map_batches under RayTaskRunner. Two options below: a minimal “parquet-in/parquet-out” helper, and a slightly more flexible factory that lets you plug different readers/writers. Minimal: parquet to parquet

from prefect import task

def ray_map_batches_task(fn, *, batch_size=None, batch_format="pandas", **map_kwargs):
    """
    Wrap a batch transform (fn) into a Prefect task that:
      - reads a parquet dataset with ray.data.read_parquet
      - applies ds.map_batches(fn, ...)
      - optionally writes parquet output
    """
    @task
    def _wrapped(read_path: str, write_path: str | None = None):
        import ray.data as rd

        ds = rd.read_parquet(read_path)
        ds_out = ds.map_batches(
            fn,
            batch_size=batch_size,
            batch_format=batch_format,
            **map_kwargs,
        )

        if write_path:
            ds_out.write_parquet(write_path)
            return {"rows": ds_out.count(), "written_to": write_path}

        # For small results only; prefer writing for large outputs
        return ds_out.take_all()

    return _wrapped

Usage

import os
from prefect import flow
from prefect_ray import RayTaskRunner
from prefect_ray.context import remote_options

# Your batch transform (Pandas or Arrow batch depending on batch_format)
def enrich(batch):
    batch["value_x2"] = batch["value"] * 2
    return batch

map_parquet = ray_map_batches_task(
    enrich,
    batch_size=1000,
    batch_format="pandas",
)

@flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS")))
def pipeline(src_parquet: str, dst_parquet: str):
    # Optional: Ray resource shaping for the Prefect task submission
    with remote_options(num_cpus=2):
        return map_parquet.submit(src_parquet, dst_parquet).result()

Flexible factory (choose reader/writer)

from typing import Callable, Optional, Any
from prefect import task

def make_map_batches_task(
    transform: Callable[[Any], Any],
    *,
    reader: Callable[..., "ray.data.Dataset"],
    writer: Optional[Callable[["ray.data.Dataset", Any], None]] = None,
    batch_size: Optional[int] = None,
    batch_format: str = "pandas",
    **map_kwargs,
):
    """
    Create a Prefect task that:
      - builds a Ray Dataset via `reader(*reader_args, **reader_kwargs)`
      - applies map_batches(transform, ...)
      - optionally persists output via `writer(dataset, *writer_args)`
    Returns small materialized results if no writer is provided.
    """
    @task
    def _task(*reader_args, _writer_args=None, **reader_kwargs):
        import ray.data as rd  # reader may be rd.read_* already

        ds = reader(*reader_args, **reader_kwargs)
        ds_out = ds.map_batches(
            transform,
            batch_size=batch_size,
            batch_format=batch_format,
            **map_kwargs,
        )

        if writer:
            writer(ds_out, *(_writer_args or ()))
            return {"rows": ds_out.count()}

        # For small datasets only; prefer writing for large outputs
        return ds_out.take_all()

    return _task

Usage

import os
import ray.data as rd
from prefect import flow
from prefect_ray import RayTaskRunner
from prefect_ray.context import remote_options

def enrich(batch):
    batch["z"] = batch["x"] + batch["y"]
    return batch

map_csv_to_parquet = make_map_batches_task(
    enrich,
    reader=rd.read_csv,                              # any ray.data reader
    writer=lambda ds, path: ds.write_parquet(path), # any ray.data writer
    batch_size=2000,
    batch_format="pandas",
)

@flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS")))
def pipeline(src_csv: str, dst_parquet: str):
    with remote_options(num_cpus=2):
        return map_csv_to_parquet.submit(src_csv, _writer_args=(dst_parquet,)).result()

Notes - Keep the dataset creation inside the task. Passing a Ray Dataset object across task boundaries can be brittle; use URIs/args and let the task read/write via ray.data. - Prefer writing results (e.g., ds.write_parquet) or returning small aggregates; avoid returning large in-memory datasets. - You can still use

remote_options(...)

to set Ray resources for the Prefect tasks that run your ray.data pipeline. - If you need Ray runtime_env (extra pip deps, env vars, etc.), set it on the RayTaskRunner via

init_kwargs={"runtime_env": {...}}

. Docs - Ray integration: prefect-ray - RayTaskRunner API: RayTaskRunner API If you tell me your exact IO format (e.g., S3 parquet in/out), I can tailor the wrapper to that and include best practices for large outputs (partitioned writes, stats, and downstream consumption).

Awesome — here are two lightweight decorators you can drop into your codebase to get “map_batches -> task” ergonomics. Option A: stackable like @task @map_batches(...) - You decorate your batch function; the result is a callable you can wrap with @task. - You pass a Ray Dataset (or choose to read/write inside the task via kwargs).

# helpers.py
from functools import partial

def map_batches(*, batch_size=None, batch_format="pandas", **map_kwargs):
    """
    Turn a batch-transform fn(batch) -> batch into a function that:
      - accepts a Ray Dataset (ds)
      - applies ds.map_batches(fn, ...)
      - optionally writes output if write_* kwargs are provided
    Use with @task on top: @task @map_batches(...)
    """
    def decorator(transform_fn):
        def runner(
            ds=None,
            # Optional IO: build/read inside the task if you don't want to pass a Dataset
            read=None,              # callable -> Dataset (e.g., partial(ray.data.read_parquet, path))
            write=None,             # callable (Dataset -> None), e.g., lambda ds: ds.write_parquet(path)
            return_dataset=False,   # if True, return Dataset handle; else return materialized small results
            take_limit=None,        # limit for take() when not writing and not returning dataset
            **_,
        ):
            import ray.data as rd

            if ds is None:
                if read is None:
                    raise ValueError("Provide either a Ray Dataset via 'ds' or a 'read' callable that returns a Dataset")
                ds = read()

            out = ds.map_batches(
                transform_fn,
                batch_size=batch_size,
                batch_format=batch_format,
                **map_kwargs,
            )

            if write is not None:
                write(out)
                return {"rows": out.count()}

            if return_dataset:
                # Only do this if you know how you'll consume the Ray Dataset downstream
                return out

            # For small results only; avoid collecting massive outputs
            if take_limit is None:
                return out.take_all()
            return out.take(take_limit)
        return runner
    return decorator

Usage

# pipeline.py
import os
import ray.data as rd
from functools import partial
from prefect import flow, task
from prefect_ray import RayTaskRunner
from prefect_ray.context import remote_options
from helpers import map_batches

@task
@map_batches(batch_size=1000, batch_format="pandas")
def enrich(batch):
    batch["value_x2"] = batch["value"] * 2
    return batch

@flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS")))
def pipeline(src_parquet: str, dst_parquet: str):
    # Option 1: let the task handle I/O via callables (recommended; no big objects passed)
    with remote_options(num_cpus=2):
        result = enrich.submit(
            read=partial(rd.read_parquet, src_parquet),
            write=lambda ds: ds.write_parquet(dst_parquet),
        ).result()
    return result

Option B: single decorator that returns a Prefect task @task_map_batches(...) - Decorate your batch fn once; it becomes a Prefect task ready to submit. ``` # helpers.py (add this too) from typing import Optional, Callable from prefect import task def task_map_batches( *, reader: Callable[..., "ray.data.Dataset"], writer: Optional[Callable[["ray.data.Dataset", ...], None]] = None, batch_size: Optional[int] = None, batch_format: str = "pandas", **map_kwargs, ): """ Wrap a batch-transform fn into a Prefect task that: - builds a Ray Dataset via

reader(*reader_args, **reader_kwargs)

- applies map_batches(transform, ...) - optionally persists output via

writer(dataset, *writer_args, **writer_kwargs)

""" def decorator(transform_fn): @task def _task( *reader_args, _writer_args=(), _writer_kwargs=None,

**reader_kwargs, ): import ray.data as rd # noqa: F401 (type hints) ds = reader(*reader_args, **reader_kwargs) out = ds.map_batches( transform_fn, batch_size=batch_size, batch_format=batch_format, **map_kwargs, ) if writer: _writer_kwargs = _writer_kwargs or {} writer(out, *_writer_args, **_writer_kwargs) return {"rows": out.count()} # For small outputs only; prefer setting a writer return out.take_all() return _task return decorator

Usage

# pipeline_task_deco.py import os import ray.data as rd from prefect import flow from prefect_ray import RayTaskRunner from prefect_ray.context import remote_options from helpers import task_map_batches @task_map_batches( reader=rd.read_parquet, writer=lambda ds, path: ds.write_parquet(path), batch_size=2000, batch_format="pandas", ) def enrich(batch): batch["value_x2"] = batch["value"] * 2 return batch @flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS"))) def pipeline(src_parquet: str, dst_parquet: str): with remote_options(num_cpus=2): return enrich.submit(src_parquet, _writer_args=(dst_parquet,)).result() ``` Notes and tips - Prefer writing results (writer=...) or returning small aggregates; avoid returning large Datasets to Prefect. - Avoid lambdas inside submit in production if you hit pickling issues; use top-level functions or functools.partial. - You can combine with remote_options to shape Ray resources for the Prefect task(s) you submit. - Prefect-Ray doesn’t ship a built-in map_batches wrapper; these are safe helpers that run entirely inside the task process where RayTaskRunner has already connected to your cluster. Docs - Ray integration: prefect-ray - RayTaskRunner API: RayTaskRunner API If you share the concrete IO format (e.g., S3 parquet → parquet), I can tailor the decorator defaults and include robust writer options (partitioned writes, filesystem args, etc.).

Short answer: use environment/job variables. Don’t use a flow parameter for this. Why env/job variables are best - The task runner is entered before your flow code runs. That means a flow parameter (Pydantic input) can’t influence the RayTaskRunner’s address in time. - Infra selection (local vs remote) is an environment concern. Job variables on the deployment or work pool reliably set env before the process starts. - Works cleanly across dev/stage/prod and with UI/CLI triggers; easy to audit/change without touching code. Recommended pattern - In code, read RAY_ADDRESS and fall back to ephemeral when unset:

import os
from prefect import flow
from prefect_ray import RayTaskRunner

@flow(task_runner=RayTaskRunner(address=os.getenv("RAY_ADDRESS")))
def my_flow(...):
    ...

- Use job variables to set RAY_ADDRESS for remote runs; leave it unset locally. prefect.yaml example

deployments:
  - name: my-flow-remote
    entrypoint: my_flow.py:my_flow
    work_pool:
      name: ray-worker-pool
      job_variables:
        env:
          RAY_ADDRESS: "<ray://ray-head.ray.svc.cluster.local:10001>"

  - name: my-flow-local
    entrypoint: my_flow.py:my_flow
    work_pool:
      name: local-pool
      job_variables:
        env: {}   # leave RAY_ADDRESS unset -> ephemeral Ray locally

Alternative toggles - Two deployments (local vs remote) is the simplest and most explicit. - You can also set RAY_ADDRESS at the work pool level so every job from that pool uses the remote cluster. What not to do - Don’t try to switch via a flow parameter; it’s evaluated after the task runner has already been started. - Don’t set os.environ["RAY_ADDRESS"] inside the flow to influence the runner; it’s too late. Docs - Ray integration: prefect-ray - RayTaskRunner API: RayTaskRunner API - Customize job variables: Customize job variables - Manage work pools: Manage work pools