pytorch

1
DOCKER_REGISTRY          ?= docker.io
2
DOCKER_ORG               ?= $(shell docker info 2>/dev/null | sed '/Username:/!d;s/.* //')
3
DOCKER_IMAGE             ?= pytorch
4
DOCKER_FULL_NAME          = $(DOCKER_REGISTRY)/$(DOCKER_ORG)/$(DOCKER_IMAGE)
5

6
ifeq ("$(DOCKER_ORG)","")
7
$(warning WARNING: No docker user found using results from whoami)
8
DOCKER_ORG                = $(shell whoami)
9
endif
10

11
CUDA_VERSION_SHORT       ?= 12.1
12
CUDA_VERSION             ?= 12.1.1
13
CUDNN_VERSION            ?= 9
14
BASE_RUNTIME              = ubuntu:22.04
15
BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-devel-ubuntu22.04
16
CMAKE_VARS               ?=
17

18
# The conda channel to use to install cudatoolkit
19
CUDA_CHANNEL              = nvidia
20
# The conda channel to use to install pytorch / torchvision
21
INSTALL_CHANNEL          ?= whl
22

23
CUDA_PATH                ?= cpu
24
ifneq ("$(CUDA_VERSION_SHORT)","cpu")
25
CUDA_PATH                = cu$(subst .,,$(CUDA_VERSION_SHORT))
26
endif
27

28
PYTHON_VERSION           ?= 3.11
29
# Match versions that start with v followed by a number, to avoid matching with tags like ciflow
30
PYTORCH_VERSION          ?= $(shell git describe --tags --always --match "v[1-9]*.*")
31
# Can be either official / dev
32
BUILD_TYPE               ?= dev
33
BUILD_PROGRESS           ?= auto
34
# Intentionally left blank
35
TRITON_VERSION           ?=
36
BUILD_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)"
44
EXTRA_DOCKER_BUILD_FLAGS ?=
45

46
BUILD                    ?= build
47
# Intentionally left blank
48
PLATFORMS_FLAG           ?=
49
PUSH_FLAG                ?=
50
USE_BUILDX               ?=
51
BUILD_PLATFORMS          ?=
52
WITH_PUSH                ?= false
53
# Setup buildx flags
54
ifneq ("$(USE_BUILDX)","")
55
BUILD                     = buildx build
56
ifneq ("$(BUILD_PLATFORMS)","")
57
PLATFORMS_FLAG            = --platform="$(BUILD_PLATFORMS)"
58
endif
59
# Only set platforms flags if using buildx
60
ifeq ("$(WITH_PUSH)","true")
61
PUSH_FLAG                 = --push
62
endif
63
endif
64

65
DOCKER_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) .
73
DOCKER_PUSH               = docker push $(DOCKER_FULL_NAME):$(DOCKER_TAG)
74

75
.PHONY: all
76
all: devel-image
77

78
.PHONY: devel-image
79
devel-image: BASE_IMAGE := $(BASE_DEVEL)
80
devel-image: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-devel
81
devel-image:
82
	$(DOCKER_BUILD)
83

84
.PHONY: devel-push
85
devel-push: BASE_IMAGE := $(BASE_DEVEL)
86
devel-push: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-devel
87
devel-push:
88
	$(DOCKER_PUSH)
89

90
ifeq ("$(CUDA_VERSION_SHORT)","cpu")
91

92
.PHONY: runtime-image
93
runtime-image: BASE_IMAGE := $(BASE_RUNTIME)
94
runtime-image: DOCKER_TAG := $(PYTORCH_VERSION)-runtime
95
runtime-image:
96
	$(DOCKER_BUILD)
97

98
.PHONY: runtime-push
99
runtime-push: BASE_IMAGE := $(BASE_RUNTIME)
100
runtime-push: DOCKER_TAG := $(PYTORCH_VERSION)-runtime
101
runtime-push:
102
	$(DOCKER_PUSH)
103

104
else
105

106
.PHONY: runtime-image
107
runtime-image: BASE_IMAGE := $(BASE_RUNTIME)
108
runtime-image: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-runtime
109
runtime-image:
110
	$(DOCKER_BUILD)
111

112
.PHONY: runtime-push
113
runtime-push: BASE_IMAGE := $(BASE_RUNTIME)
114
runtime-push: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-runtime
115
runtime-push:
116
	$(DOCKER_PUSH)
117

118
endif
119

120
.PHONY: clean
121
clean:
122
	-docker rmi -f $(shell docker images -q $(DOCKER_FULL_NAME))
123