diff --git a/docs/MULTI-NODE-TRAINING.md b/docs/MULTI-NODE-TRAINING.md index 75c1a359..1aad3a2b 100644 --- a/docs/MULTI-NODE-TRAINING.md +++ b/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. diff --git a/internal/cli/model_hostfile.go b/internal/cli/model_hostfile.go index 382d8933..a2036629 100644 --- a/internal/cli/model_hostfile.go +++ b/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 diff --git a/internal/cli/model_hostfile_test.go b/internal/cli/model_hostfile_test.go index 02267821..71f3215f 100644 --- a/internal/cli/model_hostfile_test.go +++ b/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) + } +}