Hydra: Submitit plugin breaks when using structured configs

Created on 21 Oct 2020 · 13Comments · Source: facebookresearch/hydra

🐛 Bug

Description

I get a pickling error when my main function exits when I use a structured config

Checklist

[x] I checked on the latest version of Hydra
[x] I created a minimal repro

To reproduce

# foo.py
import hydra
from hydra.core.config_store import ConfigStore
from dataclasses import dataclass
from omegaconf import OmegaConf

@dataclass
class Config:
    x: int = 42


cs = ConfigStore.instance()
cs.store(name="config", node=Config)


@hydra.main(config_name="config")
def main(cfg: Config) -> None:
    print(OmegaConf.to_yaml(cfg))


if __name__ == "__main__":
    main()

# config.yaml
defaults:
    - hydra/launcher: submitit_local

Run:

$ python foo.py -m
[2020-10-20 15:27:55,392][HYDRA] Submitit 'local' sweep output dir : multirun/2020-10-20/15-27-55
[2020-10-20 15:27:55,393][HYDRA]    #0 :
Traceback (most recent call last):
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/hydra/_internal/utils.py", line 198, in run_and_report
    return func()
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/hydra/_internal/utils.py", line 358, in <lambda>
    overrides=args.overrides,
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/hydra/_internal/hydra.py", line 136, in multirun
    return sweeper.sweep(arguments=task_overrides)
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/hydra/_internal/core_plugins/basic_sweeper.py", line 154, in sweep
    results = self.launcher.launch(batch, initial_job_idx=initial_job_idx)
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py", line 150, in launch
    return [j.results()[0] for j in jobs]
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py", line 150, in <listcomp>
    return [j.results()[0] for j in jobs]
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/submitit/core/core.py", line 286, in results
    outcome, result = self._get_outcome_and_result()
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/submitit/core/core.py", line 366, in _get_outcome_and_result
    f"Job {self.job_id} (task: {self.task_id}) with path {self.paths.result_pickle}\n"
submitit.core.utils.UncompletedJobError: Job 22319 (task: 0) with path /private/home/calebh/scratch/hydra-submitit-bug/multirun/2020-10-20/15-27-55/.submitit/22319/22319_0_result.pkl
has not produced any output (state: FINISHED)
Error stream produced:
----------------------
submitit ERROR (2020-10-20 15:27:55,943) - Could not dump error:
Can't pickle <class '__main__.Config'>: attribute lookup Config on __main__ failed

because of A temporary saved file already exists.
submitit ERROR (2020-10-20 15:27:55,943) - Submitted job triggered an exception
Traceback (most recent call last):
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/submitit/core/_submit.py", line 11, in <module>
    submitit_main()
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/submitit/core/submission.py", line 65, in submitit_main
    process_job(args.folder)
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/submitit/core/submission.py", line 58, in process_job
    raise error
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/submitit/core/submission.py", line 49, in process_job
    utils.pickle_dump(("success", result), tmppath)
  File "/private/home/calebh/miniconda3/envs/hydra-submitit-bug/lib/python3.7/site-packages/submitit/core/utils.py", line 278, in pickle_dump
    pickle.dump(obj, ofile, pickle.HIGHEST_PROTOCOL)
_pickle.PicklingError: Can't pickle <class '__main__.Config'>: attribute lookup Config on __main__ failed

Expected Behavior

Should exit without throwing. Note that the body of main does execute properly as I am able to see the print output in the log, so I'm pretty certain this is happening after main returns.

System information

Hydra Version :
hydra-core 1.0.3 hydra-submitit-launcher 1.0.0
Python version : 3.7
Virtual environment type and version : Conda 4.7.10
Operating system : Ubuntu 18.04.3

bug plugin

Source

calebho

Most helpful comment

My guess would that submitit is running a different entry point (so a different __main__) which obviously does not have this class.
not sure there is a solution for it except documentation and an example. will wait to see what @jrapin thinks.

omry on 21 Oct 2020

👍2

All 13 comments

@jrapin, any thoughts about this?

omry on 21 Oct 2020

@calebho, if you place your Config class in a different module (not __main__ is it any better) ?

omry on 21 Oct 2020

Hmm weird; moving Config to a separate module works:

# module.py
from dataclasses import dataclass
from hydra.core.config_store import ConfigStore

@dataclass
class Config:
    x: int = 42


cs = ConfigStore.instance()
cs.store(name="config", node=Config)

# foo.py
import hydra
from omegaconf import OmegaConf
from module import Config


@hydra.main(config_name="config")
def main(cfg: Config) -> None:
    print(OmegaConf.to_yaml(cfg))


if __name__ == "__main__":
    main()

calebho on 21 Oct 2020

omry on 21 Oct 2020

👍2

@ogrisel, any thoughts about this one?
The classes are defined in one __main__ and are de-serialized by another generic __main__.
I am not sure there is a solution except moving the class definition to a different module.
Interestingly, the function annotated with @hydra.main de-serializes without issues.

repro:
1.py:

from dataclasses import dataclass
import pickle
import cloudpickle
from typing import Any

@dataclass
class User:
    name : str
    age: int

def pickle_dump(obj: Any, filename: str) -> None:
    with open(filename, "wb") as ofile:
        pickle.dump(obj, ofile, pickle.HIGHEST_PROTOCOL)


def cloudpickle_dump(obj: Any, filename: str) -> None:
    with open(filename, "wb") as ofile:
        cloudpickle.dump(obj, ofile, pickle.HIGHEST_PROTOCOL)

def pickle_load(filename):
    with open(filename, "rb") as ifile:
        return pickle.load(ifile)

user = User(name="Bond", age=7)
pickle_dump(user, "user.p")
cloudpickle_dump(user, "user.cp")

print("pickle: ", pickle_load("user.p"))
print("cloudpickle: ", pickle_load("user.p"))

2.py:

import pickle

def pickle_load(filename):
    with open(filename, "rb") as ifile:
        return pickle.load(ifile)

print("cloudpickle: ", pickle_load("user.p"))
print("pickle: ", pickle_load("user.p"))

Running 1.py and then 2.py results in

$ python 2.py
Traceback (most recent call last):
  File "2.py", line 6, in <module>
    print("cloudpickle: ", pickle_load("user.p"))
  File "2.py", line 5, in pickle_load
    return pickle.load(ifile)
AttributeError: Can't get attribute 'User' on <module '__main__' from '2.py'>

omry on 21 Oct 2020

My guess would that submitit is running a different entry point (so a different __main__) which obviously does not have this class.
not sure there is a solution for it except documentation and an example. will wait to see what @jrapin thinks.

Functions/classes defined in __main__ are always the painful ones. That's actually what cloudpickle does way better than pickle. The example just above actually works with the correction that cloudpickle needs to load user.cp and not user.p.

My best guess now is that there is something more complicated going on. It actually fails when trying to pickle the result of the computation. Is the result of the function sent by the plugin to submitit supposed to contain a Config object? In that case submitit would need to pickle an object which it unpickled and was only defined in __main__ (hence a difficult case because it cannot point to a definition in a file). It's somehow 2nd order cloudpickling :s

jrapin on 21 Oct 2020

@jrapin
Yeah, that was a mistake in my repro.
Looking at the submitit code, I still see the pickle dump logic, are you sure it's not accidentally being used, like here ?

omry on 21 Oct 2020

are you sure it's not accidentally being used

What do you mean accidentally? This is through this pickle that the output of the job is transferred to the caller, so it's not accidental. And the job seems to return the output of hydra.core.utils.run_job which returns a JobReturn which does seems to contain the configuration (or am I wrong with it?), which cannot be pickled back since it is aready unpickled from a pickle of a definition in __main__.

Yes, that's a bit hard to read sorry, but the main idea is, objects pickled with definition in __main__ (the submission scripts, with opposed to having their definition in a module imported in __main__) cannot be unpickled and then pickled again, so the return type would need to be modified. That being said, it would also break the checkpointing part, which creates a pickle as well. Can the Config be cast into a dict or an OmegaConf or whatever instead?

jrapin on 21 Oct 2020

What do you mean accidentally?

I mean that I don't understand why you use pickle_dump and not cloudpickle_dump there.

Can the Config be cast into a dict or an OmegaConf or whatever instead?

I will need to have much better understanding of what is happening here to answer that. but the typing information is a critical part that should be serialized as it's used for runtime type safety.

omry on 21 Oct 2020

I mean that I don't understand why you use pickle_dump and not cloudpickle_dump there.

Good catch, I was wrong with the fact that cloudpickle could not redump inline defined functions/classes. The use of standard pickle here may have been just accidental/historical indeed since the case never arose.
I could repro and PR'ed a fix.

jrapin on 22 Oct 2020

@omry the fix is merged and released in version 1.1.3 thanks to @gwenzek

jrapin on 22 Oct 2020

Great (I hope the fix cleaned up the legacy pickle_dump function to avoid future confusion).
keeping this open as a reminder to add a test to launchers that is covering this case.

omry on 22 Oct 2020

Filed a dedicated issue.

omry on 22 Oct 2020

Was this page helpful?

0 / 5 - 0 ratings