11import inspect
2+ import functools
23import sys
34
45import sentry_sdk
1718 import ray # type: ignore[import-not-found]
1819except ImportError :
1920 raise DidNotEnable ("Ray not installed." )
20- import functools
2121
2222from 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 )
0 commit comments