Skip to content
This repository has been archived by the owner on Apr 26, 2024. It is now read-only.

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

Merged
merged 5 commits into from
Oct 25, 2022

Conversation

ahuang11
Copy link
Contributor

@ahuang11 ahuang11 commented Oct 19, 2022

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 #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

  • This pull request references any related issue by including "Closes #<ISSUE_NUMBER>"
    • 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.
  • This pull request includes tests or only affects documentation.
  • Summarized PR's changes in [CHANGELOG.md](https://github.com/{{ cookiecutter.github_organization }}/{{ cookiecutter.collection_name }}/blob/main/CHANGELOG.md)

prefect_ray/context.py Outdated Show resolved Hide resolved
README.md Outdated Show resolved Hide resolved
@ahuang11 ahuang11 requested review from zanieb and desertaxle October 24, 2022 17:00
desertaxle
desertaxle previously approved these changes Oct 25, 2022
Copy link
Member

@desertaxle desertaxle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

prefect_ray/context.py Outdated Show resolved Hide resolved
prefect_ray/context.py Outdated Show resolved Hide resolved
Co-authored-by: Michael Adkins <michael@prefect.io>
@ahuang11 ahuang11 merged commit 6089262 into main Oct 25, 2022
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

How to properly execute flows with cpu/gpu requirements
3 participants