fix(integrations/ray): Correctly pass keyword arguments to ray.remote function

sentry_sdk/integrations/ray.py

@@ -14,7 +14,7 @@
 )
 
 try:
-    import ray  # type: ignore[import-not-found]
+    import ray
 except ImportError:
     raise DidNotEnable("Ray not installed.")
 import functools
@@ -42,75 +42,83 @@ def _patch_ray_remote():
     old_remote = ray.remote
 
     @functools.wraps(old_remote)
-    def new_remote(f, *args, **kwargs):
-        # type: (Callable[..., Any], *Any, **Any) -> Callable[..., Any]
+    def new_remote(f=None, *args, **kwargs):
+        # type: (Optional[Callable[..., Any]], *Any, **Any) -> Callable[..., Any]
+
         if inspect.isclass(f):
             # Ray Actors
             # (https://docs.ray.io/en/latest/ray-core/actors.html)
             # are not supported
             # (Only Ray Tasks are supported)
-            return old_remote(f, *args, *kwargs)
-
-        def _f(*f_args, _tracing=None, **f_kwargs):
-            # type: (Any, Optional[dict[str, Any]],  Any) -> Any
-            """
-            Ray Worker
-            """
-            _check_sentry_initialized()
-
-            transaction = sentry_sdk.continue_trace(
-                _tracing or {},
-                op=OP.QUEUE_TASK_RAY,
-                name=qualname_from_function(f),
-                origin=RayIntegration.origin,
-                source=TransactionSource.TASK,
-            )
-
-            with sentry_sdk.start_transaction(transaction) as transaction:
-                try:
-                    result = 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
-
-        rv = old_remote(_f, *args, *kwargs)
-        old_remote_method = rv.remote
-
-        def _remote_method_with_header_propagation(*args, **kwargs):
-            # type: (*Any, **Any) -> Any
-            """
-            Ray Client
-            """
-            with sentry_sdk.start_span(
-                op=OP.QUEUE_SUBMIT_RAY,
-                name=qualname_from_function(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, _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
-
-        rv.remote = _remote_method_with_header_propagation
-
-        return rv
-
-    ray.remote = new_remote
+            return old_remote(f, *args, **kwargs)
+
+        def wrapper(user_f):
+            # type: (Callable[..., Any]) -> Any
+            def new_func(*f_args, _tracing=None, **f_kwargs):
+                # type: (Any, Optional[dict[str, Any]], Any) -> Any
+                _check_sentry_initialized()
+
+                transaction = sentry_sdk.continue_trace(
+                    _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
+
+            if f:
+                rv = old_remote(new_func)
+            else:
+                rv = old_remote(*args, **kwargs)(new_func)
+            old_remote_method = rv.remote
+
+            def _remote_method_with_header_propagation(*args, **kwargs):
+                # type: (*Any, **Any) -> Any
+                """
+                Ray Client
+                """
+                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, _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
+
+            rv.remote = _remote_method_with_header_propagation
+
+            return rv
+
+        if f is not None:
+            return wrapper(f)
+        else:
+            return wrapper
+
+    ray.remote = new_remote  # type: ignore[assignment]
 
 
 def _capture_exception(exc_info, **kwargs):

tests/integrations/ray/test_ray.py

@@ -59,7 +59,10 @@ def read_error_from_log(job_id):
 
 
 @pytest.mark.forked
-def test_tracing_in_ray_tasks():
+@pytest.mark.parametrize(
+    "task_options", [{}, {"num_cpus": 0, "memory": 1024 * 1024 * 10}]
+)
+def test_tracing_in_ray_tasks(task_options):
     setup_sentry()
 
     ray.init(
@@ -69,14 +72,19 @@ def test_tracing_in_ray_tasks():
         }
     )
 
-    # Setup ray task
-    @ray.remote
     def example_task():
         with sentry_sdk.start_span(op="task", name="example task step"):
             ...
 
         return sentry_sdk.get_client().transport.envelopes
 
+    # Setup ray task, calling decorator directly instead of @,
+    # to accommodate for test parametrization
+    if task_options:
+        example_task = ray.remote(**task_options)(example_task)
+    else:
+        example_task = ray.remote(example_task)
+
     with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
         worker_envelopes = ray.get(example_task.remote())