chore: import upstream snapshot with attribution

python/ray/tests/test_tracing.py

@@ -0,0 +1,287 @@
+import asyncio
+import glob
+import json
+import os
+import shutil
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.test_utils import check_call_ray
+from ray.util.tracing.setup_local_tmp_tracing import setup_tracing, spans_dir
+from ray.util.tracing.tracing_helper import _enable_tracing
+
+setup_tracing_path = "ray.util.tracing.setup_local_tmp_tracing:setup_tracing"
+
+
+@pytest.fixture()
+def cleanup_dirs():
+    """Cleanup temporary spans_dir folder at beginning and end of test."""
+    if os.path.exists(spans_dir):
+        shutil.rmtree(spans_dir)
+    os.makedirs(spans_dir)
+    yield
+    # Enable tracing only sets up tracing once per driver process.
+    # We set ray.__traced__ to False here so that each
+    # test will re-set up tracing.
+    ray.__traced__ = False
+    if os.path.exists(spans_dir):
+        shutil.rmtree(spans_dir)
+
+
+@pytest.fixture()
+def ray_start_cli_tracing(scope="function"):
+    """Start ray with tracing-startup-hook, and clean up at end of test."""
+    check_call_ray(
+        ["stop", "--force"],
+    )
+    check_call_ray(
+        ["start", "--head", "--tracing-startup-hook", setup_tracing_path],
+    )
+    yield
+    ray.shutdown()
+    check_call_ray(["stop", "--force"])
+
+
+@pytest.fixture()
+def ray_start_cli_predefined_actor_tracing(scope="function"):
+    """Start ray with tracing-startup-hook, and clean up at end of test."""
+    check_call_ray(
+        ["stop", "--force"],
+    )
+    check_call_ray(
+        ["start", "--head", "--tracing-startup-hook", setup_tracing_path],
+    )
+    yield
+    ray.shutdown()
+    check_call_ray(["stop", "--force"])
+
+
+@pytest.fixture()
+def ray_start_cli_tracing_with_client(scope="function"):
+    """Start ray with tracing startup hook and Ray Client server."""
+    check_call_ray(["stop", "--force"])
+    check_call_ray(
+        [
+            "start",
+            "--head",
+            "--ray-client-server-port=50057",
+            "--tracing-startup-hook",
+            setup_tracing_path,
+        ],
+    )
+    yield "ray://localhost:50057"
+    ray.shutdown()
+    check_call_ray(["stop", "--force"])
+
+
+@pytest.fixture()
+def ray_start_init_tracing(scope="function"):
+    """Call ray.init with tracing-startup-hook, and clean up at end of test."""
+    ray.init(_tracing_startup_hook=setup_tracing_path)
+    yield
+    ray.shutdown()
+
+
+def get_span_list():
+    """Read span files and return list of span names."""
+    span_list = []
+    for entry in glob.glob(f"{spans_dir}/**/*.txt", recursive=True):
+        with open(entry) as f:
+            for line in f.readlines():
+                span_list.append(json.loads(line))
+    return span_list
+
+
+def get_span_dict(span_list):
+    """Given a list of span names, return dictionary of span names."""
+    strip_prefix = "python.ray.tests."
+    span_names = {}
+    for span in span_list:
+        span_name = span["name"]
+        if span_name.startswith(strip_prefix):
+            span_name = span_name[len(strip_prefix) :]
+        if span_name in span_names:
+            span_names[span_name] += 1
+        else:
+            span_names[span_name] = 1
+    return span_names
+
+
+def task_helper():
+    """Run a Ray task and check the spans produced."""
+
+    @ray.remote
+    def f(value):
+        return value + 1
+
+    obj_ref = f.remote(2)
+    ray.get(obj_ref)
+
+    span_list = get_span_list()
+    assert len(span_list) == 2, span_list
+
+    # The spans could show up in a different order, so just check that
+    # all spans are as expected
+    span_names = get_span_dict(span_list)
+    assert span_names == {
+        "test_tracing.f ray.remote": 1,
+        "test_tracing.f ray.remote_worker": 1,
+    }
+
+
+def sync_actor_helper():
+    """Run a Ray sync actor and check the spans produced."""
+
+    @ray.remote
+    class Counter(object):
+        def __init__(self):
+            self.value = 0
+
+        def increment(self):
+            self.value += 1
+            return self.value
+
+    # Create an actor from this class.
+    counter = Counter.remote()
+    obj_ref = counter.increment.remote()
+    assert ray.get(obj_ref) == 1
+
+    span_list = get_span_list()
+    assert len(span_list) == 4
+
+    # The spans could show up in a different order, so just check that
+    # all spans are as expected
+    span_names = get_span_dict(span_list)
+    assert span_names == {
+        "sync_actor_helper.<locals>.Counter.__init__ ray.remote": 1,
+        "sync_actor_helper.<locals>.Counter.increment ray.remote": 1,
+        "Counter.__init__ ray.remote_worker": 1,
+        "Counter.increment ray.remote_worker": 1,
+    }, span_names
+
+
+def async_actor_helper():
+    """Run a Ray async actor and check the spans produced."""
+
+    @ray.remote
+    class AsyncActor:
+        # multiple invocation of this method can be running in
+        # the event loop at the same time
+        async def run_concurrent(self):
+            await asyncio.sleep(2)  # concurrent workload here
+
+    actor = AsyncActor.remote()
+    ray.get([actor.run_concurrent.remote() for _ in range(4)])
+
+    span_list = get_span_list()
+    assert len(span_list) == 10
+
+    # The spans could show up in a different order, so just check that
+    # all spans are as expected
+    span_names = get_span_dict(span_list)
+    assert span_names == {
+        "async_actor_helper.<locals>.AsyncActor.__init__ ray.remote": 1,
+        "async_actor_helper.<locals>.AsyncActor.run_concurrent ray.remote": 4,
+        "AsyncActor.__init__ ray.remote_worker": 1,
+        "AsyncActor.run_concurrent ray.remote_worker": 4,
+    }, span_names
+
+
+def test_tracing_task_init_workflow(cleanup_dirs, ray_start_init_tracing):
+    task_helper()
+
+
+def test_tracing_task_start_workflow(cleanup_dirs, ray_start_cli_tracing):
+    task_helper()
+
+
+def test_tracing_sync_actor_init_workflow(cleanup_dirs, ray_start_init_tracing):
+    sync_actor_helper()
+
+
+def test_tracing_sync_actor_start_workflow(cleanup_dirs, ray_start_cli_tracing):
+    sync_actor_helper()
+
+
+def test_tracing_async_actor_init_workflow(cleanup_dirs, ray_start_init_tracing):
+    async_actor_helper()
+
+
+def test_tracing_async_actor_start_workflow(cleanup_dirs, ray_start_cli_tracing):
+    async_actor_helper()
+
+
+def test_tracing_predefined_actor(cleanup_dirs, ray_start_cli_predefined_actor_tracing):
+    sync_actor_helper()
+
+
+def test_tracing_actor_client_mode(cleanup_dirs, ray_start_cli_tracing_with_client):
+    ray.init(ray_start_cli_tracing_with_client)
+    assert ray.util.client.ray.is_connected()
+    setup_tracing()
+    _enable_tracing()
+
+    @ray.remote
+    class Counter:
+        def __init__(self):
+            self.value = 0
+
+        def increment(self):
+            self.value += 1
+            return self.value
+
+    counter = Counter.remote()
+    assert ray.get(counter.increment.remote()) == 1
+
+    span_list = get_span_list()
+    span_names = get_span_dict(span_list)
+    assert span_names == {
+        "test_tracing_actor_client_mode.<locals>.Counter.__init__ ray.remote": 1,
+        "Counter.__init__ ray.remote": 1,
+        "Counter.increment ray.remote": 1,
+        "Counter.__init__ ray.remote_worker": 1,
+        "Counter.increment ray.remote_worker": 1,
+    }, span_names
+
+    actor_creation_spans = [
+        span
+        for span in span_list
+        if span["name"].endswith("Counter.__init__ ray.remote")
+    ]
+    actor_method_span = next(
+        span for span in span_list if span["name"] == "Counter.increment ray.remote"
+    )
+    actor_id = counter._actor_id.hex()
+    for span in actor_creation_spans + [actor_method_span]:
+        assert span["attributes"]["ray.actor_id"] == actor_id
+
+
+def test_wrapping(ray_start_init_tracing):
+    @ray.remote
+    def f(**_kwargs):
+        pass
+
+
+def test_deserialization_works_without_opentelemetry(ray_start_regular):
+    """
+    Test that if a function is serialized with opentelemetry, it can be
+    deserialized in a context without opentelemetry
+    """
+
+    @ray.remote
+    def f():
+        return 30
+
+    fn_bytes = ray.cloudpickle.dumps(f)
+    with patch.dict("sys.modules", opentelemetry=None):
+        # Ensure that opentelemetry cannot be imported
+        with pytest.raises(ModuleNotFoundError):
+            import opentelemetry.trace  # noqa: F401
+        ray.cloudpickle.loads(fn_bytes)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))