Docker

Run W&B inside Docker containers by configuring API keys, environment variables, and local file storage.

Docker integration

W&B can store a pointer to the Docker image that your code ran in, letting you restore a previous experiment to the exact environment it ran in. The wandb library looks for the WANDB_DOCKER environment variable to persist this state. W&B provides a few helpers that automatically set this state.

The following sections describe how to set the WANDB_DOCKER environment variable in different environments, from local development through Kubernetes-based training.

Local development

wandb docker is a command that starts a Docker container, passes in wandb environment variables, mounts your code, and ensures wandb is installed. By default, the command uses a Docker image with TensorFlow, PyTorch, Keras, and Jupyter installed. You can use the same command to start your own Docker image: wandb docker my/image:latest. The command mounts the current directory into the /app directory of the container. You can change this with the --dir flag.

Production

The wandb docker-run command is provided for production workloads. It’s a drop-in replacement for nvidia-docker that wraps the docker run command and adds your credentials and the WANDB_DOCKER environment variable to the call. If you don’t pass the --runtime flag and nvidia-docker is available on the machine, this also ensures the runtime is set to nvidia.

Kubernetes

If you run your training workloads in Kubernetes and the Kubernetes API is exposed to your pod (which is the case by default), wandb queries the API for the digest of the Docker image and automatically sets the WANDB_DOCKER environment variable.

Restore the training environment

Once the WANDB_DOCKER environment variable is set during a run, you can use it to reproduce the original training environment later.

If a run was instrumented with the WANDB_DOCKER environment variable, calling wandb restore username/project:run_id checks out a new branch restoring your code, then launches the exact Docker image used for training pre-populated with the original command.