mirror of
https://github.com/wassname/ray.git
synced 2026-07-18 12:40:56 +08:00
Move worker methods into Worker class and expose more TaskSpec fields to Python. (#796)
* Move worker methods inside worker class. Move some helper methods from actor.py into utils.py and state.py. * Add more methods exposing task spec fields to Python. * Fix linting. * Fix error. * Remove unused code in default worker.
This commit is contained in:
Robert Nishihara
committed by
Philipp Moritz
Philipp Moritz
parent
52a27be364
commit
8c8258de20
+8
-134
@@ -6,15 +6,13 @@ import cloudpickle as pickle
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import numpy as np
|
||||
import redis
|
||||
import traceback
|
||||
|
||||
import ray.local_scheduler
|
||||
import ray.signature as signature
|
||||
import ray.worker
|
||||
from ray.utils import (FunctionProperties, binary_to_hex, hex_to_binary,
|
||||
random_string)
|
||||
from ray.utils import (FunctionProperties, random_string,
|
||||
select_local_scheduler)
|
||||
|
||||
|
||||
def random_actor_id():
|
||||
@@ -102,117 +100,6 @@ def fetch_and_register_actor(actor_class_key, worker):
|
||||
# for the actor.
|
||||
|
||||
|
||||
def attempt_to_reserve_gpus(num_gpus, driver_id, local_scheduler, worker):
|
||||
"""Attempt to acquire GPUs on a particular local scheduler for an actor.
|
||||
|
||||
Args:
|
||||
num_gpus: The number of GPUs to acquire.
|
||||
driver_id: The ID of the driver responsible for creating the actor.
|
||||
local_scheduler: Information about the local scheduler.
|
||||
|
||||
Returns:
|
||||
True if the GPUs were successfully reserved and false otherwise.
|
||||
"""
|
||||
assert num_gpus != 0
|
||||
local_scheduler_id = local_scheduler["DBClientID"]
|
||||
local_scheduler_total_gpus = int(local_scheduler["NumGPUs"])
|
||||
|
||||
success = False
|
||||
|
||||
# Attempt to acquire GPU IDs atomically.
|
||||
with worker.redis_client.pipeline() as pipe:
|
||||
while True:
|
||||
try:
|
||||
# If this key is changed before the transaction below (the
|
||||
# multi/exec block), then the transaction will not take place.
|
||||
pipe.watch(local_scheduler_id)
|
||||
|
||||
# Figure out which GPUs are currently in use.
|
||||
result = worker.redis_client.hget(local_scheduler_id,
|
||||
"gpus_in_use")
|
||||
gpus_in_use = dict() if result is None else json.loads(
|
||||
result.decode("ascii"))
|
||||
num_gpus_in_use = 0
|
||||
for key in gpus_in_use:
|
||||
num_gpus_in_use += gpus_in_use[key]
|
||||
assert num_gpus_in_use <= local_scheduler_total_gpus
|
||||
|
||||
pipe.multi()
|
||||
|
||||
if local_scheduler_total_gpus - num_gpus_in_use >= num_gpus:
|
||||
# There are enough available GPUs, so try to reserve some.
|
||||
# We use the hex driver ID in hex as a dictionary key so
|
||||
# that the dictionary is JSON serializable.
|
||||
driver_id_hex = binary_to_hex(driver_id)
|
||||
if driver_id_hex not in gpus_in_use:
|
||||
gpus_in_use[driver_id_hex] = 0
|
||||
gpus_in_use[driver_id_hex] += num_gpus
|
||||
|
||||
# Stick the updated GPU IDs back in Redis
|
||||
pipe.hset(local_scheduler_id, "gpus_in_use",
|
||||
json.dumps(gpus_in_use))
|
||||
success = True
|
||||
|
||||
pipe.execute()
|
||||
# If a WatchError is not raised, then the operations should
|
||||
# have gone through atomically.
|
||||
break
|
||||
except redis.WatchError:
|
||||
# Another client must have changed the watched key between the
|
||||
# time we started WATCHing it and the pipeline's execution. We
|
||||
# should just retry.
|
||||
success = False
|
||||
continue
|
||||
|
||||
return success
|
||||
|
||||
|
||||
def select_local_scheduler(local_schedulers, num_gpus, worker):
|
||||
"""Select a local scheduler to assign this actor to.
|
||||
|
||||
Args:
|
||||
local_schedulers: A list of dictionaries of information about the local
|
||||
schedulers.
|
||||
num_gpus (int): The number of GPUs that must be reserved for this
|
||||
actor.
|
||||
|
||||
Returns:
|
||||
The ID of the local scheduler that has been chosen.
|
||||
|
||||
Raises:
|
||||
Exception: An exception is raised if no local scheduler can be found
|
||||
with sufficient resources.
|
||||
"""
|
||||
driver_id = worker.task_driver_id.id()
|
||||
|
||||
local_scheduler_id = None
|
||||
# Loop through all of the local schedulers in a random order.
|
||||
local_schedulers = np.random.permutation(local_schedulers)
|
||||
for local_scheduler in local_schedulers:
|
||||
if local_scheduler["NumCPUs"] < 1:
|
||||
continue
|
||||
if local_scheduler["NumGPUs"] < num_gpus:
|
||||
continue
|
||||
if num_gpus == 0:
|
||||
local_scheduler_id = hex_to_binary(local_scheduler["DBClientID"])
|
||||
break
|
||||
else:
|
||||
# Try to reserve enough GPUs on this local scheduler.
|
||||
success = attempt_to_reserve_gpus(num_gpus, driver_id,
|
||||
local_scheduler, worker)
|
||||
if success:
|
||||
local_scheduler_id = hex_to_binary(
|
||||
local_scheduler["DBClientID"])
|
||||
break
|
||||
|
||||
if local_scheduler_id is None:
|
||||
raise Exception("Could not find a node with enough GPUs or other "
|
||||
"resources to create this actor. The local scheduler "
|
||||
"information is {}.".format(local_schedulers))
|
||||
|
||||
return local_scheduler_id
|
||||
|
||||
|
||||
def export_actor_class(class_id, Class, actor_method_names, worker):
|
||||
if worker.mode is None:
|
||||
raise NotImplemented("TODO(pcm): Cache actors")
|
||||
@@ -255,17 +142,10 @@ def export_actor(actor_id, class_id, actor_method_names, num_cpus, num_gpus,
|
||||
num_gpus=0,
|
||||
max_calls=0))
|
||||
|
||||
# Get a list of the local schedulers from the client table.
|
||||
client_table = ray.global_state.client_table()
|
||||
local_schedulers = []
|
||||
for ip_address, clients in client_table.items():
|
||||
for client in clients:
|
||||
if (client["ClientType"] == "local_scheduler" and
|
||||
not client["Deleted"]):
|
||||
local_schedulers.append(client)
|
||||
# Select a local scheduler for the actor.
|
||||
local_scheduler_id = select_local_scheduler(local_schedulers, num_gpus,
|
||||
worker)
|
||||
local_scheduler_id = select_local_scheduler(
|
||||
worker.task_driver_id.id(), ray.global_state.local_schedulers(),
|
||||
num_gpus, worker.redis_client)
|
||||
assert local_scheduler_id is not None
|
||||
|
||||
# We must put the actor information in Redis before publishing the actor
|
||||
@@ -274,17 +154,12 @@ def export_actor(actor_id, class_id, actor_method_names, num_cpus, num_gpus,
|
||||
worker.redis_client.hmset(key, {"class_id": class_id,
|
||||
"num_gpus": num_gpus})
|
||||
|
||||
# Really we should encode this message as a flatbuffer object. However,
|
||||
# we're having trouble getting that to work. It almost works, but in Python
|
||||
# 2.7, builder.CreateString fails on byte strings that contain characters
|
||||
# outside range(128).
|
||||
|
||||
# TODO(rkn): There is actually no guarantee that the local scheduler that
|
||||
# we are publishing to has already subscribed to the actor_notifications
|
||||
# channel. Therefore, this message may be missed and the workload will
|
||||
# hang. This is a bug.
|
||||
worker.redis_client.publish("actor_notifications",
|
||||
actor_id.id() + driver_id + local_scheduler_id)
|
||||
ray.utils.publish_actor_creation(actor_id.id(), driver_id,
|
||||
local_scheduler_id, worker.redis_client)
|
||||
|
||||
|
||||
def actor(*args, **kwargs):
|
||||
@@ -319,8 +194,7 @@ def make_actor(cls, num_cpus, num_gpus):
|
||||
args = signature.extend_args(function_signature, args, kwargs)
|
||||
|
||||
function_id = get_actor_method_function_id(attr)
|
||||
object_ids = ray.worker.global_worker.submit_task(function_id, "",
|
||||
args,
|
||||
object_ids = ray.worker.global_worker.submit_task(function_id, args,
|
||||
actor_id=actor_id)
|
||||
if len(object_ids) == 1:
|
||||
return object_ids[0]
|
||||
|
||||
Reference in New Issue
Block a user