runtime_env={"pip": [...]}

ray.util.client.*

torch

torch

torch

from prefect import task, flow
     from prefect_ray.task_runners import RayTaskRunner

     @task
     def gpu_task(x):
         import torch  # defer import to worker
         return int(torch.tensor([x]).sum().item())

     @flow(task_runner=RayTaskRunner(address="auto", init_kwargs={
         "runtime_env": {
             # This ensures workers have torch; it won't fix driver/head imports.
             "pip": [
                 "torch==2.2.1"  # add extra index-url if you need CUDA wheels
             ]
         }
     }))
     def my_flow():
         futures = [gpu_task.submit(i) for i in range(5)]
         return [f.result() for f in futures]

     if __name__ == "__main__":
         my_flow()

import torch

torch.*

def gpu_task(x: torch.Tensor)

runtime_env

runtime_env

pip

RayTaskRunner

print(
        f"using prefect-ray tasks: Shouting from {socket.gethostname()}: {os.getpid()}"
    )
    shout_task.map(range(10)).wait()
    print("Shouting done")

from __future__ import annotations

import os
import socket
import time

import ray
from prefect import task


@task
def shout_task(number: int) -> str:
    time.sleep(1)
    print(f"#{number} on {socket.gethostname()}: {os.getpid()}")
    return f"#{number} on {socket.gethostname()}: {os.getpid()}"

runtime_env={"pip": [...]}

RayTaskRunner

from prefect import task

    @task
    def gpu_task(x: int):
        import torch  # defer import so only the Ray worker needs it
        t = torch.tensor([x])
        return int(t.sum().item())

from __future__ import annotations

from typing import TYPE_CHECKING
    if TYPE_CHECKING:
        import torch

runtime_env

from prefect import task

  @task
  def gpu_infer(x: int):
      import torch
      return int(torch.tensor([x]).sum().item())

from prefect import flow
  from prefect_ray.task_runners import RayTaskRunner

  @flow(task_runner=RayTaskRunner(address="ray://<head-ip>:10001", init_kwargs={
      "runtime_env": {"pip": ["torch==2.2.1"]}  # ensures workers have torch
  }))
  def my_flow():
      # Safe: importing this module does not import torch at top-level
      from .tasks_torch import gpu_infer
      futures = [gpu_infer.submit(i) for i in range(10)]
      return [f.result() for f in futures]

runtime_env={"pip": [...]}

RayTaskRunner

runtime_env={"pip": [...]}

RayTaskRunner(address="auto")

runtime_env

RayTaskRunner

runtime_env

RayTaskRunner(address="auto")

runtime_env={"pip": ["torch==...", "yourlib @ https://.../yourlib.whl"], "py_modules": ["./your_local_lib_dir"]}

ray://...

RayTaskRunner(address="auto")

RayTaskRunner(address="auto")

RayTaskRunner(address="auto")

address="auto"

runtime_env

RayTaskRunner

ray://

from prefect import flow
from prefect_ray.task_runners import RayTaskRunner

@flow(task_runner=RayTaskRunner(address="auto", init_kwargs={
    "runtime_env": {
        "pip": [
            "torch==2.2.1",
            "yourlib @ https://.../yourlib.whl"
        ],
        # optional: "py_modules": ["./your_local_lib_dir"]
    }
}))
def my_flow():
    # Your decorators and imports at module scope are fine;
    # they are resolved in the driver image/env, not by the head.
    ...

from ray.job_submission import JobSubmissionClient

     client = JobSubmissionClient("http://<kuberay-head-svc>:8265")

     client.submit_job(
         # This script/module runs your Prefect flow; inside it, use RayTaskRunner(address="auto")
         entrypoint="python -m mypkg.run_prefect_flow",
         runtime_env={
             "pip": [
                 "prefect==3.*",
                 "prefect-ray==<compatible>",
                 "torch==2.2.1",
                 "yourlib @ https://.../yourlib.whl",
             ],
             # Optional: ship your code
             # "working_dir": "git+<https://github.com/org/repo@ref#subdirectory=flows>",
             # "py_modules": ["./local_pkg_dir"],
             "env_vars": {
                 "PREFECT_API_URL": "...",
                 "PREFECT_API_KEY": "...",
                 # any other auth/config needed by your flow
             },
         },
     )

mypkg.run_prefect_flow

RayTaskRunner(address="auto")

RayTaskRunner(address="auto")

runtime_env

RayTaskRunner(address="auto")

runtime_env

ray.init(address=ray://..., runtime_env=...)

@ray.remote

shout

shout

@ray.remote

address="auto"

runtime_env

RayTaskRunner(address="auto")

RayTaskRunner(address="auto")

/tmp/ray

emptyDir

address="auto"

# your-registry/prefect-driver:latest
FROM python:3.11-slim

# System deps as needed
RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*

# Base Python deps
RUN pip install --no-cache-dir "prefect==3.*" "prefect-ray>=0.4.0" "ray>=2.9"

# Optional: add your org-wide packages or CA certs, etc.
# COPY . /app  # if you ship your code in the image
WORKDIR /app

emptyDir

/tmp/ray

3) Prefect flow: use native Ray with address="auto"
- The driver (your flow run subprocess) runs inside the `prefect-sidecar` container, which is in the same pod as a Ray worker. Because we share `/tmp/ray`, `address="auto"` finds the local Ray session. The head still orchestrates but does not import your code.

prefect worker start --pool ...

/tmp/ray

address="auto"

runtime_env

apiVersion: <http://keda.sh/v1alpha1|keda.sh/v1alpha1>
  kind: ScaledObject
  metadata:
    name: prefect-worker-autoscaler
  spec:
    scaleTargetRef:
      name: prefect-worker-deployment
    minReplicaCount: 0
    maxReplicaCount: 3
    triggers:
      - type: metrics-api
        metadata:
          targetValue: "1"
          url: <http://prefect-queue-metrics.default.svc.cluster.local/queued_runs?pool=ray-process-pool>

runtime_env

address="auto"

/tmp/ray

address="auto"

apiVersion: <http://ray.io/v1|ray.io/v1>
kind: RayCluster
metadata:
  name: raycluster
spec:
  rayVersion: "2.9.3"
  headGroupSpec:
    serviceType: ClusterIP
    rayStartParams:
      dashboard-host: "0.0.0.0"
    template:
      spec:
        containers:
          - name: ray-head
            image: rayproject/ray:2.9.3-py311
            ports:
              - containerPort: 8265  # dashboard
              - containerPort: 10001 # ray client (unused in native mode)
            volumeMounts:
              - name: ray-session
                mountPath: /tmp/ray

          - name: prefect-sidecar
            image: your-registry/prefect-driver:latest
            env:
              - name: PREFECT_API_URL
                valueFrom:
                  secretKeyRef:
                    name: prefect-creds
                    key: PREFECT_API_URL
              - name: PREFECT_API_KEY
                valueFrom:
                  secretKeyRef:
                    name: prefect-creds
                    key: PREFECT_API_KEY
              - name: PREFECT_WORK_POOL
                value: "ray-head-pool"
            resources:
              requests:
                cpu: "100m"
                memory: "256Mi"
              limits:
                cpu: "500m"
                memory: "512Mi"
            volumeMounts:
              - name: ray-session
                mountPath: /tmp/ray
            command: ["/bin/sh", "-c"]
            args:
              - |
                echo "Waiting for local Ray to be ready..."
                while [ ! -S /tmp/ray/session_latest/sockets/raylet ]; do sleep 1; done
                echo "Starting Prefect worker for pool: ${PREFECT_WORK_POOL}"
                prefect worker start --pool "${PREFECT_WORK_POOL}"

        volumes:
          - name: ray-session
            emptyDir: {}

  workerGroupSpecs:
    - groupName: compute
      replicas: 3
      minReplicas: 0
      maxReplicas: 50
      rayStartParams: {}
      template:
        spec:
          containers:
            - name: ray-worker
              image: rayproject/ray:2.9.3-py311
              resources:
                requests:
                  cpu: "2"
                  memory: "4Gi"
                limits:
                  cpu: "4"
                  memory: "8Gi"
              volumeMounts:
                - name: ray-session
                  mountPath: /tmp/ray
        volumes:
          - name: ray-session
            emptyDir: {}

from prefect import flow
from prefect_ray.task_runners import RayTaskRunner

@flow(task_runner=RayTaskRunner(address="auto", init_kwargs={
    "runtime_env": {
        "pip": ["torch==2.2.1", "yourlib @ https://.../yourlib.whl"],
        # optional: "py_modules": ["./local_pkg_dir"], "working_dir": "git+https://..."
    }
}))
def my_flow():
    ...

runtime_env

node-pool=gpu-a100

apiVersion: <http://ray.io/v1|ray.io/v1>
kind: RayCluster
metadata:
  name: raycluster
spec:
  rayVersion: "2.9.3"
  headGroupSpec:
    serviceType: ClusterIP
    template:
      spec:
        containers:
          - name: ray-head
            image: rayproject/ray:2.9.3-py311
            ports:
              - containerPort: 8265
              - containerPort: 10001
            volumeMounts:
              - name: ray-session
                mountPath: /tmp/ray
        volumes:
          - name: ray-session
            emptyDir: {}

  workerGroupSpecs:
    - groupName: gpu-a100
      replicas: 2
      minReplicas: 0
      maxReplicas: 20
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            node-pool: gpu-a100         # <- match your GPU nodes
          tolerations:
            - key: "<http://nvidia.com/gpu|nvidia.com/gpu>"
              operator: "Exists"
              effect: "NoSchedule"
          containers:
            - name: ray-worker
              # Use a GPU-capable base with CUDA + torch, or your own image
              image: <http://nvcr.io/nvidia/pytorch:24.02-py3|nvcr.io/nvidia/pytorch:24.02-py3>
              resources:
                requests:
                  cpu: "4"
                  memory: "32Gi"
                  <http://nvidia.com/gpu|nvidia.com/gpu>: "1"
                limits:
                  cpu: "8"
                  memory: "64Gi"
                  <http://nvidia.com/gpu|nvidia.com/gpu>: "1"
              volumeMounts:
                - name: ray-session
                  mountPath: /tmp/ray
          volumes:
            - name: ray-session
              emptyDir: {}

address="auto"

runtime_env

ray start --address ...

prefect work-pool set-concurrency-limit "your-pool" 2

replicas: N

from ray.job_submission import JobSubmissionClient
import os

def submit_ray_job(
    head_dashboard_url: str,
    entrypoint: str = "python -m myproj.run_flow_entrypoint",
    pip_packages: list[str] = None,
    working_dir: str | None = None,
    py_modules: list[str] | None = None,
    env: dict[str, str] = None,
):
    client = JobSubmissionClient(head_dashboard_url)

    runtime_env = {}
    if pip_packages:
        runtime_env["pip"] = pip_packages
    if working_dir:
        runtime_env["working_dir"] = working_dir
    if py_modules:
        runtime_env["py_modules"] = py_modules
    if env:
        runtime_env["env_vars"] = env

    job_id = client.submit_job(
        entrypoint=entrypoint,
        runtime_env=runtime_env,
    )
    return job_id

# Example use inside a Prefect task/flow:
# job_id = submit_ray_job(
#   head_dashboard_url="<http://raycluster-head-svc>.<ns>.svc.cluster.local:8265",
#   pip_packages=[
#     "prefect==3.*", "prefect-ray>=0.4.0", "torch==2.2.1", "yourlib @ https://.../yourlib.whl"
#   ],
#   working_dir="git+<https://github.com/org/repo@ref#subdirectory=flows>",
#   env={
#     "PREFECT_API_URL": os.environ["PREFECT_API_URL"],
#     "PREFECT_API_KEY": os.environ["PREFECT_API_KEY"],
#     # any other per-run config vars
#   },
# )

# myproj/run_flow_entrypoint.py
import os
from prefect import flow
from prefect_ray.task_runners import RayTaskRunner

def _ray_runtime_env():
    # Optional: additional worker-only deps (if different from driver)
    return {
        # "pip": ["..."],
        # "py_modules": ["./local_pkg_dir"],
    }

@flow(task_runner=RayTaskRunner(
    # In a Ray Job, the driver runs inside the cluster. Use the cluster’s address.
    # If your job image includes ray, ray.init() will be called by RayTaskRunner.
    address="<ray://raycluster-head-svc>.<ns>.svc.cluster.local:10001",
    init_kwargs={"runtime_env": _ray_runtime_env()},
))
def my_flow():
    # Safe to import heavy libs at module scope: they’re installed in the job’s driver env.
    # Submit your Ray/Prefect work here.
    ...

if __name__ == "__main__":
    my_flow()

ray.init(address="auto")

ray://<head-svc>:10001

remote_options(runtime_env=...)

@ray.remote(...)

from prefect import flow, get_run_logger
from ray.job_submission import JobSubmissionClient
import os, json

@flow
def submit_preproc_ray_job(params: dict):
    logger = get_run_logger()
    client = JobSubmissionClient("<http://raycluster-head-svc>.<ns>.svc.cluster.local:8265")

    # Driver deps for preprocessing + running the flow
    runtime_env = {
        "pip": [
            "prefect==3.*",
            "prefect-ray>=0.4.0",
            "ray>=2.9",
            # add driver-only deps here (e.g., torch) if preprocessing needs GPU locally
            # "torch==2.2.1",
            # yourlocal pkg as a wheel
            # "yourlib @ https://.../yourlib.whl",
        ],
        # Choose one of these to ship your code
        # "working_dir": "git+<https://github.com/org/repo@ref#subdirectory=flows>",
        # or package a module
        # "py_modules": ["./your_project_dir"],
        "env_vars": {
            "PREFECT_API_URL": os.environ["PREFECT_API_URL"],
            "PREFECT_API_KEY": os.environ["PREFECT_API_KEY"],
            # pass flow params if you like
            "FLOW_PARAMS_JSON": json.dumps(params),
        },
    }

    job_id = client.submit_job(
        entrypoint="python -m myproj.run_flow_entrypoint",
        runtime_env=runtime_env,
    )
    <http://logger.info|logger.info>(f"Submitted Ray Job id={job_id}")

remote_options(resources={"res_name": 0.01})

from prefect import flow, get_run_logger
from prefect.deployments import run_deployment
from ray.job_submission import JobSubmissionClient
import os

@flow
def launcher(params: dict):
    logger = get_run_logger()

    # 1) Create the child run and mark it as a subflow of this launcher
    fr = run_deployment(
        name="main-flow/main-deployment",
        parameters=params,
        as_subflow=True,  # critical for linking in the UI
    )
    <http://logger.info|logger.info>(f"Created child flow run: {fr.id}")

    # 2) Submit a Ray Job that will EXECUTE this run id
    client = JobSubmissionClient("<http://raycluster-head-svc>.<ns>.svc.cluster.local:8265")

    job_id = client.submit_job(
        entrypoint=f"prefect flow-run execute {fr.id}",
        runtime_env={
            "pip": [
                "prefect==3.*",
                "prefect-ray>=0.4.0",
                "ray>=2.9",
                # add any driver-time deps if needed for preprocessing logs, etc.
            ],
            # If your deployment uses flow.from_source(...), Prefect will load code from source
            # per the deployment config; you do NOT need to ship code here.
            "env_vars": {
                "PREFECT_API_URL": os.environ["PREFECT_API_URL"],
                "PREFECT_API_KEY": os.environ["PREFECT_API_KEY"],
            },
        },
    )
    <http://logger.info|logger.info>(f"Submitted Ray Job: {job_id}")

    # Optional: poll flow run state here if you want the launcher to wait
    return {"flow_run_id": str(fr.id), "ray_job_id": job_id}

prefect flow-run execute <id>

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

@task
def preproc():
    # heavy prep, optionally GPU via Ray tasks (set num_gpus=1 in remote_options)
    return [{"item": 1, "pip": ["torch==2.2.1"]}, {"item": 2, "pip": ["xgboost==2.0.3"]}]

@task
def unit_of_work(x: int) -> str:
    return f"processed {x}"

@flow(task_runner=RayTaskRunner(address="<ray://raycluster-head-svc>.<ns>.svc.cluster.local:10001"))
def main_flow():
    items = preproc.submit().result()

    # group by identical deps so Ray reuses envs
    from collections import defaultdict
    by_env = defaultdict(list)
    for spec in items:
        key = tuple(sorted(spec.get("pip", [])))
        by_env[key].append(spec["item"])

    results = []
    for pip_key, batch in by_env.items():
        with remote_options(runtime_env={"pip": list(pip_key)}, num_cpus=2):
            futures = [unit_of_work.submit(x) for x in batch]
            results.extend([f.result() for f in futures])

    return results

runtime_env

address="auto"

prefect flow-run execute <id>

#!/bin/sh
set -eu

# Required env:
# - FLOW_RUN_ID: the Prefect flow-run id to execute
# - RUN_SPEC_JSON: JSON with per-run Python deps (e.g., {"pip":["torch==2.2.1","yourlib @ https://.../yourlib.whl"]})
# - PREFECT_API_URL / PREFECT_API_KEY: so the run can report to Prefect
# Optional:
# - RAY_ADDRESS: if you prefer explicit address; otherwise we’ll use address="auto"
# - VENV_ROOT: path where we cache envs (default /opt/prefect-venvs)
# - WHEEL_CACHE: path for pip cache (default /opt/pip-cache)

: "${FLOW_RUN_ID:?FLOW_RUN_ID is required}"
: "${RUN_SPEC_JSON:?RUN_SPEC_JSON is required}"
: "${PREFECT_API_URL:?PREFECT_API_URL is required}"
: "${PREFECT_API_KEY:?PREFECT_API_KEY is required}"
VENV_ROOT="${VENV_ROOT:-/opt/prefect-venvs}"
WHEEL_CACHE="${WHEEL_CACHE:-/opt/pip-cache}"

echo "Waiting for local Ray to be ready..."
# KubeRay commonly exposes the raylet socket here; adjust if your Ray version differs
while [ ! -S /tmp/ray/session_latest/sockets/raylet ]; do sleep 1; done
echo "Ray local session is ready."

# Create cache directories if they don’t exist
mkdir -p "${VENV_ROOT}" "${WHEEL_CACHE}"

# Bootstrap an env and exec the run by id
exec python /app/bootstrap.py \
  --run-spec "${RUN_SPEC_JSON}" \
  --venv-root "${VENV_ROOT}" \
  --wheel-cache "${WHEEL_CACHE}" \
  --cmd "prefect flow-run execute ${FLOW_RUN_ID}"

prefect flow-run execute <id>

# flows.py
from prefect import flow, task
from prefect_kubernetes.experimental import kubernetes

@task
def preprocess(x: int) -> int:
    # trivial; replace with your real preprocessing
    return x + 1

# This entire flow runs as a single Kubernetes Job/Pod on pool-GPU,
# using image B (it does NOT inherit the entrypoint's image A).
@kubernetes(
    work_pool="pool-GPU",
    image="<http://ghcr.io/acme/ml_detection_image:B|ghcr.io/acme/ml_detection_image:B>",  # <- image B for the decorated flow
    # You can add other job variables here (e.g., cpu/memory requests, SA name, include_files)
)
@flow
def flow_decorated(data: list[int]) -> list[int]:
    # All tasks inside flow_decorated execute within the same Job/Pod (no per-task Pods)
    futures = [preprocess.submit(i) for i in data]
    return [f.result() for f in futures]

# This flow is your UI entrypoint; create a deployment for THIS flow with image A.
@flow
def entrypoint_flow(n: int = 10) -> list[int]:
    # Runs in image A (the deployment image), then delegates to flow_decorated.
    data = list(range(n))
    return flow_decorated(data)

namespace="ray-worker-only"

/tmp/ray

address="auto"

ray://<head-svc>:10001

/tmp/ray

RayTaskRunner(address="auto")

ray://…

address="auto"

# Inside the Job pod before your flow starts
ray start --address=raycluster-head-svc.<ns>.svc.cluster.local:6379 --block &
# wait for /tmp/ray/session_latest/sockets/raylet to exist
# then run your flow which does RayTaskRunner(address="auto")

# src/flow.py
import torch
from prefect import flow, task
from prefect_ray.task_runners import RayTaskRunner

@task
def gpu_inference_task(partition_id: int):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return f"Processed {partition_id} on {device}"

@flow(task_runner=RayTaskRunner(address="auto"))  # native: attach to local raylet
def heavy_execution_flow(num_partitions: int = 10000):
    futures = [gpu_inference_task.submit(i) for i in range(num_partitions)]
    return [f.result() for f in futures]

prefect flow-run execute

PREFECT__FLOW_RUN_ID

# deploy.py
from src.flow import heavy_execution_flow

if __name__ == "__main__":
    RAY_HEAD_GCS = "ray-head-svc.ray-prod.svc.cluster.local:6379"

    # Provide command as a list to avoid shell-quoting surprises
    hijacked_entrypoint = [
        "/bin/bash", "-lc",
        # 1) Start a local raylet that joins the cluster’s GCS
        # 2) Wait for the local ray session socket
        # 3) Exec the supported Prefect entrypoint
        f"ray start --address={RAY_HEAD_GCS} --block & "
        "until [ -S /tmp/ray/session_latest/sockets/raylet ]; do sleep 1; done; "
        "exec prefect flow-run execute"
    ]

    heavy_execution_flow.deploy(
        name="ml-detect-a-native-driver",
        work_pool_name="kubernetes-prod-pool",
        image="my-registry/ml_detect_a:latest",   # Torch + any driver-time deps
        job_variables={
            "command": hijacked_entrypoint
            # if needed, also set cpu/memory/gpu here or in the pool’s base template
        },
    )

from prefect import flow
from src.flow import heavy_execution_flow

if __name__ == "__main__":
    RAY_HEAD_GCS = "ray-head-svc.ray-prod.svc.cluster.local:6379"

    hijacked_entrypoint = [
        "/bin/bash", "-lc",
        f"ray start --address={RAY_HEAD_GCS} --block & "
        "until [ -S /tmp/ray/session_latest/sockets/raylet ]; do sleep 1; done; "
        "exec prefect flow-run execute"
    ]

    heavy_execution_flow.from_source(
        source="git+<https://github.com/org/repo.git@main#subdirectory=>.",
        entrypoint="src/flow.py:heavy_execution_flow",
    ).deploy(
        name="ml-detect-a-native-driver",
        work_pool_name="kubernetes-prod-pool",
        image="my-registry/ml_detect_a:latest",
        job_variables={"command": hijacked_entrypoint},
    )

prefect flow-run execute

PREFECT__FLOW_RUN_ID

/tmp/ray/session_latest/sockets/raylet

6379

address="auto"

--num-cpus=0 --num-gpus=0

--resources='{"CPU":0,"GPU":0}'

remote_options(num_cpus=..., num_gpus=..., resources=...)

{"worker_group:ml_detection": 0.01}