From 9610a27cd537efcbb3d8d82aa85f7dbf80c78b9c Mon Sep 17 00:00:00 2001 From: technillogue Date: Wed, 12 Jul 2023 16:50:31 -0400 Subject: [PATCH 1/4] initial draft Signed-off-by: technillogue --- pkg/config/config.go | 4 +++- pkg/dockerfile/generator.go | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pkg/config/config.go b/pkg/config/config.go index 95bf7e4656..9fa1365230 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -31,7 +31,9 @@ type RunItem struct { } type Build struct { - GPU bool `json:"gpu,omitempty" yaml:"gpu"` + GPU bool `json:"gpu,omitempty" yaml:"gpu"` + // UseSystemCuda? something about jax? + UseCudaBaseImage bool `json:"use_cuda_base_image,omitempty" yaml:"use_cuda_base_image"` PythonVersion string `json:"python_version,omitempty" yaml:"python_version"` PythonRequirements string `json:"python_requirements,omitempty" yaml:"python_requirements"` PythonPackages []string `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go index 99b1f93eb5..c603cf04ba 100644 --- a/pkg/dockerfile/generator.go +++ b/pkg/dockerfile/generator.go @@ -158,7 +158,7 @@ func (g *Generator) Generate(imageName string) (weightsBase string, dockerfile s return "", "", "", err } installPython := "" - if g.Config.Build.GPU { + if g.Config.Build.GPU && g.Config.Build.UseCudaBaseImage { installPython, err = g.installPythonCUDA() if err != nil { return "", "", "", err @@ -245,10 +245,10 @@ func (g *Generator) Cleanup() error { } func (g *Generator) baseImage() (string, error) { - if g.Config.Build.GPU { + if g.Config.Build.GPU && g.Config.Build.UseCudaBaseImage { return g.Config.CUDABaseImageTag() } - return "python:" + g.Config.Build.PythonVersion, nil + return "python:" + g.Config.Build.PythonVersion + "-slim", nil } func (g *Generator) preamble() string { From 244fe560a8ef9c99fafba6f5dd236004981cffa3 Mon Sep 17 00:00:00 2001 From: technillogue Date: Wed, 12 Jul 2023 16:51:38 -0400 Subject: [PATCH 2/4] fix some tests Signed-off-by: technillogue --- pkg/dockerfile/generator_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pkg/dockerfile/generator_test.go b/pkg/dockerfile/generator_test.go index 350b132028..693b6a1211 100644 --- a/pkg/dockerfile/generator_test.go +++ b/pkg/dockerfile/generator_test.go @@ -81,7 +81,7 @@ predict: predict.py:Predictor expected := `#syntax=docker/dockerfile:1.4 FROM r8.im/replicate/cog-test-weights AS weights ` + testTiniStage() + - `FROM python:3.8 + `FROM python:3.8-slim ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin @@ -152,7 +152,7 @@ predict: predict.py:Predictor expected := `#syntax=docker/dockerfile:1.4 FROM r8.im/replicate/cog-test-weights AS weights ` + testTiniStage() + - `FROM python:3.8 + `FROM python:3.8-slim ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin @@ -248,7 +248,7 @@ build: expected := `#syntax=docker/dockerfile:1.4 FROM r8.im/replicate/cog-test-weights AS weights ` + testTiniStage() + - `FROM python:3.8 + `FROM python:3.8-slim ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin @@ -428,7 +428,7 @@ predict: predict.py:Predictor expected := `#syntax=docker/dockerfile:1.4 ` + testTiniStage() + - `FROM python:3.8 + `FROM python:3.8-slim ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin From 49448ad67d378ec4d7543bc541a6ae04d6bf5d4f Mon Sep 17 00:00:00 2001 From: technillogue Date: Wed, 12 Jul 2023 17:23:59 -0400 Subject: [PATCH 3/4] set flag to be true by default and add it to the config schema Signed-off-by: technillogue --- pkg/config/config.go | 5 +++-- pkg/config/data/config_schema_v1.0.json | 5 +++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/pkg/config/config.go b/pkg/config/config.go index 9fa1365230..b971a03ab5 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -61,8 +61,9 @@ type Config struct { func DefaultConfig() *Config { return &Config{ Build: &Build{ - GPU: false, - PythonVersion: "3.8", + GPU: false, + PythonVersion: "3.8", + UseCudaBaseImage: true, }, } } diff --git a/pkg/config/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json index 7c7a48468f..1e2bc7d499 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/pkg/config/data/config_schema_v1.0.json @@ -24,6 +24,11 @@ "type": "boolean", "description": "Enable GPUs for this model. When enabled, the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image will be used, and Cog will automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you are using." }, + "use_cuda_base_image": { + "$id": "#/properties/build/properties/use_cuda_base_image", + "type": "boolean", + "description": "When enabled, the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image will be used. This is not normally necessary for torch, but may be necessary for other frameworks." + }, "python_version": { "$id": "#/properties/build/properties/python_version", "type": ["string", "number"], From bacacc986e9207b5ac370f4fa93dab863dcbbe48 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Thu, 13 Jul 2023 13:06:39 -0700 Subject: [PATCH 4/4] Fix failing tests Signed-off-by: Mattt Zmuda --- test-integration/test_integration/test_build.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test-integration/test_integration/test_build.py b/test-integration/test_integration/test_build.py index 67bd67ff4d..0260c7a0f8 100644 --- a/test-integration/test_integration/test_build.py +++ b/test-integration/test_integration/test_build.py @@ -31,7 +31,7 @@ def test_build_without_predictor(docker_image): # Deprecated. Remove for 1.0. assert len(labels["org.cogmodel.cog_version"]) > 0 assert json.loads(labels["org.cogmodel.config"]) == { - "build": {"python_version": "3.8"} + "build": {"python_version": "3.8", "use_cuda_base_image": True} } assert "org.cogmodel.openapi_schema" not in labels @@ -155,6 +155,7 @@ def test_build_gpu_model_on_cpu(tmpdir, docker_image): "gpu": True, "cuda": "11.8", "cudnn": "8", + "use_cuda_base_image": True, } } assert "run.cog.openapi_schema" not in labels @@ -167,6 +168,7 @@ def test_build_gpu_model_on_cpu(tmpdir, docker_image): "gpu": True, "cuda": "11.8", "cudnn": "8", + "use_cuda_base_image": True, } } assert "org.cogmodel.openapi_schema" not in labels