diff --git a/python/ray/_private/state.py b/python/ray/_private/state.py index f43f563a4a61..08f3e15d3421 100644 --- a/python/ray/_private/state.py +++ b/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 diff --git a/python/ray/tests/test_state_api_2.py b/python/ray/tests/test_state_api_2.py index c5f3ce4276c5..81d8d04f299b 100644 --- a/python/ray/tests/test_state_api_2.py +++ b/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