add --experimental-rerun CLI option (#2098)

examples/experimental/rerun/config.yaml

@@ -0,0 +1,8 @@
+foo: bar
+
+hydra:
+  callbacks:
+    save_job_info:
+      _target_: hydra.experimental.callbacks.PickleJobInfoCallback
+  job:
+    chdir: false

examples/experimental/rerun/my_app.py

@@ -0,0 +1,20 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import logging
+
+from omegaconf import DictConfig
+
+import hydra
+from hydra.core.hydra_config import HydraConfig
+
+log = logging.getLogger(__name__)
+
+
+@hydra.main(version_base=None, config_path=".", config_name="config")
+def my_app(cfg: DictConfig) -> None:
+    log.info(f"Output_dir={HydraConfig.get().runtime.output_dir}")
+    log.info(f"cfg.foo={cfg.foo}")
+
+
+if __name__ == "__main__":
+    my_app()

hydra/_internal/utils.py

@@ -298,6 +298,7 @@ class FakeTracebackType:
 
 
 def _run_hydra(
+    args: argparse.Namespace,
     args_parser: argparse.ArgumentParser,
     task_function: TaskFunction,
     config_path: Optional[str],
@@ -309,7 +310,6 @@ def _run_hydra(
 
     from .hydra import Hydra
 
-    args = args_parser.parse_args()
     if args.config_name is not None:
         config_name = args.config_name
 
@@ -565,6 +565,11 @@ def __repr__(self) -> str:
         help="Adds an additional config dir to the config search path",
     )
 
+    parser.add_argument(
+        "--experimental-rerun",
+        help="Rerun a job from a previous config pickle",
+    )
+
     info_choices = [
         "all",
         "config",

hydra/main.py

@@ -1,18 +1,47 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import copy
 import functools
+import pickle
+import warnings
+from pathlib import Path
 from textwrap import dedent
-from typing import Any, Callable, Optional
+from typing import Any, Callable, List, Optional
 
-from omegaconf import DictConfig
+from omegaconf import DictConfig, open_dict, read_write
 
 from . import version
 from ._internal.deprecation_warning import deprecation_warning
 from ._internal.utils import _run_hydra, get_args_parser
+from .core.hydra_config import HydraConfig
+from .core.utils import _flush_loggers, configure_log
 from .types import TaskFunction
 
 _UNSPECIFIED_: Any = object()
 
 
+def _get_rerun_conf(file_path: str, overrides: List[str]) -> DictConfig:
+    msg = "Experimental rerun CLI option, other command line args are ignored."
+    warnings.warn(msg, UserWarning)
+    file = Path(file_path)
+    if not file.exists():
+        raise ValueError(f"File {file} does not exist!")
+
+    if len(overrides) > 0:
+        msg = "Config overrides are not supported as of now."
+        warnings.warn(msg, UserWarning)
+
+    with open(str(file), "rb") as input:
+        config = pickle.load(input)  # nosec
+    configure_log(config.hydra.job_logging, config.hydra.verbose)
+    HydraConfig.instance().set_config(config)
+    task_cfg = copy.deepcopy(config)
+    with read_write(task_cfg):
+        with open_dict(task_cfg):
+            del task_cfg["hydra"]
+    assert isinstance(task_cfg, DictConfig)
+    return task_cfg
+
+
 def main(
     config_path: Optional[str] = _UNSPECIFIED_,
     config_name: Optional[str] = None,
@@ -49,15 +78,22 @@ def decorated_main(cfg_passthrough: Optional[DictConfig] = None) -> Any:
             if cfg_passthrough is not None:
                 return task_function(cfg_passthrough)
             else:
-                args = get_args_parser()
-                # no return value from run_hydra() as it may sometime actually run the task_function
-                # multiple times (--multirun)
-                _run_hydra(
-                    args_parser=args,
-                    task_function=task_function,
-                    config_path=config_path,
-                    config_name=config_name,
-                )
+                args_parser = get_args_parser()
+                args = args_parser.parse_args()
+                if args.experimental_rerun is not None:
+                    cfg = _get_rerun_conf(args.experimental_rerun, args.overrides)
+                    task_function(cfg)
+                    _flush_loggers()
+                else:
+                    # no return value from run_hydra() as it may sometime actually run the task_function
+                    # multiple times (--multirun)
+                    _run_hydra(
+                        args=args,
+                        args_parser=args_parser,
+                        task_function=task_function,
+                        config_path=config_path,
+                        config_name=config_name,
+                    )
 
         return decorated_main
 

news/1805.feature

@@ -0,0 +1 @@
+Add `--experimental-rerun` command-line option to reproduce pickled single runs

tests/test_apps/app_with_pickle_job_info_callback/my_app.py

@@ -23,6 +23,7 @@ def pickle_cfg(path: Path, obj: Any) -> Any:
     output_dir = Path(hydra_cfg.runtime.output_dir)
     pickle_cfg(Path(output_dir) / "task_cfg.pickle", cfg)
     pickle_cfg(Path(output_dir) / "hydra_cfg.pickle", hydra_cfg)
+    log.info("Running my_app")
 
     return "hello world"
 

tests/test_callbacks.py

@@ -1,5 +1,6 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 import copy
+import os
 import pickle
 import sys
 from pathlib import Path
@@ -174,3 +175,52 @@ def test_save_job_return_callback(tmpdir: Path, multirun: bool) -> None:
         with open(p, "r") as file:
             logs = file.readlines()
             assert log_msg in logs
+
+
+@mark.parametrize(
+    "warning_msg,overrides",
+    [
+        ("Experimental rerun CLI option", []),
+        ("Config overrides are not supported as of now", ["+x=1"]),
+    ],
+)
+def test_experimental_rerun(
+    tmpdir: Path, warning_msg: str, overrides: List[str]
+) -> None:
+    app_path = "tests/test_apps/app_with_pickle_job_info_callback/my_app.py"
+
+    cmd = [
+        app_path,
+        "hydra.run.dir=" + str(tmpdir),
+        "hydra.sweep.dir=" + str(tmpdir),
+        "hydra.job.chdir=False",
+        "hydra.hydra_logging.formatters.simple.format='[HYDRA] %(message)s'",
+        "hydra.job_logging.formatters.simple.format='[JOB] %(message)s'",
+    ]
+    run_python_script(cmd)
+
+    config_file = tmpdir / ".hydra" / "config.pickle"
+    log_file = tmpdir / "my_app.log"
+    assert config_file.exists()
+    assert log_file.exists()
+
+    with open(log_file, "r") as file:
+        logs = file.read().splitlines()
+        assert "[JOB] Running my_app" in logs
+
+    os.remove(str(log_file))
+    assert not log_file.exists()
+
+    # then rerun the application and verify log file is created again
+    cmd = [
+        app_path,
+        "--experimental-rerun",
+        str(config_file),
+    ]
+    cmd.extend(overrides)
+    result, err = run_python_script(cmd, allow_warnings=True)
+    assert warning_msg in err
+
+    with open(log_file, "r") as file:
+        logs = file.read().splitlines()
+        assert "[JOB] Running my_app" in logs

tests/test_examples/test_experimental.py

@@ -0,0 +1,18 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from pathlib import Path
+
+from hydra.test_utils.test_utils import run_python_script
+
+
+def test_rerun(tmpdir: Path) -> None:
+    cmd = [
+        "examples/experimental/rerun/my_app.py",
+        "hydra.run.dir=" + str(tmpdir),
+        "hydra.job.chdir=True",
+        "hydra.hydra_logging.formatters.simple.format='[HYDRA] %(message)s'",
+        "hydra.job_logging.formatters.simple.format='[JOB] %(message)s'",
+    ]
+
+    result, _err = run_python_script(cmd)
+    assert "[JOB] cfg.foo=bar" in result

tests/test_hydra.py

@@ -738,6 +738,7 @@ def test_sweep_complex_defaults(
                                     The config_path is relative to the Python file declaring @hydra.main()
                 --config-name,-cn : Overrides the config_name specified in hydra.main()
                 --config-dir,-cd : Adds an additional config dir to the config search path
+                --experimental-rerun : Rerun a job from a previous config pickle
                 --info,-i : Print Hydra information [all|config|defaults|defaults-tree|plugins|searchpath]
                 Overrides : Any key=value arguments to override config values (use dots for.nested=overrides)
                 """
@@ -795,6 +796,7 @@ def test_sweep_complex_defaults(
                                     The config_path is relative to the Python file declaring @hydra.main()
                 --config-name,-cn : Overrides the config_name specified in hydra.main()
                 --config-dir,-cd : Adds an additional config dir to the config search path
+                --experimental-rerun : Rerun a job from a previous config pickle
                 --info,-i : Print Hydra information [all|config|defaults|defaults-tree|plugins|searchpath]
                 Overrides : Any key=value arguments to override config values (use dots for.nested=overrides)
                 """

website/docs/experimental/rerun.md

@@ -0,0 +1,91 @@
+---
+id: rerun
+title: Re-run a job from previous config
+sidebar_label: Re-run 
+---
+
+import {ExampleGithubLink} from "@site/src/components/GithubLink"
+
+<ExampleGithubLink text="Example application" to="examples/experimental/rerun"/>
+
+:::caution
+This is an experimental feature. Please read through this page to understand what is supported.
+:::
+
+We use the example app linked above for demonstration. To save the configs for re-run, first use the experimental
+Hydra Callback for saving the job info:
+
+
+```yaml title="config.yaml"
+hydra:
+  callbacks:
+    save_job_info:
+      _target_: hydra.experimental.pickle_job_info_callback.PickleJobInfoCallback
+```
+
+
+
+
+```python title="Example function"
+@hydra.main(config_path=".", config_name="config")
+def my_app(cfg: DictConfig) -> None:
+    log.info(f"output_dir={HydraConfig.get().runtime.output_dir}")
+    log.info(f"cfg.foo={cfg.foo}")
+```
+
+
+Run the example app:
+```commandline
+$ python my_app.py
+[2022-03-16 14:51:30,905][hydra.experimental.pickle_job_info_callback][INFO] - Saving job configs in /Users/jieru/workspace/hydra/examples/experimental/outputs/2022-03-16/14-51-30/.hydra/config.pickle
+[2022-03-16 14:51:30,906][__main__][INFO] - Output_dir=/Users/jieru/workspace/hydra/examples/experimental/outputs/2022-03-16/14-51-30
+[2022-03-16 14:51:30,906][__main__][INFO] - cfg.foo=bar
+[2022-03-16 14:51:30,906][hydra.experimental.pickle_job_info_callback][INFO] - Saving job_return in /Users/jieru/workspace/hydra/examples/experimental/outputs/2022-03-16/14-51-30/.hydra/job_return.pickle
+```
+The Callback saves `config.pickle` in `.hydra` sub dir, this is what we will use for rerun.
+
+Now rerun the app
+```commandline
+$ OUTPUT_DIR=/Users/jieru/workspace/hydra/examples/experimental/outputs/2022-03-16/14-51-30/.hydra/
+$ python my_app.py --experimental-rerun $OUTPUT_DIR/config.pickle
+/Users/jieru/workspace/hydra/hydra/main.py:23: UserWarning: Experimental rerun CLI option.
+  warnings.warn(msg, UserWarning)
+[2022-03-16 14:59:21,666][__main__][INFO] - Output_dir=/Users/jieru/workspace/hydra/examples/experimental/outputs/2022-03-16/14-51-30
+[2022-03-16 14:59:21,666][__main__][INFO] - cfg.foo=bar
+```
+You will notice `my_app.log` is updated with the logging from the second run, but Callbacks are not called this time. Read on to learn more.
+
+
+### Important Notes
+This is an experimental feature. Please reach out if you have any question. 
+- Only single run is supported.
+- `--experimental-rerun` cannot be used with other command-line options or overrides. They will simply be ignored.
+- Rerun passes in a cfg_passthrough directly to your application, this means except for logging, no other `hydra.main` 
+functions are called (such as change working dir, or calling callbacks.) 
+- The configs are preserved and reconstructed to the best efforts. Meaning we can only guarantee that the `cfg` object 
+itself passed in by `hydra.main` stays the same across runs. However, configs are resolved lazily. Meaning we cannot 
+guarantee your application will behave the same if your application resolves configs during run time. In the following example,
+`cfg.time_now` will resolve to different value every run.
+
+<div className="row">
+<div className="col  col--5">
+
+```yaml title="config.yaml"
+time_now: ${now:%H-%M-%S}
+
+
+
+```
+
+</div>
+
+<div className="col col--7">
+
+```python title="Example function"
+@hydra.main(config_path=".", config_name="config")
+def my_app(cfg: DictConfig) -> None:
+    val = cfg.time_now
+    # the rest of the application
+```
+</div>
+</div>

website/sidebars.js

@@ -133,6 +133,7 @@ module.exports = {
         "Experimental": [
             "experimental/intro",
             "experimental/callbacks",
+            "experimental/rerun",
         ],
 
         'Developer Guide': [