platform

package
v0.1.0-rc.13 Latest Latest
Warning

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

Go to latest
Published: Aug 29, 2026 License: Apache-2.0 Imports: 15 Imported by: 0

Documentation

Overview

Package platform provides model-independent, CSP/GPU-architecture-only overrides for WorkloadRun resources. Override definitions live in YAML templates (overrides/workloadrun.yaml) that reference the same _lib/ fragments used by catalog entries, ensuring a single source of truth for CSP/GPU-specific networking, NCCL configuration, and DRA setup.

Index

Constants

View Source
const (
	AWS        = "aws"
	GCP        = "gcp"
	Azure      = "azure"
	OCI        = "oci"
	OnPrem     = "onprem"
	TogetherAI = "togetherai"
	Mistral    = "mistral"
	Forge      = "forge"
	NScale     = "nscale"
)

Platform name constants. These are the values the controller's platform detection (pkg/controller/workflow_detect.go) can return and the only values the CLI --platform flags accept. Detection returns these constants and the CLIs derive validation, help text, and error messages from Names(), so the two can never drift apart.

Variables

This section is empty.

Functions

func BaseNCCLEnvVars

func BaseNCCLEnvVars(enableMNNVL bool) []corev1.EnvVar

BaseNCCLEnvVars returns model-independent NCCL environment variables that are safe to auto-inject for any workload. These are common across all training models and NCCL tests.

Platform-specific NCCL vars (FastRak, IB settings) are injected via overrides, not here.

func BuildExecRuntime

func BuildExecRuntime(cfg RuntimeConfig) nvcrev1alpha1.DependencySpec

BuildExecRuntime creates a TrainingRuntime dependency for arbitrary command execution. Uses a simple single-replicatedJob layout with torch mlPolicy.

func BuildMPIRuntime

func BuildMPIRuntime(cfg RuntimeConfig) nvcrev1alpha1.DependencySpec

BuildMPIRuntime creates TrainingRuntime dependencies for MPI-based workloads. Generates: - A runtime with MPI mlPolicy and launcher+node replicatedJobs - Worker nodes with sshd, IPC_LOCK, readiness probe - Launcher with mpirun and SSH key setup

func BuildTorchRuntime

func BuildTorchRuntime(cfg RuntimeConfig) nvcrev1alpha1.DependencySpec

BuildTorchRuntime creates a TrainingRuntime dependency for PyTorch distributed training. Generates a runtime with torch mlPolicy and a single "node" replicatedJob.

func MergeEnvVars

func MergeEnvVars(base, user []corev1.EnvVar) []corev1.EnvVar

MergeEnvVars merges base env vars with user-provided env vars. User-provided values take precedence (override by name).

func Names

func Names() []string

Names returns every valid platform name, in the order used in flag help and error messages.

func NamesList

func NamesList() string

NamesList returns the valid platform names joined with ", " for flag help and error messages.

func ValidateFlag

func ValidateFlag(name string) error

ValidateFlag validates a --platform flag value. An empty value is allowed (platform not specified). The error text derives the valid names from Names() so the message always matches the actual set.

Types

type OverrideConfig

type OverrideConfig struct {
	EntryName     string
	NodesPerJob   int32
	GpusPerNode   int32
	MlnxPerNode   int32
	EnableMNNVL   bool
	FrameworkType string
}

OverrideConfig holds template data for rendering platform overrides. Field names match the catalog template data shape so _lib/ fragments can be rendered directly without mapping.

type RuntimeConfig

type RuntimeConfig struct {
	// EntryName is the WorkloadRun name (used as prefix for resource names).
	// Matches catalog template variable .EntryName.
	EntryName string
	// Image is the container image.
	Image string
	// NodesPerJob is the number of nodes.
	NodesPerJob int32
	// GpusPerNode is the number of GPUs per node.
	GpusPerNode int32
	// Env is the merged env vars (base NCCL + user).
	Env []corev1.EnvVar
	// Volumes are additional volumes.
	Volumes []corev1.Volume
	// VolumeMounts are additional volume mounts.
	VolumeMounts []corev1.VolumeMount
	// InitContainers are user-provided init containers.
	InitContainers []corev1.Container
	// Resources overrides GPU/memory/CPU resources.
	Resources *corev1.ResourceRequirements
	// ImagePullSecrets for container pull.
	ImagePullSecrets []corev1.LocalObjectReference
	// GangSchedulerName is the scheduler name to inject into pod specs (e.g. "kai-scheduler").
	// Empty means no gang scheduler is configured.
	GangSchedulerName string
	// GangSchedulerQueue is the queue label value for the gang scheduler.
	// Defaults to "default-queue" when GangSchedulerName is set and Queue is empty.
	GangSchedulerQueue string
}

RuntimeConfig holds all parameters needed to build a TrainingRuntime dependency.

type WorkloadRunOverride

type WorkloadRunOverride struct {
	nvcrev1alpha1.OverrideSpec

	// PreCommand contains shell lines prepended to the trainer command.
	// Applied by the WorkloadRun controller; baked into trainer.command/args.
	PreCommand []string

	// MPIArgs contains mpirun arguments prepended to the MPI launcher command.
	// Applied by the WorkloadRun controller; baked into trainer.args.
	MPIArgs []string
}

WorkloadRunOverride extends OverrideSpec with fields that are consumed by the WorkloadRun controller at build time and never stored in the Kubernetes API. Keeping them out of OverrideSpec avoids CRD schema changes.

func BuildOverrides

func BuildOverrides(cfg OverrideConfig) []WorkloadRunOverride

BuildOverrides renders the platform override templates and returns the resulting WorkloadRunOverride list.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL