Skip to content

Commit e27d743

Browse files
nevillelyhnevillelyh
authored
Separate tmp mounts for fast layers (#2182)
* So that changes in one layer does not invalidate everything else
1 parent 80d2a27 commit e27d743

File tree

4 files changed

+83
-43
lines changed

4 files changed

+83
-43
lines changed

‎pkg/dockerfile/build_tempdir.go

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -6,28 +6,16 @@ import (
66
"time"
77
)
88

9-
func BuildCogTempDir(dir string) (string, error) {
10-
rootTmp := path.Join(dir, ".cog/tmp")
9+
func BuildCogTempDir(dir string, subDir string) (string, error) {
10+
rootTmp := path.Join(dir, ".cog", "tmp", subDir)
1111
if err := os.MkdirAll(rootTmp, 0o755); err != nil {
1212
return "", err
1313
}
1414
return rootTmp, nil
1515
}
1616

1717
func BuildTempDir(dir string) (string, error) {
18-
rootTmp, err := BuildCogTempDir(dir)
19-
if err != nil {
20-
return "", err
21-
}
22-
23-
if err := os.MkdirAll(rootTmp, 0o755); err != nil {
24-
return "", err
25-
}
2618
// tmpDir ends up being something like dir/.cog/tmp/build20240620123456.000000
2719
now := time.Now().Format("20060102150405.000000")
28-
tmpDir, err := os.MkdirTemp(rootTmp, "build"+now)
29-
if err != nil {
30-
return "", err
31-
}
32-
return tmpDir, nil
20+
return BuildCogTempDir(dir, "build"+now)
3321
}

‎pkg/dockerfile/build_tempdir_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99

1010
func TestBuildCogTempDir(t *testing.T) {
1111
tmpDir := t.TempDir()
12-
cogTmpDir, err := BuildCogTempDir(tmpDir)
12+
cogTmpDir, err := BuildCogTempDir(tmpDir, "weights")
1313
require.NoError(t, err)
14-
require.Equal(t, filepath.Join(tmpDir, ".cog/tmp"), cogTmpDir)
14+
require.Equal(t, filepath.Join(tmpDir, ".cog/tmp/weights"), cogTmpDir)
1515
}

‎pkg/dockerfile/fast_generator.go

Lines changed: 75 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package dockerfile
33
import (
44
"errors"
55
"fmt"
6+
"os"
7+
"path"
68
"path/filepath"
79
"strings"
810

@@ -99,23 +101,46 @@ func (g *FastGenerator) generate() (string, error) {
99101
return "", err
100102
}
101103

102-
tmpDir, err := BuildCogTempDir(g.Dir)
104+
// Temp directories are used as bind mounts in docker build
105+
// Separate them so that changes in one layer doesn't invalidate everything else
106+
107+
// Weights layer
108+
// Includes file metadata, triggered by weights file changes
109+
tmpWeightsDir, err := BuildCogTempDir(g.Dir, "weights")
103110
if err != nil {
104111
return "", err
105112
}
106-
107-
weights, err := weights.FindFastWeights(g.Dir, tmpDir)
113+
weights, err := weights.FindFastWeights(g.Dir, tmpWeightsDir)
108114
if err != nil {
109115
return "", err
110116
}
111117

112-
aptTarFile, err := g.generateAptTarball(tmpDir)
118+
// APT layer
119+
// Includes a tarball extracted from APT packages, triggered by system_packages changes
120+
tmpAptDir, err := BuildCogTempDir(g.Dir, "apt")
121+
if err != nil {
122+
return "", err
123+
}
124+
aptTarFile, err := g.generateAptTarball(tmpAptDir)
113125
if err != nil {
114126
return "", fmt.Errorf("generate apt tarball: %w", err)
115127
}
116128

129+
// Monobase layer
130+
// Includes an ENV file, triggered by Python, Torch, or CUDA version changes
131+
tmpMonobaseDir, err := BuildCogTempDir(g.Dir, "monobase")
132+
if err != nil {
133+
return "", err
134+
}
117135
lines := []string{}
118-
lines, err = g.generateMonobase(lines, tmpDir)
136+
lines, err = g.generateMonobase(lines, tmpMonobaseDir)
137+
if err != nil {
138+
return "", err
139+
}
140+
141+
// User layer
142+
// Includes requirements.txt, triggered by python_requirements changes
143+
tmpUserDir, err := BuildCogTempDir(g.Dir, "user")
119144
if err != nil {
120145
return "", err
121146
}
@@ -125,7 +150,17 @@ func (g *FastGenerator) generate() (string, error) {
125150
return "", err
126151
}
127152

128-
lines, err = g.install(lines, weights, tmpDir, aptTarFile)
153+
lines, err = g.installApt(lines, tmpAptDir, aptTarFile)
154+
if err != nil {
155+
return "", err
156+
}
157+
158+
lines, err = g.installPython(lines, tmpUserDir)
159+
if err != nil {
160+
return "", err
161+
}
162+
163+
lines, err = g.installSrc(lines, weights)
129164
if err != nil {
130165
return "", err
131166
}
@@ -139,12 +174,8 @@ func (g *FastGenerator) generate() (string, error) {
139174
}
140175

141176
func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]string, error) {
142-
lines = append(lines, []string{
143-
"# syntax=docker/dockerfile:1-labs",
144-
"FROM r8.im/monobase:latest",
145-
}...)
146-
147-
lines = append(lines, []string{
177+
var envs []string
178+
envs = append(envs, []string{
148179
// This installs latest version of coglet
149180
"ENV R8_COG_VERSION=coglet",
150181
}...)
@@ -162,7 +193,7 @@ func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]strin
162193
return nil, fmt.Errorf("CUDNN version must be <major> only, supported versions: %s", strings.Join(g.matrix.CudnnVersions, ", "))
163194
}
164195

165-
lines = append(lines, []string{
196+
envs = append(envs, []string{
166197
"ENV R8_CUDA_VERSION=" + cudaVersion,
167198
"ENV R8_CUDNN_VERSION=" + cudnnVersion,
168199
"ENV R8_CUDA_PREFIX=https://monobase-packages.replicate.delivery/cuda",
@@ -174,7 +205,7 @@ func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]strin
174205
return nil, fmt.Errorf(
175206
"Python version must be <major>.<minor>, supported versions: %s", strings.Join(g.matrix.PythonVersions, ", "))
176207
}
177-
lines = append(lines, []string{
208+
envs = append(envs, []string{
178209
"ENV R8_PYTHON_VERSION=" + g.Config.Build.PythonVersion,
179210
}...)
180211

@@ -183,7 +214,7 @@ func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]strin
183214
if !CheckMajorMinorPatch(torchVersion) {
184215
return nil, fmt.Errorf("Torch version must be <major>.<minor>.<patch>: %s", strings.Join(g.matrix.TorchVersions, ", "))
185216
}
186-
lines = append(lines, []string{
217+
envs = append(envs, []string{
187218
"ENV R8_TORCH_VERSION=" + torchVersion,
188219
}...)
189220
}
@@ -194,19 +225,32 @@ func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]strin
194225
g.Config.Build.PythonVersion, torchVersion, g.Config.Build.CUDA)
195226
}
196227

228+
// The only input to monobase.build are these ENV vars
229+
// Write them in tmp mount for layer caching
230+
err := os.WriteFile(path.Join(tmpDir, "env.txt"), []byte(strings.Join(envs, "\n")), 0o644)
231+
if err != nil {
232+
return nil, err
233+
}
234+
197235
buildTmpMount, err := g.buildTmpMount(tmpDir)
198236
if err != nil {
199237
return nil, err
200238
}
201239

202-
return append(lines, []string{
240+
lines = append(lines, []string{
241+
"# syntax=docker/dockerfile:1-labs",
242+
"FROM r8.im/monobase:latest",
243+
}...)
244+
lines = append(lines, envs...)
245+
lines = append(lines, []string{
203246
"RUN " + strings.Join([]string{
204247
buildTmpMount,
205248
g.monobaseUsercacheMount(),
206249
APT_CACHE_MOUNT,
207250
UV_CACHE_MOUNT,
208251
}, " ") + " UV_CACHE_DIR=\"" + UV_CACHE_DIR + "\" UV_LINK_MODE=copy /opt/r8/monobase/run.sh monobase.build --mini --cache=" + MONOBASE_CACHE_PATH,
209-
}...), nil
252+
}...)
253+
return lines, nil
210254
}
211255

212256
func (g *FastGenerator) copyWeights(lines []string, weights []weights.Weight) ([]string, error) {
@@ -221,7 +265,7 @@ func (g *FastGenerator) copyWeights(lines []string, weights []weights.Weight) ([
221265
return lines, nil
222266
}
223267

224-
func (g *FastGenerator) install(lines []string, weights []weights.Weight, tmpDir string, aptTarFile string) ([]string, error) {
268+
func (g *FastGenerator) installApt(lines []string, tmpDir string, aptTarFile string) ([]string, error) {
225269
// Install apt packages
226270
buildTmpMount, err := g.buildTmpMount(tmpDir)
227271
if err != nil {
@@ -230,9 +274,16 @@ func (g *FastGenerator) install(lines []string, weights []weights.Weight, tmpDir
230274
if aptTarFile != "" {
231275
lines = append(lines, "RUN "+buildTmpMount+" tar -xf \""+filepath.Join("/buildtmp", aptTarFile)+"\" -C /")
232276
}
277+
return lines, nil
278+
}
233279

280+
func (g *FastGenerator) installPython(lines []string, tmpDir string) ([]string, error) {
234281
// Install python packages
235-
requirementsFile, err := g.pythonRequirements(tmpDir)
282+
buildTmpMount, err := g.buildTmpMount(tmpDir)
283+
if err != nil {
284+
return nil, err
285+
}
286+
requirementsFile, err := requirements.GenerateRequirements(tmpDir, g.Config)
236287
if err != nil {
237288
return nil, err
238289
}
@@ -242,6 +293,11 @@ func (g *FastGenerator) install(lines []string, weights []weights.Weight, tmpDir
242293
UV_CACHE_MOUNT,
243294
}, " ")+" 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")
244295
}
296+
return lines, nil
297+
}
298+
299+
func (g *FastGenerator) installSrc(lines []string, weights []weights.Weight) ([]string, error) {
300+
// Install /src
245301

246302
// Copy over source / without weights
247303
copyCommand := "COPY --link --exclude='.cog' "
@@ -263,10 +319,6 @@ func (g *FastGenerator) install(lines []string, weights []weights.Weight, tmpDir
263319
return lines, nil
264320
}
265321

266-
func (g *FastGenerator) pythonRequirements(tmpDir string) (string, error) {
267-
return requirements.GenerateRequirements(tmpDir, g.Config)
268-
}
269-
270322
func (g *FastGenerator) entrypoint(lines []string) ([]string, error) {
271323
return append(lines, []string{
272324
"WORKDIR /src",

‎pkg/dockerfile/fast_generator_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ func TestGenerateUVCacheMount(t *testing.T) {
8686
dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights()
8787
require.NoError(t, err)
8888
dockerfileLines := strings.Split(dockerfile, "\n")
89-
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,sharing=locked --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])
89+
require.Equal(t, "RUN --mount=type=bind,ro,source=\".cog/tmp/monobase\",target=\"/buildtmp\" --mount=type=cache,from=usercache,target=\"/var/cache/monobase\" --mount=type=cache,target=/var/cache/apt,id=apt-cache,sharing=locked --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])
9090
}
9191

9292
func TestGenerateCUDA(t *testing.T) {
@@ -157,7 +157,7 @@ func TestGeneratePythonPackages(t *testing.T) {
157157
dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights()
158158
require.NoError(t, err)
159159
dockerfileLines := strings.Split(dockerfile, "\n")
160-
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])
160+
require.Equal(t, "RUN --mount=type=bind,ro,source=\".cog/tmp/user\",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])
161161
}
162162

163163
func TestGenerateVerboseEnv(t *testing.T) {
@@ -227,7 +227,7 @@ func TestAptInstall(t *testing.T) {
227227
dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights()
228228
require.NoError(t, err)
229229
dockerfileLines := strings.Split(dockerfile, "\n")
230-
require.Equal(t, "RUN --mount=type=bind,ro,source=\".cog/tmp\",target=\"/buildtmp\" tar -xf \"/buildtmp/apt.9a881b9b9f23849475296a8cd768ea1965bc3152df7118e60c145975af6aa58a.tar.zst\" -C /", dockerfileLines[5])
230+
require.Equal(t, "RUN --mount=type=bind,ro,source=\".cog/tmp/apt\",target=\"/buildtmp\" tar -xf \"/buildtmp/apt.9a881b9b9f23849475296a8cd768ea1965bc3152df7118e60c145975af6aa58a.tar.zst\" -C /", dockerfileLines[5])
231231
}
232232

233233
func TestValidateConfigWithBuildRunItems(t *testing.T) {

0 commit comments

Comments
 (0)