Repository: facebookresearch/hydra
Pull request #289 (2019-11-14 20:49:24)

Example Sweeper plugin
Adds to #253 


Patch:
diff --git a/hydra/_internal/core_plugins/basic_launcher.py b/hydra/_internal/core_plugins/basic_launcher.py
index b4cc875623..f992c9c789 100644
--- a/hydra/_internal/core_plugins/basic_launcher.py
+++ b/hydra/_internal/core_plugins/basic_launcher.py
@@ -33,7 +33,6 @@ class BasicLauncher(Launcher):
         sweep_dir = self.config.hydra.sweep.dir
         Path(str(sweep_dir)).mkdir(parents=True, exist_ok=True)
         log.info("Launching {} jobs locally".format(len(job_overrides)))
-        log.info("Sweep output dir : {}".format(sweep_dir))
         runs = []
 
         for idx, overrides in enumerate(job_overrides):
diff --git a/hydra/_internal/core_plugins/basic_sweeper.py b/hydra/_internal/core_plugins/basic_sweeper.py
index 2d083a06c1..4f54be652d 100644
--- a/hydra/_internal/core_plugins/basic_sweeper.py
+++ b/hydra/_internal/core_plugins/basic_sweeper.py
@@ -36,6 +36,7 @@ class BasicSweeper(StepSweeper):
         :return: A list of lists of strings, each inner list is the overrides for a single job
         that should be executed.
         """
+
         lists = []
         for s in self.arguments:
             key, value = s.split("=")
diff --git a/hydra/plugins/step_sweeper.py b/hydra/plugins/step_sweeper.py
index fcc047efe5..a060d26ccc 100644
--- a/hydra/plugins/step_sweeper.py
+++ b/hydra/plugins/step_sweeper.py
@@ -5,6 +5,9 @@ A sweeper that operates on generational batches of jobs
 
 from abc import abstractmethod
 from . import Sweeper
+import logging
+
+log = logging.getLogger(__name__)
 
 
 class StepSweeper(Sweeper):
@@ -23,6 +26,8 @@ class StepSweeper(Sweeper):
     def setup(self, config, config_loader, task_function):
         from .._internal.plugins import Plugins
 
+        self.config = config
+
         self.launcher = Plugins.instantiate_launcher(
             config=config, config_loader=config_loader, task_function=task_function
         )
@@ -53,12 +58,13 @@ class StepSweeper(Sweeper):
         raise NotImplementedError()
 
     def sweep(self, arguments):
+        log.info("Sweep output dir : {}".format(self.config.hydra.sweep.dir))
+
         self.arguments = arguments
         returns = []
         while not self.is_done():
             batch = self.get_job_batch()
             results = self.launcher.launch(batch)
-            # TODO: use extend
             returns.append(results)
             self.update_results(results)
         return returns
diff --git a/plugins/examples/example_sweeper/MANIFEST.in b/plugins/examples/example_sweeper/MANIFEST.in
new file mode 100644
index 0000000000..df1000cc71
--- /dev/null
+++ b/plugins/examples/example_sweeper/MANIFEST.in
@@ -0,0 +1,3 @@
+global-exclude *.pyc
+global-exclude __pycache__
+recursive-include hydra_plugins/* *.yaml
diff --git a/plugins/examples/example_sweeper/README.md b/plugins/examples/example_sweeper/README.md
new file mode 100644
index 0000000000..8fe0ce4c15
--- /dev/null
+++ b/plugins/examples/example_sweeper/README.md
@@ -0,0 +1,32 @@
+# Hydra example Launcher plugin
+
+This plugin provides an example for how to write a custom Sweeper for Hydra.
+The provided example has a custom configuration for the sweeper that takes two parameters, foo and bar:
+```yaml
+hydra:
+  sweeper:
+    class: hydra_plugins.example_sweeper.ExampleSweeper
+    params:
+      foo: 10
+      bar: abcde
+```
+
+#### Example app using custom sweeper:
+```text
+$ python example/my_app.py -m db=mysql,postgresql
+[2019-11-14 11:42:47,941][HYDRA] ExampleSweeper (foo=10, bar=abcde) sweeping
+[2019-11-14 11:42:47,941][HYDRA] Sweep output dir : multirun/2019-11-14/11-42-47
+[2019-11-14 11:42:47,942][HYDRA] Launching 2 jobs locally
+[2019-11-14 11:42:47,942][HYDRA]        #0 : db=mysql
+db:
+  driver: mysql
+  pass: secret
+  user: omry
+
+[2019-11-14 11:42:48,011][HYDRA]        #1 : db=postgresql
+db:
+  driver: postgresql
+  pass: drowssap
+  timeout: 10
+  user: postgre_user
+```
\ No newline at end of file
diff --git a/plugins/examples/example_sweeper/example/conf/config.yaml b/plugins/examples/example_sweeper/example/conf/config.yaml
new file mode 100644
index 0000000000..730e806ad4
--- /dev/null
+++ b/plugins/examples/example_sweeper/example/conf/config.yaml
@@ -0,0 +1,3 @@
+defaults:
+  - hydra/sweeper: example
+  - db: mysql
diff --git a/plugins/examples/example_sweeper/example/conf/db/mysql.yaml b/plugins/examples/example_sweeper/example/conf/db/mysql.yaml
new file mode 100644
index 0000000000..dc46278122
--- /dev/null
+++ b/plugins/examples/example_sweeper/example/conf/db/mysql.yaml
@@ -0,0 +1,4 @@
+db:
+  driver: mysql
+  user: omry
+  pass: secret
diff --git a/plugins/examples/example_sweeper/example/conf/db/postgresql.yaml b/plugins/examples/example_sweeper/example/conf/db/postgresql.yaml
new file mode 100644
index 0000000000..862073c5d8
--- /dev/null
+++ b/plugins/examples/example_sweeper/example/conf/db/postgresql.yaml
@@ -0,0 +1,5 @@
+db:
+  driver: postgresql
+  user: postgre_user
+  pass: drowssap
+  timeout: 10
diff --git a/plugins/examples/example_sweeper/example/my_app.py b/plugins/examples/example_sweeper/example/my_app.py
new file mode 100644
index 0000000000..5b48a4e304
--- /dev/null
+++ b/plugins/examples/example_sweeper/example/my_app.py
@@ -0,0 +1,11 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import hydra
+
+
+@hydra.main(config_path="conf/config.yaml")
+def my_app(cfg):
+    print(cfg.pretty())
+
+
+if __name__ == "__main__":
+    my_app()
diff --git a/plugins/examples/example_sweeper/hydra_plugins/__init__.py b/plugins/examples/example_sweeper/hydra_plugins/__init__.py
new file mode 100644
index 0000000000..dc1c6c1ca8
--- /dev/null
+++ b/plugins/examples/example_sweeper/hydra_plugins/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
diff --git a/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/__init__.py b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/__init__.py
new file mode 100644
index 0000000000..665ce07a66
--- /dev/null
+++ b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/__init__.py
@@ -0,0 +1,4 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+from .example_sweeper import ExampleSweeper
+
+__all__ = ["ExampleSweeper"]
diff --git a/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/conf/__init__.py b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/conf/__init__.py
new file mode 100644
index 0000000000..168f9979a4
--- /dev/null
+++ b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/conf/__init__.py
@@ -0,0 +1,1 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
diff --git a/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/conf/hydra/sweeper/example.yaml b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/conf/hydra/sweeper/example.yaml
new file mode 100644
index 0000000000..eb92898eeb
--- /dev/null
+++ b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/conf/hydra/sweeper/example.yaml
@@ -0,0 +1,6 @@
+hydra:
+  sweeper:
+    class: hydra_plugins.example_sweeper.ExampleSweeper
+    params:
+      foo: 10
+      bar: abcde
\ No newline at end of file
diff --git a/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/example_sweeper.py b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/example_sweeper.py
new file mode 100644
index 0000000000..e0b0d26bb7
--- /dev/null
+++ b/plugins/examples/example_sweeper/hydra_plugins/example_sweeper/example_sweeper.py
@@ -0,0 +1,56 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import itertools
+import logging
+
+from hydra._internal.config_search_path import ConfigSearchPath
+
+# TODO: Move Plugins class outside of hydra._internal
+from hydra._internal.plugins import Plugins
+from hydra.plugins import SearchPathPlugin, Sweeper
+
+log = logging.getLogger(__name__)
+
+
+class ExampleSweeperSearchPathPlugin(SearchPathPlugin):
+    """
+    This plugin is allowing configuration files provided by the ExampleSweeper plugin to be discovered
+    and used once the ExampleSweeper plugin is installed
+    """
+
+    def manipulate_search_path(self, search_path):
+        assert isinstance(search_path, ConfigSearchPath)
+        # Appends the search path for this plugin to the end of the search path
+        search_path.append(
+            "hydra-example-sweeper", "pkg://hydra_plugins.example_sweeper.conf"
+        )
+
+
+class ExampleSweeper(Sweeper):
+    def __init__(self, foo, bar):
+        self.config = None
+        self.launcher = None
+        self.job_results = None
+        self.foo = foo
+        self.bar = bar
+
+    def setup(self, config, config_loader, task_function):
+        self.config = config
+        self.launcher = Plugins.instantiate_launcher(
+            config=config, config_loader=config_loader, task_function=task_function
+        )
+
+    def sweep(self, arguments):
+        log.info("ExampleSweeper (foo={}, bar={}) sweeping".format(self.foo, self.bar))
+        log.info("Sweep output dir : {}".format(self.config.hydra.sweep.dir))
+        # Construct list of overrides per job we want to launch
+        src_lists = []
+        for s in arguments:
+            key, value = s.split("=")
+            # for each argument, create a list. if the argument has , (aka - is a sweep), add an element for each
+            # option to that list, otherwise add a single element with the value
+            src_lists.append(["{}={}".format(key, val) for val in value.split(",")])
+
+        batch = list(itertools.product(*src_lists))
+
+        returns = [self.launcher.launch(batch)]
+        return returns
diff --git a/plugins/examples/example_sweeper/setup.py b/plugins/examples/example_sweeper/setup.py
new file mode 100644
index 0000000000..c787f1c79e
--- /dev/null
+++ b/plugins/examples/example_sweeper/setup.py
@@ -0,0 +1,31 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+from setuptools import setup, find_packages
+
+with open("README.md", "r") as fh:
+    LONG_DESC = fh.read()
+    setup(
+        name="hydra-example-sweeper",
+        version="0.1.0",
+        author="Omry Yadan",
+        author_email="omry@fb.com",
+        description="Example Hydra Sweeper plugin",
+        long_description=LONG_DESC,
+        long_description_content_type="text/markdown",
+        url="https://github.com/facebookresearch/hydra/",
+        packages=find_packages(exclude=["tests", "example"]),
+        classifiers=[
+            # Feel free to choose another license
+            "License :: OSI Approved :: MIT License",
+            # Python versions are used by the noxfile in Hydra to determine which
+            # python versions to test this plugin with
+            "Programming Language :: Python :: 2.7",
+            "Programming Language :: Python :: 3.6",
+            "Programming Language :: Python :: 3.7",
+            "Operating System :: OS Independent",
+        ],
+        install_requires=["hydra-core"],
+        # If this plugin is providing configuration files, be sure to include them in the package.
+        # See MANIFEST.in.
+        # For configurations to be discoverable at runtime, they should also be added to the search path.
+        include_package_data=True,
+    )


Test patch:
diff --git a/plugins/examples/example_sweeper/tests/__init__.py b/plugins/examples/example_sweeper/tests/__init__.py
new file mode 100644
index 0000000000..168f9979a4
--- /dev/null
+++ b/plugins/examples/example_sweeper/tests/__init__.py
@@ -0,0 +1,1 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
diff --git a/plugins/examples/example_sweeper/tests/test_example_sweeper_plugin.py b/plugins/examples/example_sweeper/tests/test_example_sweeper_plugin.py
new file mode 100644
index 0000000000..5e474ba79a
--- /dev/null
+++ b/plugins/examples/example_sweeper/tests/test_example_sweeper_plugin.py
@@ -0,0 +1,34 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from hydra_plugins.example_sweeper import ExampleSweeper
+
+from hydra._internal.plugins import Plugins
+from hydra.plugins import Sweeper
+
+# noinspection PyUnresolvedReferences
+from hydra.test_utils.test_utils import sweep_runner  # noqa: F401
+
+
+def test_discovery():
+    """
+    Tests that this plugin can be discovered via the plugins subsystem when looking for Sweeper
+    :return:
+    """
+    assert ExampleSweeper.__name__ in [x.__name__ for x in Plugins.discover(Sweeper)]
+
+
+def test_launched_jobs(sweep_runner):  # noqa: F811
+    sweep = sweep_runner(
+        calling_file=None,
+        calling_module="hydra.test_utils.a_module",
+        config_path="configs/compose.yaml",
+        overrides=["hydra/sweeper=example", "hydra/launcher=basic", "foo=1,2"],
+        strict=True,
+    )
+    with sweep:
+        job_ret = sweep.returns[0]
+        assert len(job_ret) == 2
+        assert job_ret[0].overrides == ["foo=1"]
+        assert job_ret[0].cfg == {"foo": 1, "bar": 100}
+        assert job_ret[1].overrides == ["foo=2"]
+        assert job_ret[1].cfg == {"foo": 2, "bar": 100}