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            ?= 8
14
BASE_RUNTIME              = ubuntu:22.04
15
BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_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          ?= pytorch
22

23
PYTHON_VERSION           ?= 3.10
24
# Match versions that start with v followed by a number, to avoid matching with tags like ciflow
25
PYTORCH_VERSION          ?= $(shell git describe --tags --always --match "v[1-9]*.*")
26
# Can be either official / dev
27
BUILD_TYPE               ?= dev
28
BUILD_PROGRESS           ?= auto
29
# Intentionally left blank
30
TRITON_VERSION           ?=
31
BUILD_ARGS                = --build-arg BASE_IMAGE=$(BASE_IMAGE) \
32
							--build-arg PYTHON_VERSION=$(PYTHON_VERSION) \
33
							--build-arg CUDA_VERSION=$(CUDA_VERSION) \
34
							--build-arg CUDA_CHANNEL=$(CUDA_CHANNEL) \
35
							--build-arg PYTORCH_VERSION=$(PYTORCH_VERSION) \
36
							--build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL) \
37
							--build-arg TRITON_VERSION=$(TRITON_VERSION) \
38
							--build-arg CMAKE_VARS="$(CMAKE_VARS)"
39
EXTRA_DOCKER_BUILD_FLAGS ?=
40

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

60
DOCKER_BUILD              = DOCKER_BUILDKIT=1 \
61
							docker $(BUILD) \
62
								--progress=$(BUILD_PROGRESS) \
63
								$(EXTRA_DOCKER_BUILD_FLAGS) \
64
								$(PLATFORMS_FLAG) \
65
								$(PUSH_FLAG) \
66
								--target $(BUILD_TYPE) \
67
								-t $(DOCKER_FULL_NAME):$(DOCKER_TAG) \
68
								$(BUILD_ARGS) .
69
DOCKER_PUSH               = docker push $(DOCKER_FULL_NAME):$(DOCKER_TAG)
70

71
.PHONY: all
72
all: devel-image
73

74
.PHONY: devel-image
75
devel-image: BASE_IMAGE := $(BASE_DEVEL)
76
devel-image: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-devel
77
devel-image:
78
	$(DOCKER_BUILD)
79

80
.PHONY: devel-push
81
devel-push: BASE_IMAGE := $(BASE_DEVEL)
82
devel-push: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-devel
83
devel-push:
84
	$(DOCKER_PUSH)
85

86
.PHONY: runtime-image
87
runtime-image: BASE_IMAGE := $(BASE_RUNTIME)
88
runtime-image: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-runtime
89
runtime-image:
90
	$(DOCKER_BUILD)
91

92
.PHONY: runtime-push
93
runtime-push: BASE_IMAGE := $(BASE_RUNTIME)
94
runtime-push: DOCKER_TAG := $(PYTORCH_VERSION)-cuda$(CUDA_VERSION_SHORT)-cudnn$(CUDNN_VERSION)-runtime
95
runtime-push:
96
	$(DOCKER_PUSH)
97

98
.PHONY: clean
99
clean:
100
	-docker rmi -f $(shell docker images -q $(DOCKER_FULL_NAME))
101