Repository: OpenMined/PySyft

Issue:
Memory leak in the argmax() method of AST
This code, causes a memory leak only on the first worker in the worker list passed to the `encrypt()`. In this example it is  'alice', the second worker 'bob' does not suffer from any leaks.

The reason is the `argmax()` method:

```
import torch as th
import torch.optim as optim
import syft as sy

hook = sy.TorchHook(th) #, verbose=True)
bob = sy.VirtualWorker(hook, id="bob") #, verbose=True)
alice = sy.VirtualWorker(hook, id="alice") #, verbose=True)
crypto_provider = sy.VirtualWorker(hook, id="james") #, verbose=True)

a = th.ones(1, 5)
a = a.encrypt(workers=[alice, bob], crypto_provider=crypto_provider, requires_grad=True)

for i in range(3):
    print("=" * 10 + f"Iter{i + 1}" + "=" * 10)
    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")

    a.argmax(dim=1)

    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")

print("~" * 10 + "Done" + "~" * 10)
```

Here is the output (Only Alice leaks):


```
==========Iter1==========
Alice: 1
Bob: 1
Alice: 13
Bob: 1
==========Iter2==========
Alice: 13
Bob: 1
Alice: 25
Bob: 1
==========Iter3==========
Alice: 25
Bob: 1
Alice: 37
Bob: 1
~~~~~~~~~~Done~~~~~~~~~~
```

If you comment out `a.argmax(dim=1)`, no memory leaks are observed:

```
==========Iter1==========
Alice: 1
Bob: 1
Alice: 1
Bob: 1
==========Iter2==========
Alice: 1
Bob: 1
Alice: 1
Bob: 1
==========Iter3==========
Alice: 1
Bob: 1
Alice: 1
Bob: 1
~~~~~~~~~~Done~~~~~~~~~~



```


Discussion:

Resolving patch:
diff --git a/syft/frameworks/torch/torch_attributes.py b/syft/frameworks/torch/torch_attributes.py
--- a/syft/frameworks/torch/torch_attributes.py
+++ b/syft/frameworks/torch/torch_attributes.py
@@ -1,10 +1,8 @@
+import re
 from types import ModuleType
-from typing import Union
-from typing import Callable
-from typing import Any
 
-from syft.generic.frameworks.attributes import FrameworkAttributes
 from syft.frameworks.torch.tensors.interpreters.native import TorchTensor
+from syft.generic.frameworks.attributes import FrameworkAttributes
 
 
 class TorchAttributes(FrameworkAttributes):
@@ -135,6 +133,12 @@ def __init__(self, torch: ModuleType, hook: ModuleType) -> None:
 
         # Dict {method_name: <is_inplace:bool>
         self.inplace_methods = {}
+        self._inplace_pattern = re.compile(r"(^__i(?!nit|mport|ter).+_)|^((?!^_+).+[^_])_$")
+        # Positives:
+        # __iadd__, share_
+
+        # Negatives:
+        # __init__, __import__, __iter__, __foo__, __bar_foo
 
     def is_inplace_method(self, method_name):
         """Determine if a method is inplace or not.
@@ -150,6 +154,7 @@ def is_inplace_method(self, method_name):
         try:
             return self.inplace_methods[method_name]
         except KeyError:
-            is_inplace = method_name[-1] == "_" and "__" not in method_name
+            is_inplace = True if re.search(self._inplace_pattern, method_name) else False
+
             self.inplace_methods[method_name] = is_inplace
             return is_inplace
diff --git a/syft/generic/frameworks/hook/hook_args.py b/syft/generic/frameworks/hook/hook_args.py
--- a/syft/generic/frameworks/hook/hook_args.py
+++ b/syft/generic/frameworks/hook/hook_args.py
@@ -717,6 +717,11 @@ def register_tensor(
         response_ids: List of ids where the tensor should be stored
             and each id is pop out when needed.
     """
+    # This method often leads to re-registration of tensors
+    # hence creating two copies of the same info. The older tensor
+    # is left hanging and is never deleted. De-Registering the original
+    # tensor (if-exists) before registration addresses this problem.
+    owner.de_register_obj(tensor)  # Doesn't raise Exceptions if absent on owner
     tensor.owner = owner
     try:
         tensor.id = response_ids.pop(-1)


Test patch:
diff --git a/test/torch/test_hook.py b/test/torch/test_hook.py
--- a/test/torch/test_hook.py
+++ b/test/torch/test_hook.py
@@ -1,13 +1,14 @@
 """Tests relative to verifying the hook process behaves properly."""
+import re
+
 import pytest
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 import syft
-from syft.generic.pointers.pointer_tensor import PointerTensor
-
 from syft.exceptions import RemoteObjectFoundError
+from syft.generic.pointers.pointer_tensor import PointerTensor
 
 
 def test___init__(hook):
@@ -15,6 +16,25 @@ def test___init__(hook):
     assert hook.torch.__version__ == torch.__version__
 
 
+def test_torch_inplace_method():
+    positives = ["__iadd__", "__imul__", "__idiv__", "share_", "get_", "encrypt_"]
+    negatives = [
+        "__add__",
+        "__init__",
+        "__str__",
+        "share",
+        "get",
+        "encrypt",
+        "__foo",
+        "_bar",
+        "baz__",
+    ]
+    for pos in positives:
+        assert re.search(syft.framework._inplace_pattern, pos)
+    for neg in negatives:
+        assert not re.search(syft.framework._inplace_pattern, neg)
+
+
 def test_torch_attributes():
     with pytest.raises(RuntimeError):
         syft.framework._command_guard("false_command")