Create get_clearml_task for Ignite.contrib.handlers.ClearMLLogger

ignite/contrib/handlers/clearml_logger.py

@@ -7,6 +7,8 @@
 from enum import Enum
 from typing import Any, Callable, DefaultDict, List, Mapping, Optional, Tuple, Type, Union
 
+from clearml import Task
+
 from torch.optim import Optimizer
 
 import ignite.distributed as idist
@@ -178,6 +180,20 @@ def bypass_mode(cls) -> bool:
         """
         return getattr(cls, "_bypass", bool(os.environ.get("CI")))
 
+    def get_clearml_task(self) -> Union[None, Task]:
+        """
+        Returns the ClearML Task.
+
+        Return:
+            If bypass mode state is ``True``, all outside communication is skipped,
+            and ``get_task`` will return ``None``.
+            If bypass mode state is ``False``, this method will return the ClearML Task.
+        """
+        if self.bypass_mode():
+            return None
+
+        return self._task
+
     def close(self) -> None:
         self.clearml_logger.flush()
 

tests/ignite/contrib/handlers/test_clearml_logger.py

@@ -6,6 +6,7 @@
 import clearml
 import pytest
 import torch
+from clearml import Task
 from clearml.binding.frameworks import WeightsFileHandler
 from clearml.model import Framework
 
@@ -26,6 +27,35 @@
 from ignite.handlers import Checkpoint
 
 
+def test_clearml_logger_get_clearml_task_bypass(dirname):
+
+    with pytest.warns(UserWarning, match="ClearMLSaver: running in bypass mode"):
+        ClearMLLogger.set_bypass_mode(True)
+        with ClearMLLogger(output_uri=dirname) as clearml_logger:
+            assert clearml_logger.get_clearml_task() is None
+
+
+def test_clearml_logger_get_clearml_task_not_bypass_create_task():
+
+    ClearMLLogger.set_bypass_mode(False)
+    with ClearMLLogger() as clearml_logger:
+        res_task = clearml_logger.get_clearml_task()
+        assert res_task == Task.current_task()
+        res_task.close()
+
+
+def test_clearml_logger_get_clearml_task_not_bypass_task_already_exists():
+
+    user_created_task = Task.init(project_name="experiment", task_name="experiment")
+
+    ClearMLLogger.set_bypass_mode(False)
+    with ClearMLLogger(project_name="experiment", task_name="experiment") as clearml_logger:
+        res_task = clearml_logger.get_clearml_task()
+        assert res_task == user_created_task
+        res_task.close()
+        user_created_task.close()
+
+
 def test_no_clearml():
     with patch.dict("sys.modules", {"clearml": None, "trains": None}):
         with pytest.raises(ModuleNotFoundError, match=r"This contrib module requires clearml to be installed."):