omegaconf resolution when using multirun

[ywyga]

Hi,

When using multirun, I notice that $now is resolved only once (when multirun is started) but custom resolvers are called for each job even if they are declared with use_cache=True. So I am wondering what is the general rule as to when config resolutions happen in multirun. And also whether there is a way to make $now resolve at job creation time instead of at multirun start.

Here for example:
main.py

from time import sleep

import hydra
from omegaconf import OmegaConf
import random

OmegaConf.register_new_resolver("uncached_randint", random.randint)
OmegaConf.register_new_resolver("cached_randint", random.randint, use_cache=True)

@hydra.main(config_path=".", config_name="config.yaml")
def main(cfg) -> None:
    print(OmegaConf.to_yaml(cfg, resolve=True))
    print('uncached_seed: ', cfg.uncached_seed, cfg.uncached_seed, cfg.uncached_seed)
    print('cached_seed: ', cfg.cached_seed, cfg.cached_seed, cfg.cached_seed)
    sleep(2)


if __name__ == "__main__":
    main()

config.yaml

hydra:
  sweep:
    dir: ${now:%Y-%m-%d}__${now:%H-%M-%S}
    subdir: ${now:%Y-%m-%d}__${now:%H-%M-%S}__${hydra.job.num}


a: 0
now: ${now:%Y-%m-%d}__${now:%H-%M-%S}
uncached_seed: ${uncached_randint:0,99999}
cached_seed: ${cached_randint:0,99999}

and the output after running python main.py a=0,1,2 -m

[2022-04-27 11:23:29,486][HYDRA] Launching 3 jobs locally
[2022-04-27 11:23:29,486][HYDRA]        #0 : a=0
a: 0
now: 2022-04-27__11-23-29
uncached_seed: 59181
cached_seed: 98208

uncached_seed:  86350 70087 6605
cached_seed:  98208 98208 98208
[2022-04-27 11:23:31,577][HYDRA]        #1 : a=1
a: 1
now: 2022-04-27__11-23-29
uncached_seed: 32370
cached_seed: 98274

uncached_seed:  72906 8245 6736
cached_seed:  98274 98274 98274
[2022-04-27 11:23:33,673][HYDRA]        #2 : a=2
a: 2
now: 2022-04-27__11-23-29
uncached_seed: 95071
cached_seed: 22646

uncached_seed:  11035 40775 44926
cached_seed:  22646 22646 22646

note that the fields depending on $now (and also the output directory) all have the same timestamp, while the fields depending on cached randint change between jobs.

[Jasha10]

Hi @ywyga,

I believe Hydra is doing some special maneuvering here. Without diving too deep into Hydra's internals to figure out exactly how this works, I can offer a tip: If you register your own version of the now resolver, the behavior for that resolver will be the same as for other custom resolvers. Specifically, in the use-case from your example above, the cache will not be preserved between jobs.

from datetime import datetime
OmegaConf.register_new_resolver(
    "my_now",
    lambda pattern: datetime.now().strftime(pattern),
    use_cache=True,
    replace=True,
)

Using ${my_now:...} in your config.yaml file will result in a new timestamp being cached for each job.
For reference, here is the the implementation of Hydra's now resolver (as well as of the other custom resolvers defined by Hydra).

[mattijsdp]

@Jasha10 related to this, I wonder how I can replicate the implementation of the Hydra resolvers as I would like the cache to be preserved between jobs in a multirun?

From the code I suspect this isn't possible as the hydra resolvers are created when launching the launcher and the custom resolvers are only created upon importing/running the decorated function but I'm quite new to Hydra so thought I'd check.

[odelalleau]

@mattijsdp yeah it's not really possible out-of-the-box due to the order of operations.

At a glance I suspect you could hack something with the following:

hydra:
  sweep:
    dir: "${save_cache: multirun, ${my_resolver: value1}, ${my_resolver: value2}}/${now:%Y-%m-%d}/${now:%H-%M-%S}"

where save_cache() does something like

def save_cache(s, *args, _root_):
    cache = OmegaConf.get_cache(_root_)
    # save cache to disk
    return s

and then at the start of your main() function you would do something like:

@hydra.main()
def main(cfg):
    cache = load_cache(...)  # load cache from disk (custom function you need to write)
    OmegaConf.set_cache(cfg, cache)

=> hopefully the cached value of ${my_resolver: value1} and {$my_resolver: value2} would now be available (NB: I haven't tested it).

Note however that this won't share the cache between jobs in a multirun -- this would be much trickier as multiple jobs can run in parallel so you'd need to implement some form of communication between jobs. Though if you're only running jobs sequentially, a simpler approach would be to have each job save its cache and the next one pick it up from disk with something like I wrote above (without needing the hydra.sweep.dir hack).

[mattijsdp]

@odelalleau thanks a lot for the detailed answer, I'll explore that!

Essentially I am looking for a way to run some function which returns an id that is then passed to all of the runs in the multirun regardless of whether it is run sequentially, in parallel or even across a ray cluster for example. I looked into the on_multirun_start callback but unless I make it write to some remote storage cache that also doesn't really seem to be possible.

[odelalleau]

@mattijsdp this might be launcher-dependent (not sure) but e.g. with the submitit launcher, since it relies on pickle to know what to run, it's probably possible to run your function before calling main() and store its output within the object that will be pickled. This way you could avoid manually handling your own shared storage. Sorry I can't be much more precise since I don't remember exactly what's being pickled and how it works.

[mmcdermott]

@mattijsdp I think given that the hydra now resolver runs at multirun creation, not job start, you can just use that to get a string that is persistent across all runs in the multirun then hash that to get a numeric ID, if necessary. Would that work?