From a0772427528c5376e7fc0ba9be888dbdf210dcbf Mon Sep 17 00:00:00 2001
From: Hovhannes Tamoyan <hovhannes.tamoyan@gmail.com>
Date: Thu, 27 Jun 2024 03:39:02 +0200
Subject: [PATCH] [fix] Add hydra plugin as a file in the module

---
 setup.py                                      |  5 +--
 urartu/VERSION                                |  2 +-
 urartu/commands/launch.py                     | 31 +++++++++----------
 urartu/hydra_plugins/__init__.py              |  0
 .../hydra_plugins/urartu_plugin/__init__.py   |  0
 .../hydra_plugin.py}                          |  0
 6 files changed, 17 insertions(+), 21 deletions(-)
 delete mode 100644 urartu/hydra_plugins/__init__.py
 delete mode 100644 urartu/hydra_plugins/urartu_plugin/__init__.py
 rename urartu/{hydra_plugins/urartu_plugin/urartu_plugin.py => utils/hydra_plugin.py} (100%)

diff --git a/setup.py b/setup.py
index c5ef86f..c35a109 100644
--- a/setup.py
+++ b/setup.py
@@ -37,9 +37,6 @@
     entry_points={
         "console_scripts": [
             "urartu=urartu.__init__:main",
-        ],
-        "hydra.plugins": [
-            "urartu_plugin = urartu.hydra_plugins.urartu_plugin.urartu_plugin:UrartuPlugin",
-        ],
+        ]
     },
 )
diff --git a/urartu/VERSION b/urartu/VERSION
index ed35d09..fee3c92 100644
--- a/urartu/VERSION
+++ b/urartu/VERSION
@@ -1 +1 @@
-2.0.9
\ No newline at end of file
+2.0.10
\ No newline at end of file
diff --git a/urartu/commands/launch.py b/urartu/commands/launch.py
index 75aba2e..9d8d8ec 100644
--- a/urartu/commands/launch.py
+++ b/urartu/commands/launch.py
@@ -4,16 +4,21 @@
 
 from aim import Run
 from hydra import compose, initialize
-
+from hydra.core.plugins import Plugins
 from .command import Command
-
-logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
-)
-
 from ..utils.launcher import launch, launch_on_slurm
 from ..utils.slurm import is_submitit_available
 from ..utils.registry import Registry
+from ..utils.hydra_plugin import UrartuPlugin
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+def register_my_plugin() -> None:
+    Plugins.instance().register(UrartuPlugin)
+
+
+register_my_plugin()
 
 
 @Command.register("launch")
@@ -22,9 +27,7 @@ class Launch(Command):
     Launches an action from a specific module
     """
 
-    def add_subparser(
-        self, parser: argparse._SubParsersAction
-    ) -> argparse.ArgumentParser:
+    def add_subparser(self, parser: argparse._SubParsersAction) -> argparse.ArgumentParser:
         description = """urartu: launcher"""
         subparser = parser.add_parser(
             self.name,
@@ -33,9 +36,7 @@ def add_subparser(
         )
 
         subparser.add_argument("--name", type=str, help="name of the project/module")
-        subparser.add_argument(
-            "module_args", nargs=argparse.REMAINDER, help="module arguments"
-        )
+        subparser.add_argument("module_args", nargs=argparse.REMAINDER, help="module arguments")
 
         subparser.set_defaults(fire=self._launch)
 
@@ -44,7 +45,7 @@ def add_subparser(
     def _launch(self, args: argparse.Namespace):
         module_name = re.sub(r"[^A-Za-z0-9]+", "", args.name)
         module_path = Registry.get_module_path_by_name(module_name)
- 
+
         with initialize(version_base=None, config_path="../config"):
             cfg = compose(config_name="main", overrides=args.module_args)
 
@@ -55,9 +56,7 @@ def _launch(self, args: argparse.Namespace):
         aim_run.set("cfg", cfg, strict=False)
 
         if cfg.slurm.use_slurm:
-            assert (
-                is_submitit_available()
-            ), "Please 'pip install submitit' to schedule jobs on SLURM"
+            assert is_submitit_available(), "Please 'pip install submitit' to schedule jobs on SLURM"
 
             launch_on_slurm(
                 module=module_path,
diff --git a/urartu/hydra_plugins/__init__.py b/urartu/hydra_plugins/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/urartu/hydra_plugins/urartu_plugin/__init__.py b/urartu/hydra_plugins/urartu_plugin/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/urartu/hydra_plugins/urartu_plugin/urartu_plugin.py b/urartu/utils/hydra_plugin.py
similarity index 100%
rename from urartu/hydra_plugins/urartu_plugin/urartu_plugin.py
rename to urartu/utils/hydra_plugin.py