Skip to content

Commit 97d6756

Browse files
authored
fix(Ray): Retain the original function name when patching Ray tasks (#4858)
### Description Without "@functools.wraps" added, Ray exposes Prometheus metrics with all tasks named "new_func" #### Issues * Follow up to [!4430](#4430) comments #### Reminders - Please add tests to validate your changes, and lint your code using `tox -e linters`. - Add GH Issue ID _&_ Linear ID (if applicable) - PR title should use [conventional commit](https://develop.sentry.dev/engineering-practices/commit-messages/#type) style (`feat:`, `fix:`, `ref:`, `meta:`) - For external contributors: [CONTRIBUTING.md](https://github.com/getsentry/sentry-python/blob/master/CONTRIBUTING.md), [Sentry SDK development docs](https://develop.sentry.dev/sdk/), [Discord community](https://discord.gg/Ww9hbqr)
1 parent 149a7da commit 97d6756

File tree

2 files changed

+23
-4
lines changed

2 files changed

+23
-4
lines changed

sentry_sdk/integrations/ray.py

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import inspect
2+
import functools
23
import sys
34

45
import sentry_sdk
@@ -17,7 +18,6 @@
1718
import ray # type: ignore[import-not-found]
1819
except ImportError:
1920
raise DidNotEnable("Ray not installed.")
20-
import functools
2121

2222
from typing import TYPE_CHECKING
2323

@@ -54,12 +54,13 @@ def new_remote(f=None, *args, **kwargs):
5454

5555
def wrapper(user_f):
5656
# type: (Callable[..., Any]) -> Any
57-
def new_func(*f_args, _tracing=None, **f_kwargs):
57+
@functools.wraps(user_f)
58+
def new_func(*f_args, _sentry_tracing=None, **f_kwargs):
5859
# type: (Any, Optional[dict[str, Any]], Any) -> Any
5960
_check_sentry_initialized()
6061

6162
transaction = sentry_sdk.continue_trace(
62-
_tracing or {},
63+
_sentry_tracing or {},
6364
op=OP.QUEUE_TASK_RAY,
6465
name=qualname_from_function(user_f),
6566
origin=RayIntegration.origin,
@@ -78,6 +79,19 @@ def new_func(*f_args, _tracing=None, **f_kwargs):
7879

7980
return result
8081

82+
# Patching new_func signature to add the _sentry_tracing parameter to it
83+
# Ray later inspects the signature and finds the unexpected parameter otherwise
84+
signature = inspect.signature(new_func)
85+
params = list(signature.parameters.values())
86+
params.append(
87+
inspect.Parameter(
88+
"_sentry_tracing",
89+
kind=inspect.Parameter.KEYWORD_ONLY,
90+
default=None,
91+
)
92+
)
93+
new_func.__signature__ = signature.replace(parameters=params) # type: ignore[attr-defined]
94+
8195
if f:
8296
rv = old_remote(new_func)
8397
else:
@@ -99,7 +113,9 @@ def _remote_method_with_header_propagation(*args, **kwargs):
99113
for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers()
100114
}
101115
try:
102-
result = old_remote_method(*args, **kwargs, _tracing=tracing)
116+
result = old_remote_method(
117+
*args, **kwargs, _sentry_tracing=tracing
118+
)
103119
span.set_status(SPANSTATUS.OK)
104120
except Exception:
105121
span.set_status(SPANSTATUS.INTERNAL_ERROR)

tests/integrations/ray/test_ray.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,9 @@ def example_task():
100100
else:
101101
example_task = ray.remote(example_task)
102102

103+
# Function name shouldn't be overwritten by Sentry wrapper
104+
assert example_task._function_name == "tests.integrations.ray.test_ray.example_task"
105+
103106
with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
104107
worker_envelopes = ray.get(example_task.remote())
105108

0 commit comments

Comments
 (0)