mirror of
https://github.com/wassname/ray.git
synced 2026-07-11 11:59:36 +08:00
[xray] Track ray.get calls as task dependencies (#2362)
This commit is contained in:
committed by
Robert Nishihara
parent
5b015f9a79
commit
6675361684
@@ -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)
|
||||
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user