Repository: ray-project/ray

Issue:
Passing ObjectID as a function argument in local_mode is broken
### System information
- **OS Platform and Distribution**: Ubuntu 18.04
- **Ray installed from (source or binary)**: binary
- **Ray version**: 0.8.0.dev6
- **Python version**: 3.7
- **Exact command to reproduce**: see below

### Describe the problem
The argument passing behavior with local_mode=True vs False seems to be different.
When I run the code snippet below:

```import ray
ray.init(local_mode=True)  # Replace with False to get a working example

@ray.remote
def remote_function(x):
    obj = x['a']
    return ray.get(obj)


a = ray.put(42)
d = {'a': a}
result = remote_function.remote(d)
print(ray.get(result))
```

With local_mode=False I get output `42`, as expected.
With local_mode=True I get the following error:

```
Traceback (most recent call last):
  File "/home/alex/all/projects/doom-neurobot/playground/ray_local_mode_bug.py", line 13, in <module>
    print(ray.get(result))
  File "/home/alex/miniconda3/envs/doom-rl/lib/python3.7/site-packages/ray/worker.py", line 2194, in get
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(KeyError): /home/alex/miniconda3/envs/doom-rl/bin/python /home/alex/all/projects/doom-neurobot/playground/ray_local_mode_bug.py (pid=2449, ip=10.136.109.38)
  File "/home/alex/miniconda3/envs/doom-rl/lib/python3.7/site-packages/ray/local_mode_manager.py", line 55, in execute
    results = function(*copy.deepcopy(args))
  File "/home/alex/all/projects/doom-neurobot/playground/ray_local_mode_bug.py", line 7, in remote_function
    return ray.get(obj)
  File "/home/alex/miniconda3/envs/doom-rl/lib/python3.7/site-packages/ray/local_mode_manager.py", line 105, in get_objects
    raise KeyError("Value for {} not found".format(object_id))
KeyError: 'Value for LocalModeObjectID(89f92e430883458c8107c10ed53eb35b26099831) not found'
```
It looks like the LocalObjectID instance inside `d` loses it's field `value` when it gets deep copied during the "remote" function call (currently it's `local_mode_manager.py:55`). It's hard to tell why exactly that happens, looks like a bug.






Discussion:
Indeed, this example seems to confirm that it might be a `deepcopy` problem:

```
import ray
ray.init(local_mode=True)

a = ray.put(42)

from copy import deepcopy
a_copy = deepcopy(a)
print(ray.get(a))
print(ray.get(a_copy))
```

I get `42` printed and then an error upon attempting to `ray.get(a_copy)`. Is this intended?
Related to #5853? cc @edoakes 
It is related, and I agree that local_mode should mimic normal operation as closely as possible.
Although probably there's an easier fix for this small issue.
Fixed it temporarily by implementing `__deepcopy__` on `LocalModeObjectID`:

```
    def __deepcopy__(self, memodict={}):
        obj_copy = copy.copy(self)
        content_copy = copy.deepcopy(self.__dict__, memodict)
        obj_copy.__dict__.update(content_copy)
        return obj_copy
```

Resolving patch:
diff --git a/python/ray/local_mode_manager.py b/python/ray/local_mode_manager.py
--- a/python/ray/local_mode_manager.py
+++ b/python/ray/local_mode_manager.py
@@ -5,6 +5,7 @@
 import copy
 import traceback
 
+import ray
 from ray import ObjectID
 from ray.utils import format_error_message
 from ray.exceptions import RayTaskError
@@ -20,7 +21,18 @@ class LocalModeObjectID(ObjectID):
             it equates to the object not existing in the object store. This is
             necessary because None is a valid object value.
     """
-    pass
+
+    def __copy__(self):
+        new = LocalModeObjectID(self.binary())
+        if hasattr(self, "value"):
+            new.value = self.value
+        return new
+
+    def __deepcopy__(self, memo=None):
+        new = LocalModeObjectID(self.binary())
+        if hasattr(self, "value"):
+            new.value = self.value
+        return new
 
 
 class LocalModeManager(object):
@@ -49,23 +61,37 @@ def execute(self, function, function_name, args, kwargs, num_return_vals):
         Returns:
             LocalModeObjectIDs corresponding to the function return values.
         """
-        object_ids = [
+        return_ids = [
             LocalModeObjectID.from_random() for _ in range(num_return_vals)
         ]
+        new_args = []
+        for i, arg in enumerate(args):
+            if isinstance(arg, ObjectID):
+                new_args.append(ray.get(arg))
+            else:
+                new_args.append(copy.deepcopy(arg))
+
+        new_kwargs = {}
+        for k, v in kwargs.items():
+            if isinstance(v, ObjectID):
+                new_kwargs[k] = ray.get(v)
+            else:
+                new_kwargs[k] = copy.deepcopy(v)
+
         try:
-            results = function(*copy.deepcopy(args), **copy.deepcopy(kwargs))
+            results = function(*new_args, **new_kwargs)
             if num_return_vals == 1:
-                object_ids[0].value = results
+                return_ids[0].value = results
             else:
-                for object_id, result in zip(object_ids, results):
+                for object_id, result in zip(return_ids, results):
                     object_id.value = result
         except Exception as e:
             backtrace = format_error_message(traceback.format_exc())
             task_error = RayTaskError(function_name, backtrace, e.__class__)
-            for object_id in object_ids:
+            for object_id in return_ids:
                 object_id.value = task_error
 
-        return object_ids
+        return return_ids
 
     def put_object(self, value):
         """Store an object in the emulated object store.


Test patch:
diff --git a/python/ray/tests/test_basic.py b/python/ray/tests/test_basic.py
--- a/python/ray/tests/test_basic.py
+++ b/python/ray/tests/test_basic.py
@@ -2227,6 +2227,17 @@ def function2(self):
     _ = RemoteActor2.remote()
     assert ray.get(actor1.function1.remote()) == 0
 
+    # Test passing ObjectIDs.
+    @ray.remote
+    def direct_dep(input):
+        return input
+
+    @ray.remote
+    def indirect_dep(input):
+        return ray.get(direct_dep.remote(input[0]))
+
+    assert ray.get(indirect_dep.remote(["hello"])) == "hello"
+
 
 def test_resource_constraints(shutdown_only):
     num_workers = 20