[xray] Track ray.get calls as task dependencies (#2362)

This commit is contained in:
Stephanie Wang
2018-07-27 14:59:17 -04:00
committed by Robert Nishihara
parent 5b015f9a79
commit 6675361684
20 changed files with 472 additions and 198 deletions
+2 -1
View File
@@ -96,7 +96,8 @@ class TestGlobalScheduler(unittest.TestCase):
static_resources={"CPU": 10})
# Connect to the scheduler.
local_scheduler_client = local_scheduler.LocalSchedulerClient(
local_scheduler_name, NIL_WORKER_ID, False, False)
local_scheduler_name, NIL_WORKER_ID, False, random_task_id(),
False)
self.local_scheduler_clients.append(local_scheduler_client)
self.local_scheduler_pids.append(p4)
+1 -1
View File
@@ -46,7 +46,7 @@ class TestLocalSchedulerClient(unittest.TestCase):
plasma_store_name, use_valgrind=USE_VALGRIND)
# Connect to the scheduler.
self.local_scheduler_client = local_scheduler.LocalSchedulerClient(
scheduler_name, NIL_WORKER_ID, False, False)
scheduler_name, NIL_WORKER_ID, False, random_task_id(), False)
def tearDown(self):
# Check that the processes are still alive.
+19 -12
View File
@@ -503,15 +503,6 @@ class Worker(object):
# get them until at least get_timeout_milliseconds
# milliseconds passes, then repeat.
while len(unready_ids) > 0:
for unready_id in unready_ids:
if not self.use_raylet:
self.local_scheduler_client.reconstruct_objects(
[ray.ObjectID(unready_id)], False)
# Do another fetch for objects that aren't available
# locally yet, in case they were evicted since the last
# fetch. We divide the fetch into smaller fetches so as
# to not block the manager for a prolonged period of time
# in a single call.
object_ids_to_fetch = [
plasma.ObjectID(unready_id)
for unready_id in unready_ids.keys()
@@ -525,6 +516,18 @@ class Worker(object):
for i in range(0, len(object_ids_to_fetch),
fetch_request_size):
if not self.use_raylet:
for unready_id in ray_object_ids_to_fetch[i:(
i + fetch_request_size)]:
(self.local_scheduler_client.
reconstruct_objects([unready_id], False))
# Do another fetch for objects that aren't
# available locally yet, in case they were evicted
# since the last fetch. We divide the fetch into
# smaller fetches so as to not block the manager
# for a prolonged period of time in a single call.
# This is only necessary for legacy ray since
# reconstruction and fetch are implemented by
# different processes.
self.plasma_client.fetch(object_ids_to_fetch[i:(
i + fetch_request_size)])
else:
@@ -2162,9 +2165,6 @@ def connect(info,
else:
local_scheduler_socket = info["raylet_socket_name"]
worker.local_scheduler_client = ray.local_scheduler.LocalSchedulerClient(
local_scheduler_socket, worker.worker_id, is_worker, worker.use_raylet)
# If this is a driver, set the current task ID, the task driver ID, and set
# the task index to 0.
if mode in [SCRIPT_MODE, SILENT_MODE]:
@@ -2219,6 +2219,13 @@ def connect(info,
# Set the driver's current task ID to the task ID assigned to the
# driver task.
worker.current_task_id = driver_task.task_id()
else:
# A non-driver worker begins without an assigned task.
worker.current_task_id = ray.ObjectID(NIL_ID)
worker.local_scheduler_client = ray.local_scheduler.LocalSchedulerClient(
local_scheduler_socket, worker.worker_id, is_worker,
worker.current_task_id, worker.use_raylet)
# Start the import thread
import_thread.ImportThread(worker, mode).start()