Multistage pip install (#1233)

* initial draft Signed-off-by: technillogue <[email protected]> * move install cog before install user deps Signed-off-by: technillogue <[email protected]> * fix tests Signed-off-by: technillogue <[email protected]> * try different copy dest to fix integration tests Signed-off-by: technillogue <[email protected]> * fix tests Signed-off-by: technillogue <[email protected]> * note why test is broken Signed-off-by: technillogue <[email protected]> * try to determine the actual location of site-packages in the final image Signed-off-by: technillogue <[email protected]> * check if a different version would pass the tests Signed-off-by: technillogue <[email protected]> * also update go tests Signed-off-by: technillogue <[email protected]> * clean comments Signed-off-by: technillogue <[email protected]> * notes about resolving minor versions for pyenv Signed-off-by: technillogue <[email protected]> * fmt Signed-off-by: technillogue <[email protected]> * try using copying dependencies to a known place, then symlink to whereever pyenv actually is Signed-off-by: technillogue <[email protected]> * remove dead code Signed-off-by: technillogue <[email protected]> * reword comment Signed-off-by: technillogue <[email protected]> --------- Signed-off-by: technillogue <[email protected]>
replicate · Sep 5, 2023 · caba7fd · caba7fd
1 parent d65f2c8
commit caba7fd
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 48 deletions.
diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go
@@ -95,6 +95,10 @@ func (g *Generator) SetUseCudaBaseImage(argumentValue string) {
 }
 
 func (g *Generator) GenerateBase() (string, error) {
+	pipInstallStage, err := g.pipInstallStage()
+	if err != nil {
+		return "", err
+	}
 	baseImage, err := g.baseImage()
 	if err != nil {
 		return "", err
@@ -110,28 +114,20 @@ func (g *Generator) GenerateBase() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	pipInstalls, err := g.pipInstalls()
-	if err != nil {
-		return "", err
-	}
-	installCog, err := g.installCog()
-	if err != nil {
-		return "", err
-	}
 	run, err := g.runCommands()
 	if err != nil {
 		return "", err
 	}
 
 	return strings.Join(filterEmpty([]string{
 		"#syntax=docker/dockerfile:1.4",
+		pipInstallStage,
 		"FROM " + baseImage,
 		g.preamble(),
 		g.installTini(),
 		installPython,
-		installCog,
 		aptInstalls,
-		pipInstalls,
+		g.pipInstalls(),
 		run,
 		`WORKDIR /src`,
 		`EXPOSE 5000`,
@@ -162,7 +158,10 @@ func (g *Generator) Generate(imageName string) (weightsBase string, dockerfile s
 	if err != nil {
 		return "", "", "", fmt.Errorf("Failed to generate Dockerfile for model weights files: %w", err)
 	}
-
+	pipInstallStage, err := g.pipInstallStage()
+	if err != nil {
+		return "", "", "", err
+	}
 	baseImage, err := g.baseImage()
 	if err != nil {
 		return "", "", "", err
@@ -178,29 +177,21 @@ func (g *Generator) Generate(imageName string) (weightsBase string, dockerfile s
 	if err != nil {
 		return "", "", "", err
 	}
-	pipInstalls, err := g.pipInstalls()
-	if err != nil {
-		return "", "", "", err
-	}
-	installCog, err := g.installCog()
-	if err != nil {
-		return "", "", "", err
-	}
 	runCommands, err := g.runCommands()
 	if err != nil {
 		return "", "", "", err
 	}
 
 	base := []string{
 		"#syntax=docker/dockerfile:1.4",
+		pipInstallStage,
 		fmt.Sprintf("FROM %s AS %s", imageName+"-weights", "weights"),
 		"FROM " + baseImage,
 		g.preamble(),
 		g.installTini(),
 		installPython,
-		installCog,
 		aptInstalls,
-		pipInstalls,
+		g.pipInstalls(),
 		runCommands,
 	}
 
@@ -327,6 +318,9 @@ RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get insta
 	pyenv install-latest "%s" && \
 	pyenv global $(pyenv install-latest --print "%s") && \
 	pip install "wheel<1"`, py, py), nil
+	// for sitePackagesLocation, kind of need to determine which specific version latest is (3.8 -> 3.8.17 or 3.8.18)
+	// install-latest essentially does pyenv install --list | grep $py | tail -1
+	// there are many bad options, but a symlink to $(pyenv prefix) is the least bad one
 }
 
 func (g *Generator) installCog() (string, error) {
@@ -336,28 +330,54 @@ func (g *Generator) installCog() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	lines = append(lines, fmt.Sprintf("RUN --mount=type=cache,target=/root/.cache/pip pip install %s", containerPath))
+	lines = append(lines, fmt.Sprintf("RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep %s", containerPath))
 	return strings.Join(lines, "\n"), nil
 }
 
-func (g *Generator) pipInstalls() (string, error) {
+func (g *Generator) pipInstallStage() (string, error) {
+	installCog, err := g.installCog()
+	if err != nil {
+		return "", err
+	}
 	requirements, err := g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH)
 	if err != nil {
 		return "", err
 	}
 	if strings.Trim(requirements, "") == "" {
-		return "", nil
+		return `FROM python:` + g.Config.Build.PythonVersion + ` as deps
+` + installCog, nil
 	}
 
-	lines, containerPath, err := g.writeTemp("requirements.txt", []byte(requirements))
+	copyLine, containerPath, err := g.writeTemp("requirements.txt", []byte(requirements))
 	if err != nil {
 		return "", err
 	}
-
-	lines = append(lines, "RUN --mount=type=cache,target=/root/.cache/pip pip install -r "+containerPath)
+	lines := []string{
+		// Not slim, so that we can compile wheels
+		`FROM python:` + g.Config.Build.PythonVersion + ` as deps`,
+		installCog,
+		copyLine[0],
+		"RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r " + containerPath,
+	}
 	return strings.Join(lines, "\n"), nil
 }
 
+func (g *Generator) pipInstalls() string {
+	// placing packages in workdir makes imports faster but seems to break integration tests
+	// return "COPY --from=deps --link /dep COPY --from=deps /src"
+	// ...except it's actually /root/.pyenv/versions/3.8.17/lib/python3.8/site-packages
+	py := g.Config.Build.PythonVersion
+	if g.Config.Build.GPU && g.useCudaBaseImage {
+		return strings.Join(
+			[]string{
+				"COPY --from=deps --link /dep /dep",
+				"RUN ln -s /dep/* $(pyenv prefix)/lib/python*/site-packages",
+			},
+			"\n")
+	}
+	return "COPY --from=deps --link /dep /usr/local/lib/python" + py + "/site-packages"
+}
+
 func (g *Generator) runCommands() (string, error) {
 	runCommands := g.Config.Build.Run
 

diff --git a/pkg/dockerfile/generator_test.go b/pkg/dockerfile/generator_test.go
@@ -28,7 +28,12 @@ ENTRYPOINT ["/sbin/tini", "--"]
 
 func testInstallCog(relativeTmpDir string) string {
 	return fmt.Sprintf(`COPY %s/cog-0.0.1.dev-py3-none-any.whl /tmp/cog-0.0.1.dev-py3-none-any.whl
-RUN --mount=type=cache,target=/root/.cache/pip pip install /tmp/cog-0.0.1.dev-py3-none-any.whl`, relativeTmpDir)
+RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep /tmp/cog-0.0.1.dev-py3-none-any.whl`, relativeTmpDir)
+}
+
+func testPipInstallStage(relativeTmpDir string) string {
+	return `FROM python:3.8 as deps
+` + testInstallCog(relativeTmpDir)
 }
 
 func testInstallPython(version string) string {
@@ -78,12 +83,13 @@ predict: predict.py:Predictor
 	require.NoError(t, err)
 
 	expected := `#syntax=docker/dockerfile:1.4
+` + testPipInstallStage(gen.relativeTmpDir) + `
 FROM r8.im/replicate/cog-test-weights AS weights
 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
-` + testTini() + testInstallCog(gen.relativeTmpDir) + `
+` + testTini() + `COPY --from=deps --link /dep /usr/local/lib/python3.8/site-packages
 WORKDIR /src
 EXPOSE 5000
 CMD ["python", "-m", "cog.server.http"]
@@ -108,12 +114,14 @@ predict: predict.py:Predictor
 	require.NoError(t, err)
 
 	expected := `#syntax=docker/dockerfile:1.4
+` + testPipInstallStage(gen.relativeTmpDir) + `
 FROM r8.im/replicate/cog-test-weights AS weights
 FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
 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
-` + testTini() + testInstallPython("3.8") + testInstallCog(gen.relativeTmpDir) + `
+` + testTini() + testInstallPython("3.8") + `COPY --from=deps --link /dep /dep
+RUN ln -s /dep/* $(pyenv prefix)/lib/python*/site-packages
 WORKDIR /src
 EXPOSE 5000
 CMD ["python", "-m", "cog.server.http"]
@@ -147,15 +155,16 @@ predict: predict.py:Predictor
 	require.NoError(t, err)
 
 	expected := `#syntax=docker/dockerfile:1.4
+` + testPipInstallStage(gen.relativeTmpDir) + `
+COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
+RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r /tmp/requirements.txt
 FROM r8.im/replicate/cog-test-weights AS weights
 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
-` + testTini() + testInstallCog(gen.relativeTmpDir) + `
-RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
-COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
-RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
+` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
+COPY --from=deps --link /dep /usr/local/lib/python3.8/site-packages
 RUN cowsay moo
 WORKDIR /src
 EXPOSE 5000
@@ -195,17 +204,18 @@ predict: predict.py:Predictor
 	require.NoError(t, err)
 
 	expected := `#syntax=docker/dockerfile:1.4
+` + testPipInstallStage(gen.relativeTmpDir) + `
+COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
+RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r /tmp/requirements.txt
 FROM r8.im/replicate/cog-test-weights AS weights
 FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
 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
 ` + testTini() +
-		testInstallPython("3.8") +
-		testInstallCog(gen.relativeTmpDir) + `
-RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
-COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
-RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
+		testInstallPython("3.8") + `RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
+COPY --from=deps --link /dep /dep
+RUN ln -s /dep/* $(pyenv prefix)/lib/python*/site-packages
 RUN cowsay moo
 WORKDIR /src
 EXPOSE 5000
@@ -241,13 +251,14 @@ build:
 	require.NoError(t, err)
 
 	expected := `#syntax=docker/dockerfile:1.4
+` + testPipInstallStage(gen.relativeTmpDir) + `
 FROM r8.im/replicate/cog-test-weights AS weights
 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
-` + testTini() + testInstallCog(gen.relativeTmpDir) + `
-RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
+` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
+COPY --from=deps --link /dep /usr/local/lib/python3.8/site-packages
 RUN cowsay moo
 WORKDIR /src
 EXPOSE 5000
@@ -273,7 +284,7 @@ build:
 	_, actual, _, err := gen.Generate("r8.im/replicate/cog-test")
 	require.NoError(t, err)
 	fmt.Println(actual)
-	require.Contains(t, actual, `pip install -r /tmp/requirements.txt`)
+	require.Contains(t, actual, `pip install -t /dep -r /tmp/requirements.txt`)
 }
 
 // mockFileInfo is a test type to mock os.FileInfo
@@ -345,17 +356,18 @@ COPY root-large /src/root-large`
 
 	// model copy should be run before dependency install and code copy
 	expected = `#syntax=docker/dockerfile:1.4
+` + testPipInstallStage(gen.relativeTmpDir) + `
+COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
+RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r /tmp/requirements.txt
 FROM r8.im/replicate/cog-test-weights AS weights
 FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
 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
 ` + testTini() +
-		testInstallPython("3.8") +
-		testInstallCog(gen.relativeTmpDir) + `
-RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
-COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
-RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
+		testInstallPython("3.8") + `RUN --mount=type=cache,target=/var/cache/apt apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
+COPY --from=deps --link /dep /dep
+RUN ln -s /dep/* $(pyenv prefix)/lib/python*/site-packages
 RUN cowsay moo
 COPY --from=weights --link /src/checkpoints /src/checkpoints
 COPY --from=weights --link /src/models /src/models
@@ -420,11 +432,12 @@ predict: predict.py:Predictor
 	require.NoError(t, err)
 
 	expected := `#syntax=docker/dockerfile:1.4
+` + testPipInstallStage(gen.relativeTmpDir) + `
 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
-` + testTini() + testInstallCog(gen.relativeTmpDir) + `
+` + testTini() + `COPY --from=deps --link /dep /usr/local/lib/python3.8/site-packages
 WORKDIR /src
 EXPOSE 5000
 CMD ["python", "-m", "cog.server.http"]