checkpoint

package
v0.0.0-...-017a74b Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Oct 16, 2026 License: MIT Imports: 17 Imported by: 0

Documentation

Overview

Package checkpoint publishes and resolves training checkpoints under one checkpoint root, on a POSIX filesystem reached through an afero.Fs.

Each execution writes into its own folder, <Checkpoint>/exec/<DISPATCH_ATTEMPT>/. A complete save is published by exclusively creating an empty file whose name carries the whole publication, <Checkpoint>/pub/<epoch:020d>-<step:020d>-<base32 of the save's relative path>. The latest checkpoint is the entry with the greatest (epoch, step), found by a plain read. The design is spec 0001 D16; the code is on spec 0005's simulated path and follows its D2 rules.

Store is the same log on any storage: Dir.Store for POSIX, package checkpoint/s3 for S3-compatible object storage, so this package links no S3 signer.

Index

Examples

Constants

View Source 
const (
	ExecDirName = "exec"
	PubDirName  = "pub"
)

ExecDirName and PubDirName are the two folders directly under a checkpoint root.

Variables

View Source 
var (
	// ErrMalformed reports a value that does not parse: an entry name, an epoch, a step or a
	// framework's completion marker.
	ErrMalformed = errors.New("checkpoint: malformed")
	// ErrInvalid reports an argument the operation refuses.
	ErrInvalid = errors.New("checkpoint: invalid")
	// ErrExists reports that the publication's entry already exists. It was not touched.
	ErrExists = errors.New("checkpoint: already published")
)

Functions

func Overlaps 

func Overlaps(save string, entries []Entry) bool

Overlaps reports whether any of entries names save, a path inside it, or a path containing it. Prune keeps such a folder, on any store.

func ParseEpoch 

func ParseEpoch(s string) (uint64, error)

ParseEpoch parses DISPATCH_EPOCH (spec 0003 D8): a decimal uint64 with no sign or spaces.

func ParseStep 

func ParseStep(s string) (uint64, error)

ParseStep parses a training step: a decimal uint64 with no sign or spaces.

func ValidSave 

func ValidSave(rel string) error

ValidSave refuses a slash-separated relative save path no entry could name: not clean, outside the root, the root itself, or inside pub/.

Types

type Dir 

type Dir struct {
	FS   afero.Fs
	Root string
}

Dir is a checkpoint root, <Checkpoint>, on a filesystem. Root must be absolute and must already exist: it stands for the site's storage, and durability above it is not this package's.

Example

Two executions write into their own folders and publish. The older execution publishes a save late, at a higher step, but its epoch is lower, so Resolve never chooses it. 

package main

import (
	"errors"
	"fmt"
	"path/filepath"

	"github.com/spf13/afero"

	"gitlab.com/phpboyscout/go/dispatch/checkpoint"
)

func main() {
	fsys := afero.NewMemMapFs()
	root := "/ckpt"

	if err := fsys.MkdirAll(root, 0o755); err != nil {
		fmt.Println(err)

		return
	}

	d := checkpoint.Dir{FS: fsys, Root: root}

	// DISPATCH_EPOCH for submission 1 and submission 2 of a job of generation 0.
	const first, second = 1_000_000, 2_000_000

	publish := func(attempt string, epoch, step uint64) {
		exec, err := d.ExecDir(attempt)
		if err != nil {
			fmt.Println(err)

			return
		}

		save := filepath.Join(exec, fmt.Sprintf("checkpoint-%d", step))
		if err := fsys.MkdirAll(save, 0o755); err != nil {
			fmt.Println(err)

			return
		}

		if _, err := d.Publish(epoch, step, save); err != nil {
			fmt.Println(err)
		}
	}

	publish("1.4101.0", first, 500)
	publish("2.4102.0", second, 100)
	publish("1.4101.0", first, 900) // the old run, still writing after its job ended

	latest, ok, err := d.Resolve()
	if err != nil || !ok {
		fmt.Println(ok, err)

		return
	}

	fmt.Println("resume from", d.Path(latest))

	_, err = d.Publish(second, 100, d.Path(latest))
	fmt.Println(errors.Is(err, checkpoint.ErrExists))

}

Output:
resume from /ckpt/exec/2.4102.0/checkpoint-100
true

func (Dir) Entries 

func (d Dir) Entries() ([]Entry, error)

Entries returns every entry in pub/, ascending. A name that does not parse is an error naming it, never skipped. A missing pub/ has no entries; a missing root is an error.

func (Dir) ExecDir 

func (d Dir) ExecDir(attempt string) (string, error)

ExecDir returns <Checkpoint>/exec/<attempt>/, the folder one execution writes into. The attempt is opaque (spec 0003 D8) but must be usable as one path element.

func (Dir) Path 

func (d Dir) Path(e Entry) string

Path returns the absolute path of the save an entry names.

func (Dir) Prune 

func (d Dir) Prune(keep int) (PruneResult, error)

Prune removes all but the keep latest entries, then the folders the removed entries named. The resolved entry is always among those kept, and a folder is left in place while any remaining entry names it, a path inside it, or a path containing it. A folder reached through a symbolic link is never removed; it is reported in Skipped.

func (Dir) Publish 

func (d Dir) Publish(epoch, step uint64, save string) (Entry, error)

Publish publishes the save at the absolute path save as of epoch and step, by creating its entry with O_CREAT|O_EXCL and then syncing pub/ and the root (0001 D16). A returned nil means the entry is durable. An entry that already exists is reported with ErrExists and left as it is.

func (Dir) Resolve 

func (d Dir) Resolve() (e Entry, ok bool, err error)

Resolve returns the latest publication: the entry with the greatest (epoch, step). It is a pure read, with no lock and no check that the save is still there (0001 D16). ok is false when nothing is published.

func (Dir) Store 

func (d Dir) Store() Store

Store returns d as a Store. The context is not consulted: POSIX calls through afero cannot be cancelled.

type Entry 

type Entry struct {
	Epoch uint64
	Step  uint64
	Save  string
}

Entry is one publication: a save, by its slash-separated path relative to the checkpoint root, published at an epoch and a step.

func Latest 

func Latest(entries []Entry) (Entry, bool)

Latest returns the greatest of entries sorted ascending, as Entries returns them.

func ParseEntryName 

func ParseEntryName(name string) (Entry, error)

ParseEntryName parses a file name from pub/. Anything but the exact form Name produces is refused.

func (Entry) Less 

func (e Entry) Less(o Entry) bool

Less orders entries by (epoch, step), then by save path so that equal pairs still order the same way on every read.

func (Entry) Name 

func (e Entry) Name() (string, error)

Name returns the entry's file name in pub/, refusing a save path it could not parse back.

type Folder 

type Folder struct {
	Name     string
	Location string
}

Folder is one folder directly inside a store location: its name, and its location in the store's own form.

type Framework 

type Framework int

Framework is a training framework whose saves Watch recognises, each by its folder name and the marker the framework writes once a save is complete. 

const (
	// HuggingFaceTrainer saves into checkpoint-<step>/ and writes trainer_state.json last.
	HuggingFaceTrainer Framework = iota + 1
	// PyTorchDCP saves into step-<step>/, torchtitan's layout, and renames .metadata into place
	// once every rank has written.
	PyTorchDCP
)

The frameworks Watch knows.

func ParseFramework 

func ParseFramework(s string) (Framework, error)

ParseFramework parses a framework name: "hf" or "dcp".

func (Framework) String 

func (f Framework) String() string

String returns the name ParseFramework accepts.

type PruneResult 

type PruneResult struct {
	Entries []Entry
	Folders []string
	Skipped []SkippedFolder
}

PruneResult lists what Prune removed, and the folders it left in place and why.

type SaveReader 

type SaveReader interface {
	// ReadFile returns the content of the file name inside save, wrapping fs.ErrNotExist when
	// there is none.
	ReadFile(ctx context.Context, save, name string) ([]byte, error)
	// Exists reports whether the file name inside save is present and complete.
	Exists(ctx context.Context, save, name string) (bool, error)
}

SaveReader reads the files inside a save, so one framework completion rule (0001 D16) runs on every store. A save is a location in the store's own form, as Store.ExecDir returns it.

type SkippedFolder 

type SkippedFolder struct {
	Location string
	Reason   string
}

SkippedFolder is a folder Prune would have removed but left in place, with the reason.

type Store 

type Store interface {
	// Publish creates the entry for the save at save, a location under the root in the store's
	// own form, and returns ErrExists, with the entry, when it is already there.
	Publish(ctx context.Context, epoch, step uint64, save string) (Entry, error)
	// Entries returns every entry, ascending, refusing a name that does not parse.
	Entries(ctx context.Context) ([]Entry, error)
	// Resolve returns the latest entry by a pure read; ok is false when nothing is published.
	Resolve(ctx context.Context) (e Entry, ok bool, err error)
	// Prune removes all but the keep latest entries and the folders only they named.
	Prune(ctx context.Context, keep int) (PruneResult, error)
	// Location returns where the save an entry names lives, in the store's own form.
	Location(e Entry) string
	// ExecDir returns <Checkpoint>/exec/<attempt>, in the store's own form.
	ExecDir(attempt string) (string, error)
}

Store is one checkpoint root's publication log on whichever storage holds it: a POSIX filesystem through Dir, or S3-compatible object storage through package checkpoint/s3. Both keep the same entry names and the same rules (0001 D16), so the helper drives either one.

type WatchStore 

type WatchStore interface {
	Store
	SaveReader
	// Folders returns the folders directly inside dir, in no particular order, and none when dir
	// does not exist. A root that does not exist is an error.
	Folders(ctx context.Context, dir string) ([]Folder, error)
}

WatchStore is a Store a Watcher can look into for saves: a POSIX root through Dir, or an s3:// root through package checkpoint/s3 (0001 D16).

type Watcher 

type Watcher struct {
	// Dir is the POSIX root watched when Store is nil.
	Dir Dir
	// Store, when set, is the root watched instead of Dir, such as an s3.Store.
	Store     WatchStore
	Framework Framework
	// Saves is the folder the framework writes its saves into, normally the store's
	// ExecDir(attempt), in the store's own form.
	Saves    string
	Epoch    uint64
	Interval time.Duration
	// ScanTimeout, when positive, bounds each scan, so a store that stops answering ends Watch
	// with an error instead of hanging it.
	ScanTimeout time.Duration
	// Published, if set, is called with each entry Scan creates, in step order.
	Published func(Entry)
	// contains filtered or unexported fields
}

Watcher publishes a framework's complete saves from one execution's folder, on a POSIX root or any other WatchStore (0001 D16).

func (*Watcher) Scan 

func (w *Watcher) Scan(ctx context.Context) ([]Entry, error)

Scan publishes, in step order, every complete save not yet published by this Watcher. A save whose entry already exists counts as published. An incomplete save is passed over until a later scan. It returns the entries it created.

func (*Watcher) Watch 

func (w *Watcher) Watch(ctx context.Context) error

Watch scans every Interval until ctx is cancelled or its deadline passes, then returns nil. It starts no goroutine and keeps one timer (spec 0005 D2 rules 1, 6 and 7). A scan error ends it.

Source Files

View all Source files

Directories

Path Synopsis
Package creds resolves the S3 credentials the checkpoint-publication helper signs with, from an ordered list of credential sources (spec 0001 D16).
 Package creds resolves the S3 credentials the checkpoint-publication helper signs with, from an ordered list of credential sources (spec 0001 D16).
vault
Package vault is the vault credential-source kind (spec 0001 D16): S3 credentials read from a Vault KV v2 secret with go/config-vault, through a client from go/vaultclient, after logging in with Kubernetes auth or AppRole.
 Package vault is the vault credential-source kind (spec 0001 D16): S3 credentials read from a Vault KV v2 secret with go/config-vault, through a client from go/vaultclient, after logging in with Kubernetes auth or AppRole.
Package proxy is the one rule for the HTTP proxy a checkpoint-publication client may be told to use (spec 0001 D16).
 Package proxy is the one rule for the HTTP proxy a checkpoint-publication client may be told to use (spec 0001 D16).
s3
Package s3 is the checkpoint publication log on S3-compatible object storage (spec 0001 D16).
 Package s3 is the checkpoint publication log on S3-compatible object storage (spec 0001 D16).
s3fake
Package s3fake is an in-process S3 server for tests: just the calls checkpoint publication makes (a put that honours If-None-Match: *, GET, ListObjectsV2 with prefix, delimiter and continuation, HEAD, DELETE), each refused unless its SigV4 signature verifies.
 Package s3fake is an in-process S3 server for tests: just the calls checkpoint publication makes (a put that honours If-None-Match: *, GET, ListObjectsV2 with prefix, delimiter and continuation, HEAD, DELETE), each refused unless its SigV4 signature verifies.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.