runconfig

type Compute struct {
	Workers         *int   `yaml:"workers"`
	GPUs            *int   `yaml:"gpus"`
	GPUsPerWorker   *int   `yaml:"gpus_per_worker"`
	CPUWorkers      *int   `yaml:"cpu_workers"`
	WorkloadKind    string `yaml:"workload_kind"`
	GPUResourceMode string `yaml:"gpu_resource_mode"`
	MIGProfile      string `yaml:"mig_profile"`
	CPURequest      string `yaml:"cpu_request"`
	MemoryRequest   string `yaml:"memory_request"`
	CPULimit        string `yaml:"cpu_limit"`
	MemoryLimit     string `yaml:"memory_limit"`
	HeadCPURequest  string `yaml:"head_cpu_request"`
	HeadMemRequest  string `yaml:"head_memory_request"`
	HeadCPULimit    string `yaml:"head_cpu_limit"`
	HeadMemLimit    string `yaml:"head_memory_limit"`
	WorkerCPUReq    string `yaml:"worker_cpu_request"`
	WorkerMemReq    string `yaml:"worker_memory_request"`
	WorkerCPULimit  string `yaml:"worker_cpu_limit"`
	WorkerMemLimit  string `yaml:"worker_memory_limit"`
}

type Config struct {
	SchemaVersion any    `yaml:"schema_version"`
	Name          string `yaml:"name"`
	Engine        string `yaml:"engine"`
	Entrypoint    string `yaml:"entrypoint"`
	Script        string `yaml:"script"`
	Image         string `yaml:"image"`

	Run        Run        `yaml:"run"`
	Workflow   Workflow   `yaml:"workflow"`
	Runtime    Runtime    `yaml:"runtime"`
	Compute    Compute    `yaml:"compute"`
	Policy     Policy     `yaml:"policy"`
	Storage    Storage    `yaml:"storage"`
	Profiler   Profiler   `yaml:"profiler"`
	Metrics    Metrics    `yaml:"metrics"`
	Experiment Experiment `yaml:"experiment"`
	Execution  Execution  `yaml:"execution"`
	Resilience Resilience `yaml:"resilience"`
}

type Duration struct{ time.Duration }

type Execution struct {
	Launcher            *string        `yaml:"launcher"`
	ProcessesPerNode    *int           `yaml:"processes_per_node"`
	Nodes               *int           `yaml:"nodes"`
	Metric              string         `yaml:"metric"`
	Mode                string         `yaml:"mode"`
	NumSamples          *int           `yaml:"num_samples"`
	MaxConcurrentTrials *int           `yaml:"max_concurrent_trials"`
	Configs             map[string]any `yaml:"configs"`
	AllowNCCLOverride   bool           `yaml:"allow_nccl_override"`
}

type Experiment struct {
	Project string `yaml:"project"`
	Name    string `yaml:"name"`
	Group   string `yaml:"group"`
	// Title is the pre-v0.1 experiment identity spelling. Keep accepting it so
	// existing configs migrate safely, but new configs must use Name.
	Title string `yaml:"title"`
}

type FieldInfo struct {
	Description      string
	Status           FieldStatus
	Values           []string
	DeprecatedValues []string
	Default          string
	Notes            string
}

type FieldStatus string

type ImageAsset struct {
	Name       string `yaml:"name"`
	Image      string `yaml:"image"`
	SourcePath string `yaml:"source_path"`
	MountPath  string `yaml:"mount_path"`
}

type Metrics struct {
	History []string       `yaml:"history"`
	Offload MetricsOffload `yaml:"offload"`
}

type MetricsOffload struct {
	Enabled               bool   `yaml:"enabled"`
	Runtime               string `yaml:"runtime"`
	Image                 string `yaml:"image"`
	Out                   string `yaml:"out"`
	DeliveryMode          string `yaml:"delivery_mode"`
	ADXClusterURI         string `yaml:"adx_cluster_uri"`
	ADXDatabase           string `yaml:"adx_database"`
	ADXTable              string `yaml:"adx_table"`
	ADXMapping            string `yaml:"adx_mapping"`
	ADXClientID           string `yaml:"adx_client_id"`
	ADXMaxAttempts        int    `yaml:"adx_max_attempts"`
	ADXRetryBackoff       string `yaml:"adx_retry_backoff"`
	ADXFinalStatusTimeout string `yaml:"adx_final_status_timeout"`
}

type NormalizedRDMA struct {
	Enabled bool
}

type Policy struct {
	Workspace                string            `yaml:"workspace"`
	Namespace                string            `yaml:"namespace"`
	Profile                  string            `yaml:"profile"`
	Queue                    string            `yaml:"queue"`
	Team                     string            `yaml:"team"`
	Lane                     string            `yaml:"lane"`
	GPUClass                 string            `yaml:"gpu_class"`
	Mode                     string            `yaml:"mode"`
	Topology                 string            `yaml:"topology"`
	Shape                    string            `yaml:"shape"`
	Priority                 string            `yaml:"priority"`
	PriorityTier             string            `yaml:"priority_tier"`
	WorkloadProfileSnapshot  string            `yaml:"workload_profile_snapshot"`
	WorkloadPriorityClass    string            `yaml:"workload_priority_class"`
	PodPriorityClass         string            `yaml:"pod_priority_class"`
	NodeSelector             map[string]string `yaml:"node_selector"`
	ClearNodeSelector        bool              `yaml:"clear_node_selector"`
	DisableDefaultPriorities *bool             `yaml:"disable_default_priorities"`
}

type Profiler struct {
	Mode     string `yaml:"mode"`
	Rank     string `yaml:"rank"`
	Warmup   string `yaml:"warmup"`
	Duration string `yaml:"duration"`
}

type RDMA struct {
	Enabled bool `yaml:"enabled"`
}

type Resilience struct {
	MaxRetries     int      `yaml:"max_retries"`
	RetryOn        []string `yaml:"retry_on"`
	CheckpointPath string   `yaml:"checkpoint_path"`
	BackoffInitial Duration `yaml:"backoff_initial"`
	BackoffMax     Duration `yaml:"backoff_max"`
}

type Run struct {
	Name                    string  `yaml:"name"`
	Engine                  string  `yaml:"engine"`
	Entrypoint              string  `yaml:"entrypoint"`
	Script                  string  `yaml:"script"`
	Image                   string  `yaml:"image"`
	MainScript              string  `yaml:"main_script"`
	WorkloadKind            string  `yaml:"workload_kind"`
	SmokePairs              *int    `yaml:"smoke_pairs"`
	Source                  *Source `yaml:"source"`
	TTLSecondsAfterFinished *int64  `yaml:"ttl_seconds_after_finished"`

	// WorkingDir ships a whole project directory with the run instead of the
	// entrypoint alone, so sibling modules and local packages import on
	// workers as well as the driver. Empty keeps single-file behaviour.
	WorkingDir string `yaml:"working_dir"`
	// WorkingDirExcludes are extra glob patterns to leave out of the shipped
	// project archive, on top of the built-in defaults.
	WorkingDirExcludes []string `yaml:"working_dir_excludes"`
}

type Runtime struct {
	Image      string            `yaml:"image"`
	WorkingDir string            `yaml:"working_dir"`
	Pip        []string          `yaml:"pip"`
	Env        map[string]string `yaml:"env"`
	EnvSecret  map[string]string `yaml:"env_secret"`
	EnvKV      map[string]string `yaml:"env_kv"`
	Security   Security          `yaml:"security"`
	RDMA       RDMA              `yaml:"rdma"`
}

type Security struct {
	Mode string `yaml:"mode"`
}

type Source struct {
	Image string `yaml:"image"`
	Path  string `yaml:"path"`
}

type Storage struct {
	DataPVC     string       `yaml:"data_pvc"`
	ResultPVC   string       `yaml:"result_pvc"`
	Output      string       `yaml:"output"`
	Publish     string       `yaml:"publish"`
	Volumes     []string     `yaml:"volumes"`
	Mounts      []string     `yaml:"mounts"`
	ImageAssets []ImageAsset `yaml:"image_assets"`

	// Checkpoint names the file or directory, relative to the run's
	// checkpoint directory, that this run produces as its servable model
	// (e.g. "last.safetensors"). Declaring it lets tau index the artifact
	// after a successful run so `tau serve deploy --from-finetune` and
	// `tau data model` can resolve it by run name instead of requiring the
	// researcher to pass an absolute --checkpoint path. Optional; when
	// empty, no artifact index is written and nothing else changes.
	Checkpoint string `yaml:"checkpoint"`
}

type UnknownKey struct {
	Path     string
	Children []string
}

type Workflow struct {
	File               string   `yaml:"file"`
	Script             string   `yaml:"script"`
	MainScript         string   `yaml:"main_script"`
	WorkloadKind       string   `yaml:"workload_kind"`
	ExtraScripts       []string `yaml:"extra_scripts"`
	UpstreamCheckpoint string   `yaml:"upstream_checkpoint"`
	SecretPayload      string   `yaml:"secret_payload"`
	SmokePairs         *int     `yaml:"smoke_pairs"`
}