diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go index f433e0b173..c3b093c2da 100644 --- a/pkg/cli/debug.go +++ b/pkg/cli/debug.go @@ -37,7 +37,7 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error { return err } - generator, err := dockerfile.NewGenerator(cfg, projectDir) + generator, err := dockerfile.NewGenerator(cfg, projectDir, false) if err != nil { return fmt.Errorf("Error creating Dockerfile generator: %w", err) } diff --git a/pkg/config/compatibility.go b/pkg/config/compatibility.go index fdbaaadf88..3b94ff9550 100644 --- a/pkg/config/compatibility.go +++ b/pkg/config/compatibility.go @@ -276,7 +276,7 @@ func CUDABaseImageFor(cuda string, cuDNN string) (string, error) { func tfGPUPackage(ver string, cuda string) (name string, cpuVersion string, err error) { for _, compat := range TFCompatibilityMatrix { if compat.TF == ver && version.Equal(compat.CUDA, cuda) { - name, cpuVersion, _, _, err = splitPinnedPythonRequirement(compat.TFGPUPackage) + name, cpuVersion, _, _, err = SplitPinnedPythonRequirement(compat.TFGPUPackage) return name, cpuVersion, err } } diff --git a/pkg/config/config.go b/pkg/config/config.go index f3e28eb7f6..3cf58b2f35 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -220,7 +220,7 @@ func (c *Config) cudaFromTF() (tfVersion string, tfCUDA string, tfCuDNN string, func (c *Config) pythonPackageVersion(name string) (version string, ok bool) { for _, pkg := range c.Build.pythonRequirementsContent { - pkgName, version, _, _, err := splitPinnedPythonRequirement(pkg) + pkgName, version, _, _, err := SplitPinnedPythonRequirement(pkg) if err != nil { // package is not in package==version format continue @@ -331,7 +331,11 @@ func (c *Config) PythonRequirementsForArch(goos string, goarch string, includePa includePackageNames := []string{} for _, pkg := range includePackages { - includePackageNames = append(includePackageNames, packageName(pkg)) + packageName, err := PackageName(pkg) + if err != nil { + return "", err + } + includePackageNames = append(includePackageNames, packageName) } // Include all the requirements and remove our include packages if they exist @@ -352,7 +356,7 @@ func (c *Config) PythonRequirementsForArch(goos string, goarch string, includePa } } - packageName := packageName(archPkg) + packageName, _ := PackageName(archPkg) if packageName != "" { foundIdx := -1 for i, includePkg := range includePackageNames { @@ -390,7 +394,7 @@ func (c *Config) PythonRequirementsForArch(goos string, goarch string, includePa // pythonPackageForArch takes a package==version line and // returns a package==version and index URL resolved to the correct GPU package for the given OS and architecture func (c *Config) pythonPackageForArch(pkg, goos, goarch string) (actualPackage string, findLinksList []string, extraIndexURLs []string, err error) { - name, version, findLinksList, extraIndexURLs, err := splitPinnedPythonRequirement(pkg) + name, version, findLinksList, extraIndexURLs, err := SplitPinnedPythonRequirement(pkg) if err != nil { // It's not pinned, so just return the line verbatim return pkg, []string{}, []string{}, nil @@ -562,50 +566,6 @@ Compatible cuDNN version is: %s`, c.Build.CuDNN, tfVersion, tfCuDNN) return nil } -// splitPythonPackage returns the name, version, findLinks, and extraIndexURLs from a requirements.txt line -// in the form name==version [--find-links=] [-f ] [--extra-index-url=] -func splitPinnedPythonRequirement(requirement string) (name string, version string, findLinks []string, extraIndexURLs []string, err error) { - pinnedPackageRe := regexp.MustCompile(`(?:([a-zA-Z0-9\-_]+)==([^ ]+)|--find-links=([^\s]+)|-f\s+([^\s]+)|--extra-index-url=([^\s]+))`) - - matches := pinnedPackageRe.FindAllStringSubmatch(requirement, -1) - if matches == nil { - return "", "", nil, nil, fmt.Errorf("Package %s is not in the expected format", requirement) - } - - nameFound := false - versionFound := false - - for _, match := range matches { - if match[1] != "" { - name = match[1] - nameFound = true - } - - if match[2] != "" { - version = match[2] - versionFound = true - } - - if match[3] != "" { - findLinks = append(findLinks, match[3]) - } - - if match[4] != "" { - findLinks = append(findLinks, match[4]) - } - - if match[5] != "" { - extraIndexURLs = append(extraIndexURLs, match[5]) - } - } - - if !nameFound || !versionFound { - return "", "", nil, nil, fmt.Errorf("Package name or version is missing in %s", requirement) - } - - return name, version, findLinks, extraIndexURLs, nil -} - func sliceContains(slice []string, s string) bool { for _, el := range slice { if el == s { @@ -614,11 +574,3 @@ func sliceContains(slice []string, s string) bool { } return false } - -func packageName(pipRequirement string) string { - match := PipPackageNameRegex.FindStringSubmatch(pipRequirement) - if len(match) <= 1 { - return "" - } - return match[1] -} diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 934fba5a70..9ec016cc85 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -691,7 +691,7 @@ func TestSplitPinnedPythonRequirement(t *testing.T) { } for _, tc := range testCases { - name, version, findLinks, extraIndexURLs, err := splitPinnedPythonRequirement(tc.input) + name, version, findLinks, extraIndexURLs, err := SplitPinnedPythonRequirement(tc.input) if tc.expectedError { require.Error(t, err) diff --git a/pkg/config/requirements.go b/pkg/config/requirements.go new file mode 100644 index 0000000000..1136659aa6 --- /dev/null +++ b/pkg/config/requirements.go @@ -0,0 +1,130 @@ +package config + +import ( + "bufio" + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" +) + +func GenerateRequirements(tmpDir string, config *Config) (string, error) { + // Deduplicate packages between the requirements.txt and the python packages directive. + packageNames := make(map[string]string) + + // Read the python packages configuration. + for _, requirement := range config.Build.PythonPackages { + packageName, err := PackageName(requirement) + if err != nil { + return "", err + } + packageNames[packageName] = requirement + } + + // Read the python requirements. + if config.Build.PythonRequirements != "" { + fh, err := os.Open(config.Build.PythonRequirements) + if err != nil { + return "", err + } + scanner := bufio.NewScanner(fh) + for scanner.Scan() { + requirement := scanner.Text() + packageName, err := PackageName(requirement) + if err != nil { + return "", err + } + packageNames[packageName] = requirement + } + } + + // If we don't have any packages skip further processing + if len(packageNames) == 0 { + return "", nil + } + + // Sort the package names by alphabetical order. + keys := make([]string, 0, len(packageNames)) + for k := range packageNames { + keys = append(keys, k) + } + sort.Strings(keys) + + // Render the expected contents + requirementsContent := "" + for _, k := range keys { + requirementsContent += packageNames[k] + "\n" + } + + // Check against the old requirements contents + requirementsFile := filepath.Join(tmpDir, "requirements.txt") + _, err := os.Stat(requirementsFile) + if !errors.Is(err, os.ErrNotExist) { + bytes, err := os.ReadFile(requirementsFile) + if err != nil { + return "", err + } + oldRequirementsContents := string(bytes) + if oldRequirementsContents == requirementsFile { + return requirementsFile, nil + } + } + + // Write out a new requirements file + err = os.WriteFile(requirementsFile, []byte(requirementsContent), 0o644) + if err != nil { + return "", err + } + return requirementsFile, nil +} + +// SplitPinnedPythonRequirement returns the name, version, findLinks, and extraIndexURLs from a requirements.txt line +// in the form name==version [--find-links=] [-f ] [--extra-index-url=] +func SplitPinnedPythonRequirement(requirement string) (name string, version string, findLinks []string, extraIndexURLs []string, err error) { + pinnedPackageRe := regexp.MustCompile(`(?:([a-zA-Z0-9\-_]+)==([^ ]+)|--find-links=([^\s]+)|-f\s+([^\s]+)|--extra-index-url=([^\s]+))`) + + matches := pinnedPackageRe.FindAllStringSubmatch(requirement, -1) + if matches == nil { + return "", "", nil, nil, fmt.Errorf("Package %s is not in the expected format", requirement) + } + + nameFound := false + versionFound := false + + for _, match := range matches { + if match[1] != "" { + name = match[1] + nameFound = true + } + + if match[2] != "" { + version = match[2] + versionFound = true + } + + if match[3] != "" { + findLinks = append(findLinks, match[3]) + } + + if match[4] != "" { + findLinks = append(findLinks, match[4]) + } + + if match[5] != "" { + extraIndexURLs = append(extraIndexURLs, match[5]) + } + } + + if !nameFound || !versionFound { + return "", "", nil, nil, fmt.Errorf("Package name or version is missing in %s", requirement) + } + + return name, version, findLinks, extraIndexURLs, nil +} + +func PackageName(pipRequirement string) (string, error) { + name, _, _, _, err := SplitPinnedPythonRequirement(pipRequirement) + return name, err +} diff --git a/pkg/config/requirements_test.go b/pkg/config/requirements_test.go new file mode 100644 index 0000000000..61e488319f --- /dev/null +++ b/pkg/config/requirements_test.go @@ -0,0 +1,21 @@ +package config + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGenerateRequirements(t *testing.T) { + tmpDir := t.TempDir() + build := Build{ + PythonPackages: []string{"torch==2.5.1"}, + } + config := Config{ + Build: &build, + } + requirementsFile, err := GenerateRequirements(tmpDir, &config) + require.NoError(t, err) + require.Equal(t, filepath.Join(tmpDir, "requirements.txt"), requirementsFile) +} diff --git a/pkg/docker/build.go b/pkg/docker/build.go index ebfabfa7a4..da3c8b2476 100644 --- a/pkg/docker/build.go +++ b/pkg/docker/build.go @@ -8,16 +8,22 @@ import ( "strings" "github.com/replicate/cog/pkg/config" + "github.com/replicate/cog/pkg/dockerfile" "github.com/replicate/cog/pkg/util" "github.com/replicate/cog/pkg/util/console" ) -func Build(dir, dockerfile, imageName string, secrets []string, noCache bool, progressOutput string, epoch int64) error { +func Build(dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, epoch int64) error { var args []string + userCache, err := dockerfile.UserCache() + if err != nil { + return err + } + args = append(args, - "buildx", "build", + "buildx", "build", "--build-context", "usercache="+userCache, ) if util.IsAppleSiliconMac(runtime.GOOS, runtime.GOARCH) { @@ -65,7 +71,7 @@ func Build(dir, dockerfile, imageName string, secrets []string, noCache bool, pr cmd.Dir = dir cmd.Stdout = os.Stderr // redirect stdout to stderr - build output is all messaging cmd.Stderr = os.Stderr - cmd.Stdin = strings.NewReader(dockerfile) + cmd.Stdin = strings.NewReader(dockerfileContents) console.Debug("$ " + strings.Join(cmd.Args, " ")) return cmd.Run() diff --git a/pkg/dockerfile/base.go b/pkg/dockerfile/base.go index 02b2b98fef..97b6e1a45b 100644 --- a/pkg/dockerfile/base.go +++ b/pkg/dockerfile/base.go @@ -178,7 +178,7 @@ func (g *BaseImageGenerator) GenerateDockerfile() (string, error) { return "", err } - generator, err := NewGenerator(conf, "") + generator, err := NewGenerator(conf, "", false) if err != nil { return "", err } diff --git a/pkg/dockerfile/build_tempdir.go b/pkg/dockerfile/build_tempdir.go new file mode 100644 index 0000000000..75eff0c54a --- /dev/null +++ b/pkg/dockerfile/build_tempdir.go @@ -0,0 +1,33 @@ +package dockerfile + +import ( + "os" + "path" + "time" +) + +func BuildCogTempDir(dir string) (string, error) { + rootTmp := path.Join(dir, ".cog/tmp") + if err := os.MkdirAll(rootTmp, 0o755); err != nil { + return "", err + } + return rootTmp, nil +} + +func BuildTempDir(dir string) (string, error) { + rootTmp, err := BuildCogTempDir(dir) + if err != nil { + return "", err + } + + if err := os.MkdirAll(rootTmp, 0o755); err != nil { + return "", 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 "", err + } + return tmpDir, nil +} diff --git a/pkg/dockerfile/build_tempdir_test.go b/pkg/dockerfile/build_tempdir_test.go new file mode 100644 index 0000000000..c90fcc17b3 --- /dev/null +++ b/pkg/dockerfile/build_tempdir_test.go @@ -0,0 +1,15 @@ +package dockerfile + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBuildCogTempDir(t *testing.T) { + tmpDir := t.TempDir() + cogTmpDir, err := BuildCogTempDir(tmpDir) + require.NoError(t, err) + require.Equal(t, filepath.Join(tmpDir, ".cog/tmp"), cogTmpDir) +} diff --git a/pkg/dockerfile/cog_embed.go b/pkg/dockerfile/cog_embed.go new file mode 100644 index 0000000000..3711f460e3 --- /dev/null +++ b/pkg/dockerfile/cog_embed.go @@ -0,0 +1,6 @@ +package dockerfile + +import "embed" + +//go:embed embed/*.whl +var CogEmbed embed.FS diff --git a/pkg/dockerfile/fast_generator.go b/pkg/dockerfile/fast_generator.go new file mode 100644 index 0000000000..33b1687459 --- /dev/null +++ b/pkg/dockerfile/fast_generator.go @@ -0,0 +1,272 @@ +package dockerfile + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/replicate/cog/pkg/config" + "github.com/replicate/cog/pkg/weights" +) + +const FUSE_RPC_WEIGHTS_PATH = "/srv/r8/fuse-rpc/weights" +const MONOBASE_CACHE_PATH = "/var/cache/monobase" +const APT_CACHE_MOUNT = "--mount=type=cache,target=/var/cache/apt,id=apt-cache" +const UV_CACHE_DIR = "/srv/r8/monobase/uv/cache" +const UV_CACHE_MOUNT = "--mount=type=cache,target=" + UV_CACHE_DIR + ",id=pip-cache" +const FAST_GENERATOR_NAME = "FAST_GENERATOR" + +type FastGenerator struct { + Config *config.Config + Dir string +} + +func NewFastGenerator(config *config.Config, dir string) (*FastGenerator, error) { + return &FastGenerator{ + Config: config, + Dir: dir, + }, nil +} + +func (g *FastGenerator) GenerateInitialSteps() (string, error) { + return "", errors.New("GenerateInitialSteps not supported in FastGenerator") +} + +func (g *FastGenerator) BaseImage() (string, error) { + return "", errors.New("BaseImage not supported in FastGenerator") +} + +func (g *FastGenerator) Cleanup() error { + return nil +} + +func (g *FastGenerator) GenerateDockerfileWithoutSeparateWeights() (string, error) { + return g.generate() +} + +func (g *FastGenerator) GenerateModelBase() (string, error) { + return "", errors.New("GenerateModelBase not supported in FastGenerator") +} + +func (g *FastGenerator) GenerateModelBaseWithSeparateWeights(imageName string) (weightsBase string, dockerfile string, dockerignoreContents string, err error) { + return "", "", "", errors.New("GenerateModelBaseWithSeparateWeights not supported in FastGenerator") +} + +func (g *FastGenerator) GenerateWeightsManifest() (*weights.Manifest, error) { + return nil, errors.New("GenerateWeightsManifest not supported in FastGenerator") +} + +func (g *FastGenerator) IsUsingCogBaseImage() bool { + return false +} + +func (g *FastGenerator) SetPrecompile(precompile bool) { +} + +func (g *FastGenerator) SetStrip(strip bool) { +} + +func (g *FastGenerator) SetUseCogBaseImage(useCogBaseImage bool) { +} + +func (g *FastGenerator) SetUseCogBaseImagePtr(useCogBaseImage *bool) { +} + +func (g *FastGenerator) SetUseCudaBaseImage(argumentValue string) { +} + +func (g *FastGenerator) Name() string { + return FAST_GENERATOR_NAME +} + +func (g *FastGenerator) generate() (string, error) { + tmpDir, err := BuildCogTempDir(g.Dir) + if err != nil { + return "", err + } + + weights, err := FindWeights(g.Dir, tmpDir) + if err != nil { + return "", err + } + + lines := []string{} + lines, err = g.generateMonobase(lines, tmpDir) + if err != nil { + return "", err + } + + lines, err = g.copyWeights(lines, weights) + if err != nil { + return "", err + } + + lines, err = g.install(lines, weights, tmpDir) + if err != nil { + return "", err + } + + lines, err = g.entrypoint(lines) + if err != nil { + return "", err + } + + return strings.Join(lines, "\n"), nil +} + +func (g *FastGenerator) copyCog(tmpDir string) (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 + } + path := filepath.Join(tmpDir, filename) + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return "", fmt.Errorf("Failed to write %s: %w", filename, err) + } + if err := os.WriteFile(path, data, 0o644); err != nil { + return "", fmt.Errorf("Failed to write %s: %w", filename, err) + } + return path, nil +} + +func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]string, error) { + lines = append(lines, []string{ + "# syntax=docker/dockerfile:1-labs", + "FROM r8.im/monobase:latest", + }...) + + cogPath, err := g.copyCog(tmpDir) + if err != nil { + return nil, err + } + + lines = append(lines, []string{ + "ENV R8_COG_VERSION=\"file:///buildtmp/" + filepath.Base(cogPath) + "\"", + }...) + + if g.Config.Build.GPU { + cudaVersion := g.Config.Build.CUDA + cudnnVersion := g.Config.Build.CuDNN + lines = append(lines, []string{ + "ENV R8_CUDA_VERSION=" + cudaVersion, + "ENV R8_CUDNN_VERSION=" + cudnnVersion, + "ENV R8_CUDA_PREFIX=https://monobase.replicate.delivery/cuda", + "ENV R8_CUDNN_PREFIX=https://monobase.replicate.delivery/cudnn", + }...) + } + + lines = append(lines, []string{ + "ENV R8_PYTHON_VERSION=" + g.Config.Build.PythonVersion, + }...) + + torchVersion, ok := g.Config.TorchVersion() + if ok { + lines = append(lines, []string{ + "ENV R8_TORCH_VERSION=" + torchVersion, + }...) + } + + buildTmpMount, err := g.buildTmpMount(tmpDir) + if err != nil { + return nil, err + } + + return append(lines, []string{ + "RUN " + strings.Join([]string{ + buildTmpMount, + g.monobaseUsercacheMount(), + APT_CACHE_MOUNT, + UV_CACHE_MOUNT, + }, " ") + " UV_CACHE_DIR=\"" + UV_CACHE_DIR + "\" UV_LINK_MODE=copy /opt/r8/monobase/run.sh monobase.build --mini --cache=" + MONOBASE_CACHE_PATH, + }...), nil +} + +func (g *FastGenerator) copyWeights(lines []string, weights []Weight) ([]string, error) { + if len(weights) == 0 { + return lines, nil + } + + for _, weight := range weights { + lines = append(lines, "COPY --link \""+weight.Path+"\" \""+filepath.Join(FUSE_RPC_WEIGHTS_PATH, weight.Digest)+"\"") + } + + return lines, nil +} + +func (g *FastGenerator) install(lines []string, weights []Weight, tmpDir string) ([]string, error) { + // Install apt packages + packages := g.Config.Build.SystemPackages + if len(packages) > 0 { + lines = append(lines, "RUN "+APT_CACHE_MOUNT+" apt-get update && apt-get install -qqy "+strings.Join(packages, " ")+" && rm -rf /var/lib/apt/lists/*") + } + + // Install python packages + requirementsFile, err := g.pythonRequirements(tmpDir) + if err != nil { + return nil, err + } + buildTmpMount, err := g.buildTmpMount(tmpDir) + if err != nil { + return nil, err + } + if requirementsFile != "" { + lines = append(lines, "RUN "+strings.Join([]string{ + buildTmpMount, + UV_CACHE_MOUNT, + }, " ")+" UV_CACHE_DIR=\""+UV_CACHE_DIR+"\" UV_LINK_MODE=copy UV_COMPILE_BYTECODE=0 /opt/r8/monobase/run.sh monobase.user --requirements=/buildtmp/requirements.txt") + } + + // Copy over source / without weights + copyCommand := "COPY --link " + for _, weight := range weights { + copyCommand += "--exclude='" + weight.Path + "' " + } + copyCommand += ". /src" + lines = append(lines, copyCommand) + + // Link to weights + if len(weights) > 0 { + linkCommands := []string{} + for _, weight := range weights { + linkCommands = append(linkCommands, "ln -s \""+filepath.Join(FUSE_RPC_WEIGHTS_PATH, weight.Digest)+"\" \"/src/"+weight.Path+"\"") + } + lines = append(lines, "RUN "+strings.Join(linkCommands, " && ")) + } + + return lines, nil +} + +func (g *FastGenerator) pythonRequirements(tmpDir string) (string, error) { + return config.GenerateRequirements(tmpDir, g.Config) +} + +func (g *FastGenerator) entrypoint(lines []string) ([]string, error) { + return append(lines, []string{ + "WORKDIR /src", + "ENV VERBOSE=0", + "ENTRYPOINT [\"/usr/bin/tini\", \"--\", \"/opt/r8/monobase/exec.sh\"]", + "CMD [\"python\", \"-m\", \"cog.server.http\"]", + }...), nil +} + +func (g *FastGenerator) buildTmpMount(tmpDir string) (string, error) { + relativeTmpDir, err := filepath.Rel(g.Dir, tmpDir) + if err != nil { + return "", err + } + return "--mount=type=bind,ro,source=\"" + relativeTmpDir + "\",target=\"/buildtmp\"", nil +} + +func (g *FastGenerator) monobaseUsercacheMount() string { + return "--mount=type=cache,from=usercache,target=\"" + MONOBASE_CACHE_PATH + "\"" +} diff --git a/pkg/dockerfile/fast_generator_test.go b/pkg/dockerfile/fast_generator_test.go new file mode 100644 index 0000000000..e8176b32d9 --- /dev/null +++ b/pkg/dockerfile/fast_generator_test.go @@ -0,0 +1,96 @@ +package dockerfile + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/replicate/cog/pkg/config" +) + +func TestGenerate(t *testing.T) { + dir := t.TempDir() + build := config.Build{ + PythonPackages: []string{"torch==2.5.1"}, + } + config := config.Config{ + Build: &build, + } + generator, err := NewFastGenerator(&config, dir) + require.NoError(t, err) + dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights() + require.NoError(t, err) + dockerfileLines := strings.Split(dockerfile, "\n") + require.Equal(t, "# syntax=docker/dockerfile:1-labs", dockerfileLines[0]) +} + +func TestGenerateUVCacheMount(t *testing.T) { + dir := t.TempDir() + build := config.Build{ + PythonPackages: []string{ + "torch==2.5.1", + "catboost==1.2.7", + }, + } + config := config.Config{ + Build: &build, + } + generator, err := NewFastGenerator(&config, dir) + require.NoError(t, err) + dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights() + require.NoError(t, err) + dockerfileLines := strings.Split(dockerfile, "\n") + require.Equal(t, "RUN --mount=type=bind,ro,source=\".cog/tmp\",target=\"/buildtmp\" --mount=type=cache,from=usercache,target=\"/var/cache/monobase\" --mount=type=cache,target=/var/cache/apt,id=apt-cache --mount=type=cache,target=/srv/r8/monobase/uv/cache,id=pip-cache UV_CACHE_DIR=\"/srv/r8/monobase/uv/cache\" UV_LINK_MODE=copy /opt/r8/monobase/run.sh monobase.build --mini --cache=/var/cache/monobase", dockerfileLines[4]) +} + +func TestGenerateCUDA(t *testing.T) { + dir := t.TempDir() + build := config.Build{ + GPU: true, + CUDA: "12.4", + } + config := config.Config{ + Build: &build, + } + generator, err := NewFastGenerator(&config, dir) + require.NoError(t, err) + dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights() + require.NoError(t, err) + dockerfileLines := strings.Split(dockerfile, "\n") + require.Equal(t, "ENV R8_CUDA_VERSION=12.4", dockerfileLines[3]) +} + +func TestGeneratePythonPackages(t *testing.T) { + dir := t.TempDir() + build := config.Build{ + PythonPackages: []string{ + "catboost==1.2.7", + }, + } + config := config.Config{ + Build: &build, + } + generator, err := NewFastGenerator(&config, dir) + require.NoError(t, err) + dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights() + require.NoError(t, err) + dockerfileLines := strings.Split(dockerfile, "\n") + require.Equal(t, "RUN --mount=type=bind,ro,source=\".cog/tmp\",target=\"/buildtmp\" --mount=type=cache,target=/srv/r8/monobase/uv/cache,id=pip-cache UV_CACHE_DIR=\"/srv/r8/monobase/uv/cache\" UV_LINK_MODE=copy UV_COMPILE_BYTECODE=0 /opt/r8/monobase/run.sh monobase.user --requirements=/buildtmp/requirements.txt", dockerfileLines[5]) +} + +func TestGenerateVerboseEnv(t *testing.T) { + dir := t.TempDir() + build := config.Build{ + PythonPackages: []string{"torch==2.5.1"}, + } + config := config.Config{ + Build: &build, + } + generator, err := NewFastGenerator(&config, dir) + require.NoError(t, err) + dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights() + require.NoError(t, err) + dockerfileLines := strings.Split(dockerfile, "\n") + require.Equal(t, "ENV VERBOSE=0", dockerfileLines[8]) +} diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go index 62bbbef037..ef83df8f7c 100644 --- a/pkg/dockerfile/generator.go +++ b/pkg/dockerfile/generator.go @@ -16,4 +16,5 @@ type Generator interface { GenerateWeightsManifest() (*weights.Manifest, error) GenerateDockerfileWithoutSeparateWeights() (string, error) GenerateModelBase() (string, error) + Name() string } diff --git a/pkg/dockerfile/generator_factory.go b/pkg/dockerfile/generator_factory.go index fffd6b87a0..f60d900d16 100644 --- a/pkg/dockerfile/generator_factory.go +++ b/pkg/dockerfile/generator_factory.go @@ -4,6 +4,9 @@ import ( "github.com/replicate/cog/pkg/config" ) -func NewGenerator(config *config.Config, dir string) (Generator, error) { +func NewGenerator(config *config.Config, dir string, buildFast bool) (Generator, error) { + if buildFast { + return NewFastGenerator(config, dir) + } return NewStandardGenerator(config, dir) } diff --git a/pkg/dockerfile/generator_factory_test.go b/pkg/dockerfile/generator_factory_test.go new file mode 100644 index 0000000000..363b53672c --- /dev/null +++ b/pkg/dockerfile/generator_factory_test.go @@ -0,0 +1,22 @@ +package dockerfile + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/replicate/cog/pkg/config" +) + +func TestGeneratorFactory(t *testing.T) { + dir := t.TempDir() + build := config.Build{ + PythonPackages: []string{"torch==2.5.1"}, + } + config := config.Config{ + Build: &build, + } + generator, err := NewGenerator(&config, dir, true) + require.NoError(t, err) + require.Equal(t, generator.Name(), FAST_GENERATOR_NAME) +} diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go index 8d3f076b58..92787c33f9 100644 --- a/pkg/dockerfile/standard_generator.go +++ b/pkg/dockerfile/standard_generator.go @@ -1,14 +1,12 @@ 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" @@ -17,9 +15,6 @@ import ( "github.com/replicate/cog/pkg/weights" ) -//go:embed embed/*.whl -var cogEmbed embed.FS - const DockerignoreHeader = `# generated by replicate/cog __pycache__ *.pyc @@ -46,6 +41,7 @@ const LDConfigCacheBuildCommand = "RUN find / -type f -name \"*python*.so\" -pri 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" +const STANDARD_GENERATOR_NAME = "STANDARD_GENERATOR" type StandardGenerator struct { Config *config.Config @@ -74,13 +70,7 @@ type StandardGenerator struct { } 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) + tmpDir, err := BuildTempDir(dir) if err != nil { return nil, err } @@ -317,6 +307,10 @@ func (g *StandardGenerator) BaseImage() (string, error) { return "python:" + g.Config.Build.PythonVersion + "-slim", nil } +func (g *StandardGenerator) Name() string { + return STANDARD_GENERATOR_NAME +} + func (g *StandardGenerator) preamble() string { return `ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 @@ -408,7 +402,7 @@ RUN rm -rf /usr/bin/python3 && ln -s ` + "`realpath \\`pyenv which python\\`` /u } func (g *StandardGenerator) installCog() (string, error) { - files, err := cogEmbed.ReadDir("embed") + files, err := CogEmbed.ReadDir("embed") if err != nil { return "", err } @@ -416,7 +410,7 @@ func (g *StandardGenerator) installCog() (string, error) { return "", fmt.Errorf("should only have one cog wheel embedded") } filename := files[0].Name() - data, err := cogEmbed.ReadFile("embed/" + filename) + data, err := CogEmbed.ReadFile("embed/" + filename) if err != nil { return "", err } diff --git a/pkg/dockerfile/standard_generator_test.go b/pkg/dockerfile/standard_generator_test.go index 0cb01cbb24..e4198f7f24 100644 --- a/pkg/dockerfile/standard_generator_test.go +++ b/pkg/dockerfile/standard_generator_test.go @@ -27,7 +27,7 @@ ENTRYPOINT ["/sbin/tini", "--"] } func getWheelName() string { - files, err := cogEmbed.ReadDir("embed") + files, err := CogEmbed.ReadDir("embed") if err != nil { panic(err) } @@ -314,7 +314,7 @@ build: require.NoError(t, err) require.NoError(t, conf.ValidateAndComplete(tmpDir)) - 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") @@ -753,7 +753,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") diff --git a/pkg/dockerfile/user_cache.go b/pkg/dockerfile/user_cache.go new file mode 100644 index 0000000000..f2585986d9 --- /dev/null +++ b/pkg/dockerfile/user_cache.go @@ -0,0 +1,36 @@ +package dockerfile + +import ( + "os" + "os/user" + "path" + "path/filepath" +) + +const COG_CACHE_FOLDER = ".cog_cache" + +func UserCache() (string, error) { + usr, err := user.Current() + if err != nil { + return "", err + } + + path := filepath.Join(usr.HomeDir, COG_CACHE_FOLDER) + if err := os.MkdirAll(path, 0o755); err != nil { + return "", err + } + + return path, nil +} + +func UserCacheFolder(folder string) (string, error) { + userCache, err := UserCache() + if err != nil { + return "", err + } + cacheFolder := path.Join(userCache, folder) + if err := os.MkdirAll(cacheFolder, 0o755); err != nil { + return "", err + } + return cacheFolder, nil +} diff --git a/pkg/dockerfile/user_cache_test.go b/pkg/dockerfile/user_cache_test.go new file mode 100644 index 0000000000..cddd1e50c8 --- /dev/null +++ b/pkg/dockerfile/user_cache_test.go @@ -0,0 +1,15 @@ +package dockerfile + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestUserCache(t *testing.T) { + userCache, err := UserCache() + require.NoError(t, err) + lastDirectory := filepath.Base(userCache) + require.Equal(t, COG_CACHE_FOLDER, lastDirectory) +} diff --git a/pkg/dockerfile/weights.go b/pkg/dockerfile/weights.go new file mode 100644 index 0000000000..a951a3d610 --- /dev/null +++ b/pkg/dockerfile/weights.go @@ -0,0 +1,155 @@ +package dockerfile + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "io" + "os" + "path/filepath" + "slices" + "time" +) + +var WEIGHT_FILE_EXCLUSIONS = []string{ + ".gif", + ".ipynb", + ".jpeg", + ".jpg", + ".log", + ".mp4", + ".png", + ".svg", + ".webp", +} +var WEIGHT_FILE_INCLUSIONS = []string{ + ".ckpt", + ".h5", + ".onnx", + ".pb", + ".pbtxt", + ".pt", + ".pth", + ".safetensors", + ".tflite", +} + +const WEIGHT_FILE_SIZE_EXCLUSION = 1024 * 1024 +const WEIGHT_FILE_SIZE_INCLUSION = 128 * 1024 * 1024 +const WEIGHT_FILE = "weights.json" + +type Weight struct { + Path string `json:"path"` + Digest string `json:"digest"` + Timestamp time.Time `json:"timestamp"` + Size int64 `json:"size"` +} + +func FindWeights(folder string, tmpDir string) ([]Weight, error) { + weightFile := filepath.Join(tmpDir, WEIGHT_FILE) + if _, err := os.Stat(weightFile); errors.Is(err, os.ErrNotExist) { + return findFullWeights(folder, []Weight{}, weightFile) + } + return checkWeights(folder, weightFile) +} + +func findFullWeights(folder string, weights []Weight, weightFile string) ([]Weight, error) { + err := filepath.Walk(folder, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + + relPath, err := filepath.Rel(folder, path) + if err != nil { + return err + } + + for _, weight := range weights { + if weight.Path == relPath { + return nil + } + } + + ext := filepath.Ext(path) + + if slices.Contains(WEIGHT_FILE_EXCLUSIONS, ext) || info.Size() <= WEIGHT_FILE_SIZE_EXCLUSION { + return nil + } + + if slices.Contains(WEIGHT_FILE_INCLUSIONS, ext) || info.Size() >= WEIGHT_FILE_SIZE_INCLUSION { + hash := sha256.New() + + file, err := os.Open(path) + if err != nil { + return err + } + defer file.Close() + + if _, err := io.Copy(hash, file); err != nil { + return err + } + + info, err := os.Stat(relPath) + if err != nil { + return err + } + + weights = append(weights, Weight{ + Path: relPath, + Digest: hex.EncodeToString(hash.Sum(nil)), + Timestamp: info.ModTime(), + Size: info.Size(), + }) + } + return nil + }) + if err != nil { + return nil, err + } + + jsonData, err := json.MarshalIndent(weights, "", " ") + if err != nil { + return nil, err + } + err = os.WriteFile(weightFile, jsonData, 0o644) + if err != nil { + return nil, err + } + + return weights, err +} + +func checkWeights(folder string, weightFile string) ([]Weight, error) { + var weights []Weight + + file, err := os.Open(weightFile) + if err != nil { + return nil, err + } + defer file.Close() + + decoder := json.NewDecoder(file) + err = decoder.Decode(&weights) + if err != nil { + return nil, err + } + + newWeights := []Weight{} + for _, weight := range weights { + info, err := os.Stat(weight.Path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + continue + } + return nil, err + } + + if weight.Timestamp != info.ModTime() || weight.Size != info.Size() { + continue + } + newWeights = append(newWeights, weight) + } + + return findFullWeights(folder, newWeights, weightFile) +} diff --git a/pkg/dockerfile/weights_test.go b/pkg/dockerfile/weights_test.go new file mode 100644 index 0000000000..0204545930 --- /dev/null +++ b/pkg/dockerfile/weights_test.go @@ -0,0 +1,40 @@ +package dockerfile + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestFindWeights(t *testing.T) { + folder := t.TempDir() + tmpDir := t.TempDir() + weights, err := FindWeights(folder, tmpDir) + require.NoError(t, err) + require.Empty(t, weights) +} + +func TestFindWeightsWithRemovedWeight(t *testing.T) { + folder := t.TempDir() + tmpDir := t.TempDir() + weightFile := filepath.Join(tmpDir, WEIGHT_FILE) + weights := []Weight{ + { + Path: "nonexistant_weight.h5", + Digest: "1", + Timestamp: time.Now(), + Size: 10, + }, + } + jsonData, err := json.MarshalIndent(weights, "", " ") + require.NoError(t, err) + err = os.WriteFile(weightFile, jsonData, 0o644) + require.NoError(t, err) + weights, err = FindWeights(folder, tmpDir) + require.NoError(t, err) + require.Empty(t, weights) +} diff --git a/pkg/image/build.go b/pkg/image/build.go index 7dac6b8c38..8d29efd533 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -54,7 +54,7 @@ func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache, return fmt.Errorf("Failed to build Docker image: %w", err) } } else { - generator, err := dockerfile.NewGenerator(cfg, dir) + generator, err := dockerfile.NewGenerator(cfg, dir, fastFlag) if err != nil { return fmt.Errorf("Error creating Dockerfile generator: %w", err) } @@ -244,7 +244,7 @@ func BuildBase(cfg *config.Config, dir string, useCudaBaseImage string, useCogBa imageName := config.BaseDockerImageName(dir) console.Info("Building Docker image from environment in cog.yaml...") - generator, err := dockerfile.NewGenerator(cfg, dir) + generator, err := dockerfile.NewGenerator(cfg, dir, false) if err != nil { return "", fmt.Errorf("Error creating Dockerfile generator: %w", err) } diff --git a/test-integration/test_integration/fixtures/fast-build/cog.yaml b/test-integration/test_integration/fixtures/fast-build/cog.yaml new file mode 100644 index 0000000000..62419bfb88 --- /dev/null +++ b/test-integration/test_integration/fixtures/fast-build/cog.yaml @@ -0,0 +1,9 @@ +build: + python_version: "3.12" + python_packages: + - "torch==2.5.0" + - "beautifulsoup4==4.12.3" + system_packages: + - "git" +predict: "predict.py:Predictor" +image: "test" diff --git a/test-integration/test_integration/fixtures/fast-build/predict.py b/test-integration/test_integration/fixtures/fast-build/predict.py new file mode 100644 index 0000000000..44f6992b01 --- /dev/null +++ b/test-integration/test_integration/fixtures/fast-build/predict.py @@ -0,0 +1,6 @@ +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, s: str) -> str: + return "hello " + s diff --git a/test-integration/test_integration/test_build.py b/test-integration/test_integration/test_build.py index e5a9c93df3..63ace21286 100644 --- a/test-integration/test_integration/test_build.py +++ b/test-integration/test_integration/test_build.py @@ -404,3 +404,21 @@ def test_cog_installs_apt_packages(docker_image): # Test that the build completes successfully. # If the apt-packages weren't installed the run command would fail. assert build_process.returncode == 0 + + +def test_fast_build(docker_image): + project_dir = Path(__file__).parent / "fixtures/fast-build" + weights_file = os.path.join(project_dir, "weights.h5") + with open(weights_file, "w", encoding="utf8") as handle: + handle.seek(256 * 1024 * 1024) + handle.write("\0") + + build_process = subprocess.run( + ["cog", "build", "-t", docker_image, "--x-fast"], + cwd=project_dir, + capture_output=True, + ) + + os.remove(weights_file) + + assert build_process.returncode == 0