Skip to content

Commit

Permalink
Do not run a separate python install stage
Browse files Browse the repository at this point in the history 
* Perform all the necessary installations on the
pyenv python.
  • Loading branch information
@8W9aG
8W9aG committed Dec 23, 2024
1 parent 392819d commit 9199ee5
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 84 deletions.
50 changes: 30 additions & 20 deletions pkg/dockerfile/standard_generator.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,18 +152,16 @@ func (g *StandardGenerator) GenerateInitialSteps() (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
}

if g.IsUsingCogBaseImage() {
pipInstalls, err := g.pipInstalls()
if err != nil {
return "", err
}

installCog, err := g.installCog()
if err != nil {
return "", err
}

steps := []string{
"#syntax=docker/dockerfile:1.4",
"FROM " + baseImage,
Expand All @@ -179,19 +177,15 @@ func (g *StandardGenerator) GenerateInitialSteps() (string, error) {
return joinStringsWithoutLineSpace(steps), nil
}

pipInstallStage, err := g.pipInstallStage(aptInstalls)
if err != nil {
return "", err
}

steps := []string{
"#syntax=docker/dockerfile:1.4",
pipInstallStage,
"FROM " + baseImage,
g.preamble(),
g.installTini(),
aptInstalls,
installPython,
g.copyPipPackagesFromInstallStage(),
pipInstalls,
installCog,
}
if g.precompile {
steps = append(steps, PrecompilePythonCommand)
Expand Down Expand Up @@ -431,9 +425,6 @@ func (g *StandardGenerator) installCog() (string, error) {
return "", err
}
pipInstallLine := "RUN --mount=type=cache,target=/root/.cache/pip pip install --no-cache-dir"
if !g.IsUsingCogBaseImage() {
pipInstallLine += " -t /dep"
}
pipInstallLine += " " + containerPath
// Install pydantic<2 for now, installing pydantic>2 wouldn't allow a downgrade later,
// but upgrading works fine
Expand Down Expand Up @@ -484,6 +475,25 @@ func (g *StandardGenerator) pipInstalls() (string, error) {
}, "\n"), nil
}

func (g *StandardGenerator) pipInstall() (string, error) {
copyLine, containerPath, err := g.writeTemp("requirements.txt", []byte(g.pythonRequirementsContents))
if err != nil {
return "", err
}

pipInstallLine := "RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r " + containerPath
if g.strip {
pipInstallLine += " && " + StripDebugSymbolsCommand
}

return strings.Join([]string{
copyLine[0],
CFlags,
pipInstallLine,
"ENV CFLAGS=",
}, "\n"), nil
}

func (g *StandardGenerator) pipInstallStage(aptInstalls string) (string, error) {
installCog, err := g.installCog()
if err != nil {
Expand Down
99 changes: 35 additions & 64 deletions pkg/dockerfile/standard_generator_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,20 +37,16 @@ func getWheelName() string {
return files[0].Name()
}

func testInstallCog(relativeTmpDir string, stripped bool, baseImage bool) string {
func testInstallCog(relativeTmpDir string, stripped bool) string {
wheel := getWheelName()
strippedCall := ""
if stripped {
strippedCall += " && find / -type f -name \"*python*.so\" -not -name \"*cpython*.so\" -exec strip -S {} \\;"
}
baseImageArgs := " -t /dep"
if baseImage {
baseImageArgs = ""
}
return fmt.Sprintf(`COPY %s/%s /tmp/%s
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install --no-cache-dir%s /tmp/%s 'pydantic<2'%s
ENV CFLAGS=`, relativeTmpDir, wheel, wheel, baseImageArgs, wheel, strippedCall)
RUN --mount=type=cache,target=/root/.cache/pip pip install --no-cache-dir /tmp/%s 'pydantic<2'%s
ENV CFLAGS=`, relativeTmpDir, wheel, wheel, wheel, strippedCall)
}

func testInstallPython(version string) string {
Expand Down Expand Up @@ -104,14 +100,12 @@ predict: predict.py:Predictor

expected := `#syntax=docker/dockerfile:1.4
FROM r8.im/replicate/cog-test-weights AS weights
FROM python:3.12 as deps
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
FROM python:3.12-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
ENV NVIDIA_DRIVER_CAPABILITIES=all
` + testTini() + `COPY --from=deps --link /dep /usr/local/lib/python3.12/site-packages
` + testTini() + testInstallCog(gen.relativeTmpDir, gen.strip) + `
RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
WORKDIR /src
EXPOSE 5000
Expand Down Expand Up @@ -139,18 +133,13 @@ predict: predict.py:Predictor

expected := `#syntax=docker/dockerfile:1.4
FROM r8.im/replicate/cog-test-weights AS weights
FROM python:3.12 as deps
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
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
ENV NVIDIA_DRIVER_CAPABILITIES=all
` + testTini() + testInstallPython("3.12") + "RUN rm -rf /usr/bin/python3 && ln -s `realpath \\`pyenv which python\\`` /usr/bin/python3 && chmod +x /usr/bin/python3" + `
RUN --mount=type=bind,from=deps,source=/dep,target=/dep \
cp -rf /dep/* $(pyenv prefix)/lib/python*/site-packages; \
cp -rf /dep/bin/* $(pyenv prefix)/bin; \
pyenv rehash
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
WORKDIR /src
EXPOSE 5000
Expand Down Expand Up @@ -187,19 +176,17 @@ predict: predict.py:Predictor

expected := `#syntax=docker/dockerfile:1.4
FROM r8.im/replicate/cog-test-weights AS weights
FROM python:3.12 as deps
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r /tmp/requirements.txt
ENV CFLAGS=
FROM python:3.12-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
ENV NVIDIA_DRIVER_CAPABILITIES=all
` + testTini() + `COPY --from=deps --link /dep /usr/local/lib/python3.12/site-packages
` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
ENV CFLAGS=
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
RUN cowsay moo
WORKDIR /src
Expand Down Expand Up @@ -243,24 +230,18 @@ predict: predict.py:Predictor

expected := `#syntax=docker/dockerfile:1.4
FROM r8.im/replicate/cog-test-weights AS weights
FROM python:3.12 as deps
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r /tmp/requirements.txt
ENV CFLAGS=
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
ENV NVIDIA_DRIVER_CAPABILITIES=all
` + testTini() +
testInstallPython("3.12") + "RUN rm -rf /usr/bin/python3 && ln -s `realpath \\`pyenv which python\\`` /usr/bin/python3 && chmod +x /usr/bin/python3" + `
RUN --mount=type=bind,from=deps,source=/dep,target=/dep \
cp -rf /dep/* $(pyenv prefix)/lib/python*/site-packages; \
cp -rf /dep/bin/* $(pyenv prefix)/bin; \
pyenv rehash
` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallPython("3.12") + "RUN rm -rf /usr/bin/python3 && ln -s `realpath \\`pyenv which python\\`` /usr/bin/python3 && chmod +x /usr/bin/python3" + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
ENV CFLAGS=
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
RUN cowsay moo
WORKDIR /src
Expand Down Expand Up @@ -299,15 +280,13 @@ build:

expected := `#syntax=docker/dockerfile:1.4
FROM r8.im/replicate/cog-test-weights AS weights
FROM python:3.12 as deps
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
FROM python:3.12-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
ENV NVIDIA_DRIVER_CAPABILITIES=all
` + testTini() + `COPY --from=deps --link /dep /usr/local/lib/python3.12/site-packages
` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
RUN cowsay moo
WORKDIR /src
Expand Down Expand Up @@ -335,7 +314,7 @@ build:
_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test")
require.NoError(t, err)
fmt.Println(actual)
require.Contains(t, actual, `pip install -t /dep -r /tmp/requirements.txt`)
require.Contains(t, actual, `pip install -r /tmp/requirements.txt`)
}

// mockFileInfo is a test type to mock os.FileInfo
Expand Down Expand Up @@ -409,24 +388,18 @@ COPY root-large /src/root-large`
// model copy should be run before dependency install and code copy
expected = `#syntax=docker/dockerfile:1.4
FROM r8.im/replicate/cog-test-weights AS weights
FROM python:3.12 as deps
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r /tmp/requirements.txt
ENV CFLAGS=
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
ENV NVIDIA_DRIVER_CAPABILITIES=all
` + testTini() +
testInstallPython("3.12") + `RUN rm -rf /usr/bin/python3 && ln -s ` + "`realpath \\`pyenv which python\\`` /usr/bin/python3 && chmod +x /usr/bin/python3" + `
RUN --mount=type=bind,from=deps,source=/dep,target=/dep \
cp -rf /dep/* $(pyenv prefix)/lib/python*/site-packages; \
cp -rf /dep/bin/* $(pyenv prefix)/bin; \
pyenv rehash
` + testTini() + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*` + `
` + testInstallPython("3.12") + `RUN rm -rf /usr/bin/python3 && ln -s ` + "`realpath \\`pyenv which python\\`` /usr/bin/python3 && chmod +x /usr/bin/python3" + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
ENV CFLAGS=
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
RUN cowsay moo
COPY --from=weights --link /src/checkpoints /src/checkpoints
Expand Down Expand Up @@ -493,14 +466,12 @@ predict: predict.py:Predictor
require.NoError(t, err)

expected := `#syntax=docker/dockerfile:1.4
FROM python:3.12 as deps
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
FROM python:3.12-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
ENV NVIDIA_DRIVER_CAPABILITIES=all
` + testTini() + `COPY --from=deps --link /dep /usr/local/lib/python3.12/site-packages
` + testTini() + testInstallCog(gen.relativeTmpDir, gen.strip) + `
RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
WORKDIR /src
EXPOSE 5000
Expand Down Expand Up @@ -530,7 +501,7 @@ predict: predict.py:Predictor
expected := `#syntax=docker/dockerfile:1.4
FROM r8.im/replicate/cog-test-weights AS weights
FROM r8.im/cog-base:python3.12
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
WORKDIR /src
EXPOSE 5000
CMD ["python", "-m", "cog.server.http"]
Expand Down Expand Up @@ -568,7 +539,7 @@ predict: predict.py:Predictor
FROM r8.im/replicate/cog-test-weights AS weights
FROM r8.im/cog-base:python3.12
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
Expand Down Expand Up @@ -624,7 +595,7 @@ predict: predict.py:Predictor
FROM r8.im/replicate/cog-test-weights AS weights
FROM r8.im/cog-base:cuda11.8-python3.11-torch%s
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
`+testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage())+`
`+testInstallCog(gen.relativeTmpDir, gen.strip)+`
COPY `+gen.relativeTmpDir+`/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
Expand Down Expand Up @@ -677,7 +648,7 @@ predict: predict.py:Predictor
FROM r8.im/replicate/cog-test-weights AS weights
FROM r8.im/cog-base:cuda11.8-python3.12-torch2.3.1
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt
Expand Down Expand Up @@ -729,7 +700,7 @@ predict: predict.py:Predictor
FROM r8.im/replicate/cog-test-weights AS weights
FROM r8.im/cog-base:cuda11.8-python3.12-torch2.3.1
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt && find / -type f -name "*python*.so" -not -name "*cpython*.so" -exec strip -S {} \;
Expand Down Expand Up @@ -813,7 +784,7 @@ predict: predict.py:Predictor
FROM r8.im/replicate/cog-test-weights AS weights
FROM r8.im/cog-base:cuda11.8-python3.12-torch2.3.1
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy cowsay && rm -rf /var/lib/apt/lists/*
` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
` + testInstallCog(gen.relativeTmpDir, gen.strip) + `
COPY ` + gen.relativeTmpDir + `/requirements.txt /tmp/requirements.txt
ENV CFLAGS="-O3 -funroll-loops -fno-strict-aliasing -flto -S"
RUN --mount=type=cache,target=/root/.cache/pip pip install -r /tmp/requirements.txt && find / -type f -name "*python*.so" -not -name "*cpython*.so" -exec strip -S {} \;
Expand Down

0 comments on commit 9199ee5

Please sign in to comment.