Stage multi-host launches in a verified per-user directory

docs/MULTI-NODE-TRAINING.md

@@ -62,6 +62,23 @@ launch. The binary, worker scratch, and resume staging are removed when the
 session finishes or aborts. Configure the same absolute scratch path on storage
 with enough capacity on every node; launcher state does not use `/tmp`.
 
+Uploads go to a unique temporary file in the session directory and are
+published only after their SHA-256 matches; a mismatch leaves nothing behind.
+WALDO refuses to stage the binary, worker state, or checkpoints unless the
+session directory, every ancestor, and every ancestor of its resolved path are
+owned by the remote user or root and cannot be modified by another account:
+
+- Directories owned by the remote user must not be writable by their group or
+  other users. Remove such access with `chmod go-w`. On Linux the group
+  permission bits include the POSIX ACL mask, so named ACL entries cannot grant
+  write access either.
+- Root-owned directories writable by other users, such as `/tmp`, must be
+  sticky.
+
+Staging fails and names the offending path otherwise. If another account
+created a directory on the scratch path first, an administrator must remove it
+before retrying.
+
 WALDO does not install or modify GPU drivers, CUDA, NCCL, Python, PyTorch, or
 TorchTitan. Those machine-level runtimes must be installed before training.
 

internal/cli/model_hostfile.go

@@ -470,16 +470,103 @@ func (session *hostfileSession) remoteCommandContext(ctx context.Context, argume
 	return command
 }
 
+// remoteStagingFunctions defines POSIX shell functions for staging files on
+// secondary hosts. waldo_private_dir creates an absolute directory and
+// verifies that it, every ancestor, and every ancestor of its resolved path
+// are owned by the remote user or root and cannot be modified by another
+// account. Directories owned by the remote user must not be writable by their
+// group or other users. Root-owned directories writable by others, such as
+// /tmp, must be sticky. On Linux the group permission bits report the POSIX ACL
+// mask, which bounds every named ACL entry; other systems' extended ACLs cannot
+// be evaluated from the mode and are refused. Only portable tools are used so
+// the same script runs on Linux hosts and in macOS unit tests.
+const remoteStagingFunctions = `waldo_uid=$(id -u)
+waldo_system=$(uname -s)
+waldo_check_entry() {
+	waldo_listing=$(LC_ALL=C ls -ldn -- "$1") || return 1
+	set -- "$1" $waldo_listing
+	case $4 in
+	"$waldo_uid") waldo_owner=user ;;
+	0) waldo_owner=root ;;
+	*) echo "$1 is owned by uid $4, not uid $waldo_uid or root" >&2; return 1 ;;
+	esac
+	case $2 in l*) return 0 ;; esac
+	case $2 in
+	??????????+*)
+		if [ "$waldo_system" != Linux ]; then
+			echo "$1 has an access control list that cannot be verified" >&2
+			return 1
+		fi
+		;;
+	esac
+	case $2 in ?????w*|????????w*) ;; *) return 0 ;; esac
+	if [ "$waldo_owner" = root ]; then
+		case $2 in ?????????[tT]*) return 0 ;; esac
+		echo "$1 is writable by other users and is not sticky" >&2
+		return 1
+	fi
+	echo "$1 is writable by its group or other users; remove that access with chmod go-w" >&2
+	return 1
+}
+waldo_check_ancestors() {
+	waldo_path=$1
+	while :; do
+		waldo_check_entry "$waldo_path" || return 1
+		[ "$waldo_path" = / ] && return 0
+		waldo_path=$(dirname -- "$waldo_path")
+	done
+}
+waldo_private_dir() {
+	case $1 in /*) ;; *) echo "staging directory $1 is not absolute" >&2; return 1 ;; esac
+	mkdir -p -- "$1" || return 1
+	waldo_physical=$(cd -P -- "$1" && pwd -P) || return 1
+	waldo_check_ancestors "$1" || return 1
+	waldo_check_ancestors "$waldo_physical"
+}
+waldo_sha256() {
+	if command -v sha256sum >/dev/null 2>&1; then
+		waldo_sum=$(sha256sum -- "$1")
+	else
+		waldo_sum=$(shasum -a 256 -- "$1")
+	fi
+	printf '%s\n' "${waldo_sum%% *}"
+}
+`
+
+// remoteStageScript returns a script that receives a file on standard input
+// and publishes it at target only after its SHA-256 matches digest. The upload
+// goes to a unique temporary file inside the verified private directory and
+// is removed on any failure.
+func remoteStageScript(target, digest, mode, label string) string {
+	directory := shellQuote(filepath.Dir(target))
+	sum := shellQuote(digest)
+	return "set -euf\numask 077\n" + remoteStagingFunctions +
+		"waldo_private_dir " + directory + "\n" +
+		"tmp=$(mktemp " + directory + "/.waldo-transfer.XXXXXX)\n" +
+		"trap 'rm -f -- \"$tmp\"' EXIT\n" +
+		"cat > \"$tmp\"\n" +
+		"chmod " + mode + " \"$tmp\"\n" +
+		"sum=$(waldo_sha256 \"$tmp\")\n" +
+		"if [ \"$sum\" != " + sum + " ]; then echo \"staged " + label + " SHA-256 $sum does not match \"" + sum + " >&2; exit 1; fi\n" +
+		"mv -f -- \"$tmp\" " + shellQuote(target) + "\n"
+}
+
+// remoteVerifiedFileScript succeeds only when target already exists inside a
+// verified private directory with the expected SHA-256.
+func remoteVerifiedFileScript(target, digest string) string {
+	return "set -euf\numask 077\n" + remoteStagingFunctions +
+		"waldo_private_dir " + shellQuote(filepath.Dir(target)) + "\n" +
+		"test -f " + shellQuote(target) + "\n" +
+		"test \"$(waldo_sha256 " + shellQuote(target) + ")\" = " + shellQuote(digest) + "\n"
+}
+
 func (session *hostfileSession) stageBinary(host string) error {
 	file, err := os.Open(session.binary)
 	if err != nil {
 		return err
 	}
 	defer file.Close()
-	temporary := session.remoteBinary + ".tmp"
-	remote := fmt.Sprintf("umask 077; mkdir -p %s; cat > %s; chmod 700 %s; test \"$(sha256sum %s | cut -d' ' -f1)\" = %s; mv -f %s %s",
-		shellQuote(session.remoteRoot), shellQuote(temporary), shellQuote(temporary), shellQuote(temporary), shellQuote(session.binarySHA256), shellQuote(temporary), shellQuote(session.remoteBinary))
-	command := session.remoteCommand(host, remote)
+	command := session.remoteCommand(host, remoteStageScript(session.remoteBinary, session.binarySHA256, "700", "WALDO"))
 	command.Stdin = file
 	var output strings.Builder
 	command.Stdout, command.Stderr = &output, &output
@@ -814,18 +901,14 @@ func (session *hostfileSession) stageResume(host string, artifacts []training.Ar
 	for index, artifact := range artifacts {
 		localPath := sources[index]
 		remotePath := targets[index]
-		check := fmt.Sprintf("test -f %s; test \"$(sha256sum %s | cut -d' ' -f1)\" = %s", shellQuote(remotePath), shellQuote(remotePath), shellQuote(artifact.SHA256))
-		if err := session.remoteCommand(host, check).Run(); err == nil {
+		if err := session.remoteCommand(host, remoteVerifiedFileScript(remotePath, artifact.SHA256)).Run(); err == nil {
 			continue
 		}
 		file, err := os.Open(localPath)
 		if err != nil {
 			return fmt.Errorf("stage checkpoint on %s: open %s: %w", host, artifact.Path, err)
 		}
-		temporary := remotePath + ".waldo-transfer"
-		remote := fmt.Sprintf("umask 077; mkdir -p %s; cat > %s; test \"$(sha256sum %s | cut -d' ' -f1)\" = %s; mv -f %s %s",
-			shellQuote(filepath.Dir(remotePath)), shellQuote(temporary), shellQuote(temporary), shellQuote(artifact.SHA256), shellQuote(temporary), shellQuote(remotePath))
-		command := session.remoteCommand(host, remote)
+		command := session.remoteCommand(host, remoteStageScript(remotePath, artifact.SHA256, "600", "checkpoint"))
 		command.Stdin = file
 		var output strings.Builder
 		command.Stdout, command.Stderr = &output, &output

internal/cli/model_hostfile_test.go

@@ -346,8 +346,9 @@ case "$*" in
     exit 0
     ;;
   *)
-    cat >/dev/null
-    exit 0
+    while [ "$#" -gt 0 ] && [ "$1" != -- ]; do shift; done
+    shift 2
+    exec sh -c "$1"
     ;;
 esac
 `
@@ -368,7 +369,7 @@ esac
 	}
 	t.Cleanup(func() { listenHostfileRendezvous = previousListener })
 	cluster := training.Cluster{Nodes: 2, Rendezvous: "127.0.0.1:0", RendezvousID: "session-test"}
-	scratch := t.TempDir()
+	scratch := privateTempDir(t)
 	cacheRoot := t.TempDir()
 	cache, err := lookaside.NewCache(cacheRoot, nil, lookaside.WithPersistentStorage(scratch, 1<<20))
 	if err != nil {
@@ -388,6 +389,7 @@ esac
 	if !strings.HasPrefix(session.remoteRoot, scratch+string(os.PathSeparator)) {
 		t.Fatalf("remote launch root = %q, want beneath configured scratch %q", session.remoteRoot, scratch)
 	}
+	assertFileSHA256(t, session.remoteBinary, session.binarySHA256)
 	evaluation := training.EvaluationSet{Selection: "lowest-sha256-v1", SHA256: strings.Repeat("a", 64)}
 	if err := session.publish(model.MultiNodePlan{
 		Kind: model.MultiNodePlanKind, Schema: model.MultiNodePlanSchema,
@@ -752,3 +754,201 @@ func TestHostfileRemoteCommandUsesFuzzballWrapperContract(t *testing.T) {
 		t.Fatal("Fuzzball wrapper lacks bounded graceful cancellation")
 	}
 }
+
+// installLocalSSH installs an ssh shim that runs the remote script on the
+// local host, exercising the real staging scripts with the test runner's
+// tools.
+func installLocalSSH(t *testing.T) {
+	t.Helper()
+	bin := t.TempDir()
+	script := "#!/bin/sh\nwhile [ \"$#\" -gt 0 ] && [ \"$1\" != -- ]; do shift; done\nshift 2\nexec sh -c \"$1\"\n"
+	if err := os.WriteFile(filepath.Join(bin, "ssh"), []byte(script), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	t.Setenv("PATH", bin+string(os.PathListSeparator)+os.Getenv("PATH"))
+}
+
+func newStagingSession(t *testing.T, scratch string) *hostfileSession {
+	t.Helper()
+	binary := filepath.Join(t.TempDir(), "waldo")
+	if err := os.WriteFile(binary, []byte("#!/bin/sh\nexit 0\n"), 0o700); err != nil {
+		t.Fatal(err)
+	}
+	digest, err := fileSHA256(binary)
+	if err != nil {
+		t.Fatal(err)
+	}
+	root := filepath.Join(scratch, "multinode", digest, "session-test")
+	return &hostfileSession{
+		ctx: context.Background(), cancel: func() {},
+		hostfile: trainingHostfile{Hosts: []string{"train-0", "train-1"}},
+		binary:   binary, binarySHA256: digest,
+		remoteRoot: root, remoteBinary: filepath.Join(root, "waldo"), resumeRoot: filepath.Join(root, "resume"),
+		output: io.Discard,
+	}
+}
+
+// privateTempDir returns a test directory that staging accepts. The testing
+// package creates t.TempDir with mode 0777 less the umask, which is
+// group-writable under the common umask 002.
+func privateTempDir(t *testing.T) string {
+	t.Helper()
+	directory := t.TempDir()
+	if err := os.Chmod(directory, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	return directory
+}
+
+func privateScratch(t *testing.T) string {
+	t.Helper()
+	scratch := filepath.Join(privateTempDir(t), "scratch")
+	if err := os.Mkdir(scratch, 0o700); err != nil {
+		t.Fatal(err)
+	}
+	return scratch
+}
+
+func assertFileSHA256(t *testing.T, path, want string) {
+	t.Helper()
+	got, err := fileSHA256(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got != want {
+		t.Fatalf("%s SHA-256 = %s, want %s", path, got, want)
+	}
+}
+
+func assertNoTransferFiles(t *testing.T, directory string) {
+	t.Helper()
+	entries, err := os.ReadDir(directory)
+	if err != nil && !os.IsNotExist(err) {
+		t.Fatal(err)
+	}
+	for _, entry := range entries {
+		if strings.HasPrefix(entry.Name(), ".waldo-transfer.") {
+			t.Fatalf("temporary upload %s remains in %s", entry.Name(), directory)
+		}
+	}
+}
+
+func TestHostfileStageBinaryIntoVerifiedScratch(t *testing.T) {
+	installLocalSSH(t)
+	session := newStagingSession(t, privateScratch(t))
+	for attempt := 0; attempt < 2; attempt++ {
+		if err := session.stageBinary("train-1"); err != nil {
+			t.Fatalf("stage attempt %d: %v", attempt+1, err)
+		}
+	}
+	assertFileSHA256(t, session.remoteBinary, session.binarySHA256)
+	for path, want := range map[string]os.FileMode{session.remoteBinary: 0o700, session.remoteRoot: 0o700} {
+		info, err := os.Stat(path)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if info.Mode().Perm() != want {
+			t.Fatalf("%s mode = %v, want %v", path, info.Mode().Perm(), want)
+		}
+	}
+	assertNoTransferFiles(t, session.remoteRoot)
+	session.cleanupStaging()
+	if _, err := os.Stat(session.remoteRoot); !os.IsNotExist(err) {
+		t.Fatalf("launch staging remains after cleanup: %v", err)
+	}
+}
+
+func TestHostfileStageBinaryRefusesDirectoriesOtherAccountsCanModify(t *testing.T) {
+	installLocalSSH(t)
+	for _, test := range []struct {
+		name string
+		mode os.FileMode
+	}{
+		{name: "other-writable", mode: 0o777},
+		{name: "group-writable", mode: 0o770},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			scratch := privateScratch(t)
+			if err := os.Chmod(scratch, test.mode); err != nil {
+				t.Fatal(err)
+			}
+			session := newStagingSession(t, scratch)
+			err := session.stageBinary("train-1")
+			if err == nil || !strings.Contains(err.Error(), "writable by its group or other users") {
+				t.Fatalf("stage error = %v, want writable-directory refusal", err)
+			}
+			if !strings.Contains(err.Error(), filepath.Base(scratch)) {
+				t.Fatalf("stage error %q does not name the unsafe directory", err)
+			}
+			if _, statErr := os.Stat(session.remoteBinary); !os.IsNotExist(statErr) {
+				t.Fatalf("binary was staged in an unsafe directory: %v", statErr)
+			}
+		})
+	}
+	t.Run("symlink-to-writable", func(t *testing.T) {
+		base := privateTempDir(t)
+		shared := filepath.Join(base, "shared")
+		if err := os.Mkdir(shared, 0o700); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.Chmod(shared, 0o777); err != nil {
+			t.Fatal(err)
+		}
+		link := filepath.Join(base, "scratch")
+		if err := os.Symlink(shared, link); err != nil {
+			t.Fatal(err)
+		}
+		session := newStagingSession(t, link)
+		if err := session.stageBinary("train-1"); err == nil || !strings.Contains(err.Error(), "shared is writable") {
+			t.Fatalf("stage error = %v, want refusal of the resolved writable directory", err)
+		}
+	})
+}
+
+func TestHostfileStageBinaryFailsClosedOnDigestMismatch(t *testing.T) {
+	installLocalSSH(t)
+	session := newStagingSession(t, privateScratch(t))
+	session.binarySHA256 = strings.Repeat("0", 64)
+	err := session.stageBinary("train-1")
+	if err == nil || !strings.Contains(err.Error(), "does not match") {
+		t.Fatalf("stage error = %v, want digest mismatch", err)
+	}
+	if _, statErr := os.Stat(session.remoteBinary); !os.IsNotExist(statErr) {
+		t.Fatalf("binary was published despite a digest mismatch: %v", statErr)
+	}
+	assertNoTransferFiles(t, session.remoteRoot)
+}
+
+func TestHostfileStageResumeVerifiesCheckpointsInPrivateScratch(t *testing.T) {
+	installLocalSSH(t)
+	scratch := privateScratch(t)
+	session := newStagingSession(t, scratch)
+	source := filepath.Join(t.TempDir(), "runtime.pt")
+	if err := os.WriteFile(source, []byte("synthetic checkpoint"), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	digest, err := fileSHA256(source)
+	if err != nil {
+		t.Fatal(err)
+	}
+	target := filepath.Join(session.resumeRoot, "run-42", "runtime.pt")
+	wrong := training.Artifact{Path: "runtime.pt", SHA256: strings.Repeat("0", 64)}
+	if err := session.stageResume("train-1", []training.Artifact{wrong}, []string{source}, []string{target}); err == nil || !strings.Contains(err.Error(), "does not match") {
+		t.Fatalf("stage error = %v, want digest mismatch", err)
+	}
+	if _, statErr := os.Stat(target); !os.IsNotExist(statErr) {
+		t.Fatalf("checkpoint was published despite a digest mismatch: %v", statErr)
+	}
+	assertNoTransferFiles(t, filepath.Dir(target))
+	artifact := training.Artifact{Path: "runtime.pt", SHA256: digest}
+	if err := session.stageResume("train-1", []training.Artifact{artifact}, []string{source}, []string{target}); err != nil {
+		t.Fatal(err)
+	}
+	assertFileSHA256(t, target, digest)
+	if err := os.Chmod(scratch, 0o777); err != nil {
+		t.Fatal(err)
+	}
+	if err := session.stageResume("train-1", []training.Artifact{artifact}, []string{source}, []string{target}); err == nil || !strings.Contains(err.Error(), "writable by its group or other users") {
+		t.Fatalf("stage error = %v, want an already staged checkpoint in an unsafe directory to be refused", err)
+	}
+}