Skip to content

Commit

Permalink
Fix ray tests (#3877)
Browse files Browse the repository at this point in the history
Make sure there is a transaction name
  • Loading branch information
antonpirker authored Dec 18, 2024
1 parent 60d6333 commit f6e9162
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 28 deletions.
17 changes: 12 additions & 5 deletions sentry_sdk/integrations/ray.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
from typing import Any, Optional
from sentry_sdk.utils import ExcInfo

DEFAULT_TRANSACTION_NAME = "unknown Ray function"


def _check_sentry_initialized():
# type: () -> None
Expand Down Expand Up @@ -58,18 +60,23 @@ def _f(*f_args, _tracing=None, **f_kwargs):
"""
_check_sentry_initialized()

root_span_name = qualname_from_function(f) or DEFAULT_TRANSACTION_NAME
sentry_sdk.get_current_scope().set_transaction_name(
root_span_name,
source=TRANSACTION_SOURCE_TASK,
)
with sentry_sdk.continue_trace(_tracing or {}):
with sentry_sdk.start_transaction(
with sentry_sdk.start_span(
op=OP.QUEUE_TASK_RAY,
name=qualname_from_function(f) or "unknown Ray function",
name=root_span_name,
origin=RayIntegration.origin,
source=TRANSACTION_SOURCE_TASK,
) as transaction:
) as root_span:
try:
result = f(*f_args, **f_kwargs)
transaction.set_status(SPANSTATUS.OK)
root_span.set_status(SPANSTATUS.OK)
except Exception:
transaction.set_status(SPANSTATUS.INTERNAL_ERROR)
root_span.set_status(SPANSTATUS.INTERNAL_ERROR)
exc_info = sys.exc_info()
_capture_exception(exc_info)
reraise(*exc_info)
Expand Down
50 changes: 27 additions & 23 deletions tests/integrations/ray/test_ray.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,42 +77,42 @@ def example_task():

return sentry_sdk.get_client().transport.envelopes

with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
with sentry_sdk.start_span(op="test", name="ray client root span"):
worker_envelopes = ray.get(example_task.remote())

client_envelope = sentry_sdk.get_client().transport.envelopes[0]
client_transaction = client_envelope.get_transaction_event()
assert client_transaction["transaction"] == "ray test transaction"
assert client_transaction["transaction_info"] == {"source": "custom"}
client_root_span = client_envelope.get_transaction_event()
assert client_root_span["transaction"] == "ray client root span"
assert client_root_span["transaction_info"] == {"source": "custom"}

worker_envelope = worker_envelopes[0]
worker_transaction = worker_envelope.get_transaction_event()
worker_root_span = worker_envelope.get_transaction_event()
assert (
worker_transaction["transaction"]
worker_root_span["transaction"]
== "tests.integrations.ray.test_ray.test_tracing_in_ray_tasks.<locals>.example_task"
)
assert worker_transaction["transaction_info"] == {"source": "task"}
assert worker_root_span["transaction_info"] == {"source": "task"}

(span,) = client_transaction["spans"]
(span,) = client_root_span["spans"]
assert span["op"] == "queue.submit.ray"
assert span["origin"] == "auto.queue.ray"
assert (
span["description"]
== "tests.integrations.ray.test_ray.test_tracing_in_ray_tasks.<locals>.example_task"
)
assert span["parent_span_id"] == client_transaction["contexts"]["trace"]["span_id"]
assert span["trace_id"] == client_transaction["contexts"]["trace"]["trace_id"]
assert span["parent_span_id"] == client_root_span["contexts"]["trace"]["span_id"]
assert span["trace_id"] == client_root_span["contexts"]["trace"]["trace_id"]

(span,) = worker_transaction["spans"]
(span,) = worker_root_span["spans"]
assert span["op"] == "task"
assert span["origin"] == "manual"
assert span["description"] == "example task step"
assert span["parent_span_id"] == worker_transaction["contexts"]["trace"]["span_id"]
assert span["trace_id"] == worker_transaction["contexts"]["trace"]["trace_id"]
assert span["parent_span_id"] == worker_root_span["contexts"]["trace"]["span_id"]
assert span["trace_id"] == worker_root_span["contexts"]["trace"]["trace_id"]

assert (
client_transaction["contexts"]["trace"]["trace_id"]
== worker_transaction["contexts"]["trace"]["trace_id"]
client_root_span["contexts"]["trace"]["trace_id"]
== worker_root_span["contexts"]["trace"]["trace_id"]
)


Expand All @@ -132,7 +132,7 @@ def test_errors_in_ray_tasks():
def example_task():
1 / 0

with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
with sentry_sdk.start_span(op="test", name="ray client root span"):
with pytest.raises(ZeroDivisionError):
future = example_task.remote()
ray.get(future)
Expand Down Expand Up @@ -167,22 +167,24 @@ def __init__(self):
self.n = 0

def increment(self):
with sentry_sdk.start_span(op="task", name="example actor execution"):
with sentry_sdk.start_span(
op="test", name="custom span in actor execution", only_if_parent=True
):
self.n += 1

return sentry_sdk.get_client().transport.envelopes

with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
with sentry_sdk.start_span(op="test", name="ray client root span"):
counter = Counter.remote()
worker_envelopes = ray.get(counter.increment.remote())

client_envelope = sentry_sdk.get_client().transport.envelopes[0]
client_transaction = client_envelope.get_transaction_event()
client_root_span = client_envelope.get_transaction_event()

# Spans for submitting the actor task are not created (actors are not supported yet)
assert client_transaction["spans"] == []
assert client_root_span["spans"] == []

# Transaction are not yet created when executing ray actors (actors are not supported yet)
# Root spans are not yet automatically created when executing ray actors (actors are not supported yet)
assert worker_envelopes == []


Expand All @@ -204,12 +206,14 @@ def __init__(self):
self.n = 0

def increment(self):
with sentry_sdk.start_span(op="task", name="example actor execution"):
with sentry_sdk.start_span(
op="test", name="custom span in actor execution", only_if_parent=True
):
1 / 0

return sentry_sdk.get_client().transport.envelopes

with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
with sentry_sdk.start_span(op="test", name="ray client root span"):
with pytest.raises(ZeroDivisionError):
counter = Counter.remote()
future = counter.increment.remote()
Expand Down

0 comments on commit f6e9162

Please sign in to comment.