Create pytorch-notebook docker image

[weiji14]

Another GPU-enabled docker image for running deep learning!

List of packages to include (let me know of any other relevant ones);

• pytorch
• torchgeo
• pytorch-lightning
• etc

Resolves #312

[github-actions]

👈 Try on Mybinder.org!
👈 Try on Pangeo GCP Binder!
👈 Try on Pangeo AWS Binder!

[pangeo-bot]

/condalock
Automatically locking new conda environment, building, and testing images...

[scottyhq]

Thanks for this @weiji14 ! If you add new packages (torchgeo) you have to relock the environment. The easiest way to do this is add a comment with /condalock on the first line.

[rabernat]

/condalock

[weiji14]

Thanks for this @weiji14 ! If you add new packages (torchgeo) you have to relock the environment. The easiest way to do this is add a comment with /condalock on the first line.

Ah ok, I've actually ran conda-lock locally already so there might not be stuff to update. But there might be a few other packages we might add so will do that later 😄

[scottyhq]

in the ml-notebook image we pin cudatoolkit=10,

pangeo-docker-images/ml-notebook/environment.yml

Line 7 in b5c38d0

- cudatoolkit=10

I'm guessing if both these images are offered on pangeo jupyterhubs we want cuda versions to be the same, @yuvipanda any recommendations for 10 vs 11 when it comes to node configuration?

Fine leaving it as 11 here and bumping the other image in a separate PR.

[weiji14]

Yes, it comes down to what CUDA drivers the server is using (find out using nvidia-smi). I prefer cudatoolkit=11 because it has foward and backward compatibility within a minor version (see https://docs.nvidia.com/deploy/cuda-compatibility/#forward-compatibility-title), but that's of no use if the cuda drivers on the server don't support it.

[yuvipanda]

I don't actually know enough to have an opinion here, but happy to stick to whatever upstream recommends and as long as we are consistent across the images :D If we need to upgrade the driver version on the clusters when we bump this, can do that.

[weiji14]

Ok. Just to confirm, I'm assuming that the clusters don't pull the latest pangeo-docker image? I.e. that they're pinned to a specific pangeo-docker-image version? Probably doesn't matter for this pytorch-notebook image since nobody will be using it yet, but I'd recommend updating the CUDA driver version before bumping the cudatoolkit version on ml/tensorflow-notebook, just speaking from experience working on an on-prem HPC 🙂

[scottyhq]

clusters don't pull the latest pangeo-docker image?

Actually they do currently pull the latest :) https://github.com/2i2c-org/infrastructure/search?q=ml-notebook
But I say we go ahead and merge this, and can use explicit pins on the hub if need be going forward

[weiji14]

Wow, you all like to live dangerously 😆 So looking at 2i2c-org/infrastructure#1244, it seems like the GPUs are K80s, and according to https://forums.developer.nvidia.com/t/nvidia-tesla-k80-cuda-version-support/67676/6, CUDA 11 does work but support is deprecated (i.e. there might be lots of warnings). So let's test it out then!

[rabernat]

We are not currently using ml-notebook in production in 2i2c. All production clusters use pinned tags.

Those clusters where this image turns up are not really launched yet.

[scottyhq]

One other point on cudatoolkit 10 vs 11 is that 11 brings along a few more gigabytes. I don't mind big images, but this one is currently approaching 10GB total uncompressed...

mamba create -n test --dry-run pytorch-gpu cudatoolkit=10 -> Total download: 1 GB
mamba create -n test --dry-run pytorch-gpu cudatoolkit=11 -> Total download: 3 GB

[weiji14]

Yes cudatoolkit is huge unfortunately. There are ways to trim down the docker image size as you mentioned in #188 (comment), but it's tough getting around a big binary... I think CUDA 11 will need to be used eventually though, so will need to figure out a solution 🙃

[rabernat]

This is great @weiji14 - just one small nit, then LGTM.

[rabernat]

Suggested change

	# cupy import fails unless on GPU-enabled node:
	#'cupy', #libcuda.so.1: cannot open shared object file: No such file or directory

Let's remove this comment. It's just residue from old stuff in ml-image.

[ngam]

Could I add a suggestion as well? Because of the mkl and nomkl conflicts, this is sort of worrying...

Try a test for the underlying blas to make sure things are set up correctly. Usually it should be MKL (it's pulling MKL from conda-forge):

import torch
assert torch.__config__.show()[torch.__config__.show().find('BLAS_INFO')+10:torch.__config__.show().find('BLAS_INFO')+13] == "mkl"

[weiji14]

Ok, done both in 3acef32. Let's see if the CI passes.

[weiji14]

Yep, works like a charm: tests/test_pytorch-notebook.py::test_torch_uses_mkl PASSED [100%] https://github.com/pangeo-data/pangeo-docker-images/runs/6202609085?check_suite_focus=true#step:7:41

[scottyhq]

Thanks @weiji14 for putting this together and @ngam @rabernat for the reviews, going to merge so that it'll be easier for people to test drive