Hydra: [Bug] cannot pickle 'CudnnModule' object when using Joblib

Created on 3 Dec 2020  路  8Comments  路  Source: facebookresearch/hydra

馃悰 Bug

Description

It seems that CUDNN does not play well with Hydra's Joblib pluging. See the repro below.

Checklist

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

To reproduce

import hydra
import torch.backends.cudnn as cudnn

@hydra.main(config_name="cudnn")
def main(cfg):
    cudnn.benchmark = False
    cudnn.deterministic = True


if __name__ == "__main__":
    main()

```
defaults:
- hydra/launcher: joblib


```bash
>>> python main.py -m +gpu=0,1,2,3

[2020-12-03 08:10:44,560][HYDRA] Joblib.Parallel(n_jobs=-1,backend=loky,prefer=processes,require=None,verbose=0,timeout=None,pre_dispatch=2*n_jobs,batch_size=auto,temp_folder=None,max_nbytes=None,mmap_mode=r) is launching 4 jobs
[2020-12-03 08:10:44,560][HYDRA] Launching jobs, sweep output dir : multirun/2020-12-03/08-10-44
[2020-12-03 08:10:44,560][HYDRA]        #0 : +gpu=0
[2020-12-03 08:10:44,560][HYDRA]        #1 : +gpu=1
[2020-12-03 08:10:44,560][HYDRA]        #2 : +gpu=2
[2020-12-03 08:10:44,560][HYDRA]        #3 : +gpu=3
joblib.externals.loky.process_executor._RemoteTraceback:
"""
Traceback (most recent call last):
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/loky/backend/queues.py", line 153, in _feed
    obj_ = dumps(obj, reducers=reducers)
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/loky/backend/reduction.py", line 271, in dumps
    dump(obj, buf, reducers=reducers, protocol=protocol)
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/loky/backend/reduction.py", line 264, in dump
    _LokyPickler(file, reducers=reducers, protocol=protocol).dump(obj)
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/cloudpickle/cloudpickle_fast.py", line 563, in dump
    return Pickler.dump(self, obj)
TypeError: cannot pickle 'CudnnModule' object
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra/_internal/utils.py", line 198, in run_and_report
    return func()
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra/_internal/utils.py", line 355, in <lambda>
    lambda: hydra.multirun(
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra/_internal/hydra.py", line 136, in multirun
    return sweeper.sweep(arguments=task_overrides)
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/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 "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra_plugins/hydra_joblib_launcher/joblib_launcher.py", line 45, in launch
    return _core.launch(
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra_plugins/hydra_joblib_launcher/_core.py", line 101, in launch
    runs = Parallel(**joblib_cfg)(
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/parallel.py", line 1061, in __call__
    self.retrieve()
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/parallel.py", line 940, in retrieve
    self._output.extend(job.get(timeout=self.timeout))
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/_parallel_backends.py", line 542, in wrap_future_result
    return future.result(timeout=timeout)
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/concurrent/futures/_base.py", line 439, in result
    return self.__get_result()
  File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/concurrent/futures/_base.py", line 388, in __get_result
    raise self._exception
_pickle.PicklingError: Could not pickle the task to send it to the workers.

System information

  • Hydra Version : 1.0.4
  • Python version : 3.8.5
  • Virtual environment type and version : conda 4.8.3
  • Operating system : Ubuntu 18.04.3 LTS

Additional context

This seems to be a problem with how the global CUDNN state is being set and its interaction with Joblib. This snippet works perfectly well with PyTorch and Python's distributed processing primitives.

bug

All 8 comments

This seems like a bug between joblib and Pytorch.

  1. Can you try to repro by using joblib directly?
  2. As a workaround, can you try to import cudnn lazily inside your hydra.main() function?
  3. Can you test with the submitit launcher in local mode? it also runs locally and I am curious if it has the same issue.
1. Can you try to repro by using joblib directly?
# main.py
import hydra
from joblib import Parallel, delayed
import torch.backends.cudnn as cudnn


@hydra.main(config_name="cudnn")
def main():
    cudnn.benchmark = False
    cudnn.deterministic = True


if __name__ == "__main__":
    runs = Parallel(n_jobs=4)(delayed(main)())
# cudnn.yaml

defaults:
    - hydra/launcher: joblib



md5-05f61326699d07456c006dad74424831




>     2. As a workaround, can you try to import cudnn lazily inside your `hydra.main()` function?



md5-7d73b3580e2feedaa276b1a16712d08b



cudnn.yaml

defaults:
- hydra/launcher: joblib

```shell
>>> python main.py -m +gpu=0,1,2,3

[2020-12-03 10:18:24,770][HYDRA] Joblib.Parallel(n_jobs=-1,backend=loky,prefer=processes,require=None,verbose=0,timeout=None,pre_dispatch=2*n_jobs,batch_size=auto,temp_folder=None,max_nbytes=None,mmap_mode=r) is launching 4 jobs
[2020-12-03 10:18:24,770][HYDRA] Launching jobs, sweep output dir : multirun/2020-12-03/10-18-24
[2020-12-03 10:18:24,770][HYDRA]        #0 : +gpu=0
[2020-12-03 10:18:24,770][HYDRA]        #1 : +gpu=1
[2020-12-03 10:18:24,770][HYDRA]        #2 : +gpu=2
[2020-12-03 10:18:24,770][HYDRA]        #3 : +gpu=3
3. Can you test with the submitit launcher in local mode? it also runs locally and I am curious if it has the same issue.
# main.py
import hydra
import torch.backends.cudnn as cudnn


@hydra.main(config_name="cudnn")
def main():
    cudnn.benchmark = False
    cudnn.deterministic = True


if __name__ == "__main__":
    runs = Parallel(n_jobs=4)(delayed(main)())
# cudnn.yaml

defaults:
    - hydra/launcher: submitit_local



md5-5d7f983a1fff0b23b90b75ccee93731f



Using submitit and importing lazily works as well.


Weirdly, this following code doesn't work even when you import lazily.



md5-9c4ebe2154d74b2f60d1da9fa125699f



python main.py -m gpu=0,1,2,3

[2020-12-03 10:24:53,423][HYDRA] Joblib.Parallel(n_jobs=-1,backend=loky,prefer=processes,require=None,verbose=0,timeout=None,pre_dispatch=2*n_jobs,batch_size=auto,temp_folder=None,max_nbytes=None,mmap_mode=r) is launching 4 jobs
[2020-12-03 10:24:53,423][HYDRA] Launching jobs, sweep output dir : multirun/2020-12-03/10-24-52
[2020-12-03 10:24:53,423][HYDRA] #0 : gpu=0
[2020-12-03 10:24:53,423][HYDRA] #1 : gpu=1
[2020-12-03 10:24:53,423][HYDRA] #2 : gpu=2
[2020-12-03 10:24:53,423][HYDRA] #3 : gpu=3
joblib.externals.loky.process_executor._RemoteTraceback:
"""
Traceback (most recent call last):
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/loky/backend/queues.py", line 153, in _feed
obj_ = dumps(obj, reducers=reducers)
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/loky/backend/reduction.py", line 271, in dumps
dump(obj, buf, reducers=reducers, protocol=protocol)
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/loky/backend/reduction.py", line 264, in dump
_LokyPickler(file, reducers=reducers, protocol=protocol).dump(obj)
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/externals/cloudpickle/cloudpickle_fast.py", line 563, in dump
return Pickler.dump(self, obj)
TypeError: cannot pickle 'CudnnModule' object
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra/_internal/utils.py", line 198, in run_and_report
return func()
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra/_internal/utils.py", line 355, in
lambda: hydra.multirun(
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra/_internal/hydra.py", line 136, in multirun
return sweeper.sweep(arguments=task_overrides)
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/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 "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra_plugins/hydra_joblib_launcher/joblib_launcher.py", line 45, in launch
return _core.launch(
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/hydra_plugins/hydra_joblib_launcher/_core.py", line 101, in launch
runs = Parallel(**joblib_cfg)(
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/parallel.py", line 1061, in __call__
self.retrieve()
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/parallel.py", line 940, in retrieve
self._output.extend(job.get(timeout=self.timeout))
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/site-packages/joblib/_parallel_backends.py", line 542, in wrap_future_result
return future.result(timeout=timeout)
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/concurrent/futures/_base.py", line 439, in result
return self.__get_result()
File "/home/vishwam/anaconda3/envs/hydra-torch/lib/python3.8/concurrent/futures/_base.py", line 388, in __get_result
raise self._exception
_pickle.PicklingError: Could not pickle the task to send it to the workers.
```

  1. Can you try to repro by using joblib directly?

Please take Hydra out of the equation in this one. does it still happen without @hydra.main() ?

  1. As a workaround, can you try to import cudnn lazily inside your hydra.main() function?

So importing lazily works around the issue? Is it actually functional?

  1. Can you test with the submitit launcher in local mode? it also runs locally and I am curious if it has the same issue.

So you are seeing the same behavior with submitit as well, correct?

  1. Can you try to repo by using joblib directly?
  2. As a workaround, can you try to import cudnn lazily inside your hydra.main() function?

This works:

from joblib import Parallel, delayed


def main():
    import torch.backends.cudnn as cudnn
    cudnn.benchmark = False
    cudnn.deterministic = True


if __name__ == "__main__":
    runs = Parallel(n_jobs=4)(delayed(main)() for i in range(4))

This doesn't work:

 import torch.backends.cudnn as cudnn
from joblib import Parallel, delayed


def main():
    cudnn.benchmark = False
    cudnn.deterministic = True


if __name__ == "__main__":
    runs = Parallel(n_jobs=4)(delayed(main)() for i in range(4))
  1. Can you test with the submitit launcher in local mode? it also runs locally and I am curious if it has the same issue.

Same behavior as submitit.

Minimal repro without joblib:

import torch.backends.cudnn
import cloudpickle

def foo():
    print(torch.backends.cudnn)

cloudpickle.dumps(foo)

This should probably be submitted as a bug report for pytorch and cloudpickle. it's something in between them.

@briankosw, did you an issue against pytorch/cloudpickle?

Thanks!
Closing as this is not an issue with Hydra.
I subscribed to two issues you opened and will keep an eye on the.

Was this page helpful?
0 / 5 - 0 ratings

Related issues

vict0rsch picture vict0rsch  路  4Comments

ImMrMa picture ImMrMa  路  4Comments

moskomule picture moskomule  路  5Comments

nathansomavarapu picture nathansomavarapu  路  3Comments

Khoa-NT picture Khoa-NT  路  4Comments