runtime

type Container struct {
	Name         string
	Env          []corev1ac.EnvVarApplyConfiguration
	Ports        []corev1ac.ContainerPortApplyConfiguration
	VolumeMounts []corev1ac.VolumeMountApplyConfiguration
}

type Info struct {
	// Labels and Annotations to add to the RuntimeJobTemplate.
	Labels      map[string]string
	Annotations map[string]string
	// Original policy values from the runtime.
	RuntimePolicy RuntimePolicy
	// Scheduler parameters to add to the RuntimeJobTemplate.
	Scheduler *Scheduler
	// TemplateSpec is TrainingRuntime Template object.
	// ObjApply podSpecs and this PodSets should be kept in sync by info.SyncPodSetsToTemplateSpec().
	TemplateSpec TemplateSpec
}

type InfoOption func(options *InfoOptions)

type InfoOptions struct {
	// contains filtered or unexported fields
}

type PodSet struct {
	// PodSet name is the name to identify PodSpec.
	// This typically has the name stored in each PodSpec.
	Name string
	// Ancestor is built by `trainer.kubeflow.org/trainjob-ancestor-step` label value
	// in Runtime CRDs.
	Ancestor       *string
	Count          *int32
	InitContainers []Container
	Containers     []Container
	Volumes        []corev1ac.VolumeApplyConfiguration
	Endpoints      iter.Seq[string]
	// The total PodSet requests can be calculated with
	// SinglePodRequests x Count.
	SinglePodRequests corev1.ResourceList
}

type ReconcilerBuilder func(*builder.Builder, client.Client, cache.Cache) *builder.Builder

type Runtime interface {
	NewObjects(ctx context.Context, trainJob *trainer.TrainJob) ([]runtime.ApplyConfiguration, error)
	RuntimeInfo(trainJob *trainer.TrainJob, runtimeTemplateSpec any, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy) (*Info, error)
	TrainJobStatus(ctx context.Context, trainJob *trainer.TrainJob) (*trainer.TrainJobStatus, error)
	EventHandlerRegistrars() []ReconcilerBuilder
	ValidateObjects(ctx context.Context, old, new *trainer.TrainJob) (admission.Warnings, field.ErrorList)
}

type RuntimePolicy struct {
	MLPolicySource *trainer.MLPolicySource
	PodGroupPolicy *trainer.PodGroupPolicy
}

type Scheduler struct {
	PodLabels      map[string]string
	PodAnnotations map[string]string
}

type TemplateSpec struct {
	// ObjApply is ApplyConfiguration for the TrainingRuntimes Template field.
	ObjApply any
	// PodSets is a set of Pod extracted from ObjApply.
	// This is abstract concept to represent multiple PodSpec as a unit.
	PodSets []PodSet
}