[core] Fix ray.timeline()

python/ray/_private/state.py

@@ -215,7 +215,7 @@ def profile_events(self):
         result = defaultdict(list)
         task_events = self.global_state_accessor.get_task_events()
         for i in range(len(task_events)):
-            event = common_pb2.TaskEvents.FromString(task_events[i])
+            event = gcs_pb2.TaskEvents.FromString(task_events[i])
             profile = event.profile_events
             if not profile:
                 continue

python/ray/tests/test_state_api_2.py

@@ -3,6 +3,7 @@
 import os
 import sys
 from pathlib import Path
+import tempfile
 
 from collections import defaultdict
 from ray._private.test_utils import check_call_subprocess
@@ -355,6 +356,26 @@ def test_state_api_scale_smoke(shutdown_only):
     check_call_subprocess(["python", str(full_path), "--smoke-test"])
 
 
+def test_ray_timeline(shutdown_only):
+    ray.init(num_cpus=8)
+
+    @ray.remote
+    def f():
+        pass
+
+    ray.get(f.remote())
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        filename = os.path.join(tmpdirname, "timeline.json")
+        ray.timeline(filename)
+
+        with open(filename, "r") as f:
+            dumped = json.load(f)
+        # TODO(swang): Check actual content. It doesn't seem to match the
+        # return value of chrome_tracing_dump in above tests?
+        assert len(dumped) > 0
+
+
 if __name__ == "__main__":
     import sys