From 68e3a0e0e1c5a47eea48085764112e2262df7ea3 Mon Sep 17 00:00:00 2001
From: Barak Michener <me@barakmich.com>
Date: Fri, 15 Jan 2021 15:49:38 -0800
Subject: [PATCH] [ray_client]: fix wrong reference in server_pickler (#13474)

Change-Id: Ie3d219541b1875e986e72e3ae73ece145c715acf
---
 python/ray/tests/test_actor.py                | 19 +++++++++++++++++++
 .../ray/util/client/server/server_pickler.py  |  6 +++---
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/python/ray/tests/test_actor.py b/python/ray/tests/test_actor.py
index cc252d24d..283e1a00a 100644
--- a/python/ray/tests/test_actor.py
+++ b/python/ray/tests/test_actor.py
@@ -860,6 +860,25 @@ def test_get_non_existing_named_actor(ray_start_regular_shared):
         _ = ray.get_actor("non_existing_actor")
 
 
+def test_wrapped_actor_handle(ray_start_regular_shared):
+    @ray.remote
+    class B:
+        def doit(self):
+            return 2
+
+    @ray.remote
+    class A:
+        def __init__(self):
+            self.b = B.remote()
+
+        def get_actor_ref(self):
+            return [self.b]
+
+    a = A.remote()
+    b_list = ray.get(a.get_actor_ref.remote())
+    assert ray.get(b_list[0].doit.remote()) == 2
+
+
 @pytest.mark.skip(
     "This test is just used to print the latency of creating 100 actors.")
 def test_actor_creation_latency(ray_start_regular_shared):
diff --git a/python/ray/util/client/server/server_pickler.py b/python/ray/util/client/server/server_pickler.py
index a8b66832c..b0fb46825 100644
--- a/python/ray/util/client/server/server_pickler.py
+++ b/python/ray/util/client/server/server_pickler.py
@@ -61,9 +61,9 @@ class ServerPickler(cloudpickle.CloudPickler):
             actor_id = obj._actor_id.binary()
             if actor_id not in self.server.actor_refs:
                 # We're passing back a handle, probably inside a reference.
-                self.actor_refs[actor_id] = obj
-            if actor_id not in self.actor_owners[self.client_id]:
-                self.actor_owners[self.client_id].add(actor_id)
+                self.server.actor_refs[actor_id] = obj
+            if actor_id not in self.server.actor_owners[self.client_id]:
+                self.server.actor_owners[self.client_id].add(actor_id)
             return PickleStub(
                 type="Actor",
                 client_id=self.client_id,