-
Notifications
You must be signed in to change notification settings - Fork 579
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add fast generator for cog build (#2108)
* Add fast generator * Add copy weights step * Add package installation step * Add user cache to the monobase * Resolve monobase caching into the local userspace to allow for monobase caching between builds * Add rsync copy * Check if weight files have changed before checksum * Split install into separate steps * We can get the tarballs for each of these layers instead of creating 1 big tarball. * Create requirements.txt in the build tmp directory * Fix unit tests * Fix lint * Add basic unit tests * Use UV_CACHE_DIR and mount uv cache * Remove —skip-cuda from monobase build * Monobase now handles empty CUDA env vars * Fix file not found when evaluating weights * Add UV_LINK_MODE=copy to the uv install commands * Add UV_COMPILE_BYTECODE env var * Remove verbosity from monobase exec * Fix integration test * Switch tini and exec --------- Signed-off-by: Will Sackfield <[email protected]>
- Loading branch information
Showing
26 changed files
with
913 additions
and
83 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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=<findLink>] [-f <findLink>] [--extra-index-url=<extraIndexURL>] | ||
| 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 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| package dockerfile | ||
|
|
||
| import "embed" | ||
|
|
||
| //go:embed embed/*.whl | ||
| var CogEmbed embed.FS |
Oops, something went wrong.