Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Multistage pip install #1233

Merged
merged 16 commits into from
Sep 5, 2023
Merged

Multistage pip install #1233

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
0a3664c
initial draft
technillogue Jul 26, 2023
cd72aed
move install cog before install user deps
technillogue Jul 28, 2023
e463d9c
fix tests
technillogue Jul 28, 2023
0e9464c
Merge branch 'main' into syl/pip-install-stage
technillogue Aug 7, 2023
eaa9688
try different copy dest to fix integration tests
technillogue Aug 7, 2023
59a0c16
fix tests
technillogue Aug 10, 2023
c4ae54b
note why test is broken
technillogue Aug 11, 2023
c01fac0
try to determine the actual location of site-packages in the final image
technillogue Aug 23, 2023
e5209f2
check if a different version would pass the tests
technillogue Aug 25, 2023
6affac8
also update go tests
technillogue Aug 25, 2023
050756b
clean comments
technillogue Aug 25, 2023
94ce8f0
notes about resolving minor versions for pyenv
technillogue Aug 28, 2023
c787ba7
fmt
technillogue Aug 28, 2023
caa7f6c
try using copying dependencies to a known place, then symlink to wher…
technillogue Aug 30, 2023
817134a
remove dead code
technillogue Aug 30, 2023
9c95d01
reword comment
technillogue Sep 5, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 47 additions & 27 deletions pkg/dockerfile/generator.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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`,
Expand Down Expand Up @@ -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
Expand All @@ -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,
}

Expand Down Expand Up @@ -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) {
Expand All @@ -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

Expand Down
55 changes: 34 additions & 21 deletions pkg/dockerfile/generator_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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"]
Expand All @@ -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"]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"]
Expand Down