diff --git a/pkg/dockerfile/base.go b/pkg/dockerfile/base.go index 2ad4b47677..6c38f2ca96 100644 --- a/pkg/dockerfile/base.go +++ b/pkg/dockerfile/base.go @@ -183,7 +183,7 @@ func (g *BaseImageGenerator) GenerateDockerfile() (string, error) { return "", err } useCogBaseImage := false - generator.useCogBaseImage = &useCogBaseImage + generator.SetUseCogBaseImagePtr(&useCogBaseImage) dockerfile, err := generator.generateInitialSteps() if err != nil { diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go index 2d37bb1e2e..b8b0784f9d 100644 --- a/pkg/dockerfile/generator.go +++ b/pkg/dockerfile/generator.go @@ -1,689 +1,19 @@ package dockerfile -import ( - "embed" - "fmt" - "os" - "path" - "path/filepath" - "runtime" - "strings" - "time" - - "github.com/replicate/cog/pkg/config" - "github.com/replicate/cog/pkg/util/console" - "github.com/replicate/cog/pkg/util/slices" - "github.com/replicate/cog/pkg/util/version" - "github.com/replicate/cog/pkg/weights" -) - -//go:embed embed/*.whl -var cogEmbed embed.FS - -const DockerignoreHeader = `# generated by replicate/cog -__pycache__ -*.pyc -*.pyo -*.pyd -.Python -env -pip-log.txt -pip-delete-this-directory.txt -.tox -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.log -.git -.mypy_cache -.pytest_cache -.hypothesis -` -const LDConfigCacheBuildCommand = "RUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig" -const StripDebugSymbolsCommand = "find / -type f -name \"*python*.so\" -not -name \"*cpython*.so\" -exec strip -S {} \\;" -const CFlags = "ENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"" -const PrecompilePythonCommand = "RUN find / -type f -name \"*.py[co]\" -delete && find / -type f -name \"*.py\" -exec touch -t 197001010000 {} \\; && find / -type f -name \"*.py\" -printf \"%h\\n\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0" - -type Generator struct { - Config *config.Config - Dir string - - // these are here to make this type testable - GOOS string - GOARCH string - - useCudaBaseImage bool - useCogBaseImage *bool - strip bool - precompile bool - - // absolute path to tmpDir, a directory that will be cleaned up - tmpDir string - // tmpDir relative to Dir - relativeTmpDir string - - fileWalker weights.FileWalker - - modelDirs []string - modelFiles []string - - pythonRequirementsContents string -} - -func NewGenerator(config *config.Config, dir string) (*Generator, error) { - rootTmp := path.Join(dir, ".cog/tmp") - if err := os.MkdirAll(rootTmp, 0o755); err != nil { - return nil, err - } - // tmpDir ends up being something like dir/.cog/tmp/build20240620123456.000000 - now := time.Now().Format("20060102150405.000000") - tmpDir, err := os.MkdirTemp(rootTmp, "build"+now) - if err != nil { - return nil, err - } - // tmpDir, but without dir prefix. This is the path used in the Dockerfile. - relativeTmpDir, err := filepath.Rel(dir, tmpDir) - if err != nil { - return nil, err - } - - return &Generator{ - Config: config, - Dir: dir, - GOOS: runtime.GOOS, - GOARCH: runtime.GOOS, - tmpDir: tmpDir, - relativeTmpDir: relativeTmpDir, - fileWalker: filepath.Walk, - useCudaBaseImage: true, - useCogBaseImage: nil, - strip: false, - precompile: false, - }, nil -} - -func (g *Generator) SetUseCudaBaseImage(argumentValue string) { - // "false" -> false, "true" -> true, "auto" -> true, "asdf" -> true - g.useCudaBaseImage = argumentValue != "false" -} - -func (g *Generator) SetUseCogBaseImage(useCogBaseImage bool) { - g.useCogBaseImage = new(bool) - *g.useCogBaseImage = useCogBaseImage -} - -func (g *Generator) IsUsingCogBaseImage() bool { - useCogBaseImage := g.useCogBaseImage - if useCogBaseImage != nil { - return *useCogBaseImage - } - return true -} - -func (g *Generator) SetStrip(strip bool) { - g.strip = strip -} - -func (g *Generator) SetPrecompile(precompile bool) { - g.precompile = precompile -} - -func (g *Generator) generateInitialSteps() (string, error) { - baseImage, err := g.BaseImage() - if err != nil { - return "", err - } - installPython, err := g.installPython() - if err != nil { - return "", err - } - aptInstalls, err := g.aptInstalls() - if err != nil { - return "", err - } - runCommands, err := g.runCommands() - 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, - aptInstalls, - installCog, - pipInstalls, - } - if g.precompile { - steps = append(steps, PrecompilePythonCommand) - } - steps = append(steps, runCommands) - - 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(), - installPython, - g.copyPipPackagesFromInstallStage(), - } - if g.precompile { - steps = append(steps, PrecompilePythonCommand) - } - steps = append(steps, LDConfigCacheBuildCommand, runCommands) - - return joinStringsWithoutLineSpace(steps), nil -} - -func (g *Generator) GenerateModelBase() (string, error) { - initialSteps, err := g.generateInitialSteps() - if err != nil { - return "", err - } - return strings.Join([]string{ - initialSteps, - `WORKDIR /src`, - `EXPOSE 5000`, - `CMD ["python", "-m", "cog.server.http"]`, - }, "\n"), nil -} - -// GenerateDockerfileWithoutSeparateWeights generates a Dockerfile that doesn't write model weights to a separate layer. -func (g *Generator) GenerateDockerfileWithoutSeparateWeights() (string, error) { - base, err := g.GenerateModelBase() - if err != nil { - return "", err - } - return joinStringsWithoutLineSpace([]string{ - base, - `COPY . /src`, - }), nil -} - -// GenerateModelBaseWithSeparateWeights creates the Dockerfile and .dockerignore file contents for model weights -// It returns four values: -// - weightsBase: The base image used for Dockerfile generation for model weights. -// - dockerfile: A string that represents the Dockerfile content generated by the function. -// - dockerignoreContents: A string that represents the .dockerignore content. -// - err: An error object if an error occurred during Dockerfile generation; otherwise nil. -func (g *Generator) GenerateModelBaseWithSeparateWeights(imageName string) (weightsBase string, dockerfile string, dockerignoreContents string, err error) { - weightsBase, g.modelDirs, g.modelFiles, err = g.generateForWeights() - if err != nil { - return "", "", "", fmt.Errorf("Failed to generate Dockerfile for model weights files: %w", err) - } - initialSteps, err := g.generateInitialSteps() - if err != nil { - return "", "", "", err - } - - // Inject weights base image into initial steps so we can COPY from it - base := []string{} - initialStepsLines := strings.Split(initialSteps, "\n") - for i, line := range initialStepsLines { - if strings.HasPrefix(line, "FROM ") { - base = append(base, fmt.Sprintf("FROM %s AS %s", imageName+"-weights", "weights")) - base = append(base, initialStepsLines[i:]...) - break - } else { - base = append(base, line) - } - } - - for _, p := range append(g.modelDirs, g.modelFiles...) { - base = append(base, "COPY --from=weights --link "+path.Join("/src", p)+" "+path.Join("/src", p)) - } - - base = append(base, - `WORKDIR /src`, - `EXPOSE 5000`, - `CMD ["python", "-m", "cog.server.http"]`, - `COPY . /src`, - ) - - dockerignoreContents = makeDockerignoreForWeights(g.modelDirs, g.modelFiles) - return weightsBase, joinStringsWithoutLineSpace(base), dockerignoreContents, nil -} - -func (g *Generator) generateForWeights() (string, []string, []string, error) { - modelDirs, modelFiles, err := weights.FindWeights(g.fileWalker) - if err != nil { - return "", nil, nil, err - } - // generate dockerfile to store these model weights files - dockerfileContents := `#syntax=docker/dockerfile:1.4 -FROM scratch -` - for _, p := range append(modelDirs, modelFiles...) { - dockerfileContents += fmt.Sprintf("\nCOPY %s %s", p, path.Join("/src", p)) - } - - return dockerfileContents, modelDirs, modelFiles, nil -} - -func makeDockerignoreForWeights(dirs, files []string) string { - var contents string - for _, p := range dirs { - contents += fmt.Sprintf("%[1]s\n%[1]s/**/*\n", p) - } - for _, p := range files { - contents += fmt.Sprintf("%[1]s\n", p) - } - return DockerignoreHeader + contents -} - -func (g *Generator) Cleanup() error { - if err := os.RemoveAll(g.tmpDir); err != nil { - return fmt.Errorf("Failed to clean up %s: %w", g.tmpDir, err) - } - return nil -} - -func (g *Generator) BaseImage() (string, error) { - if g.IsUsingCogBaseImage() { - baseImage, err := g.determineBaseImageName() - if err == nil || g.useCogBaseImage != nil { - return baseImage, err - } - console.Warnf("Could not find a suitable base image, continuing without base image support (%v).", err) - if g.useCogBaseImage == nil { - g.useCogBaseImage = new(bool) - *g.useCogBaseImage = false - } - } - - if g.Config.Build.GPU && g.useCudaBaseImage { - return g.Config.CUDABaseImageTag() - } - return "python:" + g.Config.Build.PythonVersion + "-slim", nil -} - -func (g *Generator) preamble() string { - return `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` -} - -func (g *Generator) installTini() string { - // Install tini as the image entrypoint to provide signal handling and process - // reaping appropriate for PID 1. - // - // N.B. If you remove/change this, consider removing/changing the `has_init` - // image label applied in image/build.go. - lines := []string{ - `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked set -eux; \ -apt-get update -qq && \ -apt-get install -qqy --no-install-recommends curl; \ -rm -rf /var/lib/apt/lists/*; \ -TINI_VERSION=v0.19.0; \ -TINI_ARCH="$(dpkg --print-architecture)"; \ -curl -sSL -o /sbin/tini "https://github.com/krallin/tini/releases/download/${TINI_VERSION}/tini-${TINI_ARCH}"; \ -chmod +x /sbin/tini`, - `ENTRYPOINT ["/sbin/tini", "--"]`, - } - return strings.Join(lines, "\n") -} - -func (g *Generator) aptInstalls() (string, error) { - packages := g.Config.Build.SystemPackages - if len(packages) == 0 { - return "", nil - } - - if g.IsUsingCogBaseImage() { - packages = slices.FilterString(packages, func(pkg string) bool { - return !slices.ContainsString(baseImageSystemPackages, pkg) - }) - } - - return "RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy " + - strings.Join(packages, " ") + - " && rm -rf /var/lib/apt/lists/*", nil -} - -func (g *Generator) installPython() (string, error) { - if g.Config.Build.GPU && g.useCudaBaseImage && !g.IsUsingCogBaseImage() { - return g.installPythonCUDA() - } - return "", nil -} - -func (g *Generator) installPythonCUDA() (string, error) { - // TODO: check that python version is valid - - py := g.Config.Build.PythonVersion - return `ENV PATH="/root/.pyenv/shims:/root/.pyenv/bin:$PATH" -RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy --no-install-recommends \ - make \ - build-essential \ - libssl-dev \ - zlib1g-dev \ - libbz2-dev \ - libreadline-dev \ - libsqlite3-dev \ - wget \ - curl \ - llvm \ - libncurses5-dev \ - libncursesw5-dev \ - xz-utils \ - tk-dev \ - libffi-dev \ - liblzma-dev \ - git \ - ca-certificates \ - && rm -rf /var/lib/apt/lists/* -` + fmt.Sprintf(` -RUN --mount=type=cache,target=/root/.cache/pip curl -s -S -L https://raw.githubusercontent.com/pyenv/pyenv-installer/master/bin/pyenv-installer | bash && \ - git clone https://github.com/momo-lab/pyenv-install-latest.git "$(pyenv root)"/plugins/pyenv-install-latest && \ - export PYTHON_CONFIGURE_OPTS='--enable-optimizations --with-lto' && \ - export PYTHON_CFLAGS='-O3' && \ - pyenv install-latest "%s" && \ - pyenv global $(pyenv install-latest --print "%s") && \ - pip install "wheel<1"`, py, py) + ` -RUN rm -rf /usr/bin/python3 && ln -s ` + "`realpath \\`pyenv which python\\`` /usr/bin/python3 && chmod +x /usr/bin/python3", 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) { - files, err := cogEmbed.ReadDir("embed") - if err != nil { - return "", err - } - if len(files) != 1 { - return "", fmt.Errorf("should only have one cog wheel embedded") - } - filename := files[0].Name() - data, err := cogEmbed.ReadFile("embed/" + filename) - if err != nil { - return "", err - } - lines, containerPath, err := g.writeTemp(filename, data) - if err != nil { - 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 - pipInstallLine += " 'pydantic<2'" - if g.strip { - pipInstallLine += " && " + StripDebugSymbolsCommand - } - lines = append(lines, CFlags, pipInstallLine, "ENV CFLAGS=") - return strings.Join(lines, "\n"), nil -} - -func (g *Generator) pipInstalls() (string, error) { - var err error - includePackages := []string{} - if torchVersion, ok := g.Config.TorchVersion(); ok { - includePackages = []string{"torch==" + torchVersion} - } - if torchvisionVersion, ok := g.Config.TorchvisionVersion(); ok { - includePackages = append(includePackages, "torchvision=="+torchvisionVersion) - } - if torchaudioVersion, ok := g.Config.TorchaudioVersion(); ok { - includePackages = append(includePackages, "torchaudio=="+torchaudioVersion) - } - g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, includePackages) - if err != nil { - return "", err - } - - if strings.Trim(g.pythonRequirementsContents, "") == "" { - return "", nil - } - - console.Debugf("Generated requirements.txt:\n%s", g.pythonRequirementsContents) - 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 -r " + containerPath - if g.strip { - pipInstallLine += " && " + StripDebugSymbolsCommand - } - return strings.Join([]string{ - copyLine[0], - CFlags, - pipInstallLine, - "ENV CFLAGS=", - }, "\n"), nil -} - -func (g *Generator) pipInstallStage(aptInstalls string) (string, error) { - installCog, err := g.installCog() - if err != nil { - return "", err - } - g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, []string{}) - if err != nil { - return "", err - } - - pipStageImage := "python:" + g.Config.Build.PythonVersion - if strings.Trim(g.pythonRequirementsContents, "") == "" { - return `FROM ` + pipStageImage + ` as deps -` + aptInstalls + ` -` + installCog, nil - } - - console.Debugf("Generated requirements.txt:\n%s", g.pythonRequirementsContents) - copyLine, containerPath, err := g.writeTemp("requirements.txt", []byte(g.pythonRequirementsContents)) - if err != nil { - return "", err - } - - // Not slim, so that we can compile wheels - fromLine := `FROM ` + pipStageImage + ` as deps` - // Sometimes, in order to run `pip install` successfully, some system packages need to be installed - // or some other change needs to happen - // this is a bodge to support that - // it will be reverted when we add custom dockerfiles - buildStageDeps := os.Getenv("COG_EXPERIMENTAL_BUILD_STAGE_DEPS") - if buildStageDeps != "" { - fromLine = fromLine + "\nRUN " + buildStageDeps - } - - pipInstallLine := "RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r " + containerPath - if g.strip { - pipInstallLine += " && " + StripDebugSymbolsCommand - } - - lines := []string{ - fromLine, - aptInstalls, - installCog, - copyLine[0], - CFlags, - pipInstallLine, - "ENV CFLAGS=", - } - return strings.Join(lines, "\n"), nil -} - -// copyPipPackagesFromInstallStage copies the Python dependencies installed in the deps stage into the main image -func (g *Generator) copyPipPackagesFromInstallStage() 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 || g.IsUsingCogBaseImage()) { - // this requires buildkit! - // we should check for buildkit and otherwise revert to symlinks or copying into /src - // we mount to avoid copying, which avoids having two copies in this layer - return ` -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 -` - } - - return "COPY --from=deps --link /dep /usr/local/lib/python" + py + "/site-packages" -} - -func (g *Generator) runCommands() (string, error) { - runCommands := g.Config.Build.Run - - // For backwards compatibility - for _, command := range g.Config.Build.PreInstall { - runCommands = append(runCommands, config.RunItem{Command: command}) - } - - lines := []string{} - for _, run := range runCommands { - command := strings.TrimSpace(run.Command) - if strings.Contains(command, "\n") { - return "", fmt.Errorf(`One of the commands in 'run' contains a new line, which won't work. You need to create a new list item in YAML prefixed with '-' for each command. - -This is the offending line: %s`, command) - } - - if len(run.Mounts) > 0 { - mounts := []string{} - for _, mount := range run.Mounts { - if mount.Type == "secret" { - secretMount := fmt.Sprintf("--mount=type=secret,id=%s,target=%s", mount.ID, mount.Target) - mounts = append(mounts, secretMount) - } - } - lines = append(lines, fmt.Sprintf("RUN %s %s", strings.Join(mounts, " "), command)) - } else { - lines = append(lines, "RUN "+command) - } - } - return strings.Join(lines, "\n"), nil -} - -// writeTemp writes a temporary file that can be used as part of the build process -// It returns the lines to add to Dockerfile to make it available and the filename it ends up as inside the container -func (g *Generator) writeTemp(filename string, contents []byte) ([]string, string, error) { - path := filepath.Join(g.tmpDir, filename) - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err) - } - if err := os.WriteFile(path, contents, 0o644); err != nil { - return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err) - } - return []string{fmt.Sprintf("COPY %s /tmp/%s", filepath.Join(g.relativeTmpDir, filename), filename)}, "/tmp/" + filename, nil -} - -func joinStringsWithoutLineSpace(chunks []string) string { - lines := []string{} - for _, chunk := range chunks { - chunkLines := strings.Split(chunk, "\n") - lines = append(lines, chunkLines...) - } - return strings.Join(filterEmpty(lines), "\n") -} - -func filterEmpty(list []string) []string { - filtered := []string{} - for _, s := range list { - if s != "" { - filtered = append(filtered, s) - } - } - return filtered -} - -func (g *Generator) GenerateWeightsManifest() (*weights.Manifest, error) { - m := weights.NewManifest() - - for _, dir := range g.modelDirs { - err := g.fileWalker(dir, func(path string, info os.FileInfo, err error) error { - if err != nil { - return err - } - if info.IsDir() { - return nil - } - - return m.AddFile(path) - }) - if err != nil { - return nil, err - } - } - - for _, path := range g.modelFiles { - err := m.AddFile(path) - if err != nil { - return nil, err - } - } - - return m, nil -} - -func (g *Generator) determineBaseImageName() (string, error) { - var changed bool - var err error - - cudaVersion := g.Config.Build.CUDA - - pythonVersion := g.Config.Build.PythonVersion - pythonVersion, changed, err = stripPatchVersion(pythonVersion) - if err != nil { - return "", err - } - if changed { - console.Warnf("Stripping patch version from Python version %s to %s", g.Config.Build.PythonVersion, pythonVersion) - } - - torchVersion, _ := g.Config.TorchVersion() - - // validate that the base image configuration exists - imageGenerator, err := NewBaseImageGenerator(cudaVersion, pythonVersion, torchVersion) - if err != nil { - return "", err - } - baseImage := BaseImageName(imageGenerator.cudaVersion, imageGenerator.pythonVersion, imageGenerator.torchVersion) - return baseImage, nil -} - -func stripPatchVersion(versionString string) (string, bool, error) { - if versionString == "" { - return "", false, nil - } - - v, err := version.NewVersion(versionString) - if err != nil { - return "", false, fmt.Errorf("Invalid version: %s", versionString) - } - - strippedVersion := fmt.Sprintf("%d.%d", v.Major, v.Minor) - changed := strippedVersion != versionString - - return strippedVersion, changed, nil +import "github.com/replicate/cog/pkg/weights" + +type Generator interface { + generateInitialSteps() (string, error) + SetUseCogBaseImage(bool) + SetUseCogBaseImagePtr(*bool) + GenerateModelBaseWithSeparateWeights(string) (string, string, string, error) + Cleanup() error + SetStrip(bool) + SetPrecompile(bool) + SetUseCudaBaseImage(string) + IsUsingCogBaseImage() bool + BaseImage() (string, error) + GenerateWeightsManifest() (*weights.Manifest, error) + GenerateDockerfileWithoutSeparateWeights() (string, error) + GenerateModelBase() (string, error) } diff --git a/pkg/dockerfile/generator_factory.go b/pkg/dockerfile/generator_factory.go new file mode 100644 index 0000000000..fffd6b87a0 --- /dev/null +++ b/pkg/dockerfile/generator_factory.go @@ -0,0 +1,9 @@ +package dockerfile + +import ( + "github.com/replicate/cog/pkg/config" +) + +func NewGenerator(config *config.Config, dir string) (Generator, error) { + return NewStandardGenerator(config, dir) +} diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go new file mode 100644 index 0000000000..489308f47b --- /dev/null +++ b/pkg/dockerfile/standard_generator.go @@ -0,0 +1,693 @@ +package dockerfile + +import ( + "embed" + "fmt" + "os" + "path" + "path/filepath" + "runtime" + "strings" + "time" + + "github.com/replicate/cog/pkg/config" + "github.com/replicate/cog/pkg/util/console" + "github.com/replicate/cog/pkg/util/slices" + "github.com/replicate/cog/pkg/util/version" + "github.com/replicate/cog/pkg/weights" +) + +//go:embed embed/*.whl +var cogEmbed embed.FS + +const DockerignoreHeader = `# generated by replicate/cog +__pycache__ +*.pyc +*.pyo +*.pyd +.Python +env +pip-log.txt +pip-delete-this-directory.txt +.tox +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.log +.git +.mypy_cache +.pytest_cache +.hypothesis +` +const LDConfigCacheBuildCommand = "RUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig" +const StripDebugSymbolsCommand = "find / -type f -name \"*python*.so\" -not -name \"*cpython*.so\" -exec strip -S {} \\;" +const CFlags = "ENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\"" +const PrecompilePythonCommand = "RUN find / -type f -name \"*.py[co]\" -delete && find / -type f -name \"*.py\" -exec touch -t 197001010000 {} \\; && find / -type f -name \"*.py\" -printf \"%h\\n\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0" + +type StandardGenerator struct { + Config *config.Config + Dir string + + // these are here to make this type testable + GOOS string + GOARCH string + + useCudaBaseImage bool + useCogBaseImage *bool + strip bool + precompile bool + + // absolute path to tmpDir, a directory that will be cleaned up + tmpDir string + // tmpDir relative to Dir + relativeTmpDir string + + fileWalker weights.FileWalker + + modelDirs []string + modelFiles []string + + pythonRequirementsContents string +} + +func NewStandardGenerator(config *config.Config, dir string) (*StandardGenerator, error) { + rootTmp := path.Join(dir, ".cog/tmp") + if err := os.MkdirAll(rootTmp, 0o755); err != nil { + return nil, err + } + // tmpDir ends up being something like dir/.cog/tmp/build20240620123456.000000 + now := time.Now().Format("20060102150405.000000") + tmpDir, err := os.MkdirTemp(rootTmp, "build"+now) + if err != nil { + return nil, err + } + // tmpDir, but without dir prefix. This is the path used in the Dockerfile. + relativeTmpDir, err := filepath.Rel(dir, tmpDir) + if err != nil { + return nil, err + } + + return &StandardGenerator{ + Config: config, + Dir: dir, + GOOS: runtime.GOOS, + GOARCH: runtime.GOOS, + tmpDir: tmpDir, + relativeTmpDir: relativeTmpDir, + fileWalker: filepath.Walk, + useCudaBaseImage: true, + useCogBaseImage: nil, + strip: false, + precompile: false, + }, nil +} + +func (g *StandardGenerator) SetUseCudaBaseImage(argumentValue string) { + // "false" -> false, "true" -> true, "auto" -> true, "asdf" -> true + g.useCudaBaseImage = argumentValue != "false" +} + +func (g *StandardGenerator) SetUseCogBaseImage(useCogBaseImage bool) { + g.useCogBaseImage = new(bool) + *g.useCogBaseImage = useCogBaseImage +} + +func (g *StandardGenerator) SetUseCogBaseImagePtr(useCogBaseImage *bool) { + g.useCogBaseImage = useCogBaseImage +} + +func (g *StandardGenerator) IsUsingCogBaseImage() bool { + useCogBaseImage := g.useCogBaseImage + if useCogBaseImage != nil { + return *useCogBaseImage + } + return true +} + +func (g *StandardGenerator) SetStrip(strip bool) { + g.strip = strip +} + +func (g *StandardGenerator) SetPrecompile(precompile bool) { + g.precompile = precompile +} + +func (g *StandardGenerator) generateInitialSteps() (string, error) { + baseImage, err := g.BaseImage() + if err != nil { + return "", err + } + installPython, err := g.installPython() + if err != nil { + return "", err + } + aptInstalls, err := g.aptInstalls() + if err != nil { + return "", err + } + runCommands, err := g.runCommands() + 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, + aptInstalls, + installCog, + pipInstalls, + } + if g.precompile { + steps = append(steps, PrecompilePythonCommand) + } + steps = append(steps, runCommands) + + 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(), + installPython, + g.copyPipPackagesFromInstallStage(), + } + if g.precompile { + steps = append(steps, PrecompilePythonCommand) + } + steps = append(steps, LDConfigCacheBuildCommand, runCommands) + + return joinStringsWithoutLineSpace(steps), nil +} + +func (g *StandardGenerator) GenerateModelBase() (string, error) { + initialSteps, err := g.generateInitialSteps() + if err != nil { + return "", err + } + return strings.Join([]string{ + initialSteps, + `WORKDIR /src`, + `EXPOSE 5000`, + `CMD ["python", "-m", "cog.server.http"]`, + }, "\n"), nil +} + +// GenerateDockerfileWithoutSeparateWeights generates a Dockerfile that doesn't write model weights to a separate layer. +func (g *StandardGenerator) GenerateDockerfileWithoutSeparateWeights() (string, error) { + base, err := g.GenerateModelBase() + if err != nil { + return "", err + } + return joinStringsWithoutLineSpace([]string{ + base, + `COPY . /src`, + }), nil +} + +// GenerateModelBaseWithSeparateWeights creates the Dockerfile and .dockerignore file contents for model weights +// It returns four values: +// - weightsBase: The base image used for Dockerfile generation for model weights. +// - dockerfile: A string that represents the Dockerfile content generated by the function. +// - dockerignoreContents: A string that represents the .dockerignore content. +// - err: An error object if an error occurred during Dockerfile generation; otherwise nil. +func (g *StandardGenerator) GenerateModelBaseWithSeparateWeights(imageName string) (weightsBase string, dockerfile string, dockerignoreContents string, err error) { + weightsBase, g.modelDirs, g.modelFiles, err = g.generateForWeights() + if err != nil { + return "", "", "", fmt.Errorf("Failed to generate Dockerfile for model weights files: %w", err) + } + initialSteps, err := g.generateInitialSteps() + if err != nil { + return "", "", "", err + } + + // Inject weights base image into initial steps so we can COPY from it + base := []string{} + initialStepsLines := strings.Split(initialSteps, "\n") + for i, line := range initialStepsLines { + if strings.HasPrefix(line, "FROM ") { + base = append(base, fmt.Sprintf("FROM %s AS %s", imageName+"-weights", "weights")) + base = append(base, initialStepsLines[i:]...) + break + } else { + base = append(base, line) + } + } + + for _, p := range append(g.modelDirs, g.modelFiles...) { + base = append(base, "COPY --from=weights --link "+path.Join("/src", p)+" "+path.Join("/src", p)) + } + + base = append(base, + `WORKDIR /src`, + `EXPOSE 5000`, + `CMD ["python", "-m", "cog.server.http"]`, + `COPY . /src`, + ) + + dockerignoreContents = makeDockerignoreForWeights(g.modelDirs, g.modelFiles) + return weightsBase, joinStringsWithoutLineSpace(base), dockerignoreContents, nil +} + +func (g *StandardGenerator) generateForWeights() (string, []string, []string, error) { + modelDirs, modelFiles, err := weights.FindWeights(g.fileWalker) + if err != nil { + return "", nil, nil, err + } + // generate dockerfile to store these model weights files + dockerfileContents := `#syntax=docker/dockerfile:1.4 +FROM scratch +` + for _, p := range append(modelDirs, modelFiles...) { + dockerfileContents += fmt.Sprintf("\nCOPY %s %s", p, path.Join("/src", p)) + } + + return dockerfileContents, modelDirs, modelFiles, nil +} + +func makeDockerignoreForWeights(dirs, files []string) string { + var contents string + for _, p := range dirs { + contents += fmt.Sprintf("%[1]s\n%[1]s/**/*\n", p) + } + for _, p := range files { + contents += fmt.Sprintf("%[1]s\n", p) + } + return DockerignoreHeader + contents +} + +func (g *StandardGenerator) Cleanup() error { + if err := os.RemoveAll(g.tmpDir); err != nil { + return fmt.Errorf("Failed to clean up %s: %w", g.tmpDir, err) + } + return nil +} + +func (g *StandardGenerator) BaseImage() (string, error) { + if g.IsUsingCogBaseImage() { + baseImage, err := g.determineBaseImageName() + if err == nil || g.useCogBaseImage != nil { + return baseImage, err + } + console.Warnf("Could not find a suitable base image, continuing without base image support (%v).", err) + if g.useCogBaseImage == nil { + g.useCogBaseImage = new(bool) + *g.useCogBaseImage = false + } + } + + if g.Config.Build.GPU && g.useCudaBaseImage { + return g.Config.CUDABaseImageTag() + } + return "python:" + g.Config.Build.PythonVersion + "-slim", nil +} + +func (g *StandardGenerator) preamble() string { + return `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` +} + +func (g *StandardGenerator) installTini() string { + // Install tini as the image entrypoint to provide signal handling and process + // reaping appropriate for PID 1. + // + // N.B. If you remove/change this, consider removing/changing the `has_init` + // image label applied in image/build.go. + lines := []string{ + `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked set -eux; \ +apt-get update -qq && \ +apt-get install -qqy --no-install-recommends curl; \ +rm -rf /var/lib/apt/lists/*; \ +TINI_VERSION=v0.19.0; \ +TINI_ARCH="$(dpkg --print-architecture)"; \ +curl -sSL -o /sbin/tini "https://github.com/krallin/tini/releases/download/${TINI_VERSION}/tini-${TINI_ARCH}"; \ +chmod +x /sbin/tini`, + `ENTRYPOINT ["/sbin/tini", "--"]`, + } + return strings.Join(lines, "\n") +} + +func (g *StandardGenerator) aptInstalls() (string, error) { + packages := g.Config.Build.SystemPackages + if len(packages) == 0 { + return "", nil + } + + if g.IsUsingCogBaseImage() { + packages = slices.FilterString(packages, func(pkg string) bool { + return !slices.ContainsString(baseImageSystemPackages, pkg) + }) + } + + return "RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy " + + strings.Join(packages, " ") + + " && rm -rf /var/lib/apt/lists/*", nil +} + +func (g *StandardGenerator) installPython() (string, error) { + if g.Config.Build.GPU && g.useCudaBaseImage && !g.IsUsingCogBaseImage() { + return g.installPythonCUDA() + } + return "", nil +} + +func (g *StandardGenerator) installPythonCUDA() (string, error) { + // TODO: check that python version is valid + + py := g.Config.Build.PythonVersion + return `ENV PATH="/root/.pyenv/shims:/root/.pyenv/bin:$PATH" +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy --no-install-recommends \ + make \ + build-essential \ + libssl-dev \ + zlib1g-dev \ + libbz2-dev \ + libreadline-dev \ + libsqlite3-dev \ + wget \ + curl \ + llvm \ + libncurses5-dev \ + libncursesw5-dev \ + xz-utils \ + tk-dev \ + libffi-dev \ + liblzma-dev \ + git \ + ca-certificates \ + && rm -rf /var/lib/apt/lists/* +` + fmt.Sprintf(` +RUN --mount=type=cache,target=/root/.cache/pip curl -s -S -L https://raw.githubusercontent.com/pyenv/pyenv-installer/master/bin/pyenv-installer | bash && \ + git clone https://github.com/momo-lab/pyenv-install-latest.git "$(pyenv root)"/plugins/pyenv-install-latest && \ + export PYTHON_CONFIGURE_OPTS='--enable-optimizations --with-lto' && \ + export PYTHON_CFLAGS='-O3' && \ + pyenv install-latest "%s" && \ + pyenv global $(pyenv install-latest --print "%s") && \ + pip install "wheel<1"`, py, py) + ` +RUN rm -rf /usr/bin/python3 && ln -s ` + "`realpath \\`pyenv which python\\`` /usr/bin/python3 && chmod +x /usr/bin/python3", 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 *StandardGenerator) installCog() (string, error) { + files, err := cogEmbed.ReadDir("embed") + if err != nil { + return "", err + } + if len(files) != 1 { + return "", fmt.Errorf("should only have one cog wheel embedded") + } + filename := files[0].Name() + data, err := cogEmbed.ReadFile("embed/" + filename) + if err != nil { + return "", err + } + lines, containerPath, err := g.writeTemp(filename, data) + if err != nil { + 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 + pipInstallLine += " 'pydantic<2'" + if g.strip { + pipInstallLine += " && " + StripDebugSymbolsCommand + } + lines = append(lines, CFlags, pipInstallLine, "ENV CFLAGS=") + return strings.Join(lines, "\n"), nil +} + +func (g *StandardGenerator) pipInstalls() (string, error) { + var err error + includePackages := []string{} + if torchVersion, ok := g.Config.TorchVersion(); ok { + includePackages = []string{"torch==" + torchVersion} + } + if torchvisionVersion, ok := g.Config.TorchvisionVersion(); ok { + includePackages = append(includePackages, "torchvision=="+torchvisionVersion) + } + if torchaudioVersion, ok := g.Config.TorchaudioVersion(); ok { + includePackages = append(includePackages, "torchaudio=="+torchaudioVersion) + } + g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, includePackages) + if err != nil { + return "", err + } + + if strings.Trim(g.pythonRequirementsContents, "") == "" { + return "", nil + } + + console.Debugf("Generated requirements.txt:\n%s", g.pythonRequirementsContents) + 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 -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 { + return "", err + } + g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, []string{}) + if err != nil { + return "", err + } + + pipStageImage := "python:" + g.Config.Build.PythonVersion + if strings.Trim(g.pythonRequirementsContents, "") == "" { + return `FROM ` + pipStageImage + ` as deps +` + aptInstalls + ` +` + installCog, nil + } + + console.Debugf("Generated requirements.txt:\n%s", g.pythonRequirementsContents) + copyLine, containerPath, err := g.writeTemp("requirements.txt", []byte(g.pythonRequirementsContents)) + if err != nil { + return "", err + } + + // Not slim, so that we can compile wheels + fromLine := `FROM ` + pipStageImage + ` as deps` + // Sometimes, in order to run `pip install` successfully, some system packages need to be installed + // or some other change needs to happen + // this is a bodge to support that + // it will be reverted when we add custom dockerfiles + buildStageDeps := os.Getenv("COG_EXPERIMENTAL_BUILD_STAGE_DEPS") + if buildStageDeps != "" { + fromLine = fromLine + "\nRUN " + buildStageDeps + } + + pipInstallLine := "RUN --mount=type=cache,target=/root/.cache/pip pip install -t /dep -r " + containerPath + if g.strip { + pipInstallLine += " && " + StripDebugSymbolsCommand + } + + lines := []string{ + fromLine, + aptInstalls, + installCog, + copyLine[0], + CFlags, + pipInstallLine, + "ENV CFLAGS=", + } + return strings.Join(lines, "\n"), nil +} + +// copyPipPackagesFromInstallStage copies the Python dependencies installed in the deps stage into the main image +func (g *StandardGenerator) copyPipPackagesFromInstallStage() 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 || g.IsUsingCogBaseImage()) { + // this requires buildkit! + // we should check for buildkit and otherwise revert to symlinks or copying into /src + // we mount to avoid copying, which avoids having two copies in this layer + return ` +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 +` + } + + return "COPY --from=deps --link /dep /usr/local/lib/python" + py + "/site-packages" +} + +func (g *StandardGenerator) runCommands() (string, error) { + runCommands := g.Config.Build.Run + + // For backwards compatibility + for _, command := range g.Config.Build.PreInstall { + runCommands = append(runCommands, config.RunItem{Command: command}) + } + + lines := []string{} + for _, run := range runCommands { + command := strings.TrimSpace(run.Command) + if strings.Contains(command, "\n") { + return "", fmt.Errorf(`One of the commands in 'run' contains a new line, which won't work. You need to create a new list item in YAML prefixed with '-' for each command. + +This is the offending line: %s`, command) + } + + if len(run.Mounts) > 0 { + mounts := []string{} + for _, mount := range run.Mounts { + if mount.Type == "secret" { + secretMount := fmt.Sprintf("--mount=type=secret,id=%s,target=%s", mount.ID, mount.Target) + mounts = append(mounts, secretMount) + } + } + lines = append(lines, fmt.Sprintf("RUN %s %s", strings.Join(mounts, " "), command)) + } else { + lines = append(lines, "RUN "+command) + } + } + return strings.Join(lines, "\n"), nil +} + +// writeTemp writes a temporary file that can be used as part of the build process +// It returns the lines to add to Dockerfile to make it available and the filename it ends up as inside the container +func (g *StandardGenerator) writeTemp(filename string, contents []byte) ([]string, string, error) { + path := filepath.Join(g.tmpDir, filename) + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err) + } + if err := os.WriteFile(path, contents, 0o644); err != nil { + return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err) + } + return []string{fmt.Sprintf("COPY %s /tmp/%s", filepath.Join(g.relativeTmpDir, filename), filename)}, "/tmp/" + filename, nil +} + +func joinStringsWithoutLineSpace(chunks []string) string { + lines := []string{} + for _, chunk := range chunks { + chunkLines := strings.Split(chunk, "\n") + lines = append(lines, chunkLines...) + } + return strings.Join(filterEmpty(lines), "\n") +} + +func filterEmpty(list []string) []string { + filtered := []string{} + for _, s := range list { + if s != "" { + filtered = append(filtered, s) + } + } + return filtered +} + +func (g *StandardGenerator) GenerateWeightsManifest() (*weights.Manifest, error) { + m := weights.NewManifest() + + for _, dir := range g.modelDirs { + err := g.fileWalker(dir, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + if info.IsDir() { + return nil + } + + return m.AddFile(path) + }) + if err != nil { + return nil, err + } + } + + for _, path := range g.modelFiles { + err := m.AddFile(path) + if err != nil { + return nil, err + } + } + + return m, nil +} + +func (g *StandardGenerator) determineBaseImageName() (string, error) { + var changed bool + var err error + + cudaVersion := g.Config.Build.CUDA + + pythonVersion := g.Config.Build.PythonVersion + pythonVersion, changed, err = stripPatchVersion(pythonVersion) + if err != nil { + return "", err + } + if changed { + console.Warnf("Stripping patch version from Python version %s to %s", g.Config.Build.PythonVersion, pythonVersion) + } + + torchVersion, _ := g.Config.TorchVersion() + + // validate that the base image configuration exists + imageGenerator, err := NewBaseImageGenerator(cudaVersion, pythonVersion, torchVersion) + if err != nil { + return "", err + } + baseImage := BaseImageName(imageGenerator.cudaVersion, imageGenerator.pythonVersion, imageGenerator.torchVersion) + return baseImage, nil +} + +func stripPatchVersion(versionString string) (string, bool, error) { + if versionString == "" { + return "", false, nil + } + + v, err := version.NewVersion(versionString) + if err != nil { + return "", false, fmt.Errorf("Invalid version: %s", versionString) + } + + strippedVersion := fmt.Sprintf("%d.%d", v.Major, v.Minor) + changed := strippedVersion != versionString + + return strippedVersion, changed, nil +} diff --git a/pkg/dockerfile/generator_test.go b/pkg/dockerfile/standard_generator_test.go similarity index 97% rename from pkg/dockerfile/generator_test.go rename to pkg/dockerfile/standard_generator_test.go index 6ad4b5ee82..e82c557936 100644 --- a/pkg/dockerfile/generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -96,7 +96,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(false) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -131,7 +131,7 @@ predict: predict.py:Predictor `)) require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(false) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -179,7 +179,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(false) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -235,7 +235,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(false) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -291,7 +291,7 @@ build: require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(false) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -383,7 +383,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(false) @@ -486,7 +486,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(false) actual, err := gen.GenerateDockerfileWithoutSeparateWeights() @@ -521,7 +521,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(true) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -558,7 +558,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(true) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -609,7 +609,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(true) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -667,7 +667,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(true) _, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test") @@ -718,7 +718,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(true) gen.SetStrip(true) @@ -801,7 +801,7 @@ predict: predict.py:Predictor require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete("")) - gen, err := NewGenerator(conf, tmpDir) + gen, err := NewStandardGenerator(conf, tmpDir) require.NoError(t, err) gen.SetUseCogBaseImage(true) gen.SetStrip(true)