Use weak references to resolve circular references and add test

Author: goanpeca

ignite/engine/engine.py

@@ -328,7 +328,7 @@ def execute_something():
 
         try:
             _check_signature(handler, "handler", self, *(event_args + args), **kwargs)
-            self._event_handlers[event_name].append((handler, (self,) + args, kwargs))
+            self._event_handlers[event_name].append((handler, (weakref.ref(self),) + args, kwargs))
         except ValueError:
             _check_signature(handler, "handler", *(event_args + args), **kwargs)
             self._event_handlers[event_name].append((handler, args, kwargs))
@@ -432,7 +432,15 @@ def _fire_event(self, event_name: Any, *event_args: Any, **event_kwargs: Any) ->
         self.last_event_name = event_name
         for func, args, kwargs in self._event_handlers[event_name]:
             kwargs.update(event_kwargs)
-            first, others = ((args[0],), args[1:]) if (args and args[0] == self) else ((), args)
+            if args and isinstance(args[0], weakref.ref):
+                resolved_engine = args[0]()
+                if resolved_engine is None:
+                    raise RuntimeError("Engine reference not resolved. Cannot execute event handler.")
+                first, others = ((resolved_engine,), args[1:])
+            else:
+                # metrics do not provide engine when registered
+                first, others = (tuple(), args)  # type: ignore[assignment]
+
             func(*first, *(event_args + others), **kwargs)
 
     def fire_event(self, event_name: Any) -> None:

tests/ignite/engine/test_memory_leaks.py

@@ -0,0 +1,58 @@
+import weakref
+
+import pytest
+
+from ignite.engine import Engine, Events
+
+
+class TestEngineMemoryLeak:
+    """See: https://github.com/pytorch/ignite/issues/3438"""
+
+    ENGINE_WEAK_REFS = set()
+
+    def do_train(self, cls, with_handler) -> None:
+        engine = cls(lambda e, b: None)
+
+        if with_handler:
+
+            @engine.on(Events.EPOCH_STARTED)
+            def handler(engine) -> None:
+                pass
+
+        engine.run(range(5), max_epochs=5)
+        self.ENGINE_WEAK_REFS.add(weakref.ref(engine))
+
+    @pytest.mark.parametrize(
+        "with_handler",
+        [
+            (True,),
+            (False,),
+        ],
+    )
+    def test_memory_leak(self, with_handler):
+        num_iters = 5
+        counter = 0
+
+        class EngineForTests(Engine):
+
+            def __del__(self):
+                nonlocal counter
+                counter += 1
+
+        for i in range(num_iters):
+            self.do_train(EngineForTests, with_handler)
+            for weak_engine_ref in self.ENGINE_WEAK_REFS:
+                engine = weak_engine_ref()
+                assert engine is None
+
+            print(counter)
+            if with_handler:
+                assert counter == i + 1
+            else:
+                assert counter == 0
+
+        print(counter)
+        if with_handler:
+            assert counter == i + 1
+        else:
+            assert counter == 0