Skip to content

Commit ca6e84d

Browse files
authored
Add development_image flag to gpu_version section (#11)
1 parent 038eef3 commit ca6e84d

File tree

2 files changed

+6
-1
lines changed

2 files changed

+6
-1
lines changed

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ gpu: <bool> # whether your model server will use a GPU (default: false)
107107
gpu_version: # gpu version (only applicable for python model type, required if gpu is used)
108108
cuda: <string> # cuda version (tested with 10.0, 10.1, 10.2, 11.0, 11.1) (required)
109109
cudnn: <string> # cudnn version (tested with 7 and 8) (required)
110+
development_image: <bool> # whether to use the dev cuda image or not (default: false)
110111

111112
# dependencies
112113
dependencies:

nucleus/generate.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,8 @@ def validate_config(config: dict):
177177
raise CortexModelServerBuilder("gpu_version: 'cuda' field must be specified")
178178
if "cudnn" not in config["gpu_version"]:
179179
raise CortexModelServerBuilder("gpu_version: 'cudnn' field must be specified")
180+
if "development_image" not in config["gpu_version"]:
181+
config["gpu_version"]["development_image"] = False
180182

181183

182184
def build_handler_dockerfile(config: dict, path_to_config: str, dev_env: bool) -> str:
@@ -193,8 +195,10 @@ def build_handler_dockerfile(config: dict, path_to_config: str, dev_env: bool) -
193195
and config["gpu_version"]
194196
and config["gpu_version"]["cuda"] not in ["", None]
195197
and config["gpu_version"]["cudnn"] not in ["", None]
198+
and config["gpu_version"]["development_image"] not in ["", None]
196199
):
197-
base_image = f"nvidia/cuda:{config['gpu_version']['cuda']}-cudnn{config['gpu_version']['cudnn']}-runtime-ubuntu18.04"
200+
image_type = "devel" if config["gpu_version"]["development_image"] else "runtime"
201+
base_image = f"nvidia/cuda:{config['gpu_version']['cuda']}-cudnn{config['gpu_version']['cudnn']}-{image_type}-ubuntu18.04"
198202
cortex_image_type = "python-handler-gpu"
199203
if config["type"] == "tensorflow":
200204
cortex_image_type = "tensorflow-handler"

0 commit comments

Comments
 (0)