Possible Improvement to Rerun

[jgbos]

I had an idea on how to improve the new experimental rerun functionality that might help solve other problems (for me, the PyTorch Lightning issues).

What if we changed the signature of Callback to include the task function and pickle the config and task function together? E.g.,

class PickleJobInfoCallback(Callback):
    ...
    def on_job_start(self, task_fn: TaskFunction, config: DictConfig, **kwargs: Any) -> None:
       # pickle task_fn and config

Then have a main function in hydra that can execute from the pickled file:

$ python -m hydra.rerun <path to pickled file>

I think this is very similar to how submitit works (for reference).

[Jasha10]

Thanks for the suggestion @jgbos!

Regarding the on_job_start API:

For backward compatibility, perhaps we could make task_fn into a keyword arg:

class PickleJobInfoCallback(Callback):
    ...
    def on_job_start(self, config: DictConfig, *, task_fn: TaskFunction, **kwargs: Any) -> None:
       # pickle task_fn and config

That way, older user-defined callbacks with the signature on_job_start(self, config: DictConfig, **kwargs: Any) would still work, as the task_fn keyword would be slurped up by kwargs.

Regarding hydra.rerun:

I'm not intimately familiar with the interaction between hydra/pytorch-lightning/ddp.
Is the goal to run the task function multiple times from different python processes? Could the proposed hydra.rerun functionality be supplied by some other library besides hydra (e.g. by pytorch-lightning)? If so, we could get away with simply updating the signature of on_job_start and not have to worry about standardizing the pickle-rerun functionality in the Hydra library.

[jgbos]

Awesome, I think this will go a long ways towards reproducibility and supporting frameworks like pytorch-lightning that want to launch the task function multiple times. For pytorch-lightning they distinguish between processes by setting environment variables like the LOCAL_RANK.

The first thought I had about doing the pickling was "how much work does it take to standardize this process." I think it could be a powerful tool within Hydra but understand the headaches that come along with doing that. Certainly isn't an issue for me to build that functionality into each library.

[Jasha10]

I've run into a snag in implementing this. See my comment on #2361.