PrefectHQ / prefect-ray

Prefect integrations with Ray
https://prefecthq.github.io/prefect-ray/
Apache License 2.0
63 stars 5 forks source link

Add resources context manager to set CPUs/GPUs for task #54

Closed ahuang11 closed 1 year ago

ahuang11 commented 2 years ago

This PR allows users to specify required resources, like CPUs and GPUs, for a task.

To accomplish this, we collect resources kwargs from a user through a context manager (see example below), and then pass this to ray.remote internally in submit as suggested by Michael..

Closes https://github.com/PrefectHQ/prefect-ray/issues/44

Example

Use 4 CPUs and 2 GPUs for the process task:

from prefect import flow, task
from prefect_ray.task_runners import RayTaskRunner
from prefect_ray.context import resources
@task
def process(x):
    return x + 1
@flow(task_runner=RayTaskRunner())
def my_flow():
    # equivalent to setting @ray.remote(num_cpus=4, num_gpus=2)
    with resources(num_cpus=4, num_gpus=2):
        process.submit(42)

Checklist