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
- func ApplyGangSchedulerToDependencies(deps []nvcrev1alpha1.DependencySpec, gs *nvcrev1alpha1.GangSchedulerSpec) error
- func BaseNCCLEnvVars(enableMNNVL bool) []corev1.EnvVar
- func BuildExecRuntime(cfg RuntimeConfig) nvcrev1alpha1.DependencySpec
- func BuildMPIRuntime(cfg RuntimeConfig) nvcrev1alpha1.DependencySpec
- func BuildTorchRuntime(cfg RuntimeConfig) nvcrev1alpha1.DependencySpec
- func MergeEnvVars(base, user []corev1.EnvVar) []corev1.EnvVar
- func Names() []string
- func NamesList() string
- func ValidateFlag(name string) error
- type OverrideConfig
- type RuntimeConfig
- type WorkloadRunOverride
Constants ¶
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 ApplyGangSchedulerToDependencies ¶ added in v0.2.0
func ApplyGangSchedulerToDependencies(deps []nvcrev1alpha1.DependencySpec, gs *nvcrev1alpha1.GangSchedulerSpec) error
ApplyGangSchedulerToDependencies rewrites every TrainingRuntime dependency in place so each of its replicatedJobs runs under the configured gang scheduler. For each replicatedJob it sets schedulerName on the pod spec and the kai.scheduler/queue label on the job template metadata, matching what BuildTorchRuntime and BuildMPIRuntime already emit for a WorkloadRun.
It is a no-op when gs is nil, so a Certification that does not ask for gang scheduling renders byte-identically to before. It is also a no-op when the scheduler name is empty, matching applyGangScheduler: the CRD requires a non-empty schedulerName, but nvcrectl renders straight from a file without consulting the API server, so without this guard a typo would render pod templates pinned to an empty scheduler and carrying a stray queue label.
Callers must invoke this after overrides are resolved. The scheduler name is overwritten unconditionally rather than filled in only when absent, because some catalog entries hardcode schedulerName: default-scheduler and the whole point of the field is to replace it.
func BaseNCCLEnvVars ¶
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, and cfg.Env - Launcher with mpirun, SSH key setup, and cfg.Env
cfg.Env goes on both containers as container-level env, the same way BuildTorchRuntime emits it (issue #68: it used to be dropped here, so spec.env behaved differently between the two frameworks). On the launcher the variables reach mpirun directly; on the worker they reach the sshd process, but sshd gives each SSH session a fresh environment, so ranks launched through it may not inherit them. A variable that must reach the ranks should be passed as "-x NAME=value" in mpiArgs, which forwards it through mpirun itself.
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 ¶
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 ¶
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.
Source Files