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 submitas suggested by Michael..
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
[x] This pull request references any related issue by including "Closes #"
If no issue exists and your change is not a small fix, please [create an issue](https://github.com/{{ cookiecutter.github_organization }}/{{ cookiecutter.collection_name }}/issues/new/choose) first.
[x] This pull request includes tests or only affects documentation.
[x] Summarized PR's changes in [CHANGELOG.md](https://github.com/{{ cookiecutter.github_organization }}/{{ cookiecutter.collection_name }}/blob/main/CHANGELOG.md)
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 insubmit
as suggested by Michael..Closes https://github.com/PrefectHQ/prefect-ray/issues/44
Example
Use 4 CPUs and 2 GPUs for the
process
task:Checklist