Repository: facebookresearch/CompilerGym

Issue:
Either explicitly prevent copy operators on environments, or add support for them
## 🚀 Feature

As of v0.1.9, the shallow copy operator works on `CompilerEnv` instances but the deep copy operator fails:

```py
In [1]: import compiler_gym

In [2]: env = compiler_gym.make("llvm-v0")

In [3]: import copy

In [4]: copy.copy(env)
Out[4]: <compiler_gym.envs.llvm.llvm_env.LlvmEnv at 0x7f908027bc40>

In [5]: copy.deepcopy(env)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
...
TypeError: cannot pickle '_thread.lock' object
```

## Pitch

Either cause both `copy()` and `deepcopy()` to raise an error, or use `env.fork()` to add support for them.


Discussion:

Resolving patch:
diff --git a/compiler_gym/envs/compiler_env.py b/compiler_gym/envs/compiler_env.py
--- a/compiler_gym/envs/compiler_env.py
+++ b/compiler_gym/envs/compiler_env.py
@@ -1322,3 +1322,12 @@ def send_params(self, *params: Iterable[Tuple[str, str]]) -> List[str]:
             )
 
         return list(reply.reply)
+
+    def __copy__(self) -> "CompilerEnv":
+        raise TypeError(
+            "CompilerEnv instances do not support shallow copies. Use deepcopy()"
+        )
+
+    def __deepcopy__(self, memo) -> "CompilerEnv":
+        del memo  # unused
+        return self.fork()


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -26,6 +26,16 @@ py_test(
     ],
 )
 
+py_test(
+    name = "env_copy_test",
+    srcs = ["env_copy_test.py"],
+    deps = [
+        "//compiler_gym/envs",
+        "//tests:test_main",
+        "//tests/pytest_plugins:llvm",
+    ],
+)
+
 py_test(
     name = "make_test",
     timeout = "short",
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -28,6 +28,17 @@ cg_py_test(
     tests::test_main
 )
 
+cg_py_test(
+  NAME
+    env_copy_test
+  SRCS
+    "env_copy_test.py"
+  DEPS
+    compiler_gym::envs::envs
+    tests::pytest_plugins::llvm
+    tests::test_main
+)
+
 cg_py_test(
   NAME
     make_test
diff --git a/tests/env_copy_test.py b/tests/env_copy_test.py
new file mode 100644
--- /dev/null
+++ b/tests/env_copy_test.py
@@ -0,0 +1,35 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""Tests for the copy() and deepcopy() operators on CompilerEnv."""
+from copy import copy, deepcopy
+
+import pytest
+
+from compiler_gym.envs.llvm import LlvmEnv
+from tests.test_main import main
+
+pytest_plugins = ["tests.pytest_plugins.llvm"]
+
+
+def test_forbidden_shallow_copy(env: LlvmEnv):
+    """Test that shallow copy operator is explicitly forbidden."""
+    with pytest.raises(
+        TypeError,
+        match=r"^CompilerEnv instances do not support shallow copies. Use deepcopy\(\)",
+    ):
+        copy(env)
+
+
+def test_deep_copy(env: LlvmEnv):
+    """Test that deep copy creates an independent copy."""
+    env.reset()
+    with deepcopy(env) as cpy:
+        assert cpy.state == env.state
+        env.step(env.action_space.sample())
+        assert cpy.state != env.state
+
+
+if __name__ == "__main__":
+    main()