Skip to main content

prefect_ray.context

Contexts to manage Ray clusters and tasks.

Functions

remote_options

remote_options(**new_remote_options: Dict[str, Any]) -> Generator[None, Dict[str, Any], None]
Context manager to add keyword arguments to Ray @remote calls for task runs. If contexts are nested, new options are merged with options in the outer context. If a key is present in both, the new option will be used. Examples: 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 remote_options

@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 remote_options(num_cpus=4, num_gpus=2):
        process.submit(42)

Classes

RemoteOptionsContext

The context for Ray remote_options management. Attributes:
  • current_remote_options: A set of current remote_options in the context.
Methods:

get

get(cls) -> 'RemoteOptionsContext'
Return an empty RemoteOptionsContext instead of None if no context exists.