try using copying dependencies to a known place, then symlink to wher…

…eever pyenv actually is Signed-off-by: technillogue <[email protected]>
replicate · Aug 30, 2023 · 01a08b7 · 01a08b7
1 parent 2111ad5
commit 01a08b7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go
@@ -386,8 +386,16 @@ 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
-	target := g.sitePackagesLocation()
-	return "COPY --from=deps --link /dep " + target
+	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) {

diff --git a/pkg/dockerfile/generator_test.go b/pkg/dockerfile/generator_test.go
@@ -120,7 +120,8 @@ 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") + `COPY --from=deps --link /dep /root/.pyenv/versions/3.8.18/lib/python3.8/site-packages
+` + 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"]
@@ -213,7 +214,8 @@ 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") + `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 /root/.pyenv/versions/3.8.18/lib/python3.8/site-packages
+COPY --from=deps --link /dep /dep
+RUN ln -s /dep/* $(pyenv prefix)/lib/python*/site-packages
 RUN cowsay moo
 WORKDIR /src
 EXPOSE 5000
@@ -364,7 +366,8 @@ 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") + `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 /root/.pyenv/versions/3.8.18/lib/python3.8/site-packages
+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