1
1
import inspect
2
+ import functools
2
3
import sys
3
4
4
5
import sentry_sdk
17
18
import ray # type: ignore[import-not-found]
18
19
except ImportError :
19
20
raise DidNotEnable ("Ray not installed." )
20
- import functools
21
21
22
22
from typing import TYPE_CHECKING
23
23
@@ -54,12 +54,13 @@ def new_remote(f=None, *args, **kwargs):
54
54
55
55
def wrapper (user_f ):
56
56
# 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 ):
58
59
# type: (Any, Optional[dict[str, Any]], Any) -> Any
59
60
_check_sentry_initialized ()
60
61
61
62
transaction = sentry_sdk .continue_trace (
62
- _tracing or {},
63
+ _sentry_tracing or {},
63
64
op = OP .QUEUE_TASK_RAY ,
64
65
name = qualname_from_function (user_f ),
65
66
origin = RayIntegration .origin ,
@@ -78,6 +79,19 @@ def new_func(*f_args, _tracing=None, **f_kwargs):
78
79
79
80
return result
80
81
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
+
81
95
if f :
82
96
rv = old_remote (new_func )
83
97
else :
@@ -99,7 +113,9 @@ def _remote_method_with_header_propagation(*args, **kwargs):
99
113
for k , v in sentry_sdk .get_current_scope ().iter_trace_propagation_headers ()
100
114
}
101
115
try :
102
- result = old_remote_method (* args , ** kwargs , _tracing = tracing )
116
+ result = old_remote_method (
117
+ * args , ** kwargs , _sentry_tracing = tracing
118
+ )
103
119
span .set_status (SPANSTATUS .OK )
104
120
except Exception :
105
121
span .set_status (SPANSTATUS .INTERNAL_ERROR )
0 commit comments