diff --git a/sentry_sdk/integrations/ray.py b/sentry_sdk/integrations/ray.py index 07ad7c100c..080d0437ef 100644 --- a/sentry_sdk/integrations/ray.py +++ b/sentry_sdk/integrations/ray.py @@ -3,11 +3,9 @@ import sys import sentry_sdk -from sentry_sdk.consts import OP, SPANSTATUS +from sentry_sdk.consts import OP from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version from sentry_sdk.traces import SegmentNameSource -from sentry_sdk.tracing import TransactionSource -from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import ( event_from_exception, logger, @@ -91,52 +89,26 @@ def new_func( ) -> "Any": _check_sentry_initialized() - span_streaming = has_span_streaming_enabled( - sentry_sdk.get_client().options - ) - if span_streaming: - sentry_sdk.traces.continue_trace(_sentry_tracing or {}) - - function_name = qualname_from_function(user_f) - with sentry_sdk.traces.start_span( - name="unknown Ray task" - if function_name is None - else function_name, - attributes={ - "sentry.op": OP.QUEUE_TASK_RAY, - "sentry.origin": RayIntegration.origin, - "sentry.segment.name.source": SegmentNameSource.TASK, - }, - parent_span=None, - ): - try: - result = user_f(*f_args, **f_kwargs) - except Exception: - exc_info = sys.exc_info() - _capture_exception(exc_info) - reraise(*exc_info) - - return result - else: - transaction = sentry_sdk.continue_trace( - _sentry_tracing or {}, - op=OP.QUEUE_TASK_RAY, - name=qualname_from_function(user_f), - origin=RayIntegration.origin, - source=TransactionSource.TASK, - ) - - with sentry_sdk.start_transaction(transaction) as transaction: - try: - result = user_f(*f_args, **f_kwargs) - transaction.set_status(SPANSTATUS.OK) - except Exception: - transaction.set_status(SPANSTATUS.INTERNAL_ERROR) - exc_info = sys.exc_info() - _capture_exception(exc_info) - reraise(*exc_info) - - return result + sentry_sdk.traces.continue_trace(_sentry_tracing or {}) + + function_name = qualname_from_function(user_f) + with sentry_sdk.traces.start_span( + name="unknown Ray task" if function_name is None else function_name, + attributes={ + "sentry.op": OP.QUEUE_TASK_RAY, + "sentry.origin": RayIntegration.origin, + "sentry.segment.name.source": SegmentNameSource.TASK, + }, + parent_span=None, + ): + try: + result = user_f(*f_args, **f_kwargs) + except Exception: + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + + return result _insert_sentry_tracing_in_signature(new_func) @@ -152,73 +124,45 @@ def _remote_method_with_header_propagation( """ Ray Client """ - span_streaming = has_span_streaming_enabled( - sentry_sdk.get_client().options - ) - if span_streaming: - function_name = qualname_from_function(user_f) - - if sentry_sdk.traces.get_current_span() is None: - tracing = { - k: v - for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers() - } - try: - result = old_remote_method( - *args, **kwargs, _sentry_tracing=tracing - ) - except Exception: - exc_info = sys.exc_info() - _capture_exception(exc_info) - reraise(*exc_info) - - return result - - with sentry_sdk.traces.start_span( - name="unknown Ray task" - if function_name is None - else function_name, - attributes={ - "sentry.op": OP.QUEUE_SUBMIT_RAY, - "sentry.origin": RayIntegration.origin, - }, - ): - tracing = { - k: v - for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers() - } - try: - result = old_remote_method( - *args, **kwargs, _sentry_tracing=tracing - ) - except Exception: - exc_info = sys.exc_info() - _capture_exception(exc_info) - reraise(*exc_info) - - return result - else: - with sentry_sdk.start_span( - op=OP.QUEUE_SUBMIT_RAY, - name=qualname_from_function(user_f), - origin=RayIntegration.origin, - ) as span: - tracing = { - k: v - for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers() - } - try: - result = old_remote_method( - *args, **kwargs, _sentry_tracing=tracing - ) - span.set_status(SPANSTATUS.OK) - except Exception: - span.set_status(SPANSTATUS.INTERNAL_ERROR) - exc_info = sys.exc_info() - _capture_exception(exc_info) - reraise(*exc_info) - - return result + function_name = qualname_from_function(user_f) + + if sentry_sdk.traces.get_current_span() is None: + tracing = { + k: v + for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers() + } + try: + result = old_remote_method( + *args, **kwargs, _sentry_tracing=tracing + ) + except Exception: + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + + return result + + with sentry_sdk.traces.start_span( + name="unknown Ray task" if function_name is None else function_name, + attributes={ + "sentry.op": OP.QUEUE_SUBMIT_RAY, + "sentry.origin": RayIntegration.origin, + }, + ): + tracing = { + k: v + for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers() + } + try: + result = old_remote_method( + *args, **kwargs, _sentry_tracing=tracing + ) + except Exception: + exc_info = sys.exc_info() + _capture_exception(exc_info) + reraise(*exc_info) + + return result rv.remote = _remote_method_with_header_propagation diff --git a/tests/integrations/ray/test_ray.py b/tests/integrations/ray/test_ray.py index bef084fc5c..05a1e39d8c 100644 --- a/tests/integrations/ray/test_ray.py +++ b/tests/integrations/ray/test_ray.py @@ -39,19 +39,14 @@ def setup_sentry_with_logging_transport(): setup_sentry(transport=RayLoggingTransport()) -def setup_sentry_with_logging_transport_and_span_streaming(): - setup_sentry(span_streaming=True, transport=RayLoggingTransport()) - - -def setup_sentry(span_streaming=False, transport=None): - if span_streaming: - sentry_sdk._span_batcher.SpanBatcher.MAX_BEFORE_FLUSH = 1 +def setup_sentry(transport=None): + sentry_sdk._span_batcher.SpanBatcher.MAX_BEFORE_FLUSH = 1 sentry_sdk.init( integrations=[RayIntegration()], transport=RayTestTransport() if transport is None else transport, traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) @@ -141,34 +136,22 @@ def read_spans_from_log(job_id, ray_temp_dir, min_spans=1, timeout=10): time.sleep(0.1) -def example_task(span_streaming: bool): - if span_streaming: - with sentry_sdk.traces.start_span( - name="example task step", - attributes={ - "sentry.op": "task", - }, - ): - ... - else: - with sentry_sdk.start_span(op="task", name="example task step"): - ... - - return sentry_sdk.get_client().transport.envelopes +def example_task(): + with sentry_sdk.traces.start_span( + name="example task step", + attributes={ + "sentry.op": "task", + }, + ): + ... # RayIntegration must leave variadic keyword arguments at the end -def example_task_with_kwargs(span_streaming: bool, **kwargs): - if span_streaming: - with sentry_sdk.traces.start_span( - name="example task step", attributes={"sentry.op": "task"} - ): - ... - else: - with sentry_sdk.start_span(op="task", name="example task step"): - ... - - return sentry_sdk.get_client().transport.envelopes +def example_task_with_kwargs(**kwargs): + with sentry_sdk.traces.start_span( + name="example task step", attributes={"sentry.op": "task"} + ): + ... @pytest.mark.parametrize( @@ -178,14 +161,16 @@ def example_task_with_kwargs(span_streaming: bool, **kwargs): "task", [example_task, example_task_with_kwargs], ) -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_tracing_in_ray_tasks(task_options, task, span_streaming): +def test_tracing_in_ray_tasks( + task_options, + task, +): sentry_sdk.init( integrations=[RayIntegration()], disabled_integrations=[StdlibIntegration], transport=RayTestTransport(), traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) # Setup ray task, calling decorator directly instead of @, @@ -200,121 +185,67 @@ def test_tracing_in_ray_tasks(task_options, task, span_streaming): example_task._function_name == f"tests.integrations.ray.test_ray.{task.__name__}" ) + ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}") + os.makedirs(ray_temp_dir, exist_ok=True) - if span_streaming: - ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}") - os.makedirs(ray_temp_dir, exist_ok=True) - - try: - ray.init( - runtime_env={ - "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming, - "working_dir": "./", - }, - _temp_dir=ray_temp_dir, - ) - - with sentry_sdk.traces.start_span( - name="ray test parent", attributes={"sentry.op": "task"} - ): - future = example_task.remote(span_streaming) - ray.get(future) - - job_id = future.job_id().hex() - worker_spans = read_spans_from_log(job_id, ray_temp_dir, min_spans=2) - finally: - if os.path.exists(ray_temp_dir): - shutil.rmtree(ray_temp_dir, ignore_errors=True) - - sentry_sdk.flush() - client_envelope = sentry_sdk.get_client().transport.envelopes[0] - client_spans = [ - span - for item in client_envelope.items - for span in item.payload.json["items"] - ] - assert client_spans[1]["name"] == "ray test parent" - - assert ( - worker_spans[1]["name"] - == f"tests.integrations.ray.test_ray.{task.__name__}" - ) - assert ( - worker_spans[1]["attributes"]["sentry.segment.name.source"]["value"] - == SegmentNameSource.TASK - ) - - span = client_spans[0] - assert span["attributes"]["sentry.op"]["value"] == "queue.submit.ray" - assert span["attributes"]["sentry.origin"]["value"] == "auto.queue.ray" - assert span["name"] == f"tests.integrations.ray.test_ray.{task.__name__}" - assert span["parent_span_id"] == client_spans[1]["span_id"] - assert span["trace_id"] == client_spans[1]["trace_id"] - - span = worker_spans[0] - assert span["attributes"]["sentry.op"]["value"] == "task" - assert span["attributes"]["sentry.origin"]["value"] == "manual" - assert span["name"] == "example task step" - assert span["parent_span_id"] == worker_spans[1]["span_id"] - assert span["trace_id"] == worker_spans[1]["trace_id"] - - assert client_spans[1]["trace_id"] == worker_spans[1]["trace_id"] - else: + try: ray.init( runtime_env={ - "worker_process_setup_hook": setup_sentry, + "worker_process_setup_hook": setup_sentry_with_logging_transport, "working_dir": "./", - } + }, + _temp_dir=ray_temp_dir, ) - with sentry_sdk.start_transaction(op="task", name="ray test transaction"): - worker_envelopes = ray.get(example_task.remote(span_streaming)) + with sentry_sdk.traces.start_span( + name="ray test parent", attributes={"sentry.op": "task"} + ): + future = example_task.remote() + ray.get(future) - 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"} + job_id = future.job_id().hex() + worker_spans = read_spans_from_log(job_id, ray_temp_dir, min_spans=2) + finally: + if os.path.exists(ray_temp_dir): + shutil.rmtree(ray_temp_dir, ignore_errors=True) - worker_envelope = worker_envelopes[0] - worker_transaction = worker_envelope.get_transaction_event() - assert ( - worker_transaction["transaction"] - == f"tests.integrations.ray.test_ray.{task.__name__}" - ) - assert worker_transaction["transaction_info"] == {"source": "task"} + sentry_sdk.flush() + client_envelope = sentry_sdk.get_client().transport.envelopes[0] + client_spans = [ + span for item in client_envelope.items for span in item.payload.json["items"] + ] + assert client_spans[1]["name"] == "ray test parent" - (span,) = client_transaction["spans"] - assert span["op"] == "queue.submit.ray" - assert span["origin"] == "auto.queue.ray" - assert span["description"] == f"tests.integrations.ray.test_ray.{task.__name__}" - assert ( - span["parent_span_id"] == client_transaction["contexts"]["trace"]["span_id"] - ) - assert span["trace_id"] == client_transaction["contexts"]["trace"]["trace_id"] + assert worker_spans[1]["name"] == f"tests.integrations.ray.test_ray.{task.__name__}" + assert ( + worker_spans[1]["attributes"]["sentry.segment.name.source"]["value"] + == SegmentNameSource.TASK + ) - (span,) = worker_transaction["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"] + span = client_spans[0] + assert span["attributes"]["sentry.op"]["value"] == "queue.submit.ray" + assert span["attributes"]["sentry.origin"]["value"] == "auto.queue.ray" + assert span["name"] == f"tests.integrations.ray.test_ray.{task.__name__}" + assert span["parent_span_id"] == client_spans[1]["span_id"] + assert span["trace_id"] == client_spans[1]["trace_id"] - assert ( - client_transaction["contexts"]["trace"]["trace_id"] - == worker_transaction["contexts"]["trace"]["trace_id"] - ) + span = worker_spans[0] + assert span["attributes"]["sentry.op"]["value"] == "task" + assert span["attributes"]["sentry.origin"]["value"] == "manual" + assert span["name"] == "example task step" + assert span["parent_span_id"] == worker_spans[1]["span_id"] + assert span["trace_id"] == worker_spans[1]["trace_id"] + + assert client_spans[1]["trace_id"] == worker_spans[1]["trace_id"] -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_errors_in_ray_tasks(span_streaming): +def test_errors_in_ray_tasks(): sentry_sdk.init( integrations=[RayIntegration()], disabled_integrations=[StdlibIntegration], transport=RayTestTransport(), traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}") @@ -323,9 +254,7 @@ def test_errors_in_ray_tasks(span_streaming): try: ray.init( runtime_env={ - "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming - if span_streaming - else setup_sentry_with_logging_transport, + "worker_process_setup_hook": setup_sentry_with_logging_transport, "working_dir": "./", }, _temp_dir=ray_temp_dir, @@ -336,18 +265,12 @@ def test_errors_in_ray_tasks(span_streaming): def example_task(): 1 / 0 - if span_streaming: - with sentry_sdk.traces.start_span( - name="ray test parent", attributes={"sentry.op": "task"} - ): - with pytest.raises(ZeroDivisionError): - future = example_task.remote() - ray.get(future) - else: - with sentry_sdk.start_transaction(op="task", name="ray test transaction"): - with pytest.raises(ZeroDivisionError): - future = example_task.remote() - ray.get(future) + with sentry_sdk.traces.start_span( + name="ray test parent", attributes={"sentry.op": "task"} + ): + with pytest.raises(ZeroDivisionError): + future = example_task.remote() + ray.get(future) job_id = future.job_id().hex() error = read_error_from_log(job_id, ray_temp_dir) @@ -367,14 +290,15 @@ def example_task(): # Arbitrary keyword argument to test all decorator paths @pytest.mark.parametrize("remote_kwargs", [{}, {"namespace": "actors"}]) -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_tracing_in_ray_actors(remote_kwargs, span_streaming): +def test_tracing_in_ray_actors( + remote_kwargs, +): sentry_sdk.init( integrations=[RayIntegration()], disabled_integrations=[StdlibIntegration], transport=RayTestTransport(), traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) # Setup ray actor @@ -386,18 +310,10 @@ def __init__(self): self.n = 0 def increment(self): - if span_streaming: - with sentry_sdk.traces.start_span( - name="example actor execution", attributes={"sentry.op": "task"} - ): - self.n += 1 - else: - with sentry_sdk.start_span( - op="task", name="example actor execution" - ): - self.n += 1 - - return sentry_sdk.get_client().transport.envelopes + with sentry_sdk.traces.start_span( + name="example actor execution", attributes={"sentry.op": "task"} + ): + self.n += 1 else: @ray.remote @@ -406,90 +322,58 @@ def __init__(self): self.n = 0 def increment(self): - if span_streaming: - with sentry_sdk.traces.start_span( - name="example actor execution", attributes={"sentry.op": "task"} - ): - self.n += 1 - else: - with sentry_sdk.start_span( - op="task", name="example actor execution" - ): - self.n += 1 - - return sentry_sdk.get_client().transport.envelopes - - if span_streaming: - ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}") - os.makedirs(ray_temp_dir, exist_ok=True) + with sentry_sdk.traces.start_span( + name="example actor execution", attributes={"sentry.op": "task"} + ): + self.n += 1 - try: - ray.init( - runtime_env={ - "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming, - "working_dir": "./", - }, - _temp_dir=ray_temp_dir, - ) - - with sentry_sdk.traces.start_span( - name="ray test parent", attributes={"sentry.op": "task"} - ): - counter = Counter.remote() - future = counter.increment.remote() - ray.get(future) + ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}") + os.makedirs(ray_temp_dir, exist_ok=True) - job_id = future.job_id().hex() - worker_spans = read_spans_from_log(job_id, ray_temp_dir) - finally: - if os.path.exists(ray_temp_dir): - shutil.rmtree(ray_temp_dir, ignore_errors=True) - - sentry_sdk.flush() - client_envelope = sentry_sdk.get_client().transport.envelopes[0] - client_spans = [ - span - for item in client_envelope.items - for span in item.payload.json["items"] - ] - - # Spans for submitting the actor task are not created (actors are not supported yet) - # Only the manual "example actor execution" span is recorded. - assert len(client_spans) == 1 - - # Transaction are not yet created when executing ray actors (actors are not supported yet) - # Only the manual "example actor execution" span is recorded. - assert len(worker_spans) == 1 - else: + try: ray.init( runtime_env={ - "worker_process_setup_hook": setup_sentry, + "worker_process_setup_hook": setup_sentry_with_logging_transport, "working_dir": "./", - } + }, + _temp_dir=ray_temp_dir, ) - with sentry_sdk.start_transaction(op="task", name="ray test transaction"): + with sentry_sdk.traces.start_span( + name="ray test parent", attributes={"sentry.op": "task"} + ): counter = Counter.remote() - worker_envelopes = ray.get(counter.increment.remote()) + future = counter.increment.remote() + ray.get(future) - client_envelope = sentry_sdk.get_client().transport.envelopes[0] - client_transaction = client_envelope.get_transaction_event() + job_id = future.job_id().hex() + worker_spans = read_spans_from_log(job_id, ray_temp_dir) + finally: + if os.path.exists(ray_temp_dir): + shutil.rmtree(ray_temp_dir, ignore_errors=True) + + sentry_sdk.flush() + client_envelope = sentry_sdk.get_client().transport.envelopes[0] + client_spans = [ + span for item in client_envelope.items for span in item.payload.json["items"] + ] - # Spans for submitting the actor task are not created (actors are not supported yet) - assert client_transaction["spans"] == [] + # Spans for submitting the actor task are not created (actors are not supported yet) + # Only the manual "example actor execution" span is recorded. + assert len(client_spans) == 1 - # Transaction are not yet created when executing ray actors (actors are not supported yet) - assert worker_envelopes == [] + # Spans are not yet created when executing ray actors (actors are not supported yet) + # Only the manual "example actor execution" span is recorded. + assert len(worker_spans) == 1 -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_errors_in_ray_actors(span_streaming): +def test_errors_in_ray_actors(): sentry_sdk.init( integrations=[RayIntegration()], disabled_integrations=[StdlibIntegration], transport=RayLoggingTransport(), traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", ) ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}") @@ -498,9 +382,7 @@ def test_errors_in_ray_actors(span_streaming): try: ray.init( runtime_env={ - "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming - if span_streaming - else setup_sentry_with_logging_transport, + "worker_process_setup_hook": setup_sentry_with_logging_transport, "working_dir": "./", }, _temp_dir=ray_temp_dir, @@ -513,33 +395,20 @@ def __init__(self): self.n = 0 def increment(self): - if span_streaming: - with sentry_sdk.traces.start_span( - name="example actor execution", attributes={"sentry.op": "task"} - ): - 1 / 0 - else: - with sentry_sdk.start_span( - op="task", name="example actor execution" - ): - 1 / 0 + with sentry_sdk.traces.start_span( + name="example actor execution", attributes={"sentry.op": "task"} + ): + 1 / 0 return sentry_sdk.get_client().transport.envelopes - if span_streaming: - with sentry_sdk.traces.start_span( - name="ray test parent", attributes={"sentry.op": "task"} - ): - with pytest.raises(ZeroDivisionError): - counter = Counter.remote() - future = counter.increment.remote() - ray.get(future) - else: - with sentry_sdk.start_transaction(op="task", name="ray test transaction"): - with pytest.raises(ZeroDivisionError): - counter = Counter.remote() - future = counter.increment.remote() - ray.get(future) + with sentry_sdk.traces.start_span( + name="ray test parent", attributes={"sentry.op": "task"} + ): + with pytest.raises(ZeroDivisionError): + counter = Counter.remote() + future = counter.increment.remote() + ray.get(future) job_id = future.job_id().hex() error = read_error_from_log(job_id, ray_temp_dir)