Dynamic resource example in Ray

dynamic_resource_example.py

import time
import ray
from ray.tests.cluster_utils import Cluster

# Create a cluster with some worker nodes with 1 CPU each, to force colocated
# tasks to run one at a time.
cluster = Cluster(initialize_head=True, connect=True, head_node_args={'num_cpus': 0})
num_nodes = 3
for _ in range(num_nodes):
    cluster.add_node(num_cpus=1)

# Define a custom resource to control scheduling.
CUSTOM_RESOURCE = 'my_custom_resource'

# Can run on any node.
@ray.remote
def scheduler_task():
    time.sleep(1)
    # Adds the custom resource to the local node.
    ray.experimental.set_resource(CUSTOM_RESOURCE, 1)

# Can only run on a node where CUSTOM_RESOURCE is available.
@ray.remote(resources={CUSTOM_RESOURCE: 1})
def colocated_task():
    time.sleep(1)

# Launch the tasks. The colocated_task will not be able to run until the
# scheduler_task adds the custom resource.
tasks = [scheduler_task.remote(), colocated_task.remote()]
start = time.time()
ray.get(tasks)
print("Took", time.time() - start, "seconds")  # Should take >2s.