Skip to content

feat: Allow to set gpu and worker scheduling params for ray remote tasks#6323

Open
ntkathole wants to merge 1 commit intofeast-dev:masterfrom
ntkathole:ray_gpu
Open

feat: Allow to set gpu and worker scheduling params for ray remote tasks#6323
ntkathole wants to merge 1 commit intofeast-dev:masterfrom
ntkathole:ray_gpu

Conversation

@ntkathole
Copy link
Copy Markdown
Member

@ntkathole ntkathole commented Apr 23, 2026

What this PR does / why we need it:

This PR adds first-class GPU support and a general worker_task_options escape
hatch for any Ray .options() scheduling parameter, wired consistently
through all execution modes (local, remote, KubeRay).

Example config

batch_engine:
    type: ray.engine
    num_gpus: 1
    gpu_batch_format: numpy
    worker_task_options:
        num_cpus: 4
        memory: 8589934592        # 8 GB
        accelerator_type: "A100"  # pin to A100 on mixed GPU pools (KubeRay)
        max_retries: 5
        runtime_env:
            pip: ["cudf-cu12==24.10.0"]
            env_vars: {CUDA_VISIBLE_DEVICES: "0"}

Open in Devin Review

@ntkathole ntkathole self-assigned this Apr 23, 2026
@ntkathole ntkathole requested a review from a team as a code owner April 23, 2026 17:21
@ntkathole ntkathole changed the title feat: Allow to set gpu and worker scheduling params for ray feat: Allow to set gpu and worker scheduling params for ray remote tasks Apr 23, 2026
devin-ai-integration[bot]

This comment was marked as resolved.

devin-ai-integration[bot]

This comment was marked as resolved.

devin-ai-integration[bot]

This comment was marked as resolved.

Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant