Repository: mit-ll-responsible-ai/hydra-zen
Pull request #313 (2022-10-05 19:53:15)

Allow dictionary for overrides in `launch`
Currently `launch` requires overrides to be a list of strings of type `param=value`. 

This PR allows the user to provide a dictionary as an override: `{param: val, ...}`. Where `param` is a string and `val` can be one of `int`, `float`, `str`, and `bool` along with with two special types:
  - `hydra_list`: A list type supported by Hydra, e.g., `hydra_list([1,2,3]) -> "[1,2,3]"`
  - `multirun`: A list that will be converted into a Hydra multirun value, e.g., `multirun([1,2,3]) -> "1,2,3"`

## Override Types

```python

override_primitive_examples = {
    "str": "\(1,2\)",   # --> "\(1,2\)" (no change)
    "NULL": None,   # --> "null"
    "int": 10,  #  --> "10"
    "float" 3.14,  # --> "3.14"
    "bool": True,    # --> "true"
    "list": hydra_list([1, 2, 3]),  # --> "[1, 2, 3]"
    "multirun_list": multirun([1, 2, 3])  # --> "1,2,3"
}
```

## Example of `multirun`

The use of the `multirun` type is what really makes this PR special.  With the current `overrides` structure a user must create the string that contains the list of values to multirun overview.  

```python
import random
from hydra_zen import launch, instantiate, make_config

values_for_experiment = [random.uniform(0, 1) for i in range(10)]
ov = []

jobs = launch(
    make_config(),
    instantiate,
    overrides=[
         "+param=" + ",".join([str(i) for i in values_for_experiment])
    ],
    multirun=True
)
```

This PR allows the user to provide the multirun value without having to create the string:

```python
import random
from hydra_zen import launch, instantiate, make_config, multirun

values_for_experiment = [random.uniform(0, 1) for i in range(10)]

jobs = launch(
    make_config(),
    instantiate,
    overrides={
         "+param": multirun(values_for_experiment)
    },
    multirun=True
)
```




TODO:
  - [ ] Update this PR description
  - [ ] Add `null` support in supported types
  - [ ] Test configs are the same for list of strings and dictionary
  - [ ] Update docstrings
  - [ ] Update docs


Patch:
diff --git a/docs/source/explanation/type_refinement.rst b/docs/source/explanation/type_refinement.rst
index 22a67bae..ef0b9668 100644
--- a/docs/source/explanation/type_refinement.rst
+++ b/docs/source/explanation/type_refinement.rst
@@ -1,8 +1,8 @@
 .. _type-support:
 
-##############################################
-hydra-zen's Automatic Type Refinement Behavior
-##############################################
+############################################################
+Understanding hydra-zen's Automatic Type Refinement Behavior
+############################################################
 
 hydra-zen's :ref:`config-creation functions <create-config>` will automatically 
 refine user-supplied type annotations so that they are made compatible with Hydra's 
diff --git a/docs/source/how_to/pytorch_lightning.rst b/docs/source/how_to/pytorch_lightning.rst
index f5ec7bab..fe3d16aa 100644
--- a/docs/source/how_to/pytorch_lightning.rst
+++ b/docs/source/how_to/pytorch_lightning.rst
@@ -222,7 +222,7 @@ and run the following code.
 
    >>> from hydra_zen import launch
    >>> from experiment import ExperimentConfig, task_function
-   >>> (jobs,) = launch(  # type: ignore
+   >>> (jobs,) = launch(
    ...     ExperimentConfig,
    ...     task_function,
    ...     overrides=[
diff --git a/src/hydra_zen/__init__.py b/src/hydra_zen/__init__.py
index 67e1099d..b3d6a747 100644
--- a/src/hydra_zen/__init__.py
+++ b/src/hydra_zen/__init__.py
@@ -9,7 +9,7 @@ from ._hydra_overloads import (
     save_as_yaml,
     to_yaml,
 )
-from ._launch import launch
+from ._launch import hydra_list, launch, multirun
 from .structured_configs import (
     ZenField,
     builds,
@@ -41,6 +41,8 @@ __all__ = [
     "is_partial_builds",
     "uses_zen_processing",
     "zen",
+    "hydra_list",
+    "multirun",
     "store",
     "ZenStore",
 ]
diff --git a/src/hydra_zen/_launch.py b/src/hydra_zen/_launch.py
index b4f31c34..56ed8f51 100644
--- a/src/hydra_zen/_launch.py
+++ b/src/hydra_zen/_launch.py
@@ -1,8 +1,24 @@
 # Copyright (c) 2023 Massachusetts Institute of Technology
 # SPDX-License-Identifier: MIT
 import warnings
+from collections import UserList
 from dataclasses import fields, is_dataclass
-from typing import Any, Callable, List, Mapping, Optional, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
 from hydra import initialize
 from hydra._internal.callbacks import Callbacks
@@ -12,16 +28,106 @@ from hydra.core.utils import JobReturn, run_job
 from hydra.plugins.sweeper import Sweeper
 from hydra.types import HydraContext, RunMode
 from omegaconf import DictConfig, ListConfig, OmegaConf
+from typing_extensions import Literal, TypeAlias
 
 from hydra_zen._compatibility import SUPPORTS_VERSION_BASE
 from hydra_zen._hydra_overloads import instantiate
 from hydra_zen.typing._implementations import DataClass, InstOrType
 
+T = TypeVar("T", bound=Any)
+HydraPrimitives: TypeAlias = Union[None, int, float, bool, str, Dict[str, str]]
+
+if TYPE_CHECKING:  # pragma: no cover
+    # branching needed to deal with pyright type-completeness complaints
+    TUserList: TypeAlias = UserList[Any]
+else:
+    TUserList = UserList
+
 
 class _NotSet:  # pragma: no cover
     pass
 
 
+T1 = TypeVar("T1", bound=HydraPrimitives)
+
+
+class hydra_list(TUserList, Generic[T1]):
+    """Signals that a sequence is provided as a single configured value (i.e. it is not
+    to be iterated over during a multirun)"""
+
+    pass
+
+
+T2 = TypeVar("T2", bound=Union[HydraPrimitives, hydra_list[HydraPrimitives]])
+
+
+class multirun(TUserList, Generic[T2]):
+    """Signals that a sequence is to be iterated over in a multirun"""
+
+    pass
+
+
+def _safe_name(x: Any) -> str:
+    return getattr(x, "__name__", str(x))
+
+
+def value_check(
+    name: str,
+    value: T,
+    type_: Union[type, Tuple[type, ...]],
+) -> T:
+    """
+    For internal use only.
+
+    Used to check the type of `value`. Numerical types can also be bound-checked.
+
+    Examples
+    --------
+    >>> value_check("x", 1, type_=str)
+    TypeError: `x` must be of type(s) `str`, got 1 (type: int)
+
+    Raises
+    ------
+    TypeError"""
+    # check internal params
+    assert isinstance(name, str), name
+
+    if not isinstance(value, type_):
+        raise TypeError(
+            f"`{name}` must be of type(s) "
+            f"`{_safe_name(type_)}`, got {value} (type: {_safe_name(type(value))})"
+        )
+
+    return cast(T, value)
+
+
+OverrideValues: TypeAlias = Union[
+    HydraPrimitives,
+    multirun[Union[HydraPrimitives, hydra_list[HydraPrimitives]]],
+    hydra_list[HydraPrimitives],
+]
+OverrideDict: TypeAlias = Mapping[str, OverrideValues]
+
+
+def _process_dict_overrides(overrides: OverrideDict) -> List[str]:
+    """Convert dict overrides to a list of Hydra CLI compatible args"""
+    launch_overrides = []
+    for k, v in overrides.items():
+        if v is None:
+            v = "null"
+
+        value_check(
+            k,
+            v,
+            type_=(int, float, bool, str, dict, multirun, hydra_list),
+        )
+        if isinstance(v, multirun):
+            v = ",".join(str(item) for item in v)
+
+        launch_overrides.append(f"{k}={v}")
+    return launch_overrides
+
+
 def _store_config(
     cfg: Union[DataClass, Type[DataClass], DictConfig, ListConfig, Mapping[Any, Any]],
     config_name: str = "hydra_launch",
@@ -55,18 +161,52 @@ def _store_config(
     return config_name
 
 
+@overload
+def launch(
+    config: Union[InstOrType[DataClass], Mapping[str, Any]],
+    task_function: Callable[[Any], Any],
+    overrides: Optional[Union[OverrideDict, List[str]]] = ...,
+    multirun: Literal[False] = ...,
+    version_base: Optional[Union[str, Type[_NotSet]]] = ...,
+    to_dictconfig: bool = ...,
+    config_name: str = ...,
+    job_name: str = ...,
+    with_log_configuration: bool = ...,
+    **override_kwargs: OverrideValues,
+) -> JobReturn:
+    ...
+
+
+@overload
+def launch(
+    config: Union[InstOrType[DataClass], Mapping[str, Any]],
+    task_function: Callable[[Any], Any],
+    overrides: Optional[Union[OverrideDict, List[str]]] = ...,
+    multirun: Literal[True] = ...,
+    version_base: Optional[Union[str, Type[_NotSet]]] = ...,
+    to_dictconfig: bool = ...,
+    config_name: str = ...,
+    job_name: str = ...,
+    with_log_configuration: bool = ...,
+    **override_kwargs: OverrideValues,
+) -> Any:
+    ...
+
+
 def launch(
     config: Union[InstOrType[DataClass], Mapping[str, Any]],
     task_function: Callable[[Any], Any],
-    overrides: Optional[List[str]] = None,
+    overrides: Optional[Union[OverrideDict, List[str]]] = None,
     multirun: bool = False,
     version_base: Optional[Union[str, Type[_NotSet]]] = _NotSet,
     to_dictconfig: bool = False,
     config_name: str = "zen_launch",
     job_name: str = "zen_launch",
     with_log_configuration: bool = True,
+    **override_kwargs: OverrideValues,
 ) -> Union[JobReturn, Any]:
-    r"""Launch a Hydra job using a Python-based interface.
+    r"""
+    Launches a Hydra job from a Python function rather than a CLI.
 
     `launch` is designed to closely match the interface of the standard Hydra CLI.
     For example, launching a Hydra job from the CLI via::
@@ -86,22 +226,23 @@ def launch(
         The function that Hydra will execute. Its input will be ``config``, which
         has been modified via the specified ``overrides``
 
-    overrides : Optional[List[str]]
+    overrides : Optional[Union[OverrideMapping, List[str]]] (default: None)
         If provided, sets/overrides values in ``config``. See [1]_ and [2]_
         for a detailed discussion of the "grammar" supported by ``overrides``.
 
     multirun : bool (default: False)
         Launch a Hydra multi-run ([3]_).
 
-    version_base : Optional[str], optional (default=_NotSet)
+    version_base : Optional[str], optional (default=not-specified)
         Available starting with Hydra 1.2.0.
         - If the `version_base parameter` is not specified, Hydra 1.x will use defaults compatible with version 1.1. Also in this case, a warning is issued to indicate an explicit version_base is preferred.
         - If the `version_base parameter` is `None`, then the defaults are chosen for the current minor Hydra version. For example for Hydra 1.2, then would imply `config_path=None` and `hydra.job.chdir=False`.
         - If the `version_base` parameter is an explicit version string like "1.1", then the defaults appropriate to that version are used.
 
-    to_dictconfig: bool (default: False)
-        If ``True``, convert a ``dataclasses.dataclass`` to a ``omegaconf.DictConfig``. Note, this
-        will remove Hydra's cabability for validation with structured configurations.
+    to_dictconfig : bool (default: False)
+        If ``True``, convert a ``dataclasses.dataclass`` to a ``omegaconf.DictConfig``.
+        Note, this will remove Hydra's cabability for validation with structured
+        configurations.
 
     config_name : str (default: "zen_launch")
         Name of the stored configuration in Hydra's ConfigStore API.
@@ -109,13 +250,20 @@ def launch(
     job_name : str (default: "zen_launch")
 
     with_log_configuration : bool (default: True)
-        If ``True``, enables the configuration of the logging subsystem from the loaded config.
+        If ``True``, enables the configuration of the logging subsystem from the loaded
+        config.
+
+    **override_kwargs : OverrideValues
+        Keyword arguments to override existing configuration values.  Note, this only
+        works when the configuration value name is a valid Python identifier; e.g.,
+        this does not support adding (`+param`) values.
 
     Returns
     -------
-    result : JobReturn | Any
+    result : hydra.core.utils.JobReturn | Any
         If ``multirun is False``:
-            A ``JobReturn`` object storing the results of the Hydra experiment via the following attributes
+            A ``JobReturn`` object storing the results of the Hydra experiment via the
+            following attributes
                 - ``cfg``: Reflects ``config``
                 - ``overrides``: Reflects ``overrides``
                 - ``return_value``: The return value of the task function
@@ -159,7 +307,7 @@ def launch(
     Now, let's use `launch` to run this task function via Hydra, using particular configured
     values (or, "overrides") for ``a`` and ``b``.
 
-    >>> job_out = launch(Conf, task_fn, overrides=["a=1", "b='foo'"])
+    >>> job_out = launch(Conf, task_fn, a=1, b='foo')
     a: 1
     b: foo
 
@@ -188,7 +336,8 @@ def launch(
     >>> (outputs,) = launch(
     ...     Conf,
     ...     task_fn,
-    ...     overrides=["a=1,2,3", "b='bar'"],
+    ...     a="1,2,3",
+    ...     b="bar",
     ...     multirun=True,
     ... )
     [2021-10-19 17:50:07,334][HYDRA] Launching 3 jobs locally
@@ -215,6 +364,22 @@ def launch(
     ['multirun/2021-10-19/17-50-07\\0',
     'multirun/2021-10-19/17-50-07\\1',
     'multirun/2021-10-19/17-50-07\\2']
+
+    **Launching with quoted overrides**
+
+    Some of the Hydra CLI override syntax cannot be specified as keyword arguments. In such cases we can instead provide a list or a dict with quoted overrides.
+
+    >>> job_out = launch(Conf, task_fn, a=1, b="foo", overrides={"+c": 22})
+    a: 1
+    b: foo
+    c: 22
+    >>> job_out.overrides  # the overrides that we provides
+    ['a=1', 'b=foo', '+c=22']
+
+    >>> launch(Conf, task_fn, overrides=["a=1", "b='foo'", "+c=22"])
+    a: 1
+    b: foo
+    c: 22
     """
 
     # used for check below
@@ -232,6 +397,15 @@ def launch(
     else:
         config_name = _store_config(config, config_name)
 
+    # allow user to provide a dictionary of override values
+    # instead of just a list of strings
+    overrides = overrides if overrides is not None else []
+    if isinstance(overrides, Mapping):
+        overrides = _process_dict_overrides(overrides)
+
+    override_kwargs_list = _process_dict_overrides(override_kwargs)
+    overrides += override_kwargs_list
+
     # Initializes Hydra and add the config_path to the config search path
     with initialize(
         config_path=None,
@@ -240,7 +414,7 @@ def launch(
             {}
             if (not SUPPORTS_VERSION_BASE or version_base is _NotSet)
             else {"version_base": version_base}
-        )
+        ),
     ):
 
         # taken from hydra.compose with support for MULTIRUN
@@ -250,7 +424,7 @@ def launch(
         # Load configuration
         cfg = gh.hydra.compose_config(
             config_name=config_name,
-            overrides=overrides if overrides is not None else [],
+            overrides=overrides,
             run_mode=RunMode.RUN if not multirun else RunMode.MULTIRUN,
             from_shell=False,
             with_log_configuration=with_log_configuration,
@@ -303,11 +477,13 @@ def launch(
             and _num_dataclass_fields_after < _num_dataclass_fields
         ):
             warnings.warn(
-                "Your dataclass-based config was mutated by this run. If you just executed with a "
-                "`hydra/launcher` that utilizes cloudpickle (e.g., hydra-submitit-launcher), there is a known "
-                "issue with dataclasses (see: https://github.com/cloudpipe/cloudpickle/issues/386). You will have "
-                "to restart your interactive environment ro run `launch` again. To avoid this issue you can use the "
-                "`launch` option: `to_dictconfig=True`."
+                "Your dataclass-based config was mutated by this run. If you just "
+                "executed with a `hydra/launcher` that utilizes cloudpickle (e.g., "
+                "hydra-submitit-launcher), there is a known issue with dataclasses "
+                "(see: https://github.com/cloudpipe/cloudpickle/issues/386). You will "
+                "have to restart your interactive environment ro run `launch` again. "
+                "To avoid this issue you can use the `launch` option: "
+                "`to_dictconfig=True`."
             )
 
     return job


Test patch:
diff --git a/tests/annotations/declarations.py b/tests/annotations/declarations.py
index c98bf164..931c2f64 100644
--- a/tests/annotations/declarations.py
+++ b/tests/annotations/declarations.py
@@ -30,6 +30,7 @@ from typing import (
     Union,
 )
 
+from hydra.core.utils import JobReturn
 from omegaconf import MISSING, DictConfig, ListConfig
 from typing_extensions import Literal, assert_type
 
@@ -1054,7 +1055,15 @@ def check_launch():
         pass
 
     f(Xonf)
-    launch(Xonf, f)
+
+    job1 = launch(Xonf(), f)
+    assert_type(job1, JobReturn)
+
+    job2 = launch(Xonf, f, multirun=False)
+    assert_type(job2, JobReturn)
+
+    job3 = launch(Xonf, f, multirun=True)
+    assert_type(job3, Any)
 
 
 def check_instantiate():
diff --git a/tests/test_launch/test_dict_overrides.py b/tests/test_launch/test_dict_overrides.py
new file mode 100644
index 00000000..bc222411
--- /dev/null
+++ b/tests/test_launch/test_dict_overrides.py
@@ -0,0 +1,202 @@
+# Copyright (c) 2023 Massachusetts Institute of Technology
+# SPDX-License-Identifier: MIT
+
+import string
+from pathlib import Path
+
+import hypothesis.strategies as st
+import pytest
+from hypothesis import given, settings
+from omegaconf import ListConfig
+
+from hydra_zen import (
+    hydra_list,
+    instantiate,
+    launch,
+    load_from_yaml,
+    make_config,
+    multirun,
+)
+from hydra_zen._launch import value_check
+
+any_types = st.from_type(type)
+
+
+def everything_except(excluded_types):
+    return (
+        st.from_type(type)
+        .flatmap(st.from_type)
+        .filter(lambda x: not isinstance(x, excluded_types))
+    )
+
+
+@settings(max_examples=10)
+@given(
+    name=st.sampled_from(["name_a", "name_b"]),
+    target_type=st.shared(any_types, key="target_type"),
+    value=st.shared(any_types, key="target_type").flatmap(everything_except),
+)
+def test_type_catches_bad_type(name, target_type, value):
+    with pytest.raises(TypeError, match=rf"`{name}` must be of type\(s\) .*"):
+        value_check(name, value=value, type_=target_type)
+
+
+@given(
+    target_type=st.shared(any_types, key="target_type"),
+    value=st.shared(any_types, key="target_type").flatmap(st.from_type),
+)
+def test_type_passes_valid_type(target_type, value):
+    value_check("dummy", value=value, type_=target_type)
+
+
+@pytest.mark.usefixtures("cleandir")
+@settings(max_examples=10, deadline=None)
+@given(
+    none_=st.none(),
+    int_=st.integers(),
+    bool_=st.booleans(),
+    float_=st.floats(-10, 10),
+    list_=st.lists(st.integers()),
+    str_=st.text(alphabet=string.ascii_lowercase).filter(
+        lambda x: x != "true" and x != "false"
+    ),
+    mrun=st.lists(
+        st.booleans() | st.lists(st.integers()),
+        min_size=2,
+        max_size=5,
+    ),
+)
+def test_overrides_roundtrip(
+    none_,
+    int_,
+    bool_,
+    float_,
+    str_,
+    list_,
+    mrun,
+):
+    overrides = {
+        "+none_": none_,
+        "+int_": int_,
+        "+float_": float_,
+        "+str_": str_,
+        "+bool_": bool_,
+        "+list_": hydra_list(list_),
+        "+mrun": multirun(mrun),
+    }
+    (jobs,) = launch(make_config(), instantiate, overrides, multirun=True)
+
+    assert len(jobs) == len(mrun)
+    for i, job in enumerate(jobs):
+        assert job.return_value.none_ == none_
+        assert job.return_value.int_ == int_
+        assert job.return_value.bool_ == bool_
+        assert job.return_value.float_ == float_
+        assert job.return_value.str_ == str_
+        assert job.return_value.list_ == list_
+        assert job.return_value.mrun == mrun[i]
+
+
+@pytest.mark.usefixtures("cleandir")
+@settings(max_examples=10, deadline=None)
+@given(
+    none_=st.none(),
+    int_=st.integers(),
+    bool_=st.booleans(),
+    float_=st.floats(-10, 10),
+    list_=st.lists(st.integers()),
+    str_=st.text(alphabet=string.ascii_lowercase).filter(
+        lambda x: x != "true" and x != "false"
+    ),
+    mrun=st.lists(
+        st.booleans() | st.lists(st.integers()),
+        min_size=2,
+        max_size=5,
+    ),
+)
+def test_overrides_kwargs_roundtrip(
+    none_,
+    int_,
+    bool_,
+    float_,
+    str_,
+    list_,
+    mrun,
+):
+    overrides = {
+        "+none_": none_,
+        "+int_": int_,
+        "+float_": float_,
+        "+str_": str_,
+        "+bool_": bool_,
+        "+list_": hydra_list(list_),
+        "+mrun": multirun(mrun),
+    }
+    jobs = launch(make_config(), instantiate, ["+foo=1"], multirun=True, **overrides)
+    jobs1 = launch(make_config(), instantiate, {"+foo": 2}, multirun=True, **overrides)
+
+    assert len(jobs[0]) == len(mrun)
+    assert len(jobs1[0]) == len(mrun)
+    for i, (job, job1) in enumerate(zip(jobs[0], jobs1[0])):
+        assert job.return_value.foo == 1 and job1.return_value.foo == 2
+        assert job.return_value.none_ == job1.return_value.none_ == none_
+        assert job.return_value.int_ == job1.return_value.int_ == int_
+        assert job.return_value.bool_ == job1.return_value.bool_ == bool_
+        assert job.return_value.float_ == job1.return_value.float_ == float_
+        assert job.return_value.str_ == job1.return_value.str_ == str_
+        assert job.return_value.list_ == job1.return_value.list_ == list_
+        assert job.return_value.mrun == job1.return_value.mrun == mrun[i]
+
+
+@pytest.mark.usefixtures("cleandir")
+@settings(max_examples=10, deadline=None)
+@given(
+    int_=st.integers(),
+    bool_=st.booleans(),
+    float_=st.floats(-10, 10),
+    list_=st.lists(st.integers()),
+    str_=st.text(alphabet=string.ascii_lowercase).filter(
+        lambda x: x != "true" and x != "false"
+    ),
+    mrun=st.lists(
+        st.booleans() | st.lists(st.integers()),
+        min_size=2,
+        max_size=5,
+    ),
+)
+def test_overrides_yaml(
+    int_,
+    bool_,
+    float_,
+    str_,
+    list_,
+    mrun,
+):
+    overrides = {
+        "+int_": int_,
+        "+float_": float_,
+        "+str_": str_,
+        "+bool_": bool_,
+        "+list_": hydra_list(list_),
+        "+mrun": multirun(mrun),
+    }
+    (jobs,) = launch(make_config(), instantiate, overrides, multirun=True)
+
+    assert len(jobs) == len(mrun)
+    for i, job in enumerate(jobs):
+        overrides_yamls = list(Path(job.working_dir).glob("**/overrides.yaml"))
+        assert len(overrides_yamls) == 1
+        overrides_yaml = overrides_yamls[0]
+
+        yaml = load_from_yaml(overrides_yaml)
+        assert isinstance(yaml, ListConfig)
+        assert len(yaml) == 6
+        for param_val in yaml:
+            param, val = param_val.split("=")
+            assert param in list(overrides.keys())
+
+            override_val = overrides[param]
+            if isinstance(override_val, multirun):
+                assert val == str(mrun[i]).strip().replace(" ", "")
+            else:
+                assert val == str(override_val).strip().replace(" ", "")