Skip to content

Commit 236d5fd

Browse files
rclareyIsaac
rclarey
and
Isaac
authored
[SSH] Add versioning to agent dependency directory names (#7188)
## Changes Install agent dependencies (`uv`, `npm`, `ug`) under versioned directories so old dependencies are not used when dependency pins are updated. ## Why This is needed for dedicated cluster support, since dedicated clusters might remaining running for a long time and we want to be able to update dependencies without requiring customers to restart their clusters ## Tests updated/added unit tests Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 1d5ca29 commit 236d5fd

2 files changed

Lines changed: 81 additions & 29 deletions

File tree

‎experimental/ssh/internal/client/agentshim.go‎

Lines changed: 37 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -50,15 +50,18 @@ const (
5050
ugExecutable = "ug"
5151

5252
// pins the version of uv downloaded
53-
uvBaseURL = "https://github.com/astral-sh/uv/releases/download/0.12.18/"
53+
uvVersion = "0.12.18"
54+
uvBaseURL = "https://github.com/astral-sh/uv/releases/download/" + uvVersion + "/"
5455

5556
// SHA-256 of each pinned uv release tarball
5657
uvX64LinuxChecksum = "89eadd7c76fc063887959510d5ba0ab1264dfd5f1143b925ddb73021a40acf16"
5758
uvArm64LinuxChecksum = "afb6291f3f0a6b4521fc67b947822506c41dde5b60d2189dd8f3695b2ac8c9e7"
5859

59-
// overrides for uv tool installs so they are isolated to a directory we control
60-
uvToolDir = agentDepsDir + "/uv_tools"
61-
uvToolBinDir = agentDepsDir + "/uv_tools_bin"
60+
// overrides for uv tool installs so they are isolated to a directory we control.
61+
// Keyed by ugCommit so bumping the pin installs it afresh on a long-lived cluster
62+
// rather than reusing the previous commit's install.
63+
uvToolDir = agentDepsDir + "/uv_tools-" + ugCommit
64+
uvToolBinDir = agentDepsDir + "/uv_tools_bin-" + ugCommit
6265

6366
// pins the version of node/npm downloaded
6467
nodeVersion = "v24.21.0"
@@ -464,10 +467,17 @@ func bootstrapToolchain(ctx context.Context, home string) (string, error) {
464467
return ugPath, nil
465468
}
466469

467-
func toolchainReady() bool {
468-
_, ugErr := exec.LookPath(ugExecutable)
470+
// toolchainReady reports whether the pinned Unity Gateway CLI and some npm are
471+
// installed. npm (and uv) shipped with the image are used as-is, but ug must be the
472+
// pinned commit, so it is checked at its versioned path rather than looked up on PATH.
473+
func toolchainReady(home string) bool {
469474
_, npmErr := exec.LookPath(npmExecutable)
470-
return ugErr == nil && npmErr == nil
475+
return ugInstalled(home) && npmErr == nil
476+
}
477+
478+
func ugInstalled(home string) bool {
479+
_, err := os.Stat(filepath.Join(home, uvToolBinDir, ugExecutable))
480+
return err == nil
471481
}
472482

473483
func ensureToolchain(ctx context.Context, home, arch string) error {
@@ -478,11 +488,10 @@ func ensureToolchain(ctx context.Context, home, arch string) error {
478488
uvToolBinDirAbsolute := filepath.Join(home, uvToolBinDir)
479489
prependPath(ctx, uvToolBinDirAbsolute)
480490
for _, spec := range []archiveSpec{uvSpec, nodeSpec} {
481-
binDir := filepath.Join(home, agentDepsDir, spec.dirName, spec.binDirName)
482-
prependPath(ctx, binDir)
491+
prependPath(ctx, spec.binDir(home))
483492
}
484493

485-
if toolchainReady() {
494+
if toolchainReady(home) {
486495
return nil
487496
}
488497
unlock, err := acquireSetupLock(ctx, home)
@@ -491,7 +500,7 @@ func ensureToolchain(ctx context.Context, home, arch string) error {
491500
}
492501
defer unlock()
493502
// Another client may have finished the install while we waited for the lock.
494-
if toolchainReady() {
503+
if toolchainReady(home) {
495504
return nil
496505
}
497506

@@ -506,7 +515,7 @@ func ensureToolchain(ctx context.Context, home, arch string) error {
506515
}
507516
}
508517

509-
if _, err := exec.LookPath(ugExecutable); err != nil {
518+
if !ugInstalled(home) {
510519
if err := runStep(ctx, "Installing Unity Gateway CLI", func(out io.Writer) error {
511520
env := os.Environ()
512521
env = append(env, "UV_TOOL_DIR="+uvToolDirAbsolute, "UV_TOOL_BIN_DIR="+uvToolBinDirAbsolute)
@@ -535,11 +544,22 @@ type archiveSpec struct {
535544
archiveURL string
536545
binaryName string
537546
dirName string
538-
// sub-directory inside dirName where the binary is located, "" if it's the same as dirName
547+
// installed under dirName-version, so a version bump installs afresh on a
548+
// long-lived cluster instead of reusing the previous version
549+
version string
550+
// sub-directory inside the install dir where the binary is located, "" if it's the install dir itself
539551
binDirName string
540552
checksum string
541553
}
542554

555+
func (s archiveSpec) installDirName() string {
556+
return s.dirName + "-" + s.version
557+
}
558+
559+
func (s archiveSpec) binDir(home string) string {
560+
return filepath.Join(home, agentDepsDir, s.installDirName(), s.binDirName)
561+
}
562+
543563
// ensure the specified binary is installed, installing it if it's not already
544564
// present. Extraction output is written to out. Returns the path to the directory
545565
// containing the installed binary
@@ -549,9 +569,8 @@ func ensureBinary(ctx context.Context, home string, out io.Writer, spec archiveS
549569
}
550570

551571
depsRoot := filepath.Join(home, agentDepsDir)
552-
depsDir := filepath.Join(depsRoot, spec.dirName)
553-
binDir := filepath.Join(depsDir, spec.binDirName)
554-
if _, err := os.Stat(filepath.Join(binDir, spec.binaryName)); err == nil {
572+
depsDir := filepath.Join(depsRoot, spec.installDirName())
573+
if _, err := os.Stat(filepath.Join(spec.binDir(home), spec.binaryName)); err == nil {
555574
return nil
556575
}
557576

@@ -588,7 +607,7 @@ func ensureBinary(ctx context.Context, home string, out io.Writer, spec archiveS
588607
}
589608

590609
func uvArchiveSpec(goarch string) archiveSpec {
591-
spec := archiveSpec{dirName: "uv", binaryName: "uv"}
610+
spec := archiveSpec{dirName: "uv", version: uvVersion, binaryName: "uv"}
592611
switch goarch {
593612
case "amd64":
594613
spec.archiveURL = uvBaseURL + "uv-x86_64-unknown-linux-gnu.tar.gz"
@@ -749,7 +768,7 @@ func downloadVerified(ctx context.Context, spec archiveSpec) (string, error) {
749768
}
750769

751770
func nodeArchiveSpec(goarch string) archiveSpec {
752-
spec := archiveSpec{binDirName: "bin", dirName: "node", binaryName: "npm"}
771+
spec := archiveSpec{binDirName: "bin", dirName: "node", version: nodeVersion, binaryName: "npm"}
753772
switch goarch {
754773
case "amd64":
755774
spec.archiveURL = nodeBaseURL + "node-" + nodeVersion + "-linux-x64.tar.xz"

‎experimental/ssh/internal/client/agentshim_test.go‎

Lines changed: 44 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,15 @@ func TestEnsureBinary(t *testing.T) {
190190
defer srv.Close()
191191

192192
home := t.TempDir()
193+
// Older installs: an unversioned one from a CLI before versioned installs, and a
194+
// previous version. Neither counts as the pinned version being installed.
195+
depsRoot := filepath.Join(home, agentDepsDir)
196+
for _, dir := range []string{"node", "node-v1"} {
197+
writeFakeExec(t, filepath.Join(depsRoot, dir, "bin"), "npm", ":")
198+
}
199+
for _, dir := range []string{"uv", "uv-1"} {
200+
writeFakeExec(t, filepath.Join(depsRoot, dir), "uv", ":")
201+
}
193202
for _, c := range []struct {
194203
spec archiveSpec
195204
wantBinDir string
@@ -200,16 +209,18 @@ func TestEnsureBinary(t *testing.T) {
200209
binaryName: "npm",
201210
checksum: hex.EncodeToString(nodeSum[:]),
202211
dirName: "node",
212+
version: "v2",
203213
binDirName: "bin",
204-
}, wantBinDir: filepath.Join(home, agentDepsDir, "node", "bin"),
214+
}, wantBinDir: filepath.Join(home, agentDepsDir, "node-v2", "bin"),
205215
},
206216
{
207217
spec: archiveSpec{
208218
archiveURL: srv.URL + "/uv.tar.gz",
209219
binaryName: "uv",
210220
checksum: hex.EncodeToString(uvSum[:]),
211221
dirName: "uv",
212-
}, wantBinDir: filepath.Join(home, agentDepsDir, "uv"),
222+
version: "2",
223+
}, wantBinDir: filepath.Join(home, agentDepsDir, "uv-2"),
213224
},
214225
} {
215226
t.Run("downloads, extracts, and returns the bin dir ("+c.spec.binaryName+")", func(t *testing.T) {
@@ -255,8 +266,8 @@ func TestEnsureToolchain(t *testing.T) {
255266
t.Run(fmt.Sprintf("no-op when the toolchain is already present (staleLock=%v)", staleLock), func(t *testing.T) {
256267
home := t.TempDir()
257268
writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
258-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
259-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
269+
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
270+
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
260271
if staleLock {
261272
// left by a session killed mid-install
262273
writeFakeExec(t, filepath.Join(home, agentRootDir), setupLockName, "")
@@ -274,8 +285,8 @@ func TestEnsureToolchain(t *testing.T) {
274285
home := t.TempDir()
275286
// uv and npm are present; ug is not, so only ug gets installed and no
276287
// download happens
277-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
278-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
288+
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
289+
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
279290
marker := filepath.Join(t.TempDir(), "uv.log")
280291
t.Setenv("FAKE_UV_LOG", marker)
281292
t.Setenv("PATH", t.TempDir())
@@ -290,11 +301,33 @@ func TestEnsureToolchain(t *testing.T) {
290301
assert.NoFileExists(t, filepath.Join(home, agentRootDir, setupLockName), "the setup lock is released once setup finishes")
291302
})
292303

304+
t.Run("replaces a Unity Gateway CLI from another commit with the pinned one", func(t *testing.T) {
305+
home := t.TempDir()
306+
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
307+
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
308+
// Installs from an older pin and from a CLI before installs were versioned. Both
309+
// put a ug on PATH, but neither is the pinned commit.
310+
depsRoot := filepath.Join(home, agentDepsDir)
311+
oldBinDirs := []string{filepath.Join(depsRoot, "uv_tools_bin"), filepath.Join(depsRoot, "uv_tools_bin-oldcommit")}
312+
for _, dir := range oldBinDirs {
313+
writeFakeExec(t, dir, ugExecutable, ":")
314+
}
315+
marker := filepath.Join(t.TempDir(), "uv.log")
316+
t.Setenv("FAKE_UV_LOG", marker)
317+
t.Setenv("PATH", strings.Join(oldBinDirs, string(os.PathListSeparator)))
318+
319+
require.NoError(t, ensureToolchain(cmdio.MockDiscard(t.Context()), home, "amd64"))
320+
321+
out, err := os.ReadFile(marker)
322+
require.NoError(t, err)
323+
assert.Contains(t, string(out), "tool install git+https://github.com/"+ugRepo+"@"+ugCommit)
324+
})
325+
293326
t.Run("recheck after acquiring the lock skips a redundant install", func(t *testing.T) {
294327
home := t.TempDir()
295328
// uv is present and records if it ever installs; it must not be called once a
296329
// concurrent client finishes the toolchain while we wait for the lock.
297-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
330+
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", fakeUvRecorder)
298331
marker := filepath.Join(t.TempDir(), "uv.log")
299332
t.Setenv("FAKE_UV_LOG", marker)
300333
t.Setenv("PATH", t.TempDir())
@@ -310,7 +343,7 @@ func TestEnsureToolchain(t *testing.T) {
310343
// The other client finishes: publish ug+npm into the PATH dir, then release
311344
// the lock so the waiting ensureToolchain can proceed.
312345
writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
313-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
346+
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), "npm", ":")
314347
unlock()
315348

316349
require.NoError(t, <-done)
@@ -635,7 +668,7 @@ func TestBootstrapToolchain(t *testing.T) {
635668
t.Run("returns the path to the Unity Gateway CLI", func(t *testing.T) {
636669
home := t.TempDir()
637670
writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
638-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), npmExecutable, ":")
671+
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), npmExecutable, ":")
639672
t.Setenv("PATH", t.TempDir())
640673

641674
ugPath, err := bootstrapToolchain(cmdio.MockDiscard(t.Context()), home)
@@ -646,8 +679,8 @@ func TestBootstrapToolchain(t *testing.T) {
646679
t.Run("reports UNITY_GATEWAY_CLI_NOT_FOUND when the install leaves no CLI on PATH", func(t *testing.T) {
647680
home := t.TempDir()
648681
// uv "succeeds" without installing anything, so the post-setup lookup fails.
649-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", ":")
650-
writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), npmExecutable, ":")
682+
writeFakeExec(t, uvArchiveSpec("amd64").binDir(home), "uv", ":")
683+
writeFakeExec(t, nodeArchiveSpec("amd64").binDir(home), npmExecutable, ":")
651684
t.Setenv("PATH", t.TempDir())
652685

653686
_, err := bootstrapToolchain(cmdio.MockDiscard(t.Context()), home)

0 commit comments

Comments
 (0)