_ray_trace_ctx fix follow-up (ray-project#18950)

* sanity check

* add test case

* fix assert

* refactor

* check kwargs instead of _kwargs

* format

python/ray/tests/client_test_utils.py

@@ -18,3 +18,20 @@ async def wait(self, should_wait=True):
  await self.ready_event.wait()
 
  return SignalActor
+
+
+# See test_client::test_wrapped_actor_creation for details on usage of
+# run_wrapped_actor_creation and SomeClass.
+def run_wrapped_actor_creation():
+ import ray
+ RemoteClass = ray.remote(SomeClass)
+ handle = RemoteClass.remote()
+ return ray.get(handle.ready.remote())
+
+
+class SomeClass:
+ def __init__(self):
+ pass
+
+ def ready(self):
+ return 1

python/ray/tests/test_client.py

@@ -10,6 +10,7 @@
 
 import ray.util.client.server.server as ray_client_server
 from ray.tests.client_test_utils import create_remote_signal_actor
+from ray.tests.client_test_utils import run_wrapped_actor_creation
 from ray.util.client.common import ClientObjectRef
 from ray.util.client.ray_client_helpers import connect_to_client_or_not
 from ray.util.client.ray_client_helpers import ray_start_client_server
@@ -707,5 +708,37 @@ def test_object_ref_cleanup():
  assert result == ""
 
 
+@pytest.mark.parametrize(
+ "call_ray_start",
+ ["ray start --head --ray-client-server-port 25552 --port 0"],
+ indirect=True)
+def test_wrapped_actor_creation(call_ray_start):
+ """
+ When the client schedules an actor, the server will load a separate
+ copy of the actor class if it's defined in a separate file. This
+ means that modifications to the client's copy of the actor class
+ aren't propagated to the server. Currently, tracing logic modifies
+ the signatures of actor methods to pass around metadata when ray.remote
+ is applied to an actor class. However, if a user does something like:
+
+ class SomeActor:
+ def __init__(self):
+ pass
+
+ def decorate_actor():
+ RemoteActor = ray.remote(SomeActor)
+ ...
+
+ Then the SomeActor class will have its signatures modified on the client
+ side, but not on the server side, since ray.remote was applied inside of
+ the function instead of directly on the actor. Note if it were directly
+ applied to the actor then the signature would be modified when the server
+ imports the class.
+ """
+ import ray
+ ray.init("ray:https://localhost:25552")
+ run_wrapped_actor_creation()
+
+
 if __name__ == "__main__":
  sys.exit(pytest.main(["-v", __file__]))

python/ray/util/tracing/tracing_helper.py

@@ -290,6 +290,8 @@ def _invocation_remote_span(
  # If tracing feature flag is not on, perform a no-op.
  # Tracing doesn't work for cross lang yet.
  if not is_tracing_enabled() or self._is_cross_language:
+ if kwargs is not None:
+ assert "_ray_trace_ctx" not in kwargs
  return method(self, args, kwargs, *_args, **_kwargs)
 
  assert "_ray_trace_ctx" not in kwargs
@@ -365,9 +367,7 @@ def _invocation_actor_class_remote_span(
 
  # If tracing feature flag is not on, perform a no-op
  if not is_tracing_enabled():
- if not self.__ray_metadata__.is_cross_language:
- # Remove _ray_trace_ctx from kwargs if tracing disabled
- kwargs.pop("_ray_trace_ctx", None)
+ assert "_ray_trace_ctx" not in kwargs
  return method(self, args, kwargs, *_args, **_kwargs)
 
  class_name = self.__ray_metadata__.class_name
@@ -405,6 +405,8 @@ def _start_span(
  # If tracing feature flag is not on, perform a no-op
  if (not is_tracing_enabled()
  or self._actor_ref()._ray_is_cross_language):
+ if kwargs is not None:
+ assert "_ray_trace_ctx" not in kwargs
  return method(self, args, kwargs, *_args, **_kwargs)
 
  class_name = (self._actor_ref()