Merge branch 'master' into expose-fblogger

vfdev-5 · web-flow · commit b7894a97981a · 2025-10-14T08:18:05.000+01:00
diff --git a/.github/workflows/gpu-hvd-tests.yml b/.github/workflows/gpu-hvd-tests.yml
@@ -24,7 +24,7 @@ jobs:
         pytorch-channel: [pytorch]
       fail-fast: false
     env:
-      DOCKER_IMAGE: "pytorch/conda-builder:cuda12.1"
+      DOCKER_IMAGE: "pytorch/almalinux-builder:cuda12.8"
       REPOSITORY: ${{ github.repository }}
       PR_NUMBER: ${{ github.event.pull_request.number }}
     runs-on: linux.8xlarge.nvidia.gpu
@@ -113,6 +113,10 @@ jobs:
           pip install -r requirements-dev.txt
           pip install -e .
 
+          # Upgrade pyOpenSSL to avoid issue:
+          # AttributeError: module 'lib' has no attribute 'X509_V_FLAG_NOTIFY_POLICY'. Did you mean: 'X509_V_FLAG_EXPLICIT_POLICY'?
+          pip install -U pyOpenSSL
+
           EOF
           )
 
diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml
@@ -24,7 +24,7 @@ jobs:
         pytorch-channel: [pytorch, pytorch-nightly]
       fail-fast: false
     env:
-      DOCKER_IMAGE: "pytorch/almalinux-builder:cuda12.4"
+      DOCKER_IMAGE: "pytorch/almalinux-builder:cuda12.8"
       REPOSITORY: ${{ github.repository }}
       PR_NUMBER: ${{ github.event.pull_request.number }}
     runs-on: linux.g4dn.12xlarge.nvidia.gpu
@@ -113,6 +113,10 @@ jobs:
           pip install -r requirements-dev.txt
           pip install -e .
 
+          # Upgrade pyOpenSSL to avoid issue:
+          # AttributeError: module 'lib' has no attribute 'X509_V_FLAG_NOTIFY_POLICY'. Did you mean: 'X509_V_FLAG_EXPLICIT_POLICY'?
+          pip install -U pyOpenSSL
+
           EOF
           )
 
diff --git a/.github/workflows/mps-tests.yml b/.github/workflows/mps-tests.yml
@@ -36,7 +36,7 @@ jobs:
         pytorch-channel: ["pytorch"]
         skip-distrib-tests: [1]
       fail-fast: false
-    runs-on: ["macos-m1-stable"]
+    runs-on: ["macos-14"]
     timeout-minutes: 60
 
     steps:
@@ -71,7 +71,7 @@ jobs:
       - name: Install uv
         shell: bash -l {0}
         run: |
-          conda shell.bash hook
+          eval "$(conda shell.bash hook)"
           conda activate $CONDA_ENV
           pip install uv
 
@@ -85,23 +85,23 @@ jobs:
         if: ${{ matrix.pytorch-channel == 'pytorch' }}
         shell: bash -l {0}
         run: |
-          conda shell.bash hook
+          eval "$(conda shell.bash hook)"
           conda activate $CONDA_ENV
           uv pip install torch torchvision
 
       - name: Install PyTorch (nightly)
         if: ${{ matrix.pytorch-channel == 'pytorch-nightly' }}
         shell: bash -l {0}
         run: |
-          conda shell.bash hook
+          eval "$(conda shell.bash hook)"
           conda activate $CONDA_ENV
           uv pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
 
       - name: Install dependencies
         shell: bash -l {0}
         working-directory: ${{ github.repository }}
         run: |
-          conda shell.bash hook
+          eval "$(conda shell.bash hook)"
           conda activate $CONDA_ENV
           # TODO: We add set -xe to explicitly fail the CI if one of the commands is failing.
           # Somehow the step is passing even if a subcommand failed
@@ -129,7 +129,7 @@ jobs:
         shell: bash -l {0}
         working-directory: ${{ github.repository }}
         run: |
-          conda shell.bash hook
+          eval "$(conda shell.bash hook)"
           conda activate $CONDA_ENV
           SKIP_DISTRIB_TESTS=${{ matrix.skip-distrib-tests }} bash tests/run_cpu_tests.sh
 
@@ -144,6 +144,6 @@ jobs:
         shell: bash -l {0}
         working-directory: ${{ github.repository }}
         run: |
-          conda shell.bash hook
+          eval "$(conda shell.bash hook)"
           conda activate $CONDA_ENV
           python examples/mnist/mnist.py --epochs=1
diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
@@ -182,9 +182,7 @@ jobs:
           set -e
           echo "Testing Caltech101 dataset availability..."
           if python -c "import torchvision; torchvision.datasets.Caltech101(root='./data', download=True)"; then
-            echo "Caltech101 dataset downloaded successfully. Please remove this workaround and restore dataset check."
-            exit 1
-            # python examples/super_resolution/main.py --upscale_factor 3 --crop_size 180 --batch_size 4 --test_batch_size 100 --n_epochs 1 --lr 0.001 --threads 2 --debug
+            python examples/super_resolution/main.py --upscale_factor 3 --crop_size 180 --batch_size 4 --test_batch_size 100 --n_epochs 1 --lr 0.001 --threads 2 --debug
           else
             echo "Caltech101 dataset failed to download. Skipping SR example test."
           fi
diff --git a/docs/source/handlers.rst b/docs/source/handlers.rst
@@ -11,6 +11,7 @@ Complete list of generic handlers
     :toctree: generated
 
     checkpoint.Checkpoint
+    checkpoint.CheckpointEvents
     DiskSaver
     checkpoint.ModelCheckpoint
     ema_handler.EMAHandler
diff --git a/ignite/engine/engine.py b/ignite/engine/engine.py
@@ -249,6 +249,17 @@ class TBPTT_Events(EventEnum):
         # we need to update state attributes associated with new custom events
         self.state._update_attrs()
 
+    def has_registered_events(self, event: Any) -> bool:
+        """Check whether engine has a registered event.
+
+        Args:
+            event: Event to check for registration.
+
+        Returns:
+            bool: True if the event is registered, False otherwise.
+        """
+        return event in self._allowed_events
+
     def _handler_wrapper(self, handler: Callable, event_name: Any, event_filter: Callable) -> Callable:
         # signature of the following wrapper will be inspected during registering to check if engine is necessary
         # we have to build a wrapper with relevant signature : solution is functools.wraps
@@ -328,7 +339,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 +443,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:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
@@ -21,10 +21,21 @@
 
 import ignite.distributed as idist
 from ignite.base import Serializable
-from ignite.engine import Engine, Events
+from ignite.engine import Engine, Events, EventEnum
 from ignite.utils import _tree_apply2, _tree_map
 
-__all__ = ["Checkpoint", "DiskSaver", "ModelCheckpoint", "BaseSaveHandler"]
+__all__ = ["Checkpoint", "DiskSaver", "ModelCheckpoint", "BaseSaveHandler", "CheckpointEvents"]
+
+
+class CheckpointEvents(EventEnum):
+    """Events fired by :class:`~ignite.handlers.checkpoint.Checkpoint`
+
+    - SAVED_CHECKPOINT : triggered when checkpoint handler has saved objects
+
+    .. versionadded:: 0.5.3
+    """
+
+    SAVED_CHECKPOINT = "saved_checkpoint"
 
 
 class BaseSaveHandler(metaclass=ABCMeta):
@@ -264,6 +275,29 @@ class Checkpoint(Serializable):
                 to_save, save_handler=DiskSaver('/tmp/models', create_dir=True, **kwargs), n_saved=2
             )
 
+        Respond to checkpoint events:
+
+        .. code-block:: python
+
+            from ignite.handlers import Checkpoint
+            from ignite.engine import Engine, Events
+
+            checkpoint_handler = Checkpoint(
+                {'model': model, 'optimizer': optimizer},
+                save_dir,
+                n_saved=2
+            )
+
+            @trainer.on(Checkpoint.SAVED_CHECKPOINT)
+            def on_checkpoint_saved(engine):
+                print(f"Checkpoint saved at epoch {engine.state.epoch}")
+
+            trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler)
+
+    Attributes:
+        SAVED_CHECKPOINT: Alias of ``SAVED_CHECKPOINT`` from
+            :class:`~ignite.handlers.checkpoint.CheckpointEvents`.
+
     .. versionchanged:: 0.4.3
 
         - Checkpoint can save model with same filename.
@@ -274,8 +308,13 @@ class Checkpoint(Serializable):
         - `score_name` can be used to define `score_function` automatically without providing `score_function`.
         - `save_handler` automatically saves to disk if path to directory is provided.
         - `save_on_rank` saves objects on this rank in a distributed configuration.
+
+    .. versionchanged:: 0.5.3
+
+        - Added ``SAVED_CHECKPOINT`` class attribute.
     """
 
+    SAVED_CHECKPOINT = CheckpointEvents.SAVED_CHECKPOINT
     Item = NamedTuple("Item", [("priority", int), ("filename", str)])
     _state_dict_all_req_keys = ("_saved",)
 
@@ -400,6 +439,8 @@ def _compare_fn(self, new: Union[int, float]) -> bool:
             return new > self._saved[0].priority
 
     def __call__(self, engine: Engine) -> None:
+        if not engine.has_registered_events(CheckpointEvents.SAVED_CHECKPOINT):
+            engine.register_events(*CheckpointEvents)
         global_step = None
         if self.global_step_transform is not None:
             global_step = self.global_step_transform(engine, engine.last_event_name)
@@ -460,11 +501,11 @@ def __call__(self, engine: Engine) -> None:
             if self.include_self:
                 # Now that we've updated _saved, we can add our own state_dict.
                 checkpoint["checkpointer"] = self.state_dict()
-
             try:
                 self.save_handler(checkpoint, filename, metadata)
             except TypeError:
                 self.save_handler(checkpoint, filename)
+            engine.fire_event(CheckpointEvents.SAVED_CHECKPOINT)
 
     def _setup_checkpoint(self) -> Dict[str, Any]:
         if self.to_save is not None:
diff --git a/ignite/handlers/visdom_logger.py b/ignite/handlers/visdom_logger.py
@@ -1,7 +1,7 @@
 """Visdom logger and its helper handlers."""
 
 import os
-from typing import Any, Callable, cast, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -179,7 +179,7 @@ def __init__(
                 )
 
         if server is None:
-            server = cast(str, os.environ.get("VISDOM_SERVER_URL", "localhost"))
+            server = os.environ.get("VISDOM_SERVER_URL", "localhost")
 
         if port is None:
             port = int(os.environ.get("VISDOM_PORT", 8097))
diff --git a/tests/ignite/distributed/test_auto.py b/tests/ignite/distributed/test_auto.py
@@ -311,7 +311,7 @@ def test_dist_proxy_sampler():
         DistributedProxySampler(None)
 
     with pytest.raises(TypeError, match=r"Argument sampler should have length"):
-        DistributedProxySampler(Sampler([1]))
+        DistributedProxySampler(Sampler())
 
     with pytest.raises(TypeError, match=r"Argument sampler must not be a distributed sampler already"):
         DistributedProxySampler(DistributedSampler(sampler, num_replicas=num_replicas, rank=0))
diff --git a/tests/ignite/engine/test_custom_events.py b/tests/ignite/engine/test_custom_events.py
@@ -46,6 +46,24 @@ def process_func(engine, batch):
         assert h.called
 
 
+def test_has_registered_events_custom():
+    """Test has_registered_events with custom events."""
+
+    class TestEvents(EventEnum):
+        CUSTOM_EVENT = "custom_event"
+
+    engine = Engine(lambda e, b: None)
+
+    # Custom event not registered yet
+    assert not engine.has_registered_events(TestEvents.CUSTOM_EVENT)
+
+    # Register custom event
+    engine.register_events(TestEvents.CUSTOM_EVENT)
+
+    # Now should return True
+    assert engine.has_registered_events(TestEvents.CUSTOM_EVENT)
+
+
 def test_custom_events_asserts():
     # Dummy engine
     engine = Engine(lambda engine, batch: 0)
diff --git a/tests/ignite/engine/test_engine.py b/tests/ignite/engine/test_engine.py
@@ -1174,6 +1174,18 @@ def update_fn(engine, batch):
 
         assert mem_consumption1 == mem_consumption2
 
+    def test_has_registered_events_builtin(self):
+        """Test has_registered_events with built-in events."""
+        engine = Engine(lambda e, b: None)
+
+        # Built-in events should be registered by default
+        assert engine.has_registered_events(Events.STARTED)
+        assert engine.has_registered_events(Events.COMPLETED)
+        assert engine.has_registered_events(Events.ITERATION_COMPLETED)
+
+        # Non-existent event should return False
+        assert not engine.has_registered_events("non_existent_event")
+
     def test_engine_no_data_asserts(self):
         trainer = Engine(lambda e, b: None)
 
diff --git a/tests/ignite/engine/test_memory_leaks.py b/tests/ignite/engine/test_memory_leaks.py
@@ -0,0 +1,44 @@
+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
+
+            assert counter == i + 1
+
+        assert counter == i + 1
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
