[xray] Implement timeline and profiling API. (#2306)

* Add profile table and store profiling information there. * Code for dumping timeline. * Improve color scheme. * Push timeline events on driver only for raylet. * Improvements to profiling and timeline visualization * Some linting * Small fix. * Linting * Propagate node IP address through profiling events. * Fix test. * object_id.hex() should return byte string in python 2. * Include gcs.fbs in node_manager.fbs. * Remove flatbuffer definition duplication. * Decode to unicode in Python 3 and bytes in Python 2. * Minor * Submit profile events in a batch. Revert some CMake changes. * Fix * Workaround test failure. * Fix linting * Linting * Don't return anything from chrome_tracing_dump when filename is provided. * Remove some redundancy from profile table. * Linting * Move TODOs out of docstring. * Minor
2026-08-18 12:20:14 +08:00 · 2018-07-04 23:23:48 -07:00
parent 8e687cbc98
commit b90e551b41
27 changed files with 777 additions and 147 deletions
@@ -562,7 +562,7 @@ class Worker(object):
        Returns:
            The return object IDs for this task.
        """
-        with log_span("ray:submit_task", worker=self):
+        with profile("submit_task", worker=self):
            check_main_thread()
            if actor_id is None:
                assert actor_handle_id is None
@@ -867,7 +867,7 @@ class Worker(object):

        # Get task arguments from the object store.
        try:
-            with log_span("ray:task:get_arguments", worker=self):
+            with profile("task:deserialize_arguments", worker=self):
                arguments = self._get_arguments_for_execution(
                    function_name, args)
        except (RayGetError, RayGetArgumentError) as e:
@@ -882,7 +882,7 @@ class Worker(object):

        # Execute the task.
        try:
-            with log_span("ray:task:execute", worker=self):
+            with profile("task:execute", worker=self):
                if task.actor_id().id() == NIL_ACTOR_ID:
                    outputs = function_executor(*arguments)
                else:
@@ -901,7 +901,7 @@ class Worker(object):

        # Store the outputs in the local object store.
        try:
-            with log_span("ray:task:store_outputs", worker=self):
+            with profile("task:store_outputs", worker=self):
                # If this is an actor task, then the last object ID returned by
                # the task is a dummy output, not returned by the function
                # itself. Decrement to get the correct number of return values.
@@ -976,7 +976,7 @@ class Worker(object):
        # Wait until the function to be executed has actually been registered
        # on this worker. We will push warnings to the user if we spend too
        # long in this loop.
-        with log_span("ray:wait_for_function", worker=self):
+        with profile("wait_for_function", worker=self):
            self._wait_for_function(function_id, driver_id)

        # Execute the task.
@@ -984,22 +984,26 @@ class Worker(object):
        # warning to the user if we are waiting too long to acquire the lock
        # because that may indicate that the system is hanging, and it'd be
        # good to know where the system is hanging.
-        log(event_type="ray:acquire_lock", kind=LOG_SPAN_START, worker=self)
        with self.lock:
-            log(event_type="ray:acquire_lock", kind=LOG_SPAN_END, worker=self)

            function_name = (self.function_execution_info[driver_id][
                function_id.id()]).function_name
-            contents = {
-                "function_name": function_name,
-                "task_id": task.task_id().hex(),
-                "worker_id": binary_to_hex(self.worker_id)
-            }
-            with log_span("ray:task", contents=contents, worker=self):
+            if not self.use_raylet:
+                extra_data = {
+                    "function_name": function_name,
+                    "task_id": task.task_id().hex(),
+                    "worker_id": binary_to_hex(self.worker_id)
+                }
+            else:
+                extra_data = {
+                    "name": function_name,
+                    "task_id": task.task_id().hex()
+                }
+            with profile("task", extra_data=extra_data, worker=self):
                self._process_task(task)

        # Push all of the log events to the global state store.
-        flush_log()
+        flush_profile_data()

        # Increase the task execution counter.
        self.num_task_executions[driver_id][function_id.id()] += 1
@@ -1017,7 +1021,7 @@ class Worker(object):
        Returns:
            A task from the local scheduler.
        """
-        with log_span("ray:get_task", worker=self):
+        with profile("get_task", worker=self):
            task = self.local_scheduler_client.get_task()

        # Automatically restrict the GPUs available to this task.
@@ -1103,7 +1107,7 @@ def _webui_url_helper(client):
        The URL of the web UI as a string.
    """
    result = client.hmget("webui", "url")[0]
-    return result.decode("ascii") if result is not None else result
+    return ray.utils.decode(result) if result is not None else result


 def get_webui_url():
@@ -1194,9 +1198,9 @@ def error_info(worker=global_worker):
        if error_applies_to_driver(error_key, worker=worker):
            error_contents = worker.redis_client.hgetall(error_key)
            error_contents = {
-                "type": error_contents[b"type"].decode("ascii"),
-                "message": error_contents[b"message"].decode("ascii"),
-                "data": error_contents[b"data"].decode("ascii")
+                "type": ray.utils.decode(error_contents[b"type"]),
+                "message": ray.utils.decode(error_contents[b"message"]),
+                "data": ray.utils.decode(error_contents[b"data"])
            }
            errors.append(error_contents)

@@ -1296,13 +1300,14 @@ def get_address_info_from_redis_helper(redis_address,
            assert b"ray_client_id" in info
            assert b"node_ip_address" in info
            assert b"client_type" in info
-            client_node_ip_address = info[b"node_ip_address"].decode("ascii")
+            client_node_ip_address = ray.utils.decode(info[b"node_ip_address"])
            if (client_node_ip_address == node_ip_address or
                (client_node_ip_address == "127.0.0.1"
                 and redis_ip_address == ray.services.get_node_ip_address())):
-                if info[b"client_type"].decode("ascii") == "plasma_manager":
+                if ray.utils.decode(info[b"client_type"]) == "plasma_manager":
                    plasma_managers.append(info)
-                elif info[b"client_type"].decode("ascii") == "local_scheduler":
+                elif (ray.utils.decode(
+                        info[b"client_type"]) == "local_scheduler"):
                    local_schedulers.append(info)
        # Make sure that we got at least one plasma manager and local
        # scheduler.
@@ -1311,16 +1316,16 @@ def get_address_info_from_redis_helper(redis_address,
        # Build the address information.
        object_store_addresses = []
        for manager in plasma_managers:
-            address = manager[b"manager_address"].decode("ascii")
+            address = ray.utils.decode(manager[b"manager_address"])
            port = services.get_port(address)
            object_store_addresses.append(
                services.ObjectStoreAddress(
-                    name=manager[b"store_socket_name"].decode("ascii"),
-                    manager_name=manager[b"manager_socket_name"].decode(
-                        "ascii"),
+                    name=ray.utils.decode(manager[b"store_socket_name"]),
+                    manager_name=ray.utils.decode(
+                        manager[b"manager_socket_name"]),
                    manager_port=port))
        scheduler_names = [
-            scheduler[b"local_scheduler_socket_name"].decode("ascii")
+            ray.utils.decode(scheduler[b"local_scheduler_socket_name"])
            for scheduler in local_schedulers
        ]
        client_info = {
@@ -1343,8 +1348,8 @@ def get_address_info_from_redis_helper(redis_address,
        for client_message in clients:
            client = ray.gcs_utils.ClientTableData.GetRootAsClientTableData(
                client_message, 0)
-            client_node_ip_address = client.NodeManagerAddress().decode(
-                "ascii")
+            client_node_ip_address = ray.utils.decode(
+                client.NodeManagerAddress())
            if (client_node_ip_address == node_ip_address or
                (client_node_ip_address == "127.0.0.1"
                 and redis_ip_address == ray.services.get_node_ip_address())):
@@ -1352,12 +1357,12 @@ def get_address_info_from_redis_helper(redis_address,

        object_store_addresses = [
            services.ObjectStoreAddress(
-                name=raylet.ObjectStoreSocketName().decode("ascii"),
+                name=ray.utils.decode(raylet.ObjectStoreSocketName()),
                manager_name=None,
                manager_port=None) for raylet in raylets
        ]
        raylet_socket_names = [
-            raylet.RayletSocketName().decode("ascii") for raylet in raylets
+            ray.utils.decode(raylet.RayletSocketName()) for raylet in raylets
        ]
        return {
            "node_ip_address": node_ip_address,
@@ -1807,6 +1812,21 @@ def custom_excepthook(type, value, tb):
 sys.excepthook = custom_excepthook


+def _flush_profile_events(worker):
+    """Drivers run this as a thread to flush profile data in the background."""
+    # Note(rkn): This is run on a background thread in the driver. It uses the
+    # local scheduler client. This should be ok because it doesn't read from
+    # the local scheduler client and we have the GIL here. However, if either
+    # of those things changes, then we could run into issues.
+    try:
+        while True:
+            time.sleep(1)
+            flush_profile_data(worker=worker)
+    except AttributeError:
+        # This is to suppress errors that occur at shutdown.
+        pass
+
+
 def print_error_messages_raylet(worker):
    """Print error messages in the background on the driver.

@@ -1858,7 +1878,7 @@ def print_error_messages_raylet(worker):
            if job_id not in [worker.task_driver_id.id(), NIL_JOB_ID]:
                continue

-            error_message = error_data.ErrorMessage().decode("ascii")
+            error_message = ray.utils.decode(error_data.ErrorMessage())

            if error_message not in old_error_messages:
                logger.error(error_message)
@@ -1900,8 +1920,8 @@ def print_error_messages(worker):
        error_keys = worker.redis_client.lrange("ErrorKeys", 0, -1)
        for error_key in error_keys:
            if error_applies_to_driver(error_key, worker=worker):
-                error_message = worker.redis_client.hget(
-                    error_key, "message").decode("ascii")
+                error_message = ray.utils.decode(
+                    worker.redis_client.hget(error_key, "message"))
                if error_message not in old_error_messages:
                    logger.error(error_message)
                    old_error_messages.add(error_message)
@@ -1915,8 +1935,8 @@ def print_error_messages(worker):
                for error_key in worker.redis_client.lrange(
                        "ErrorKeys", num_errors_received, -1):
                    if error_applies_to_driver(error_key, worker=worker):
-                        error_message = worker.redis_client.hget(
-                            error_key, "message").decode("ascii")
+                        error_message = ray.utils.decode(
+                            worker.redis_client.hget(error_key, "message"))
                        if error_message not in old_error_messages:
                            logger.error(error_message)
                            old_error_messages.add(error_message)
@@ -1939,9 +1959,9 @@ def fetch_and_register_remote_function(key, worker=global_worker):
         "module", "resources", "max_calls"
     ])
    function_id = ray.ObjectID(function_id_str)
-    function_name = function_name.decode("ascii")
+    function_name = ray.utils.decode(function_name)
    max_calls = int(max_calls)
-    module = module.decode("ascii")
+    module = ray.utils.decode(module)

    # This is a placeholder in case the function can't be unpickled. This will
    # be overwritten if the function is successfully registered.
@@ -2031,15 +2051,18 @@ def import_thread(worker, mode):
            # Handle the driver case first.
            if mode != WORKER_MODE:
                if key.startswith(b"FunctionsToRun"):
-                    fetch_and_execute_function_to_run(key, worker=worker)
+                    with profile("fetch_and_run_function", worker=worker):
+                        fetch_and_execute_function_to_run(key, worker=worker)
                # Continue because FunctionsToRun are the only things that the
                # driver should import.
                continue

            if key.startswith(b"RemoteFunction"):
-                fetch_and_register_remote_function(key, worker=worker)
+                with profile("register_remote_function", worker=worker):
+                    fetch_and_register_remote_function(key, worker=worker)
            elif key.startswith(b"FunctionsToRun"):
-                fetch_and_execute_function_to_run(key, worker=worker)
+                with profile("fetch_and_run_function", worker=worker):
+                    fetch_and_execute_function_to_run(key, worker=worker)
            elif key.startswith(b"ActorClass"):
                # Keep track of the fact that this actor class has been
                # exported so that we know it is safe to turn this worker into
@@ -2063,9 +2086,8 @@ def import_thread(worker, mode):
                    # Handle the driver case first.
                    if mode != WORKER_MODE:
                        if key.startswith(b"FunctionsToRun"):
-                            with log_span(
-                                    "ray:import_function_to_run",
-                                    worker=worker):
+                            with profile(
+                                    "fetch_and_run_function", worker=worker):
                                fetch_and_execute_function_to_run(
                                    key, worker=worker)
                        # Continue because FunctionsToRun are the only things
@@ -2073,13 +2095,12 @@ def import_thread(worker, mode):
                        continue

                    if key.startswith(b"RemoteFunction"):
-                        with log_span(
-                                "ray:import_remote_function", worker=worker):
+                        with profile(
+                                "register_remote_function", worker=worker):
                            fetch_and_register_remote_function(
                                key, worker=worker)
                    elif key.startswith(b"FunctionsToRun"):
-                        with log_span(
-                                "ray:import_function_to_run", worker=worker):
+                        with profile("fetch_and_run_function", worker=worker):
                            fetch_and_execute_function_to_run(
                                key, worker=worker)
                    elif key.startswith(b"ActorClass"):
@@ -2333,6 +2354,13 @@ def connect(info,
        t.daemon = True
        t.start()

+    if mode in [SCRIPT_MODE, SILENT_MODE] and worker.use_raylet:
+        t = threading.Thread(target=_flush_profile_events, args=(worker, ))
+        # Making the thread a daemon causes it to exit when the main thread
+        # exits.
+        t.daemon = True
+        t.start()
+
    if mode in [SCRIPT_MODE, SILENT_MODE]:
        # Add the directory containing the script that is running to the Python
        # paths of the workers. Also add the current directory. Note that this
@@ -2526,7 +2554,8 @@ class RayLogSpan(object):

    def __enter__(self):
        """Log the beginning of a span event."""
-        log(event_type=self.event_type,
+        _log(
+            event_type=self.event_type,
            contents=self.contents,
            kind=LOG_SPAN_START,
            worker=self.worker)
@@ -2534,11 +2563,13 @@ class RayLogSpan(object):
    def __exit__(self, type, value, tb):
        """Log the end of a span event. Log any exception that occurred."""
        if type is None:
-            log(event_type=self.event_type,
+            _log(
+                event_type=self.event_type,
                kind=LOG_SPAN_END,
                worker=self.worker)
        else:
-            log(event_type=self.event_type,
+            _log(
+                event_type=self.event_type,
                contents={
                    "type": str(type),
                    "value": value,
@@ -2548,19 +2579,109 @@ class RayLogSpan(object):
                worker=self.worker)


-def log_span(event_type, contents=None, worker=global_worker):
-    return RayLogSpan(event_type, contents=contents, worker=worker)
+class RayLogSpanRaylet(object):
+    """An object used to enable logging a span of events with a with statement.
+
+    Attributes:
+        event_type (str): The type of the event being logged.
+        contents: Additional information to log.
+    """
+
+    def __init__(self, event_type, extra_data=None, worker=global_worker):
+        """Initialize a RayLogSpan object."""
+        self.event_type = event_type
+        self.extra_data = extra_data if extra_data is not None else {}
+        self.worker = worker
+
+    def set_attribute(self, key, value):
+        """Add a key-value pair to the extra_data dict.
+
+        This can be used to add attributes that are not available when
+        ray.profile was called.
+
+        Args:
+            key: The attribute name.
+            value: The attribute value.
+        """
+        if not isinstance(key, str) or not isinstance(value, str):
+            raise ValueError("The extra_data argument must be a "
+                             "dictionary mapping strings to strings.")
+        self.extra_data[key] = value
+
+    def __enter__(self):
+        """Log the beginning of a span event.
+
+        Returns:
+            The object itself is returned so that if the block is opened using
+                "with ray.profile(...) as prof:", we can call
+                "prof.set_attribute" inside the block.
+        """
+        self.start_time = time.time()
+        return self
+
+    def __exit__(self, type, value, tb):
+        """Log the end of a span event. Log any exception that occurred."""
+        for key, value in self.extra_data.items():
+            if not isinstance(key, str) or not isinstance(value, str):
+                raise ValueError("The extra_data argument must be a "
+                                 "dictionary mapping strings to strings.")
+
+        event = {
+            "event_type": self.event_type,
+            "start_time": self.start_time,
+            "end_time": time.time(),
+            "extra_data": json.dumps(self.extra_data),
+        }
+
+        if type is not None:
+            event["extra_data"] = json.dumps({
+                "type": str(type),
+                "value": str(value),
+                "traceback": str(traceback.format_exc()),
+            })
+
+        self.worker.events.append(event)


-def log_event(event_type, contents=None, worker=global_worker):
-    log(event_type, kind=LOG_POINT, contents=contents, worker=worker)
+def profile(event_type, extra_data=None, worker=global_worker):
+    """Profile a span of time so that it appears in the timeline visualization.
+
+    This function can be used as follows (both on the driver or within a task).
+
+        with ray.profile("custom event", extra_data={'key': 'value'}):
+            # Do some computation here.
+
+    Optionally, a dictionary can be passed as the "extra_data" argument, and
+    it can have keys "name" and "cname" if you want to override the default
+    timeline display text and box color. Other values will appear at the bottom
+    of the chrome tracing GUI when you click on the box corresponding to this
+    profile span.
+
+    Args:
+        event_type: A string describing the type of the event.
+        extra_data: This must be a dictionary mapping strings to strings. This
+            data will be added to the json objects that are used to populate
+            the timeline, so if you want to set a particular color, you can
+            simply set the "cname" attribute to an appropriate color.
+            Similarly, if you set the "name" attribute, then that will set the
+            text displayed on the box in the timeline.
+
+    Returns:
+        An object that can profile a span of time via a "with" statement.
+    """
+    if not worker.use_raylet:
+        return RayLogSpan(event_type, contents=extra_data, worker=worker)
+    else:
+        return RayLogSpanRaylet(
+            event_type, extra_data=extra_data, worker=worker)


-def log(event_type, kind, contents=None, worker=global_worker):
+def _log(event_type, kind, contents=None, worker=global_worker):
    """Log an event to the global state store.

    This adds the event to a buffer of events locally. The buffer can be
-    flushed and written to the global state store by calling flush_log().
+    flushed and written to the global state store by calling
+    flush_profile_data().

    Args:
        event_type (str): The type of the event.
@@ -2571,6 +2692,9 @@ def log(event_type, kind, contents=None, worker=global_worker):
            time, and it is LOG_SPAN_END if we are finishing logging a span of
            time.
    """
+    if worker.use_raylet:
+        raise Exception(
+            "This method is not supported in the raylet code path.")
    # TODO(rkn): This code currently takes around half a microsecond. Since we
    # call it tens of times per task, this adds up. We will need to redo the
    # logging code, perhaps in C.
@@ -2584,13 +2708,32 @@ def log(event_type, kind, contents=None, worker=global_worker):
        worker.events.append((time.time(), event_type, kind, contents))


-def flush_log(worker=global_worker):
-    """Send the logged worker events to the global state store."""
-    event_log_key = b"event_log:" + worker.worker_id
-    event_log_value = json.dumps(worker.events)
+# TODO(rkn): Support calling this function in the middle of a task, and also
+# call this periodically in the background from the driver.
+def flush_profile_data(worker=global_worker):
+    """Push the logged profiling data to the global control store.
+
+    By default, profiling information for a given task won't appear in the
+    timeline until after the task has completed. For very long-running tasks,
+    we may want profiling information to appear more quickly. In such cases,
+    this function can be called. Note that as an alternative, we could start
+    a thread in the background on workers that calls this automatically.
+    """
    if not worker.use_raylet:
+        event_log_key = b"event_log:" + worker.worker_id
+        event_log_value = json.dumps(worker.events)
        worker.local_scheduler_client.log_event(event_log_key, event_log_value,
                                                time.time())
+    else:
+        if worker.mode == WORKER_MODE:
+            component_type = "worker"
+        else:
+            component_type = "driver"
+
+        worker.local_scheduler_client.push_profile_events(
+            component_type, ray.ObjectID(worker.worker_id),
+            worker.node_ip_address, worker.events)
+
    worker.events = []


@@ -2611,7 +2754,7 @@ def get(object_ids, worker=global_worker):
        A Python object or a list of Python objects.
    """
    worker.check_connected()
-    with log_span("ray:get", worker=worker):
+    with profile("ray.get", worker=worker):
        check_main_thread()

        if worker.mode == PYTHON_MODE:
@@ -2644,7 +2787,7 @@ def put(value, worker=global_worker):
        The object ID assigned to this value.
    """
    worker.check_connected()
-    with log_span("ray:put", worker=worker):
+    with profile("ray.put", worker=worker):
        check_main_thread()

        if worker.mode == PYTHON_MODE:
@@ -2702,7 +2845,7 @@ def wait(object_ids, num_returns=1, timeout=None, worker=global_worker):
                                    type(object_id)))

    worker.check_connected()
-    with log_span("ray:wait", worker=worker):
+    with profile("ray.wait", worker=worker):
        check_main_thread()

        # When Ray is run in PYTHON_MODE, all functions are run immediately,