1DOCKER_REGISTRY ?= docker.io 2DOCKER_ORG ?= $(shell docker info 2>/dev/null | sed '/Username:/!d;s/.* //') 3DOCKER_IMAGE ?= pytorch 4DOCKER_FULL_NAME = $(DOCKER_REGISTRY)/$(DOCKER_ORG)/$(DOCKER_IMAGE) 5 6ifeq ("$(DOCKER_ORG)","") 7$(warning WARNING: No docker user found using results from whoami) 8DOCKER_ORG = $(shell whoami) 9endif 10 11CUDA_VERSION_SHORT ?= 12.1 12CUDA_VERSION ?= 12.1.1 13CUDNN_VERSION ?= 9 14BASE_RUNTIME = ubuntu:22.04 15BASE_DEVEL = nvidia/cuda:$(CUDA_VERSION)-devel-ubuntu22.04 16CMAKE_VARS ?= 17 18# The conda channel to use to install cudatoolkit 19CUDA_CHANNEL = nvidia 20# The conda channel to use to install pytorch / torchvision 21INSTALL_CHANNEL ?= whl 22 23CUDA_PATH ?= cpu 24ifneq ("$(CUDA_VERSION_SHORT)","cpu") 25CUDA_PATH = cu$(subst .,,$(CUDA_VERSION_SHORT)) 26endif 27 28PYTHON_VERSION ?= 3.11 29# Match versions that start with v followed by a number, to avoid matching with tags like ciflow 30PYTORCH_VERSION ?= $(shell git describe --tags --always --match "v[1-9]*.*") 31# Can be either official / dev 32BUILD_TYPE ?= dev 33BUILD_PROGRESS ?= auto 34# Intentionally left blank 35TRITON_VERSION ?= 36BUILD_ARGS = --build-arg BASE_IMAGE=$(BASE_IMAGE) \ 37 --build-arg PYTHON_VERSION=$(PYTHON_VERSION) \ 38 --build-arg CUDA_VERSION=$(CUDA_VERSION) \ 39 --build-arg CUDA_PATH=$(CUDA_PATH) \ 40 --build-arg PYTORCH_VERSION=$(PYTORCH_VERSION) \ 41 --build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL) \ 42 --build-arg TRITON_VERSION=$(TRITON_VERSION) \ 43 --build-arg CMAKE_VARS="$(CMAKE_VARS)" 44EXTRA_DOCKER_BUILD_FLAGS ?= 45 46BUILD ?= build 47# Intentionally left blank 48PLATFORMS_FLAG ?= 49PUSH_FLAG ?= 50USE_BUILDX ?= 51BUILD_PLATFORMS ?= 52WITH_PUSH ?= false 53# Setup buildx flags 54ifneq ("$(USE_BUILDX)","") 55BUILD = buildx build 56ifneq ("$(BUILD_PLATFORMS)","") 57PLATFORMS_FLAG = --platform="$(BUILD_PLATFORMS)" 58endif 59# Only set platforms flags if using buildx 60ifeq ("$(WITH_PUSH)","true") 61PUSH_FLAG = --push 62endif 63endif 64 65DOCKER_BUILD = docker $(BUILD) \ 66 --progress=$(BUILD_PROGRESS) \ 67 $(EXTRA_DOCKER_BUILD_FLAGS) \ 68 $(PLATFORMS_FLAG) \ 69 $(PUSH_FLAG) \ 70 --target $(BUILD_TYPE) \ 71 -t $(DOCKER_FULL_NAME):$(DOCKER_TAG) \ 72 $(BUILD_ARGS) . 73DOCKER_PUSH = docker push $(DOCKER_FULL_NAME):$(DOCKER_TAG) 74 75.PHONY: all 76all: devel-image 77 78.PHONY: devel-image 79devel-image: BASE_IMAGE := $(BASE_DEVEL) 80devel-image: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-devel 81devel-image: 82 $(DOCKER_BUILD) 83 84.PHONY: devel-push 85devel-push: BASE_IMAGE := $(BASE_DEVEL) 86devel-push: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-devel 87devel-push: 88 $(DOCKER_PUSH) 89 90ifeq ("$(CUDA_VERSION_SHORT)","cpu") 91 92.PHONY: runtime-image 93runtime-image: BASE_IMAGE := $(BASE_RUNTIME) 94runtime-image: DOCKER_TAG := $(PYTORCH_VERSION)-runtime 95runtime-image: 96 $(DOCKER_BUILD) 97 98.PHONY: runtime-push 99runtime-push: BASE_IMAGE := $(BASE_RUNTIME) 100runtime-push: DOCKER_TAG := $(PYTORCH_VERSION)-runtime 101runtime-push: 102 $(DOCKER_PUSH) 103 104else 105 106.PHONY: runtime-image 107runtime-image: BASE_IMAGE := $(BASE_RUNTIME) 108runtime-image: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-runtime 109runtime-image: 110 $(DOCKER_BUILD) 111 112.PHONY: runtime-push 113runtime-push: BASE_IMAGE := $(BASE_RUNTIME) 114runtime-push: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-runtime 115runtime-push: 116 $(DOCKER_PUSH) 117 118endif 119 120.PHONY: clean 121clean: 122 -docker rmi -f $(shell docker images -q $(DOCKER_FULL_NAME)) 123