mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
Implemenet MPI Plugin for OpenMPI (#2493)
* Implemenet MPI Plugin for OpenMPI Signed-off-by: Yuki Iwai <[email protected]> * Directory pass the JobSetApplyconfiguration to RuntimeInfo Signed-off-by: Yuki Iwai <[email protected]> * Make repeated string as constants Signed-off-by: Yuki Iwai <[email protected]> * Use numNodes=1 as default mpi_distributed ClusterTrainingRuntime Signed-off-by: Yuki Iwai <[email protected]> * Remove unused errors Signed-off-by: Yuki Iwai <[email protected]> * Rename runLauncherAsWorker with runLauncherAsNode Signed-off-by: Yuki Iwai <[email protected]> * Fix unintended constants usage for ModelMountPath Signed-off-by: Yuki Iwai <[email protected]> * Rename SecretDataComparer with MPISecretDataComparer Signed-off-by: Yuki Iwai <[email protected]> * Add TODO STAEEMENT to deprecated env wrappers. Signed-off-by: Yuki Iwai <[email protected]> --------- Signed-off-by: Yuki Iwai <[email protected]>
This commit is contained in:
Yuki Iwai
@@ -11,35 +11,47 @@ spec:
|
||||
mpi:
|
||||
numProcPerNode: 1
|
||||
mpiImplementation: OpenMPI
|
||||
sshAuthMountPath: /root/.ssh
|
||||
sshAuthMountPath: /home/mpiuser/.ssh
|
||||
template:
|
||||
spec:
|
||||
network:
|
||||
publishNotReadyAddresses: false
|
||||
successPolicy:
|
||||
operator: All
|
||||
targetReplicatedJobs:
|
||||
- launcher
|
||||
replicatedJobs:
|
||||
- name: launcher
|
||||
template:
|
||||
spec:
|
||||
template:
|
||||
spec:
|
||||
# TODO (andreyvelich): Change the command with mpirun.
|
||||
containers:
|
||||
- name: launcher
|
||||
image: busybox
|
||||
image: mpioperator/mpi-pi:openmpi
|
||||
securityContext:
|
||||
runAsUser: 1000
|
||||
command:
|
||||
- /bin/sh
|
||||
- -c
|
||||
- "echo 'launcher runs for 15 seconds' && sleep 15"
|
||||
- mpirun
|
||||
args:
|
||||
- /home/mpiuser/pi
|
||||
- name: trainer-node
|
||||
dependsOn:
|
||||
- name: launcher
|
||||
status: Ready
|
||||
template:
|
||||
spec:
|
||||
template:
|
||||
spec:
|
||||
containers:
|
||||
- name: trainer
|
||||
image: busybox
|
||||
- name: trainer-node
|
||||
image: mpioperator/mpi-pi:openmpi
|
||||
securityContext:
|
||||
runAsUser: 1000
|
||||
command:
|
||||
- /bin/sh
|
||||
- -c
|
||||
- "echo 'trainer node runs for 15 seconds' && sleep 15"
|
||||
- /usr/sbin/sshd
|
||||
args:
|
||||
- -De
|
||||
- -f
|
||||
- /home/mpiuser/.sshd_config
|
||||
readinessProbe:
|
||||
tcpSocket:
|
||||
port: 2222
|
||||
initialDelaySeconds: 5
|
||||
|
||||
+45
-10
@@ -17,38 +17,50 @@ limitations under the License.
|
||||
package apply
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
|
||||
"k8s.io/apimachinery/pkg/runtime"
|
||||
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
||||
"k8s.io/utils/ptr"
|
||||
"sigs.k8s.io/controller-runtime/pkg/client"
|
||||
)
|
||||
|
||||
func UpsertEnvVar(envVars *[]corev1ac.EnvVarApplyConfiguration, envVar ...*corev1ac.EnvVarApplyConfiguration) {
|
||||
var (
|
||||
errorRequestedFieldPathNotFound = errors.New("requested field path not found")
|
||||
)
|
||||
|
||||
func UpsertEnvVar(envVars *[]corev1ac.EnvVarApplyConfiguration, envVar ...corev1ac.EnvVarApplyConfiguration) {
|
||||
for _, e := range envVar {
|
||||
upsert(envVars, *e, byEnvVarName)
|
||||
upsert(envVars, e, byEnvVarName)
|
||||
}
|
||||
}
|
||||
|
||||
func UpsertEnvVars(envVars *[]corev1ac.EnvVarApplyConfiguration, upEnvVars []corev1ac.EnvVarApplyConfiguration) {
|
||||
func UpsertEnvVars(envVars *[]corev1ac.EnvVarApplyConfiguration, upEnvVars ...corev1ac.EnvVarApplyConfiguration) {
|
||||
for _, e := range upEnvVars {
|
||||
upsert(envVars, e, byEnvVarName)
|
||||
}
|
||||
}
|
||||
|
||||
func UpsertPort(ports *[]corev1ac.ContainerPortApplyConfiguration, port ...*corev1ac.ContainerPortApplyConfiguration) {
|
||||
func UpsertPort(ports *[]corev1ac.ContainerPortApplyConfiguration, port ...corev1ac.ContainerPortApplyConfiguration) {
|
||||
for _, p := range port {
|
||||
upsert(ports, *p, byContainerPortOrName)
|
||||
upsert(ports, p, byContainerPortOrName)
|
||||
}
|
||||
}
|
||||
|
||||
func UpsertVolumes(volumes *[]corev1ac.VolumeApplyConfiguration, upVolumes []corev1ac.VolumeApplyConfiguration) {
|
||||
func UpsertVolumes(volumes *[]corev1ac.VolumeApplyConfiguration, upVolumes ...corev1ac.VolumeApplyConfiguration) {
|
||||
for _, v := range upVolumes {
|
||||
upsert(volumes, v, byVolumeName)
|
||||
}
|
||||
}
|
||||
|
||||
func UpsertVolumeMounts(mounts *[]corev1ac.VolumeMountApplyConfiguration, upMounts []corev1ac.VolumeMountApplyConfiguration) {
|
||||
func UpsertVolumeMounts(mounts *[]corev1ac.VolumeMountApplyConfiguration, upMounts ...corev1ac.VolumeMountApplyConfiguration) {
|
||||
for _, m := range upMounts {
|
||||
upsert(mounts, m, byVolumeMountName)
|
||||
upsert(mounts, m, byVolumeMountPath)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -64,8 +76,8 @@ func byVolumeName(a, b corev1ac.VolumeApplyConfiguration) bool {
|
||||
return ptr.Equal(a.Name, b.Name)
|
||||
}
|
||||
|
||||
func byVolumeMountName(a, b corev1ac.VolumeMountApplyConfiguration) bool {
|
||||
return ptr.Equal(a.Name, b.Name)
|
||||
func byVolumeMountPath(a, b corev1ac.VolumeMountApplyConfiguration) bool {
|
||||
return ptr.Equal(a.MountPath, b.MountPath)
|
||||
}
|
||||
|
||||
type compare[T any] func(T, T) bool
|
||||
@@ -121,3 +133,26 @@ func EnvVars(e ...corev1.EnvVar) []corev1ac.EnvVarApplyConfiguration {
|
||||
}
|
||||
return envs
|
||||
}
|
||||
|
||||
func FromTypedObjWithFields[A any](typed client.Object, fields ...string) (*A, error) {
|
||||
u, err := runtime.DefaultUnstructuredConverter.ToUnstructured(typed)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
uObj, ok, err := unstructured.NestedFieldCopy(u, fields...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%w: '.%s'", errorRequestedFieldPathNotFound, strings.Join(fields, "."))
|
||||
}
|
||||
raw, err := json.Marshal(uObj)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var objApply *A
|
||||
if err = json.Unmarshal(raw, &objApply); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return objApply, nil
|
||||
}
|
||||
|
||||
@@ -32,6 +32,12 @@ const (
|
||||
// ContainerDatasetInitializer is the container name for the dataset initializer.
|
||||
ContainerDatasetInitializer string = "dataset-initializer"
|
||||
|
||||
// DatasetMountPath is the volumeMount path for dataset.
|
||||
DatasetMountPath string = "/workspace/dataset"
|
||||
|
||||
// ModelMountPath is the volumeMount path for model.
|
||||
ModelMountPath string = "/workspace/model"
|
||||
|
||||
// PodGroupKind is the Kind name for the PodGroup.
|
||||
PodGroupKind string = "PodGroup"
|
||||
|
||||
@@ -110,7 +116,21 @@ const (
|
||||
|
||||
// Distributed envs for mpirun.
|
||||
// Values for OpenMPI implementation.
|
||||
|
||||
// OpenMPIEnvHostFileLocation is the OpenMPI default hostfile env key.
|
||||
OpenMPIEnvHostFileLocation string = "OMPI_MCA_orte_default_hostfile"
|
||||
|
||||
// OpenMPIEnvKeyRSHArgs is the env key for the OpenMPI rsh arguments.
|
||||
OpenMPIEnvKeyRSHArgs string = "OMPI_MCA_plm_rsh_args"
|
||||
|
||||
// OpenMPIEnvDefaultValueRSHArgs is the default env valur for the OpenMPI rsh arguments.
|
||||
OpenMPIEnvDefaultValueRSHArgs string = "-o ConnectionAttempts=10"
|
||||
|
||||
// OpenMPIEnvKeepFQDNHostNames is the env key for OpenMPI if FQDN should be kept.
|
||||
OpenMPIEnvKeepFQDNHostNames string = "OMPI_MCA_orte_keep_fqdn_hostnames"
|
||||
|
||||
// OpenMPIEnvDefaultSlots is the OpenMPI default number of slots env key.
|
||||
OpenMPIEnvDefaultSlots string = "OMPI_MCA_orte_set_default_slots"
|
||||
)
|
||||
|
||||
var (
|
||||
|
||||
@@ -26,9 +26,11 @@ import (
|
||||
"k8s.io/apimachinery/pkg/api/resource"
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
"k8s.io/apimachinery/pkg/runtime"
|
||||
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
||||
schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
testingutil "github.com/kubeflow/trainer/pkg/util/testing"
|
||||
)
|
||||
|
||||
@@ -47,6 +49,12 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
|
||||
"succeeded to build PodGroup and JobSet with NumNodes from the Runtime and container from the Trainer.": {
|
||||
clusterTrainingRuntime: testingutil.MakeClusterTrainingRuntimeWrapper("test-runtime").RuntimeSpec(
|
||||
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeClusterTrainingRuntimeWrapper("test-runtime").Spec).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper("", "").
|
||||
DependsOn(constants.JobInitializer, jobsetv1alpha2.DependencyComplete, constants.JobTrainerNode).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
ContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
WithMLPolicy(
|
||||
testingutil.MakeMLPolicyWrapper().
|
||||
@@ -70,6 +78,10 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
|
||||
wantObjs: []runtime.Object{
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
ContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
DependsOn(constants.JobInitializer, jobsetv1alpha2.DependencyComplete, constants.JobTrainerNode).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Parallelism(constants.JobInitializer, 1).
|
||||
Completions(constants.JobInitializer, 1).
|
||||
NumNodes(100).
|
||||
ContainerTrainer("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
|
||||
Suspend(true).
|
||||
|
||||
@@ -21,11 +21,18 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/kubeflow/trainer/pkg/apply"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
"k8s.io/apimachinery/pkg/runtime/schema"
|
||||
"k8s.io/apimachinery/pkg/util/validation/field"
|
||||
resourcehelpers "k8s.io/component-helpers/resource"
|
||||
"k8s.io/utils/ptr"
|
||||
"sigs.k8s.io/controller-runtime/pkg/client"
|
||||
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
|
||||
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
||||
jobsetv1alpha2ac "sigs.k8s.io/jobset/client-go/applyconfiguration/jobset/v1alpha2"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/runtime"
|
||||
@@ -98,31 +105,86 @@ func (r *TrainingRuntime) buildObjects(
|
||||
// The JobSetTemplateSpec annotations are overridden by the TrainJob Annotations (.spec.annotations).
|
||||
propagationAnnotations[k] = v
|
||||
}
|
||||
jobSetSpecApply, err := apply.FromTypedObjWithFields[jobsetv1alpha2ac.JobSetSpecApplyConfiguration](&jobsetv1alpha2.JobSet{
|
||||
TypeMeta: metav1.TypeMeta{
|
||||
APIVersion: jobsetv1alpha2.GroupVersion.String(),
|
||||
Kind: "JobSet",
|
||||
},
|
||||
Spec: jobSetTemplateSpec.Spec,
|
||||
}, "spec")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
opts := []runtime.InfoOption{
|
||||
runtime.WithLabels(propagationLabels),
|
||||
runtime.WithAnnotations(propagationAnnotations),
|
||||
runtime.WithMLPolicy(mlPolicy),
|
||||
runtime.WithPodGroupPolicy(podGroupPolicy),
|
||||
runtime.WithTemplateSpec(jobSetSpecApply),
|
||||
runtime.WithPodSetSyncer(syncPodSets),
|
||||
}
|
||||
|
||||
for _, rJob := range jobSetTemplateSpec.Spec.ReplicatedJobs {
|
||||
// By default every ReplicatedJob has only 1 replica.
|
||||
opts = append(opts, runtime.WithPodSpecReplicas(rJob.Name, 1, rJob.Template.Spec.Template.Spec))
|
||||
for i, rJob := range jobSetSpecApply.ReplicatedJobs {
|
||||
// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
|
||||
// REF: https://github.com/kubeflow/trainer/issues/2318
|
||||
count := ptr.Deref(rJob.Template.Spec.Parallelism, 1)
|
||||
if *rJob.Name == constants.JobTrainerNode {
|
||||
count = ptr.Deref(mlPolicy.NumNodes, 1)
|
||||
}
|
||||
opts = append(opts, runtime.WithPodSpecReplicas(
|
||||
*rJob.Name,
|
||||
count,
|
||||
resourcehelpers.PodRequests(&corev1.Pod{Spec: *jobSetTemplateSpec.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.DeepCopy()}, resourcehelpers.PodResourcesOptions{}),
|
||||
rJob.Template.Spec.Template.Spec),
|
||||
)
|
||||
}
|
||||
|
||||
info := runtime.NewInfo(opts...)
|
||||
|
||||
if err := r.framework.RunEnforceMLPolicyPlugins(info, trainJob); err != nil {
|
||||
if err = r.framework.RunEnforceMLPolicyPlugins(info, trainJob); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := r.framework.RunEnforcePodGroupPolicyPlugins(info, trainJob); err != nil {
|
||||
if err = r.framework.RunEnforcePodGroupPolicyPlugins(info, trainJob); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err = r.framework.RunPodNetworkPlugins(info, trainJob); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return r.framework.RunComponentBuilderPlugins(ctx, info, trainJob)
|
||||
}
|
||||
|
||||
func syncPodSets(info *runtime.Info) {
|
||||
jsSpec, ok := runtime.TemplateSpecApply[jobsetv1alpha2ac.JobSetSpecApplyConfiguration](info)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for psIdx, ps := range info.TemplateSpec.PodSets {
|
||||
if ps.CountForNonTrainer != nil {
|
||||
jsSpec.ReplicatedJobs[psIdx].Template.Spec.Parallelism = ps.CountForNonTrainer
|
||||
jsSpec.ReplicatedJobs[psIdx].Template.Spec.Completions = ps.CountForNonTrainer
|
||||
}
|
||||
apply.UpsertVolumes(&jsSpec.ReplicatedJobs[psIdx].Template.Spec.Template.Spec.Volumes, ps.Volumes...)
|
||||
for containerIdx, container := range ps.Containers {
|
||||
apply.UpsertEnvVar(
|
||||
&jsSpec.ReplicatedJobs[psIdx].Template.Spec.Template.Spec.Containers[containerIdx].Env,
|
||||
container.Env...,
|
||||
)
|
||||
apply.UpsertPort(
|
||||
&jsSpec.ReplicatedJobs[psIdx].Template.Spec.Template.Spec.Containers[containerIdx].Ports,
|
||||
container.Ports...,
|
||||
)
|
||||
apply.UpsertVolumeMounts(
|
||||
&jsSpec.ReplicatedJobs[psIdx].Template.Spec.Template.Spec.Containers[containerIdx].VolumeMounts,
|
||||
container.VolumeMounts...,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *TrainingRuntime) TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error) {
|
||||
return r.framework.RunTerminalConditionPlugins(ctx, trainJob)
|
||||
}
|
||||
|
||||
@@ -28,11 +28,14 @@ import (
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
"k8s.io/apimachinery/pkg/runtime"
|
||||
"k8s.io/apimachinery/pkg/util/intstr"
|
||||
"k8s.io/klog/v2/ktesting"
|
||||
"k8s.io/utils/ptr"
|
||||
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
||||
schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
jobsetplugin "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset"
|
||||
jobsetplgconsts "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset/constants"
|
||||
testingutil "github.com/kubeflow/trainer/pkg/util/testing"
|
||||
)
|
||||
|
||||
@@ -45,6 +48,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
trainingRuntime *trainer.TrainingRuntime
|
||||
trainJob *trainer.TrainJob
|
||||
ObjCmpOpts []cmp.Option
|
||||
wantObjs []runtime.Object
|
||||
wantError error
|
||||
}{
|
||||
@@ -80,7 +84,10 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
wantObjs: []runtime.Object{
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
ContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
Parallelism(constants.JobInitializer, 1).
|
||||
Completions(constants.JobInitializer, 1).
|
||||
NumNodes(30).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
Suspend(true).
|
||||
Label("conflictLabel", "override").
|
||||
@@ -146,6 +153,9 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
wantObjs: []runtime.Object{
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
NumNodes(100).
|
||||
Parallelism(constants.JobInitializer, 1).
|
||||
Completions(constants.JobInitializer, 1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
ContainerTrainer("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
|
||||
ContainerTrainerEnv(
|
||||
[]corev1.EnvVar{
|
||||
@@ -218,12 +228,15 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
wantObjs: []runtime.Object{
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
NumNodes(100).
|
||||
Parallelism(constants.JobInitializer, 1).
|
||||
Completions(constants.JobInitializer, 1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
ContainerDatasetModelInitializer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
ContainerDatasetInitializerEnv(
|
||||
[]corev1.EnvVar{
|
||||
{
|
||||
Name: jobsetplugin.InitializerEnvStorageUri,
|
||||
Name: jobsetplgconsts.InitializerEnvStorageUri,
|
||||
Value: "hf://trainjob-dataset",
|
||||
},
|
||||
{
|
||||
@@ -246,7 +259,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
ContainerModelInitializerEnv(
|
||||
[]corev1.EnvVar{
|
||||
{
|
||||
Name: jobsetplugin.InitializerEnvStorageUri,
|
||||
Name: jobsetplgconsts.InitializerEnvStorageUri,
|
||||
Value: "hf://trainjob-model",
|
||||
},
|
||||
{
|
||||
@@ -280,6 +293,12 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
TorchPolicy("auto", nil).
|
||||
Obj(),
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper("", "").
|
||||
DependsOn(constants.JobInitializer, jobsetv1alpha2.DependencyComplete, constants.JobTrainerNode).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
Obj(),
|
||||
).Obj(),
|
||||
@@ -296,6 +315,10 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
wantObjs: []runtime.Object{
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
NumNodes(30).
|
||||
DependsOn(constants.JobInitializer, jobsetv1alpha2.DependencyComplete, constants.JobTrainerNode).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Parallelism(constants.JobInitializer, 1).
|
||||
Completions(constants.JobInitializer, 1).
|
||||
ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
|
||||
ContainerTrainerEnv(
|
||||
@@ -378,6 +401,9 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
wantObjs: []runtime.Object{
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
NumNodes(100).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Parallelism(constants.JobInitializer, 1).
|
||||
Completions(constants.JobInitializer, 1).
|
||||
ContainerTrainer("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
|
||||
ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
|
||||
ContainerTrainerEnv(
|
||||
@@ -424,6 +450,149 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
Obj(),
|
||||
},
|
||||
},
|
||||
"succeeded to build JobSet with OpenMPI values from the TrainJob": {
|
||||
ObjCmpOpts: cmp.Options{
|
||||
cmp.Comparer(testingutil.MPISecretDataComparer),
|
||||
},
|
||||
trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
|
||||
testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
|
||||
WithMLPolicy(
|
||||
testingutil.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(false)).
|
||||
Obj(),
|
||||
).
|
||||
LauncherReplica().
|
||||
ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
Obj(),
|
||||
).
|
||||
Obj(),
|
||||
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
|
||||
UID("uid").
|
||||
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
|
||||
Trainer(
|
||||
testingutil.MakeTrainJobTrainerWrapper().
|
||||
NumNodes(2).
|
||||
NumProcPerNode(intstr.FromInt32(8)).
|
||||
Obj(),
|
||||
).
|
||||
Obj(),
|
||||
wantObjs: []runtime.Object{
|
||||
testingutil.MakeConfigMapWrapper(fmt.Sprintf("test-job%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
|
||||
WithData(map[string]string{
|
||||
constants.MPIHostfileName: `test-job-trainer-node-0-0.test-job slots=8
|
||||
test-job-trainer-node-0-1.test-job slots=8
|
||||
`,
|
||||
}).
|
||||
Obj(),
|
||||
testingutil.MakeSecretWrapper(fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix), metav1.NamespaceDefault).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
|
||||
WithImmutable(true).
|
||||
WithData(map[string][]byte{
|
||||
corev1.SSHAuthPrivateKey: []byte("EXIST"),
|
||||
constants.MPISSHPublicKey: []byte("EXIST"),
|
||||
}).
|
||||
WithType(corev1.SecretTypeSSHAuth).
|
||||
Obj(),
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
NumNodes(2).
|
||||
LauncherReplica().
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Parallelism(constants.JobInitializer, 1).
|
||||
Completions(constants.JobInitializer, 1).
|
||||
Parallelism(constants.JobLauncher, 1).
|
||||
Completions(constants.JobLauncher, 1).
|
||||
Volumes(constants.JobLauncher,
|
||||
corev1.Volume{
|
||||
Name: constants.MPISSHAuthVolumeName,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
Secret: &corev1.SecretVolumeSource{
|
||||
SecretName: fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix),
|
||||
Items: []corev1.KeyToPath{
|
||||
{
|
||||
Key: corev1.SSHAuthPrivateKey,
|
||||
Path: constants.MPISSHPrivateKeyFile,
|
||||
},
|
||||
{
|
||||
Key: constants.MPISSHPublicKey,
|
||||
Path: constants.MPISSHPublicKeyFile,
|
||||
},
|
||||
{
|
||||
Key: constants.MPISSHPublicKey,
|
||||
Path: constants.MPISSHAuthorizedKeys,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
corev1.Volume{
|
||||
Name: constants.MPIHostfileVolumeName,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
ConfigMap: &corev1.ConfigMapVolumeSource{
|
||||
LocalObjectReference: corev1.LocalObjectReference{
|
||||
Name: fmt.Sprintf("test-job%s", constants.MPIHostfileConfigMapSuffix),
|
||||
},
|
||||
Items: []corev1.KeyToPath{{
|
||||
Key: constants.MPIHostfileName,
|
||||
Path: constants.MPIHostfileName,
|
||||
Mode: ptr.To[int32](0444),
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
).
|
||||
Volumes(constants.JobTrainerNode,
|
||||
corev1.Volume{
|
||||
Name: constants.MPISSHAuthVolumeName,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
Secret: &corev1.SecretVolumeSource{
|
||||
SecretName: fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix),
|
||||
Items: []corev1.KeyToPath{
|
||||
{
|
||||
Key: corev1.SSHAuthPrivateKey,
|
||||
Path: constants.MPISSHPrivateKeyFile,
|
||||
},
|
||||
{
|
||||
Key: constants.MPISSHPublicKey,
|
||||
Path: constants.MPISSHPublicKeyFile,
|
||||
},
|
||||
{
|
||||
Key: constants.MPISSHPublicKey,
|
||||
Path: constants.MPISSHAuthorizedKeys,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
).
|
||||
VolumeMounts(constants.JobLauncher, constants.ContainerLauncher,
|
||||
corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
|
||||
corev1.VolumeMount{Name: constants.MPIHostfileVolumeName, MountPath: constants.MPIHostfileDir},
|
||||
).
|
||||
Env(constants.JobLauncher, constants.ContainerLauncher,
|
||||
corev1.EnvVar{
|
||||
Name: constants.OpenMPIEnvHostFileLocation,
|
||||
Value: fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName),
|
||||
},
|
||||
corev1.EnvVar{
|
||||
Name: constants.OpenMPIEnvKeepFQDNHostNames,
|
||||
Value: "true",
|
||||
},
|
||||
corev1.EnvVar{
|
||||
Name: constants.OpenMPIEnvDefaultSlots,
|
||||
Value: "8",
|
||||
},
|
||||
corev1.EnvVar{
|
||||
Name: constants.OpenMPIEnvKeyRSHArgs,
|
||||
Value: constants.OpenMPIEnvDefaultValueRSHArgs,
|
||||
},
|
||||
).
|
||||
ContainerTrainer("test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
|
||||
Obj(),
|
||||
},
|
||||
},
|
||||
// Failed test cases.
|
||||
"missing trainingRuntime resource": {
|
||||
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job-3").
|
||||
@@ -444,12 +613,12 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
cmpopts.SortSlices(func(a, b corev1.EnvVar) bool {
|
||||
return a.Name < b.Name
|
||||
}),
|
||||
cmpopts.EquateEmpty(),
|
||||
cmpopts.SortMaps(func(a, b string) bool { return a < b }),
|
||||
}
|
||||
for name, tc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
_, ctx := ktesting.NewTestContext(t)
|
||||
var cancel func()
|
||||
ctx, cancel = context.WithCancel(ctx)
|
||||
t.Cleanup(cancel)
|
||||
clientBuilder := testingutil.NewClientBuilder()
|
||||
if tc.trainingRuntime != nil {
|
||||
@@ -472,7 +641,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
|
||||
t.Errorf("Pipeline built unrecognizable objects: %v", err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tc.wantObjs, resultObjs, cmpOpts...); len(diff) != 0 {
|
||||
if diff := cmp.Diff(tc.wantObjs, resultObjs, append(cmpOpts, tc.ObjCmpOpts...)...); len(diff) != 0 {
|
||||
t.Errorf("Unexpected objects (-want,+got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
|
||||
@@ -40,6 +40,7 @@ type Framework struct {
|
||||
enforcePodGroupPolicyPlugins []framework.EnforcePodGroupPolicyPlugin
|
||||
customValidationPlugins []framework.CustomValidationPlugin
|
||||
watchExtensionPlugins []framework.WatchExtensionPlugin
|
||||
podNetworkPlugins []framework.PodNetworkPlugin
|
||||
componentBuilderPlugins []framework.ComponentBuilderPlugin
|
||||
terminalConditionPlugins []framework.TerminalConditionPlugin
|
||||
}
|
||||
@@ -68,6 +69,9 @@ func New(ctx context.Context, c client.Client, r fwkplugins.Registry, indexer cl
|
||||
if p, ok := plugin.(framework.WatchExtensionPlugin); ok {
|
||||
f.watchExtensionPlugins = append(f.watchExtensionPlugins, p)
|
||||
}
|
||||
if p, ok := plugin.(framework.PodNetworkPlugin); ok {
|
||||
f.podNetworkPlugins = append(f.podNetworkPlugins, p)
|
||||
}
|
||||
if p, ok := plugin.(framework.ComponentBuilderPlugin); ok {
|
||||
f.componentBuilderPlugins = append(f.componentBuilderPlugins, p)
|
||||
}
|
||||
@@ -112,6 +116,15 @@ func (f *Framework) RunCustomValidationPlugins(oldObj, newObj *trainer.TrainJob)
|
||||
return aggregatedWarnings, aggregatedErrors
|
||||
}
|
||||
|
||||
func (f *Framework) RunPodNetworkPlugins(info *runtime.Info, trainJob *trainer.TrainJob) error {
|
||||
for _, plugin := range f.podNetworkPlugins {
|
||||
if err := plugin.IdentifyPodNetwork(info, trainJob); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *Framework) RunComponentBuilderPlugins(ctx context.Context, info *runtime.Info, trainJob *trainer.TrainJob) ([]any, error) {
|
||||
var objs []any
|
||||
for _, plugin := range f.componentBuilderPlugins {
|
||||
|
||||
@@ -28,10 +28,14 @@ import (
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
apiruntime "k8s.io/apimachinery/pkg/runtime"
|
||||
"k8s.io/apimachinery/pkg/util/validation/field"
|
||||
batchv1ac "k8s.io/client-go/applyconfigurations/batch/v1"
|
||||
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
||||
"k8s.io/klog/v2/ktesting"
|
||||
"k8s.io/utils/ptr"
|
||||
"sigs.k8s.io/controller-runtime/pkg/client"
|
||||
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
|
||||
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
||||
jobsetv1alpha2ac "sigs.k8s.io/jobset/client-go/applyconfiguration/jobset/v1alpha2"
|
||||
jobsetconsts "sigs.k8s.io/jobset/pkg/constants"
|
||||
schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"
|
||||
|
||||
@@ -42,6 +46,7 @@ import (
|
||||
fwkplugins "github.com/kubeflow/trainer/pkg/runtime/framework/plugins"
|
||||
"github.com/kubeflow/trainer/pkg/runtime/framework/plugins/coscheduling"
|
||||
"github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset"
|
||||
jobsetplgconsts "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset/constants"
|
||||
"github.com/kubeflow/trainer/pkg/runtime/framework/plugins/mpi"
|
||||
"github.com/kubeflow/trainer/pkg/runtime/framework/plugins/plainml"
|
||||
"github.com/kubeflow/trainer/pkg/runtime/framework/plugins/torch"
|
||||
@@ -88,6 +93,9 @@ func TestNew(t *testing.T) {
|
||||
&jobset.JobSet{},
|
||||
&mpi.MPI{},
|
||||
},
|
||||
podNetworkPlugins: []framework.PodNetworkPlugin{
|
||||
&jobset.JobSet{},
|
||||
},
|
||||
componentBuilderPlugins: []framework.ComponentBuilderPlugin{
|
||||
&coscheduling.CoScheduling{},
|
||||
&jobset.JobSet{},
|
||||
@@ -416,6 +424,112 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
|
||||
Trainer: runtime.Trainer{
|
||||
NumNodes: ptr.To[int32](10),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobInitializer,
|
||||
Containers: []runtime.Container{
|
||||
{
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
},
|
||||
},
|
||||
{
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
||||
),
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: []runtime.Container{{
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobInitializer).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerDatasetInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
),
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerModelInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
),
|
||||
).
|
||||
WithVolumes(
|
||||
corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer)),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerTrainer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
),
|
||||
).
|
||||
WithVolumes(
|
||||
corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer)),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
TotalRequests: map[string]runtime.TotalResourceRequest{
|
||||
constants.JobInitializer: {
|
||||
@@ -455,6 +569,127 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
|
||||
Trainer: runtime.Trainer{
|
||||
NumNodes: ptr.To[int32](100),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithReplicas(1).
|
||||
WithName(constants.JobInitializer).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithLabels(map[string]string{
|
||||
schedulerpluginsv1alpha1.PodGroupLabel: "test-job",
|
||||
}).
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerDatasetInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
),
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerModelInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
),
|
||||
).
|
||||
WithVolumes(
|
||||
corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer)),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithReplicas(1).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(100).
|
||||
WithCompletions(100).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithLabels(map[string]string{
|
||||
schedulerpluginsv1alpha1.PodGroupLabel: "test-job",
|
||||
}).
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerTrainer).
|
||||
WithImage("test:trainjob").
|
||||
WithCommand("trainjob").
|
||||
WithArgs("trainjob").
|
||||
WithResources(corev1ac.ResourceRequirements().
|
||||
WithRequests(resRequests)).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
),
|
||||
).
|
||||
WithVolumes(
|
||||
corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer)),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobInitializer,
|
||||
Containers: []runtime.Container{
|
||||
{
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
},
|
||||
},
|
||||
{
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
||||
),
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: []runtime.Container{{
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
PodLabels: map[string]string{schedulerpluginsv1alpha1.PodGroupLabel: "test-job"},
|
||||
TotalRequests: map[string]runtime.TotalResourceRequest{
|
||||
@@ -480,6 +715,7 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "test-job", "uid").
|
||||
Obj(),
|
||||
testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "test-job", "uid").
|
||||
NumNodes(100).
|
||||
PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
|
||||
@@ -501,9 +737,6 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
|
||||
t.Cleanup(cancel)
|
||||
|
||||
clientBuilder := testingutil.NewClientBuilder()
|
||||
if tc.trainingRuntime != nil {
|
||||
clientBuilder.WithObjects(tc.trainingRuntime)
|
||||
}
|
||||
c := clientBuilder.Build()
|
||||
|
||||
fwk, err := New(ctx, c, tc.registry, testingutil.AsIndex(clientBuilder))
|
||||
@@ -686,3 +919,112 @@ func TestTerminalConditionPlugins(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPodNetworkPlugins(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
registry fwkplugins.Registry
|
||||
runtimeInfo *runtime.Info
|
||||
trainJob *trainer.TrainJob
|
||||
wantError error
|
||||
wantRuntimeInfo *runtime.Info
|
||||
}{
|
||||
"Pod network is calculated by jobset plugin": {
|
||||
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
|
||||
Obj(),
|
||||
registry: fwkplugins.NewRegistry(),
|
||||
runtimeInfo: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: testingutil.MakeMLPolicyWrapper().
|
||||
WithNumNodes(2).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerTrainer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
wantRuntimeInfo: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: testingutil.MakeMLPolicyWrapper().
|
||||
WithNumNodes(2).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("test-job-trainer-node-0-0.test-job")
|
||||
yield("test-job-trainer-node-0-1.test-job")
|
||||
},
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerTrainer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
"am empty registry": {
|
||||
trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").Obj(),
|
||||
runtimeInfo: &runtime.Info{},
|
||||
wantRuntimeInfo: &runtime.Info{},
|
||||
},
|
||||
}
|
||||
for name, tc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
_, ctx := ktesting.NewTestContext(t)
|
||||
var cancel func()
|
||||
ctx, cancel = context.WithCancel(ctx)
|
||||
t.Cleanup(cancel)
|
||||
cliBuilder := testingutil.NewClientBuilder()
|
||||
fwk, err := New(ctx, cliBuilder.Build(), tc.registry, testingutil.AsIndex(cliBuilder))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
err = fwk.RunPodNetworkPlugins(tc.runtimeInfo, tc.trainJob)
|
||||
if diff := cmp.Diff(tc.wantError, err); len(diff) != 0 {
|
||||
t.Errorf("Unexpected error (-want,+got):\n%s", diff)
|
||||
}
|
||||
if diff := cmp.Diff(tc.wantRuntimeInfo, tc.runtimeInfo, testingutil.PodSetEndpointsCmpOpts); len(diff) != 0 {
|
||||
t.Errorf("Unexpected runtimeInfo (-want,+got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -51,6 +51,11 @@ type EnforceMLPolicyPlugin interface {
|
||||
EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) error
|
||||
}
|
||||
|
||||
type PodNetworkPlugin interface {
|
||||
Plugin
|
||||
IdentifyPodNetwork(info *runtime.Info, trainJob *trainer.TrainJob) error
|
||||
}
|
||||
|
||||
type ComponentBuilderPlugin interface {
|
||||
Plugin
|
||||
Build(ctx context.Context, info *runtime.Info, trainJob *trainer.TrainJob) ([]any, error)
|
||||
|
||||
@@ -18,12 +18,14 @@ package jobset
|
||||
|
||||
import (
|
||||
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
||||
"k8s.io/utils/ptr"
|
||||
jobsetv1alpha2ac "sigs.k8s.io/jobset/client-go/applyconfiguration/jobset/v1alpha2"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/apply"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
"github.com/kubeflow/trainer/pkg/runtime"
|
||||
jobsetplgconsts "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset/constants"
|
||||
)
|
||||
|
||||
type Builder struct {
|
||||
@@ -40,17 +42,21 @@ func NewBuilder(jobSet *jobsetv1alpha2ac.JobSetApplyConfiguration) *Builder {
|
||||
func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
|
||||
for i, rJob := range b.Spec.ReplicatedJobs {
|
||||
if *rJob.Name == constants.JobInitializer {
|
||||
// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
|
||||
// REF: https://github.com/kubeflow/trainer/issues/2318
|
||||
b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
|
||||
|
||||
for j, container := range rJob.Template.Spec.Template.Spec.Containers {
|
||||
// Update values for the dataset initializer container.
|
||||
if *container.Name == constants.ContainerDatasetInitializer && trainJob.Spec.DatasetConfig != nil {
|
||||
env := &b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Env
|
||||
// Update the dataset initializer envs.
|
||||
if storageUri := trainJob.Spec.DatasetConfig.StorageUri; storageUri != nil {
|
||||
apply.UpsertEnvVar(env, corev1ac.EnvVar().
|
||||
WithName(InitializerEnvStorageUri).
|
||||
apply.UpsertEnvVar(env, *corev1ac.EnvVar().
|
||||
WithName(jobsetplgconsts.InitializerEnvStorageUri).
|
||||
WithValue(*storageUri))
|
||||
}
|
||||
apply.UpsertEnvVars(env, apply.EnvVars(trainJob.Spec.DatasetConfig.Env...))
|
||||
apply.UpsertEnvVars(env, apply.EnvVars(trainJob.Spec.DatasetConfig.Env...)...)
|
||||
// Update the dataset initializer secret reference.
|
||||
if trainJob.Spec.DatasetConfig.SecretRef != nil {
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].
|
||||
@@ -67,11 +73,11 @@ func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
|
||||
// Update the model initializer envs.
|
||||
env := &b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Env
|
||||
if storageUri := trainJob.Spec.ModelConfig.Input.StorageUri; storageUri != nil {
|
||||
apply.UpsertEnvVar(env, corev1ac.EnvVar().
|
||||
WithName(InitializerEnvStorageUri).
|
||||
apply.UpsertEnvVar(env, *corev1ac.EnvVar().
|
||||
WithName(jobsetplgconsts.InitializerEnvStorageUri).
|
||||
WithValue(*storageUri))
|
||||
}
|
||||
apply.UpsertEnvVars(env, apply.EnvVars(trainJob.Spec.ModelConfig.Input.Env...))
|
||||
apply.UpsertEnvVars(env, apply.EnvVars(trainJob.Spec.ModelConfig.Input.Env...)...)
|
||||
// Update the model initializer secret reference.
|
||||
if trainJob.Spec.ModelConfig.Input.SecretRef != nil {
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].
|
||||
@@ -87,32 +93,12 @@ func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
|
||||
}
|
||||
|
||||
// Launcher updates JobSet values for the launcher Job.
|
||||
func (b *Builder) Launcher(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
|
||||
func (b *Builder) Launcher() *Builder {
|
||||
for i, rJob := range b.Spec.ReplicatedJobs {
|
||||
if *rJob.Name == constants.JobLauncher {
|
||||
|
||||
// Update the volumes for the Trainer Job.
|
||||
apply.UpsertVolumes(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Volumes, info.Trainer.Volumes)
|
||||
|
||||
// Update values for the launcher container.
|
||||
for j, container := range rJob.Template.Spec.Template.Spec.Containers {
|
||||
if *container.Name == constants.ContainerLauncher {
|
||||
// Update values from the Info object.
|
||||
if env := info.Trainer.Env; env != nil {
|
||||
// Update JobSet envs from the Info.
|
||||
apply.UpsertEnvVars(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Env, env)
|
||||
}
|
||||
|
||||
// Update the launcher container port.
|
||||
if port := info.Trainer.ContainerPort; port != nil {
|
||||
apply.UpsertPort(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Ports, port)
|
||||
}
|
||||
// Update the launcher container volume mounts.
|
||||
if mounts := info.Trainer.VolumeMounts; mounts != nil {
|
||||
apply.UpsertVolumeMounts(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].VolumeMounts, mounts)
|
||||
}
|
||||
}
|
||||
}
|
||||
// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
|
||||
// REF: https://github.com/kubeflow/trainer/issues/2318
|
||||
b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
|
||||
}
|
||||
}
|
||||
return b
|
||||
@@ -122,12 +108,64 @@ func (b *Builder) Launcher(info *runtime.Info, trainJob *trainer.TrainJob) *Buil
|
||||
func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
|
||||
for i, rJob := range b.Spec.ReplicatedJobs {
|
||||
if *rJob.Name == constants.JobTrainerNode {
|
||||
// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
|
||||
// REF: https://github.com/kubeflow/trainer/issues/2318
|
||||
b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
|
||||
// Update the Parallelism and Completions values for the Trainer Job.
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.RuntimePolicy.MLPolicy.NumNodes
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.RuntimePolicy.MLPolicy.NumNodes
|
||||
|
||||
// Update values for the Trainer container.
|
||||
for j, container := range rJob.Template.Spec.Template.Spec.Containers {
|
||||
if *container.Name == constants.ContainerTrainer {
|
||||
// Update values from the TrainJob trainer.
|
||||
if jobTrainer := trainJob.Spec.Trainer; jobTrainer != nil {
|
||||
if image := jobTrainer.Image; image != nil {
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Image = image
|
||||
}
|
||||
if command := jobTrainer.Command; command != nil {
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Command = command
|
||||
}
|
||||
if args := jobTrainer.Args; args != nil {
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Args = args
|
||||
}
|
||||
if resourcesPerNode := jobTrainer.ResourcesPerNode; resourcesPerNode != nil &&
|
||||
(resourcesPerNode.Limits != nil || resourcesPerNode.Requests != nil) {
|
||||
requirements := corev1ac.ResourceRequirements()
|
||||
if limits := resourcesPerNode.Limits; limits != nil {
|
||||
requirements.WithLimits(limits)
|
||||
}
|
||||
if requests := resourcesPerNode.Requests; requests != nil {
|
||||
requirements.WithRequests(requests)
|
||||
}
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].
|
||||
WithResources(requirements)
|
||||
}
|
||||
apply.UpsertEnvVars(
|
||||
&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Env,
|
||||
apply.EnvVars(jobTrainer.Env...)...,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// DeprecatedTrainer updates JobSet values for the trainer Job.
|
||||
func (b *Builder) DeprecatedTrainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
|
||||
for i, rJob := range b.Spec.ReplicatedJobs {
|
||||
if *rJob.Name == constants.JobTrainerNode {
|
||||
// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
|
||||
// REF: https://github.com/kubeflow/trainer/issues/2318
|
||||
b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
|
||||
// Update the Parallelism and Completions values for the Trainer Job.
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.Trainer.NumNodes
|
||||
b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.Trainer.NumNodes
|
||||
|
||||
// Update the volumes for the Trainer Job.
|
||||
apply.UpsertVolumes(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Volumes, info.Trainer.Volumes)
|
||||
apply.UpsertVolumes(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Volumes, info.Trainer.Volumes...)
|
||||
|
||||
// Update values for the Trainer container.
|
||||
for j, container := range rJob.Template.Spec.Template.Spec.Containers {
|
||||
@@ -159,15 +197,15 @@ func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Build
|
||||
// Update values from the Info object.
|
||||
if env := info.Trainer.Env; env != nil {
|
||||
// Update JobSet envs from the Info.
|
||||
apply.UpsertEnvVars(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Env, env)
|
||||
apply.UpsertEnvVars(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Env, env...)
|
||||
}
|
||||
// Update the Trainer container port.
|
||||
if port := info.Trainer.ContainerPort; port != nil {
|
||||
apply.UpsertPort(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Ports, port)
|
||||
apply.UpsertPort(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].Ports, *port)
|
||||
}
|
||||
// Update the Trainer container volume mounts.
|
||||
if mounts := info.Trainer.VolumeMounts; mounts != nil {
|
||||
apply.UpsertVolumeMounts(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].VolumeMounts, mounts)
|
||||
if mounts := info.Trainer.VolumeMounts; len(mounts) != 0 {
|
||||
apply.UpsertVolumeMounts(&b.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[j].VolumeMounts, mounts...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,42 +0,0 @@
|
||||
package jobset
|
||||
|
||||
import (
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
)
|
||||
|
||||
const (
|
||||
|
||||
// VolumeNameInitializer is the name for the initializer Pod's Volume and VolumeMount.
|
||||
// TODO (andreyvelich): Add validation to check that initializer Pod has the correct volume.
|
||||
VolumeNameInitializer string = "initializer"
|
||||
|
||||
// InitializerEnvStorageUri is the env name for the initializer storage uri.
|
||||
InitializerEnvStorageUri string = "STORAGE_URI"
|
||||
)
|
||||
|
||||
var (
|
||||
|
||||
// VolumeMountModelInitializer is the volume mount for the model initializer container.
|
||||
// TODO (andreyvelich): Add validation to check that initializer ReplicatedJob has the following volumes.
|
||||
VolumeMountModelInitializer = corev1.VolumeMount{
|
||||
Name: VolumeNameInitializer,
|
||||
MountPath: "/workspace/model",
|
||||
}
|
||||
|
||||
// VolumeMountModelInitializer is the volume mount for the dataset initializer container.
|
||||
VolumeMountDatasetInitializer = corev1.VolumeMount{
|
||||
Name: VolumeNameInitializer,
|
||||
MountPath: "/workspace/dataset",
|
||||
}
|
||||
|
||||
// VolumeInitializer is the volume for the initializer ReplicatedJob.
|
||||
// TODO (andreyvelich): We should make VolumeSource configurable.
|
||||
VolumeInitializer = corev1.Volume{
|
||||
Name: VolumeNameInitializer,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
|
||||
ClaimName: VolumeNameInitializer,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,26 @@
|
||||
/*
|
||||
Copyright 2025 The Kubeflow Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package constants
|
||||
|
||||
const (
|
||||
// VolumeNameInitializer is the name for the initializer Pod's Volume and VolumeMount.
|
||||
// TODO (andreyvelich): Add validation to check that initializer Pod has the correct volume.
|
||||
VolumeNameInitializer string = "initializer"
|
||||
|
||||
// InitializerEnvStorageUri is the env name for the initializer storage uri.
|
||||
InitializerEnvStorageUri string = "STORAGE_URI"
|
||||
)
|
||||
@@ -18,7 +18,6 @@ package jobset
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"maps"
|
||||
|
||||
@@ -26,7 +25,6 @@ import (
|
||||
apierrors "k8s.io/apimachinery/pkg/api/errors"
|
||||
"k8s.io/apimachinery/pkg/api/meta"
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
|
||||
apiruntime "k8s.io/apimachinery/pkg/runtime"
|
||||
"k8s.io/apimachinery/pkg/runtime/schema"
|
||||
metav1ac "k8s.io/client-go/applyconfigurations/meta/v1"
|
||||
@@ -52,6 +50,7 @@ type JobSet struct {
|
||||
}
|
||||
|
||||
var _ framework.WatchExtensionPlugin = (*JobSet)(nil)
|
||||
var _ framework.PodNetworkPlugin = (*JobSet)(nil)
|
||||
var _ framework.ComponentBuilderPlugin = (*JobSet)(nil)
|
||||
var _ framework.TerminalConditionPlugin = (*JobSet)(nil)
|
||||
|
||||
@@ -87,10 +86,47 @@ func (j *JobSet) ReconcilerBuilders() []runtime.ReconcilerBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
func (j *JobSet) IdentifyPodNetwork(info *runtime.Info, trainJob *trainer.TrainJob) error {
|
||||
if info == nil || trainJob == nil {
|
||||
return nil
|
||||
}
|
||||
spec, ok := runtime.TemplateSpecApply[jobsetv1alpha2ac.JobSetSpecApplyConfiguration](info)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
subDomain := trainJob.Name
|
||||
if jobSetNet := spec.Network; jobSetNet != nil && jobSetNet.Subdomain != nil {
|
||||
subDomain = *jobSetNet.Subdomain
|
||||
}
|
||||
for rJobIdx, rJob := range spec.ReplicatedJobs {
|
||||
// TODO: Support multiple replicas for replicated Jobs.
|
||||
// REF: https://github.com/kubeflow/trainer/issues/2318
|
||||
podCount := info.TemplateSpec.PodSets[rJobIdx].CountForNonTrainer
|
||||
if *rJob.Name == constants.JobTrainerNode {
|
||||
podCount = info.RuntimePolicy.MLPolicy.NumNodes
|
||||
}
|
||||
rJobReplicas := 1
|
||||
info.TemplateSpec.PodSets[rJobIdx].Endpoints = func(yield func(string) bool) {
|
||||
for podIdx := range ptr.Deref(podCount, 1) {
|
||||
endpoint := fmt.Sprintf("%s-%s-%d-%d.%s", trainJob.Name, *rJob.Name, rJobReplicas-1, podIdx, subDomain)
|
||||
if !yield(endpoint) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
info.SyncPodSetsToTemplateSpec()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *trainer.TrainJob) ([]any, error) {
|
||||
if info == nil || trainJob == nil {
|
||||
return nil, fmt.Errorf("runtime info or object is missing")
|
||||
}
|
||||
jobSetSpec, ok := runtime.TemplateSpecApply[jobsetv1alpha2ac.JobSetSpecApplyConfiguration](info)
|
||||
if !ok {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Do not update the JobSet if it already exists and is not suspended
|
||||
oldJobSet := &jobsetv1alpha2.JobSet{}
|
||||
@@ -106,55 +142,32 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Get the runtime as unstructured from the TrainJob ref
|
||||
runtimeJobTemplate := &unstructured.Unstructured{}
|
||||
runtimeJobTemplate.SetAPIVersion(trainer.GroupVersion.String())
|
||||
|
||||
if kind := trainJob.Spec.RuntimeRef.Kind; kind != nil {
|
||||
runtimeJobTemplate.SetKind(*kind)
|
||||
} else {
|
||||
runtimeJobTemplate.SetKind(trainer.ClusterTrainingRuntimeKind)
|
||||
}
|
||||
|
||||
key := client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}
|
||||
if runtimeJobTemplate.GetKind() == trainer.TrainingRuntimeKind {
|
||||
key.Namespace = trainJob.Namespace
|
||||
}
|
||||
|
||||
if err := j.client.Get(ctx, key, runtimeJobTemplate); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Populate the JobSet template spec apply configuration
|
||||
jobSetTemplateSpec := &jobsetv1alpha2ac.JobSetSpecApplyConfiguration{}
|
||||
if jobSetSpec, ok, err := unstructured.NestedFieldCopy(runtimeJobTemplate.Object, "spec", "template", "spec"); err != nil {
|
||||
return nil, err
|
||||
} else if !ok {
|
||||
return nil, fmt.Errorf("trainJob runtime %s does not have a spec.template.spec field", trainJob.Spec.RuntimeRef.Name)
|
||||
} else {
|
||||
if raw, err := json.Marshal(jobSetSpec); err != nil {
|
||||
return nil, err
|
||||
} else if err := json.Unmarshal(raw, jobSetTemplateSpec); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
// Init the JobSet apply configuration from the runtime template spec
|
||||
jobSetBuilder := NewBuilder(jobsetv1alpha2ac.JobSet(trainJob.Name, trainJob.Namespace).
|
||||
WithLabels(maps.Clone(info.Labels)).
|
||||
WithAnnotations(maps.Clone(info.Annotations)).
|
||||
WithSpec(jobSetTemplateSpec))
|
||||
WithSpec(jobSetSpec))
|
||||
|
||||
// TODO (andreyvelich): Add support for the PodSpecOverride.
|
||||
// TODO (andreyvelich): Refactor the builder with wrappers for PodSpec.
|
||||
// Apply the runtime info
|
||||
jobSet := jobSetBuilder.
|
||||
Initializer(trainJob).
|
||||
Launcher(info, trainJob).
|
||||
Trainer(info, trainJob).
|
||||
PodLabels(info.PodLabels).
|
||||
Suspend(trainJob.Spec.Suspend).
|
||||
Build()
|
||||
// TODO: Once we remove deprecated runtime.Info.Trainer, we should remove JobSet Builder with DeprecatedTrainer().
|
||||
var jobSet *jobsetv1alpha2ac.JobSetApplyConfiguration
|
||||
if info.RuntimePolicy.MLPolicy != nil && info.RuntimePolicy.MLPolicy.MPI != nil {
|
||||
jobSet = jobSetBuilder.
|
||||
Initializer(trainJob).
|
||||
Launcher().
|
||||
Trainer(info, trainJob).
|
||||
PodLabels(info.PodLabels).
|
||||
Suspend(trainJob.Spec.Suspend).
|
||||
Build()
|
||||
} else {
|
||||
jobSet = jobSetBuilder.
|
||||
Initializer(trainJob).
|
||||
DeprecatedTrainer(info, trainJob).
|
||||
PodLabels(info.PodLabels).
|
||||
Suspend(trainJob.Spec.Suspend).
|
||||
Build()
|
||||
}
|
||||
|
||||
// Set the TrainJob as owner
|
||||
jobSet.WithOwnerReferences(metav1ac.OwnerReference().
|
||||
|
||||
@@ -0,0 +1,331 @@
|
||||
/*
|
||||
Copyright 2025 The Kubeflow Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package jobset
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/google/go-cmp/cmp/cmpopts"
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
batchv1ac "k8s.io/client-go/applyconfigurations/batch/v1"
|
||||
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
||||
"k8s.io/klog/v2/ktesting"
|
||||
"k8s.io/utils/ptr"
|
||||
jobsetv1alpha2ac "sigs.k8s.io/jobset/client-go/applyconfiguration/jobset/v1alpha2"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
"github.com/kubeflow/trainer/pkg/runtime"
|
||||
"github.com/kubeflow/trainer/pkg/runtime/framework"
|
||||
utiltesting "github.com/kubeflow/trainer/pkg/util/testing"
|
||||
)
|
||||
|
||||
// TODO: Add tests for all Interfaces.
|
||||
// REF: https://github.com/kubeflow/trainer/issues/2468
|
||||
|
||||
func TestJobSet(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
trainJob *trainer.TrainJob
|
||||
info *runtime.Info
|
||||
wantInfo *runtime.Info
|
||||
wantError error
|
||||
}{
|
||||
"no action when info is nil": {
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
Obj(),
|
||||
},
|
||||
"no action when trainJob is not nil": {
|
||||
info: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
},
|
||||
},
|
||||
"no action when template.spec is not JobSet": {
|
||||
info: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
ObjApply: batchv1ac.JobSpec(),
|
||||
},
|
||||
},
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
Obj(),
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
ObjApply: batchv1ac.JobSpec(),
|
||||
},
|
||||
},
|
||||
},
|
||||
"trainer numNodes is respected rather than parallelism when replicatedJob name is trainer-node": {
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
Obj(),
|
||||
info: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(2).
|
||||
MPIPolicy(nil, ptr.To(trainer.MPIImplementationOpenMPI), nil, nil).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Containers: make([]runtime.Container, 2),
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobLauncher).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName("sidecar"),
|
||||
corev1ac.Container().WithName(constants.ContainerLauncher),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName(constants.ContainerTrainer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(2).
|
||||
MPIPolicy(nil, ptr.To(trainer.MPIImplementationOpenMPI), nil, nil).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Containers: make([]runtime.Container, 2),
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-0-0.trainJob")
|
||||
yield("trainJob-trainer-node-0-1.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobLauncher).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName("sidecar"),
|
||||
corev1ac.Container().WithName(constants.ContainerLauncher),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName(constants.ContainerTrainer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
"subDomain in jobSetSpec is used to endpoint": {
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
Obj(),
|
||||
info: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithNetwork(jobsetv1alpha2ac.Network().
|
||||
WithSubdomain("kubeflow.org")).
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobLauncher).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName(constants.ContainerLauncher),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName(constants.ContainerTrainer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.kubeflow.org")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: make([]runtime.Container, 1),
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-0-0.kubeflow.org")
|
||||
},
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithNetwork(jobsetv1alpha2ac.Network().
|
||||
WithSubdomain("kubeflow.org")).
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobLauncher).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName(constants.ContainerLauncher),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithParallelism(1).
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().WithName(constants.ContainerTrainer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for name, tc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
_, ctx := ktesting.NewTestContext(t)
|
||||
var cancel func()
|
||||
ctx, cancel = context.WithCancel(ctx)
|
||||
t.Cleanup(cancel)
|
||||
cli := utiltesting.NewClientBuilder().Build()
|
||||
p, err := New(ctx, cli, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to initialize JobSet plugin: %v", err)
|
||||
}
|
||||
err = p.(framework.PodNetworkPlugin).IdentifyPodNetwork(tc.info, tc.trainJob)
|
||||
if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
|
||||
t.Errorf("Unexpected error (-want,+got):\n%s", diff)
|
||||
}
|
||||
if diff := cmp.Diff(tc.wantInfo, tc.info,
|
||||
cmpopts.SortSlices(func(a, b string) bool { return a < b }),
|
||||
cmpopts.SortMaps(func(a, b string) bool { return a < b }),
|
||||
utiltesting.PodSetEndpointsCmpOpts,
|
||||
); len(diff) != 0 {
|
||||
t.Errorf("Unexpected Info from IdentifyPodNetwork (-want,+got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
"strconv"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
apiruntime "k8s.io/apimachinery/pkg/runtime"
|
||||
"k8s.io/apimachinery/pkg/util/validation/field"
|
||||
@@ -46,6 +47,8 @@ import (
|
||||
"github.com/kubeflow/trainer/pkg/runtime/framework"
|
||||
)
|
||||
|
||||
// TODO : Support MPICH and IntelMPI implementations.
|
||||
|
||||
type MPI struct {
|
||||
client client.Client
|
||||
scheme *apiruntime.Scheme
|
||||
@@ -73,6 +76,10 @@ func (m *MPI) Name() string {
|
||||
}
|
||||
|
||||
// TODO: Need to implement validations for MPI Policy.
|
||||
// TODO (andreyvelich): Add validation to check that TrainJob doesn't have MPI envs.
|
||||
// TODO (andreyvelich): We should validate that envs from different plugins don't conflict with each other.
|
||||
// Ref: https://github.com/kubeflow/trainer/pull/2308#discussion_r1823229940
|
||||
|
||||
func (m *MPI) Validate(oldObj, newObj *trainer.TrainJob) (admission.Warnings, field.ErrorList) {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -83,85 +90,113 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
|
||||
}
|
||||
|
||||
// TrainJob contains the actual information for the Trainer.
|
||||
numNodes := info.RuntimePolicy.MLPolicy.NumNodes
|
||||
if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
|
||||
numNodes = trainJob.Spec.Trainer.NumNodes
|
||||
info.RuntimePolicy.MLPolicy.NumNodes = trainJob.Spec.Trainer.NumNodes
|
||||
}
|
||||
info.Trainer.NumNodes = numNodes
|
||||
|
||||
numProcPerNode := strconv.Itoa(int(*info.RuntimePolicy.MLPolicy.MPI.NumProcPerNode))
|
||||
if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumProcPerNode != nil {
|
||||
numProcPerNode = (*trainJob.Spec.Trainer.NumProcPerNode).String()
|
||||
info.RuntimePolicy.MLPolicy.MPI.NumProcPerNode = ptr.To(int32(trainJob.Spec.Trainer.NumProcPerNode.IntValue()))
|
||||
}
|
||||
info.Trainer.NumProcPerNode = numProcPerNode
|
||||
|
||||
// Add Secret and ConfigMap volumes to the Info object
|
||||
info.Trainer.Volumes = []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(trainJob.Name+constants.MPISSHAuthSecretSuffix).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
)),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(trainJob.Name + constants.MPIHostfileConfigMapSuffix)),
|
||||
for psIdx, ps := range info.TemplateSpec.PodSets {
|
||||
if ps.Name != constants.JobTrainerNode && ps.Name != constants.JobLauncher {
|
||||
continue
|
||||
}
|
||||
apply.UpsertVolumes(
|
||||
&info.TemplateSpec.PodSets[psIdx].Volumes,
|
||||
[]corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("%s%s", trainJob.Name, constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
}...,
|
||||
)
|
||||
if ps.Name == constants.JobLauncher {
|
||||
apply.UpsertVolumes(
|
||||
&info.TemplateSpec.PodSets[psIdx].Volumes,
|
||||
[]corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("%s%s", trainJob.Name, constants.MPIHostfileConfigMapSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPIHostfileName).
|
||||
WithPath(constants.MPIHostfileName).
|
||||
WithMode(0444),
|
||||
),
|
||||
),
|
||||
}...,
|
||||
)
|
||||
}
|
||||
for cIdx, container := range ps.Containers {
|
||||
if container.Name != constants.JobLauncher && container.Name != constants.JobTrainerNode {
|
||||
continue
|
||||
}
|
||||
apply.UpsertVolumeMounts(
|
||||
&info.TemplateSpec.PodSets[psIdx].Containers[cIdx].VolumeMounts,
|
||||
[]corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithMountPath(*info.RuntimePolicy.MLPolicy.MPI.SSHAuthMountPath),
|
||||
}...,
|
||||
)
|
||||
if ps.Name == constants.JobLauncher && container.Name == constants.ContainerLauncher {
|
||||
apply.UpsertVolumeMounts(
|
||||
&info.TemplateSpec.PodSets[psIdx].Containers[cIdx].VolumeMounts,
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithMountPath(constants.MPIHostfileDir),
|
||||
)
|
||||
switch *info.RuntimePolicy.MLPolicy.MPI.MPIImplementation {
|
||||
case trainer.MPIImplementationOpenMPI:
|
||||
apply.UpsertEnvVars(
|
||||
&info.TemplateSpec.PodSets[psIdx].Containers[cIdx].Env,
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.OpenMPIEnvHostFileLocation).
|
||||
WithValue(fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName)),
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.OpenMPIEnvKeepFQDNHostNames).
|
||||
WithValue("true"),
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.OpenMPIEnvDefaultSlots).
|
||||
WithValue(strconv.Itoa(int(*info.RuntimePolicy.MLPolicy.MPI.NumProcPerNode))),
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.OpenMPIEnvKeyRSHArgs).
|
||||
WithValue(constants.OpenMPIEnvDefaultValueRSHArgs),
|
||||
)
|
||||
default:
|
||||
return fmt.Errorf("MPI implementation for %v doesn't supported", info.RuntimePolicy.MLPolicy.MPI.MPIImplementation)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info.Trainer.VolumeMounts = []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithMountPath(*info.RuntimePolicy.MLPolicy.MPI.SSHAuthMountPath),
|
||||
*corev1ac.VolumeMount().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithMountPath(constants.MPIHostfileDir),
|
||||
}
|
||||
|
||||
// Update envs for Info object.
|
||||
// TODO (andreyvelich): Add validation to check that TrainJob doesn't have MPI envs.
|
||||
// TODO (andreyvelich): We should validate that envs from different plugins don't conflict with each other.
|
||||
// Ref: https://github.com/kubeflow/trainer/pull/2308#discussion_r1823229940
|
||||
// TODO (andreyvelich): Support other MPI implementations.
|
||||
|
||||
if trainJob.Spec.Trainer != nil {
|
||||
info.Trainer.Env = apply.EnvVars(trainJob.Spec.Trainer.Env...)
|
||||
}
|
||||
|
||||
switch *info.RuntimePolicy.MLPolicy.MPI.MPIImplementation {
|
||||
case trainer.MPIImplementationOpenMPI:
|
||||
apply.UpsertEnvVar(&info.Trainer.Env, corev1ac.EnvVar().
|
||||
WithName(constants.OpenMPIEnvHostFileLocation).
|
||||
WithValue(fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName)))
|
||||
default:
|
||||
return fmt.Errorf("MPI implementation for %v doesn't supported", info.RuntimePolicy.MLPolicy.MPI.MPIImplementation)
|
||||
}
|
||||
|
||||
// Add container port for the headless service.
|
||||
info.Trainer.ContainerPort = corev1ac.ContainerPort().
|
||||
WithContainerPort(constants.ContainerTrainerPort)
|
||||
|
||||
// Update total Pod requests for the PodGroupPolicy plugin.
|
||||
for rName := range info.TotalRequests {
|
||||
// For other Jobs like the Initializer, replica is always equal to 1.
|
||||
// TODO (andreyvelich): Add support for total requests from the TrainJob's ResourcesPerNode.
|
||||
if rName == constants.JobTrainerNode {
|
||||
info.TotalRequests[rName] = runtime.TotalResourceRequest{
|
||||
Replicas: ptr.Deref(numNodes, constants.DefaultJobReplicas),
|
||||
Replicas: ptr.Deref(info.RuntimePolicy.MLPolicy.NumNodes, constants.DefaultJobReplicas),
|
||||
PodRequests: info.TotalRequests[rName].PodRequests,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info.SyncPodSetsToTemplateSpec()
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -176,90 +211,92 @@ func (m *MPI) ReconcilerBuilders() []runtime.ReconcilerBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MPI) Build(_ context.Context, info *runtime.Info, trainJob *trainer.TrainJob) ([]any, error) {
|
||||
func (m *MPI) Build(ctx context.Context, info *runtime.Info, trainJob *trainer.TrainJob) ([]any, error) {
|
||||
if info == nil || info.RuntimePolicy.MLPolicy == nil || info.RuntimePolicy.MLPolicy.MPI == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
secret, err := m.buildSSHAuthSecret(trainJob)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to build Secret with SSH auth keys. Error: %v", err)
|
||||
}
|
||||
var objects []any
|
||||
|
||||
configMap, err := m.buildHostFileConfigMap(info, trainJob)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to build ConfigMap with hostfile. Error: %v", err)
|
||||
// SSHAuthSecret is immutable.
|
||||
if err := m.client.Get(ctx, client.ObjectKey{Name: sshAuthSecretName(trainJob.Name), Namespace: trainJob.Namespace}, &corev1.Secret{}); err != nil {
|
||||
if client.IgnoreNotFound(err) != nil {
|
||||
return nil, err
|
||||
}
|
||||
secret, err := m.buildSSHAuthSecret(trainJob)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to build SSH Auth secret: %w", err)
|
||||
}
|
||||
objects = append(objects, secret)
|
||||
}
|
||||
|
||||
return []any{secret, configMap}, nil
|
||||
return append(objects, m.buildHostFileConfigMap(info, trainJob)), nil
|
||||
}
|
||||
|
||||
func (m *MPI) buildSSHAuthSecret(trainJob *trainer.TrainJob) (*corev1ac.SecretApplyConfiguration, error) {
|
||||
// Generate SSH private and public keys.
|
||||
privateKey, err := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to generate private SSH key, Error: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
privateDER, err := x509.MarshalECPrivateKey(privateKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to covert private SSH key to DER format. Error: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
privatePEM := pem.EncodeToMemory(&pem.Block{
|
||||
Type: "EC PRIVATE KEY",
|
||||
Bytes: privateDER,
|
||||
})
|
||||
|
||||
publicKey, err := ssh.NewPublicKey(&privateKey.PublicKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to generate public SSH key. Error: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Create Secret to store ssh keys.
|
||||
secret := corev1ac.Secret(trainJob.Name+constants.MPISSHAuthSecretSuffix, trainJob.Namespace).
|
||||
return corev1ac.Secret(sshAuthSecretName(trainJob.Name), trainJob.Namespace).
|
||||
WithType(corev1.SecretTypeSSHAuth).
|
||||
WithData(map[string][]byte{
|
||||
corev1.SSHAuthPrivateKey: privatePEM,
|
||||
constants.MPISSHPublicKey: ssh.MarshalAuthorizedKey(publicKey),
|
||||
})
|
||||
|
||||
secret.WithOwnerReferences(metav1ac.OwnerReference().
|
||||
WithAPIVersion(trainer.GroupVersion.String()).
|
||||
WithKind(trainer.TrainJobKind).
|
||||
WithName(trainJob.Name).
|
||||
WithUID(trainJob.UID).
|
||||
WithController(true).
|
||||
WithBlockOwnerDeletion(true))
|
||||
|
||||
return secret, nil
|
||||
}).
|
||||
WithImmutable(true).
|
||||
WithOwnerReferences(metav1ac.OwnerReference().
|
||||
WithAPIVersion(trainer.GroupVersion.String()).
|
||||
WithKind(trainer.TrainJobKind).
|
||||
WithName(trainJob.Name).
|
||||
WithUID(trainJob.UID).
|
||||
WithController(true).
|
||||
WithBlockOwnerDeletion(true)), nil
|
||||
}
|
||||
|
||||
func (m *MPI) buildHostFileConfigMap(info *runtime.Info, trainJob *trainer.TrainJob) (*corev1ac.ConfigMapApplyConfiguration, error) {
|
||||
// Generate hostfile for the MPI communication.
|
||||
var hostfile bytes.Buffer
|
||||
// TODO (andreyvelich): Support other MPI implementations.
|
||||
for i := range *info.Trainer.NumNodes {
|
||||
func sshAuthSecretName(trainJobName string) string {
|
||||
return fmt.Sprintf("%s%s", trainJobName, constants.MPISSHAuthSecretSuffix)
|
||||
}
|
||||
|
||||
func (m *MPI) buildHostFileConfigMap(info *runtime.Info, trainJob *trainer.TrainJob) *corev1ac.ConfigMapApplyConfiguration {
|
||||
var hostFile bytes.Buffer
|
||||
runLauncherAsNode := ptr.Deref(info.RuntimePolicy.MLPolicy.MPI.RunLauncherAsNode, false)
|
||||
slots := ptr.Deref(info.RuntimePolicy.MLPolicy.MPI.NumProcPerNode, 1)
|
||||
for _, ps := range info.TemplateSpec.PodSets {
|
||||
if !isTrainerNode(runLauncherAsNode, ps) {
|
||||
continue
|
||||
}
|
||||
switch *info.RuntimePolicy.MLPolicy.MPI.MPIImplementation {
|
||||
case trainer.MPIImplementationOpenMPI:
|
||||
// hostfile.WriteString(fmt.Sprintf("%s-%s-0-%s.%s.%s.svc slots=%s\n", trainJob.Name, constants.JobTrainerNode, i, trainJob.Name, trainJob.Namespace, info.NumProcPerNode))
|
||||
hostfile.WriteString(fmt.Sprintf("%s-%s-0-%d.%s slots=%s\n", trainJob.Name, constants.JobTrainerNode, i, info.NumProcPerNode, trainJob.Name))
|
||||
for e := range ps.Endpoints {
|
||||
hostFile.WriteString(fmt.Sprintf("%s slots=%d\n", e, slots))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Create ConfigMap to store hostfile.
|
||||
configMap := corev1ac.ConfigMap(trainJob.Name+constants.MPIHostfileConfigMapSuffix, trainJob.Namespace).
|
||||
return corev1ac.ConfigMap(fmt.Sprintf("%s%s", trainJob.Name, constants.MPIHostfileConfigMapSuffix), trainJob.Namespace).
|
||||
WithData(map[string]string{
|
||||
constants.MPIHostfileName: hostfile.String(),
|
||||
})
|
||||
constants.MPIHostfileName: hostFile.String(),
|
||||
}).
|
||||
WithOwnerReferences(metav1ac.OwnerReference().
|
||||
WithAPIVersion(trainer.GroupVersion.String()).
|
||||
WithKind(trainer.TrainJobKind).
|
||||
WithName(trainJob.Name).
|
||||
WithUID(trainJob.UID).
|
||||
WithController(true).
|
||||
WithBlockOwnerDeletion(true))
|
||||
}
|
||||
|
||||
configMap.WithOwnerReferences(metav1ac.OwnerReference().
|
||||
WithAPIVersion(trainer.GroupVersion.String()).
|
||||
WithKind(trainer.TrainJobKind).
|
||||
WithName(trainJob.Name).
|
||||
WithUID(trainJob.UID).
|
||||
WithController(true).
|
||||
WithBlockOwnerDeletion(true))
|
||||
|
||||
return configMap, nil
|
||||
func isTrainerNode(runLauncherAsNode bool, ps runtime.PodSet) bool {
|
||||
return (runLauncherAsNode && ps.Name == constants.JobLauncher) || ps.Name == constants.JobTrainerNode
|
||||
}
|
||||
|
||||
@@ -17,11 +17,13 @@ limitations under the License.
|
||||
package mpi
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
gocmp "github.com/google/go-cmp/cmp"
|
||||
"github.com/google/go-cmp/cmp/cmpopts"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
@@ -32,6 +34,8 @@ import (
|
||||
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
||||
"k8s.io/klog/v2/ktesting"
|
||||
"k8s.io/utils/ptr"
|
||||
"sigs.k8s.io/controller-runtime/pkg/client"
|
||||
"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/runtime"
|
||||
@@ -40,16 +44,18 @@ import (
|
||||
)
|
||||
|
||||
func TestMPI(t *testing.T) {
|
||||
objCmpOpts := []cmp.Option{
|
||||
cmpopts.SortSlices(func(a, b apiruntime.Object) bool {
|
||||
return a.GetObjectKind().GroupVersionKind().String() < b.GetObjectKind().GroupVersionKind().String()
|
||||
objCmpOpts := []gocmp.Option{
|
||||
cmpopts.SortSlices(func(a, b apiruntime.Object) int {
|
||||
return cmp.Compare(a.GetObjectKind().GroupVersionKind().String(), b.GetObjectKind().GroupVersionKind().String())
|
||||
}),
|
||||
cmp.Comparer(secretDataComparer),
|
||||
gocmp.Comparer(utiltesting.MPISecretDataComparer),
|
||||
}
|
||||
errorGetSSHAuthSecretFromAPI := errors.New("failed to get SSH Auth Secret from API during Build")
|
||||
|
||||
cases := map[string]struct {
|
||||
info *runtime.Info
|
||||
trainJob *trainer.TrainJob
|
||||
objs []client.Object
|
||||
wantInfo *runtime.Info
|
||||
wantObjs []apiruntime.Object
|
||||
wantMLPolicyError error
|
||||
@@ -57,32 +63,61 @@ func TestMPI(t *testing.T) {
|
||||
}{
|
||||
"no action when info is nil": {},
|
||||
"no action when mlPolicy is nil": {
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
),
|
||||
wantInfo: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
),
|
||||
info: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
},
|
||||
},
|
||||
"no action when mlPolicy mpi is null": {
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithMLPolicy(utiltesting.MakeMLPolicyWrapper().
|
||||
Obj()),
|
||||
),
|
||||
wantInfo: runtime.NewInfo(
|
||||
runtime.WithMLPolicy(utiltesting.MakeMLPolicyWrapper().
|
||||
Obj()),
|
||||
),
|
||||
info: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().Obj(),
|
||||
},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().Obj(),
|
||||
},
|
||||
},
|
||||
},
|
||||
"trainJob numNodes is respected rather than mlPolicy one": {
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithMLPolicy(
|
||||
utiltesting.MakeMLPolicyWrapper().
|
||||
info: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-1.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), nil).
|
||||
Obj(),
|
||||
),
|
||||
),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
TotalRequests: make(map[string]runtime.TotalResourceRequest),
|
||||
},
|
||||
},
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
UID("trainJob").
|
||||
Trainer(
|
||||
@@ -93,72 +128,150 @@ func TestMPI(t *testing.T) {
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPIHostfileName).
|
||||
WithPath(constants.MPIHostfileName).
|
||||
WithMode(0444),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-1.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
WithNumNodes(2).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), nil).
|
||||
Obj(),
|
||||
},
|
||||
Trainer: runtime.Trainer{
|
||||
NumNodes: ptr.To[int32](2),
|
||||
NumProcPerNode: "1",
|
||||
Env: []corev1ac.EnvVarApplyConfiguration{{
|
||||
Name: ptr.To(constants.OpenMPIEnvHostFileLocation),
|
||||
Value: ptr.To(fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName)),
|
||||
}},
|
||||
ContainerPort: &corev1ac.ContainerPortApplyConfiguration{
|
||||
ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
|
||||
},
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().WithKey(corev1.SSHAuthPrivateKey).WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().WithKey(constants.MPISSHPublicKey).WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().WithKey(constants.MPISSHPublicKey).WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob-%s", constants.MPIHostfileVolumeName))),
|
||||
},
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().WithName(constants.MPISSHAuthVolumeName).WithMountPath("/root/.ssh"),
|
||||
*corev1ac.VolumeMount().WithName(constants.MPIHostfileVolumeName).WithMountPath(constants.MPIHostfileDir),
|
||||
},
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{TotalRequests: make(map[string]runtime.TotalResourceRequest)},
|
||||
},
|
||||
wantObjs: []apiruntime.Object{
|
||||
utiltesting.MakeSecretWrapper(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix), metav1.NamespaceDefault).
|
||||
WithImmutable(true).
|
||||
WithType(corev1.SecretTypeSSHAuth).
|
||||
WithData(map[string][]byte{
|
||||
constants.MPISSHPublicKey: []byte("EXIST"),
|
||||
corev1.SSHAuthPrivateKey: []byte("EXIST"),
|
||||
}).
|
||||
OwnerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "trainJob", "trainJob").
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
|
||||
WithData(map[string]string{
|
||||
constants.MPIHostfileName: `trainJob-trainer-node-0-0.1 slots=trainJob
|
||||
trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
constants.MPIHostfileName: `trainJob-trainer-node-1-0.trainJob slots=1
|
||||
trainJob-trainer-node-1-1.trainJob slots=1
|
||||
`,
|
||||
}).
|
||||
OwnerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "trainJob", "trainJob").
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
},
|
||||
},
|
||||
"trainJob numProcPerNode is respected rather than mlPolicy one": {
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithMLPolicy(
|
||||
utiltesting.MakeMLPolicyWrapper().
|
||||
info: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), nil).
|
||||
Obj(),
|
||||
),
|
||||
),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
TotalRequests: make(map[string]runtime.TotalResourceRequest),
|
||||
},
|
||||
},
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
UID("trainJob").
|
||||
Trainer(
|
||||
@@ -170,59 +283,93 @@ trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPIHostfileName).
|
||||
WithPath(constants.MPIHostfileName).
|
||||
WithMode(0444),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), nil).
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](2), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), nil).
|
||||
Obj(),
|
||||
},
|
||||
Trainer: runtime.Trainer{
|
||||
NumNodes: ptr.To[int32](1),
|
||||
NumProcPerNode: "2",
|
||||
Env: []corev1ac.EnvVarApplyConfiguration{{
|
||||
Name: ptr.To(constants.OpenMPIEnvHostFileLocation),
|
||||
Value: ptr.To(fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName)),
|
||||
}},
|
||||
ContainerPort: &corev1ac.ContainerPortApplyConfiguration{
|
||||
ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
|
||||
},
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().WithKey(corev1.SSHAuthPrivateKey).WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().WithKey(constants.MPISSHPublicKey).WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().WithKey(constants.MPISSHPublicKey).WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob-%s", constants.MPIHostfileVolumeName))),
|
||||
},
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().WithName(constants.MPISSHAuthVolumeName).WithMountPath("/root/.ssh"),
|
||||
*corev1ac.VolumeMount().WithName(constants.MPIHostfileVolumeName).WithMountPath(constants.MPIHostfileDir),
|
||||
},
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{TotalRequests: make(map[string]runtime.TotalResourceRequest)},
|
||||
},
|
||||
wantObjs: []apiruntime.Object{
|
||||
utiltesting.MakeSecretWrapper(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix), metav1.NamespaceDefault).
|
||||
WithImmutable(true).
|
||||
WithType(corev1.SecretTypeSSHAuth).
|
||||
WithData(map[string][]byte{
|
||||
constants.MPISSHPublicKey: []byte("EXIST"),
|
||||
corev1.SSHAuthPrivateKey: []byte("EXIST"),
|
||||
}).
|
||||
OwnerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "trainJob", "trainJob").
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
|
||||
WithData(map[string]string{
|
||||
constants.MPIHostfileName: `trainJob-trainer-node-0-0.2 slots=trainJob
|
||||
constants.MPIHostfileName: `trainJob-trainer-node-1-0.trainJob slots=2
|
||||
`,
|
||||
}).
|
||||
OwnerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "trainJob", "trainJob").
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
},
|
||||
},
|
||||
@@ -236,6 +383,14 @@ trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), nil).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{{
|
||||
Name: constants.JobLauncher,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
}},
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{TotalRequests: map[string]runtime.TotalResourceRequest{
|
||||
constants.JobTrainerNode: {
|
||||
Replicas: 1,
|
||||
@@ -258,40 +413,46 @@ trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
Annotations: make(map[string]string),
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
WithNumNodes(2).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), nil).
|
||||
Obj(),
|
||||
},
|
||||
Trainer: runtime.Trainer{
|
||||
NumNodes: ptr.To[int32](2),
|
||||
NumProcPerNode: "1",
|
||||
Env: []corev1ac.EnvVarApplyConfiguration{{
|
||||
Name: ptr.To(constants.OpenMPIEnvHostFileLocation),
|
||||
Value: ptr.To(fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName)),
|
||||
}},
|
||||
ContainerPort: &corev1ac.ContainerPortApplyConfiguration{
|
||||
ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
|
||||
},
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().WithKey(corev1.SSHAuthPrivateKey).WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().WithKey(constants.MPISSHPublicKey).WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().WithKey(constants.MPISSHPublicKey).WithPath(constants.MPISSHAuthorizedKeys),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{{
|
||||
Name: constants.JobLauncher,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob-%s", constants.MPIHostfileVolumeName))),
|
||||
},
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
|
||||
*corev1ac.VolumeMount().WithName(constants.MPISSHAuthVolumeName).WithMountPath("/root/.ssh"),
|
||||
*corev1ac.VolumeMount().WithName(constants.MPIHostfileVolumeName).WithMountPath(constants.MPIHostfileDir),
|
||||
},
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPIHostfileName).
|
||||
WithPath(constants.MPIHostfileName).
|
||||
WithMode(0444),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
}},
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{TotalRequests: map[string]runtime.TotalResourceRequest{
|
||||
constants.JobTrainerNode: {
|
||||
@@ -305,23 +466,325 @@ trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
},
|
||||
wantObjs: []apiruntime.Object{
|
||||
utiltesting.MakeSecretWrapper(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix), metav1.NamespaceDefault).
|
||||
WithImmutable(true).
|
||||
WithType(corev1.SecretTypeSSHAuth).
|
||||
WithData(map[string][]byte{
|
||||
constants.MPISSHPublicKey: []byte("EXIST"),
|
||||
corev1.SSHAuthPrivateKey: []byte("EXIST"),
|
||||
}).
|
||||
OwnerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "trainJob", "trainJob").
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
|
||||
WithData(map[string]string{
|
||||
constants.MPIHostfileName: `trainJob-trainer-node-0-0.1 slots=trainJob
|
||||
trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
`,
|
||||
constants.MPIHostfileName: "",
|
||||
}).
|
||||
OwnerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "trainJob", "trainJob").
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
},
|
||||
},
|
||||
"runLauncherAsNode is true": {
|
||||
info: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(true)).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
TotalRequests: make(map[string]runtime.TotalResourceRequest),
|
||||
},
|
||||
},
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
UID("trainJob").
|
||||
Trainer(
|
||||
utiltesting.MakeTrainJobTrainerWrapper().
|
||||
NumNodes(1).
|
||||
Obj()).
|
||||
Obj(),
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(true)).
|
||||
Obj(),
|
||||
},
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPIHostfileName).
|
||||
WithPath(constants.MPIHostfileName).
|
||||
WithMode(0444),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-trainer-node-1-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{TotalRequests: make(map[string]runtime.TotalResourceRequest)},
|
||||
},
|
||||
wantObjs: []apiruntime.Object{
|
||||
utiltesting.MakeSecretWrapper(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix), metav1.NamespaceDefault).
|
||||
WithImmutable(true).
|
||||
WithType(corev1.SecretTypeSSHAuth).
|
||||
WithData(map[string][]byte{
|
||||
constants.MPISSHPublicKey: []byte("EXIST"),
|
||||
corev1.SSHAuthPrivateKey: []byte("EXIST"),
|
||||
}).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
|
||||
WithData(map[string]string{
|
||||
constants.MPIHostfileName: `trainJob-launcher-0-0.trainJob slots=1
|
||||
trainJob-trainer-node-1-0.trainJob slots=1
|
||||
`,
|
||||
}).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
},
|
||||
},
|
||||
"sshAuth secret already has existed in the cluster": {
|
||||
objs: []client.Object{
|
||||
utiltesting.MakeSecretWrapper(sshAuthSecretName("trainJob"), metav1.NamespaceDefault).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
WithImmutable(true).
|
||||
Obj(),
|
||||
},
|
||||
info: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(true)).
|
||||
Obj(),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
TotalRequests: make(map[string]runtime.TotalResourceRequest),
|
||||
},
|
||||
},
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
UID("trainJob").
|
||||
Obj(),
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPIHostfileName).
|
||||
WithPath(constants.MPIHostfileName).
|
||||
WithMode(0444),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(true)).
|
||||
Obj(),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{TotalRequests: make(map[string]runtime.TotalResourceRequest)},
|
||||
},
|
||||
wantObjs: []apiruntime.Object{
|
||||
utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
|
||||
WithData(map[string]string{
|
||||
constants.MPIHostfileName: `trainJob-launcher-0-0.trainJob slots=1
|
||||
`,
|
||||
}).
|
||||
ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
|
||||
Obj(),
|
||||
},
|
||||
},
|
||||
"failed to get sshAuth secret due to API error": {
|
||||
info: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(true)).
|
||||
Obj(),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
TotalRequests: make(map[string]runtime.TotalResourceRequest),
|
||||
},
|
||||
},
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
|
||||
UID("trainJob").
|
||||
Obj(),
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
TemplateSpec: runtime.TemplateSpec{
|
||||
PodSets: []runtime.PodSet{
|
||||
{
|
||||
Name: constants.JobLauncher,
|
||||
Volumes: []corev1ac.VolumeApplyConfiguration{
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPISSHAuthVolumeName).
|
||||
WithSecret(corev1ac.SecretVolumeSource().
|
||||
WithSecretName(fmt.Sprintf("trainJob%s", constants.MPISSHAuthSecretSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(corev1.SSHAuthPrivateKey).
|
||||
WithPath(constants.MPISSHPrivateKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHPublicKeyFile),
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPISSHPublicKey).
|
||||
WithPath(constants.MPISSHAuthorizedKeys),
|
||||
),
|
||||
),
|
||||
*corev1ac.Volume().
|
||||
WithName(constants.MPIHostfileVolumeName).
|
||||
WithConfigMap(corev1ac.ConfigMapVolumeSource().
|
||||
WithName(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix)).
|
||||
WithItems(
|
||||
corev1ac.KeyToPath().
|
||||
WithKey(constants.MPIHostfileName).
|
||||
WithPath(constants.MPIHostfileName).
|
||||
WithMode(0444),
|
||||
),
|
||||
),
|
||||
},
|
||||
Endpoints: func(yield func(string) bool) {
|
||||
yield("trainJob-launcher-0-0.trainJob")
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(1).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(true)).
|
||||
Obj(),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{TotalRequests: make(map[string]runtime.TotalResourceRequest)},
|
||||
},
|
||||
wantBuildError: errorGetSSHAuthSecretFromAPI,
|
||||
},
|
||||
}
|
||||
for name, tc := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
@@ -329,24 +792,35 @@ trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
var cancel func()
|
||||
ctx, cancel = context.WithCancel(ctx)
|
||||
t.Cleanup(cancel)
|
||||
cli := utiltesting.NewClientBuilder().Build()
|
||||
b := utiltesting.NewClientBuilder().WithObjects(tc.objs...)
|
||||
b.WithInterceptorFuncs(interceptor.Funcs{
|
||||
Get: func(ctx context.Context, client client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
|
||||
if _, ok := obj.(*corev1.Secret); ok && errors.Is(tc.wantBuildError, errorGetSSHAuthSecretFromAPI) {
|
||||
return errorGetSSHAuthSecretFromAPI
|
||||
}
|
||||
return client.Get(ctx, key, obj, opts...)
|
||||
},
|
||||
})
|
||||
cli := b.Build()
|
||||
|
||||
p, err := New(ctx, cli, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to initialize MPI plugin: %v", err)
|
||||
}
|
||||
err = p.(framework.EnforceMLPolicyPlugin).EnforceMLPolicy(tc.info, tc.trainJob)
|
||||
if diff := cmp.Diff(tc.wantMLPolicyError, err, cmpopts.EquateErrors()); len(diff) != 0 {
|
||||
if diff := gocmp.Diff(tc.wantMLPolicyError, err, cmpopts.EquateErrors()); len(diff) != 0 {
|
||||
t.Errorf("Unexpected error from EnforceMLPolicy (-want, +got): %s", diff)
|
||||
}
|
||||
if diff := cmp.Diff(tc.wantInfo, tc.info,
|
||||
if diff := gocmp.Diff(tc.wantInfo, tc.info,
|
||||
cmpopts.SortSlices(func(a, b string) bool { return a < b }),
|
||||
cmpopts.SortMaps(func(a, b int) bool { return a < b }),
|
||||
utiltesting.PodSetEndpointsCmpOpts,
|
||||
); len(diff) != 0 {
|
||||
t.Errorf("Unexpected info from EnforceMLPolicy (-want, +got): %s", diff)
|
||||
}
|
||||
var objs []any
|
||||
objs, err = p.(framework.ComponentBuilderPlugin).Build(ctx, tc.info, tc.trainJob)
|
||||
if diff := cmp.Diff(tc.wantBuildError, err, cmpopts.EquateErrors()); len(diff) != 0 {
|
||||
if diff := gocmp.Diff(tc.wantBuildError, err, cmpopts.EquateErrors()); len(diff) != 0 {
|
||||
t.Errorf("Unexpected error from Build (-want, +got): %s", diff)
|
||||
}
|
||||
var typedObjs []apiruntime.Object
|
||||
@@ -354,21 +828,9 @@ trainJob-trainer-node-0-1.1 slots=trainJob
|
||||
if err != nil {
|
||||
t.Errorf("Failed to convert object: %v", err)
|
||||
}
|
||||
if diff := cmp.Diff(typedObjs, tc.wantObjs, objCmpOpts...); len(diff) != 0 {
|
||||
if diff := gocmp.Diff(tc.wantObjs, typedObjs, objCmpOpts...); len(diff) != 0 {
|
||||
t.Errorf("Unexpected objects from Build (-want, +got): %s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func secretDataComparer(a, b map[string][]byte) bool {
|
||||
isKeysEqual := true
|
||||
if (a != nil && b != nil) &&
|
||||
((len(a[constants.MPISSHPublicKey]) > 0) != (len(b[constants.MPISSHPublicKey]) > 0) ||
|
||||
(len(a[corev1.SSHAuthPrivateKey]) > 0) != (len(b[corev1.SSHAuthPrivateKey]) > 0)) {
|
||||
isKeysEqual = false
|
||||
}
|
||||
return isKeysEqual && cmp.Equal(a, b, cmpopts.IgnoreMapEntries(func(k string, _ []byte) bool {
|
||||
return k == constants.MPISSHPublicKey || k == corev1.SSHAuthPrivateKey
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -58,7 +58,7 @@ func (p *PlainML) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob
|
||||
|
||||
// Add envs from the TrainJob.
|
||||
if trainJob.Spec.Trainer != nil {
|
||||
apply.UpsertEnvVars(&info.Trainer.Env, apply.EnvVars(trainJob.Spec.Trainer.Env...))
|
||||
apply.UpsertEnvVars(&info.Trainer.Env, apply.EnvVars(trainJob.Spec.Trainer.Env...)...)
|
||||
}
|
||||
|
||||
// Update total Pod requests for the PodGroupPolicy plugin.
|
||||
|
||||
@@ -45,52 +45,52 @@ func TestPlainML(t *testing.T) {
|
||||
}{
|
||||
"no action when info is null": {},
|
||||
"no action when mlPolicy is null": {
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
),
|
||||
wantInfo: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
),
|
||||
info: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
},
|
||||
},
|
||||
"no action when mlPolicy torch is not null": {
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
runtime.WithMLPolicy(
|
||||
utiltesting.MakeMLPolicyWrapper().
|
||||
info: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(100).
|
||||
TorchPolicy("auto", nil).
|
||||
Obj(),
|
||||
),
|
||||
),
|
||||
wantInfo: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
runtime.WithMLPolicy(
|
||||
utiltesting.MakeMLPolicyWrapper().
|
||||
},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(100).
|
||||
TorchPolicy("auto", nil).
|
||||
Obj(),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
"no action when mlPolicy mpi is not null": {
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
runtime.WithMLPolicy(
|
||||
utiltesting.MakeMLPolicyWrapper().
|
||||
info: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(100).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), nil, ptr.To(false)).
|
||||
Obj(),
|
||||
),
|
||||
),
|
||||
wantInfo: runtime.NewInfo(
|
||||
runtime.WithLabels(map[string]string{"key": "value"}),
|
||||
runtime.WithMLPolicy(
|
||||
utiltesting.MakeMLPolicyWrapper().
|
||||
},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: map[string]string{"key": "value"},
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().
|
||||
WithNumNodes(100).
|
||||
MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), nil, ptr.To(false)).
|
||||
Obj(),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
"trainJob trainer numNodes are respected rather than runtimeInfo": {
|
||||
trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
|
||||
@@ -98,11 +98,16 @@ func TestPlainML(t *testing.T) {
|
||||
utiltesting.MakeTrainJobTrainerWrapper().NumNodes(200).Obj(),
|
||||
).
|
||||
Obj(),
|
||||
info: runtime.NewInfo(
|
||||
runtime.WithMLPolicy(
|
||||
utiltesting.MakeMLPolicyWrapper().WithNumNodes(100).Obj(),
|
||||
),
|
||||
),
|
||||
info: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
RuntimePolicy: runtime.RuntimePolicy{
|
||||
MLPolicy: utiltesting.MakeMLPolicyWrapper().WithNumNodes(100).Obj(),
|
||||
},
|
||||
Scheduler: &runtime.Scheduler{
|
||||
TotalRequests: make(map[string]runtime.TotalResourceRequest),
|
||||
},
|
||||
},
|
||||
wantInfo: &runtime.Info{
|
||||
Labels: make(map[string]string),
|
||||
Annotations: make(map[string]string),
|
||||
|
||||
@@ -82,21 +82,21 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob)
|
||||
}
|
||||
|
||||
apply.UpsertEnvVar(&info.Trainer.Env,
|
||||
corev1ac.EnvVar().
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.TorchEnvNumNodes).
|
||||
WithValue(fmt.Sprintf("%d", ptr.Deref(numNodes, 1))),
|
||||
corev1ac.EnvVar().
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.TorchEnvNumProcPerNode).
|
||||
WithValue(numProcPerNode.String()),
|
||||
corev1ac.EnvVar().
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.TorchEnvNodeRank).
|
||||
WithValueFrom(corev1ac.EnvVarSource().
|
||||
WithFieldRef(corev1ac.ObjectFieldSelector().
|
||||
WithFieldPath(constants.JobCompletionIndexFieldPath))),
|
||||
corev1ac.EnvVar().
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.TorchEnvMasterAddr).
|
||||
WithValue(fmt.Sprintf("%s-%s-0-0.%s", trainJob.Name, constants.JobTrainerNode, trainJob.Name)),
|
||||
corev1ac.EnvVar().
|
||||
*corev1ac.EnvVar().
|
||||
WithName(constants.TorchEnvMasterPort).
|
||||
WithValue(fmt.Sprintf("%d", constants.ContainerTrainerPort)),
|
||||
)
|
||||
|
||||
+96
-11
@@ -17,13 +17,19 @@ limitations under the License.
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"iter"
|
||||
"maps"
|
||||
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
||||
resourcehelpers "k8s.io/component-helpers/resource"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
)
|
||||
|
||||
var (
|
||||
defaultPodSetsSyncer = func(*Info) {}
|
||||
syncPodSets = defaultPodSetsSyncer
|
||||
)
|
||||
|
||||
type Info struct {
|
||||
@@ -36,6 +42,9 @@ type Info struct {
|
||||
Trainer
|
||||
// Scheduler parameters to add to the RuntimeJobTemplate.
|
||||
*Scheduler
|
||||
// TemplateSpec is TrainingRuntime Template object.
|
||||
// ObjApply podSpecs and this PodSets should be kept in sync by info.SyncPodSetsToTemplateSpec().
|
||||
TemplateSpec TemplateSpec
|
||||
}
|
||||
|
||||
type RuntimePolicy struct {
|
||||
@@ -43,6 +52,32 @@ type RuntimePolicy struct {
|
||||
PodGroupPolicy *trainer.PodGroupPolicy
|
||||
}
|
||||
|
||||
type TemplateSpec struct {
|
||||
// ObjApply is ApplyConfiguration for the TrainingRuntimes Template field.
|
||||
ObjApply any
|
||||
// PodSets is a set of Pod extracted from ObjApply.
|
||||
PodSets []PodSet
|
||||
}
|
||||
|
||||
type PodSet struct {
|
||||
Name string
|
||||
// If Name is trainer-node, CountForNonTrainer is null.
|
||||
// For Trainer, PodSet Count should be stored in Info.RuntimePolicy.MLPolicy.NumNodes.
|
||||
CountForNonTrainer *int32
|
||||
Containers []Container
|
||||
Volumes []corev1ac.VolumeApplyConfiguration
|
||||
Endpoints iter.Seq[string]
|
||||
}
|
||||
|
||||
type Container struct {
|
||||
Name string
|
||||
Env []corev1ac.EnvVarApplyConfiguration
|
||||
Ports []corev1ac.ContainerPortApplyConfiguration
|
||||
VolumeMounts []corev1ac.VolumeMountApplyConfiguration
|
||||
}
|
||||
|
||||
// DEPRECATED: Replace all Trainer usage with RuntimePolicy and PodSet.
|
||||
|
||||
type Trainer struct {
|
||||
NumNodes *int32
|
||||
NumProcPerNode string
|
||||
@@ -58,6 +93,8 @@ type Scheduler struct {
|
||||
TotalRequests map[string]TotalResourceRequest
|
||||
}
|
||||
|
||||
// DEPRECATED: Replace all TotalResourceRequest usage with PodSet.
|
||||
|
||||
type TotalResourceRequest struct {
|
||||
Replicas int32
|
||||
PodRequests corev1.ResourceList
|
||||
@@ -68,16 +105,21 @@ type InfoOptions struct {
|
||||
annotations map[string]string
|
||||
runtimePolicy RuntimePolicy
|
||||
podSpecReplicas []podSpecReplica
|
||||
templateSpec TemplateSpec
|
||||
}
|
||||
|
||||
type InfoOption func(options *InfoOptions)
|
||||
|
||||
var defaultOptions = InfoOptions{}
|
||||
|
||||
// DEPRECATED: Replace all podSpecReplica usage with PodSet
|
||||
// once we remove TotalResourceRequest.
|
||||
|
||||
type podSpecReplica struct {
|
||||
replicas int32
|
||||
name string
|
||||
podSpec corev1.PodSpec
|
||||
count int32
|
||||
name string
|
||||
podSpecApply *corev1ac.PodSpecApplyConfiguration
|
||||
singlePodRequests corev1.ResourceList
|
||||
}
|
||||
|
||||
func WithLabels(labels map[string]string) InfoOption {
|
||||
@@ -104,16 +146,34 @@ func WithPodGroupPolicy(pgPolicy *trainer.PodGroupPolicy) InfoOption {
|
||||
}
|
||||
}
|
||||
|
||||
func WithPodSpecReplicas(replicaName string, replicas int32, podSpec corev1.PodSpec) InfoOption {
|
||||
// DEPRECATED: Replace WithPodSpecReplicas with WithTemplateSpec
|
||||
// once we remove TotalResourceRequest.
|
||||
|
||||
func WithPodSpecReplicas(
|
||||
replicaName string, count int32, singlePodRequest corev1.ResourceList, podSpecApply *corev1ac.PodSpecApplyConfiguration,
|
||||
) InfoOption {
|
||||
return func(o *InfoOptions) {
|
||||
o.podSpecReplicas = append(o.podSpecReplicas, podSpecReplica{
|
||||
name: replicaName,
|
||||
replicas: replicas,
|
||||
podSpec: *podSpec.DeepCopy(),
|
||||
name: replicaName,
|
||||
count: max(count, 1),
|
||||
podSpecApply: podSpecApply,
|
||||
singlePodRequests: singlePodRequest,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func WithTemplateSpec(objApply any) InfoOption {
|
||||
return func(o *InfoOptions) {
|
||||
o.templateSpec.ObjApply = objApply
|
||||
}
|
||||
}
|
||||
|
||||
func WithPodSetSyncer(syncer func(*Info)) InfoOption {
|
||||
return func(o *InfoOptions) {
|
||||
syncPodSets = syncer
|
||||
}
|
||||
}
|
||||
|
||||
func NewInfo(opts ...InfoOption) *Info {
|
||||
options := defaultOptions
|
||||
for _, opt := range opts {
|
||||
@@ -127,13 +187,30 @@ func NewInfo(opts ...InfoOption) *Info {
|
||||
Scheduler: &Scheduler{
|
||||
TotalRequests: make(map[string]TotalResourceRequest, len(options.podSpecReplicas)),
|
||||
},
|
||||
TemplateSpec: options.templateSpec,
|
||||
}
|
||||
|
||||
for _, spec := range options.podSpecReplicas {
|
||||
info.TotalRequests[spec.name] = TotalResourceRequest{
|
||||
Replicas: spec.replicas,
|
||||
PodRequests: resourcehelpers.PodRequests(&corev1.Pod{Spec: spec.podSpec}, resourcehelpers.PodResourcesOptions{}),
|
||||
Replicas: spec.count,
|
||||
PodRequests: spec.singlePodRequests,
|
||||
}
|
||||
ps := PodSet{
|
||||
Name: spec.name,
|
||||
Volumes: spec.podSpecApply.Volumes,
|
||||
}
|
||||
if spec.name != constants.JobTrainerNode {
|
||||
ps.CountForNonTrainer = &spec.count
|
||||
}
|
||||
for _, container := range spec.podSpecApply.Containers {
|
||||
ps.Containers = append(ps.Containers, Container{
|
||||
Name: *container.Name,
|
||||
Env: container.Env,
|
||||
Ports: container.Ports,
|
||||
VolumeMounts: container.VolumeMounts,
|
||||
})
|
||||
}
|
||||
info.TemplateSpec.PodSets = append(info.TemplateSpec.PodSets, ps)
|
||||
}
|
||||
if options.labels != nil {
|
||||
info.Labels = options.labels
|
||||
@@ -141,6 +218,14 @@ func NewInfo(opts ...InfoOption) *Info {
|
||||
if options.annotations != nil {
|
||||
info.Annotations = options.annotations
|
||||
}
|
||||
|
||||
return info
|
||||
}
|
||||
|
||||
func (i *Info) SyncPodSetsToTemplateSpec() {
|
||||
syncPodSets(i)
|
||||
}
|
||||
|
||||
func TemplateSpecApply[A any](info *Info) (*A, bool) {
|
||||
spec, ok := info.TemplateSpec.ObjApply.(*A)
|
||||
return spec, ok
|
||||
}
|
||||
|
||||
+244
-20
@@ -23,13 +23,17 @@ import (
|
||||
"github.com/google/go-cmp/cmp/cmpopts"
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
"k8s.io/apimachinery/pkg/api/resource"
|
||||
batchv1ac "k8s.io/client-go/applyconfigurations/batch/v1"
|
||||
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
||||
resourcehelpers "k8s.io/component-helpers/resource"
|
||||
"k8s.io/utils/ptr"
|
||||
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
jobsetplgconsts "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset/constants"
|
||||
jobsetv1alpha2ac "sigs.k8s.io/jobset/client-go/applyconfiguration/jobset/v1alpha2"
|
||||
)
|
||||
|
||||
func TestNewInfo(t *testing.T) {
|
||||
|
||||
cases := map[string]struct {
|
||||
infoOpts []InfoOption
|
||||
wantInfo *Info
|
||||
@@ -42,40 +46,159 @@ func TestNewInfo(t *testing.T) {
|
||||
WithAnnotations(map[string]string{
|
||||
"annotationKey": "annotationValue",
|
||||
}),
|
||||
WithPodSpecReplicas(constants.JobInitializer, 1, corev1.PodSpec{
|
||||
InitContainers: []corev1.Container{{
|
||||
Resources: corev1.ResourceRequirements{
|
||||
Requests: corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("5"),
|
||||
},
|
||||
},
|
||||
RestartPolicy: ptr.To(corev1.ContainerRestartPolicyAlways),
|
||||
}},
|
||||
Containers: []corev1.Container{{
|
||||
WithPodSpecReplicas(constants.JobInitializer, 1, resourcehelpers.PodRequests(&corev1.Pod{
|
||||
Spec: corev1.PodSpec{Containers: []corev1.Container{{
|
||||
Name: constants.ContainerModelInitializer,
|
||||
Resources: corev1.ResourceRequirements{
|
||||
Requests: corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("10"),
|
||||
},
|
||||
},
|
||||
}},
|
||||
}),
|
||||
WithPodSpecReplicas(constants.JobTrainerNode, 10, corev1.PodSpec{
|
||||
InitContainers: []corev1.Container{{
|
||||
}}}, InitContainers: []corev1.Container{{
|
||||
RestartPolicy: ptr.To(corev1.ContainerRestartPolicyAlways),
|
||||
Resources: corev1.ResourceRequirements{
|
||||
Requests: corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("5"),
|
||||
},
|
||||
}}},
|
||||
},
|
||||
}, resourcehelpers.PodResourcesOptions{}), corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerModelInitializer).
|
||||
WithResources(corev1ac.ResourceRequirements().
|
||||
WithRequests(corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("10"),
|
||||
})),
|
||||
).
|
||||
WithInitContainers(
|
||||
corev1ac.Container().
|
||||
WithRestartPolicy(corev1.ContainerRestartPolicyAlways).
|
||||
WithResources(corev1ac.ResourceRequirements().
|
||||
WithRequests(corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("5"),
|
||||
})),
|
||||
),
|
||||
),
|
||||
WithPodSpecReplicas(constants.JobTrainerNode, 10, resourcehelpers.PodRequests(&corev1.Pod{
|
||||
Spec: corev1.PodSpec{Containers: []corev1.Container{{
|
||||
Name: constants.ContainerTrainer,
|
||||
Resources: corev1.ResourceRequirements{
|
||||
Requests: corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("15"),
|
||||
},
|
||||
},
|
||||
}}, InitContainers: []corev1.Container{{
|
||||
Name: "preparation",
|
||||
RestartPolicy: ptr.To(corev1.ContainerRestartPolicyAlways),
|
||||
}},
|
||||
Containers: []corev1.Container{{
|
||||
Resources: corev1.ResourceRequirements{
|
||||
Requests: corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("25"),
|
||||
},
|
||||
},
|
||||
}},
|
||||
}),
|
||||
}}},
|
||||
}, resourcehelpers.PodResourcesOptions{}), corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerTrainer).
|
||||
WithResources(corev1ac.ResourceRequirements().
|
||||
WithRequests(corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("25"),
|
||||
})).
|
||||
WithEnv(
|
||||
corev1ac.EnvVar().WithName("TEST").WithValue("TEST"),
|
||||
).
|
||||
WithPorts(corev1ac.ContainerPort().
|
||||
WithName("http").
|
||||
WithProtocol(corev1.ProtocolTCP).
|
||||
WithContainerPort(8080),
|
||||
).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName("TEST").
|
||||
WithMountPath("/var").
|
||||
WithReadOnly(true),
|
||||
),
|
||||
).
|
||||
WithInitContainers(corev1ac.Container().
|
||||
WithName("preparation").
|
||||
WithResources(corev1ac.ResourceRequirements().
|
||||
WithRequests(corev1.ResourceList{
|
||||
corev1.ResourceCPU: resource.MustParse("15"),
|
||||
})).
|
||||
WithRestartPolicy(corev1.ContainerRestartPolicyAlways),
|
||||
)),
|
||||
WithTemplateSpec(
|
||||
jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobInitializer).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerDatasetInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
).
|
||||
WithResources(corev1ac.ResourceRequirements()),
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerModelInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
).
|
||||
WithResources(corev1ac.ResourceRequirements()),
|
||||
).
|
||||
WithVolumes(corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerTrainer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
).
|
||||
WithResources(corev1ac.ResourceRequirements()),
|
||||
).
|
||||
WithVolumes(corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
wantInfo: &Info{
|
||||
Labels: map[string]string{
|
||||
@@ -100,6 +223,107 @@ func TestNewInfo(t *testing.T) {
|
||||
},
|
||||
},
|
||||
},
|
||||
TemplateSpec: TemplateSpec{
|
||||
PodSets: []PodSet{
|
||||
{
|
||||
Name: constants.JobInitializer,
|
||||
CountForNonTrainer: ptr.To[int32](1),
|
||||
Containers: []Container{{
|
||||
Name: constants.ContainerModelInitializer,
|
||||
}},
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
Containers: []Container{{
|
||||
Name: constants.ContainerTrainer,
|
||||
Env: []corev1ac.EnvVarApplyConfiguration{{
|
||||
Name: ptr.To("TEST"),
|
||||
Value: ptr.To("TEST"),
|
||||
}},
|
||||
Ports: []corev1ac.ContainerPortApplyConfiguration{{
|
||||
Name: ptr.To("http"),
|
||||
Protocol: ptr.To(corev1.ProtocolTCP),
|
||||
ContainerPort: ptr.To[int32](8080),
|
||||
}},
|
||||
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{{
|
||||
Name: ptr.To("TEST"),
|
||||
ReadOnly: ptr.To(true),
|
||||
MountPath: ptr.To("/var"),
|
||||
}},
|
||||
}},
|
||||
},
|
||||
},
|
||||
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
||||
WithReplicatedJobs(
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobInitializer).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerDatasetInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
).
|
||||
WithResources(corev1ac.ResourceRequirements()),
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerModelInitializer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
).
|
||||
WithResources(corev1ac.ResourceRequirements()),
|
||||
).
|
||||
WithVolumes(corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
jobsetv1alpha2ac.ReplicatedJob().
|
||||
WithName(constants.JobTrainerNode).
|
||||
WithTemplate(batchv1ac.JobTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(batchv1ac.JobSpec().
|
||||
WithTemplate(corev1ac.PodTemplateSpec().
|
||||
WithLabels(nil).
|
||||
WithSpec(corev1ac.PodSpec().
|
||||
WithContainers(
|
||||
corev1ac.Container().
|
||||
WithName(constants.ContainerTrainer).
|
||||
WithVolumeMounts(
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.DatasetMountPath),
|
||||
corev1ac.VolumeMount().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithMountPath(constants.ModelMountPath),
|
||||
).
|
||||
WithResources(corev1ac.ResourceRequirements()),
|
||||
).
|
||||
WithVolumes(corev1ac.Volume().
|
||||
WithName(jobsetplgconsts.VolumeNameInitializer).
|
||||
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
||||
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
"all arguments are not specified": {
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
/*
|
||||
Copyright 2025 The Kubeflow Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package testing
|
||||
|
||||
import (
|
||||
"iter"
|
||||
"slices"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/google/go-cmp/cmp/cmpopts"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
corev1 "k8s.io/api/core/v1"
|
||||
)
|
||||
|
||||
var (
|
||||
PodSetEndpointsCmpOpts = cmp.Transformer("Seq", func(a iter.Seq[string]) []string { return slices.Collect(a) })
|
||||
)
|
||||
|
||||
func MPISecretDataComparer(a, b map[string][]byte) bool {
|
||||
isKeysEqual := true
|
||||
if (a != nil && b != nil) &&
|
||||
((len(a[constants.MPISSHPublicKey]) > 0) != (len(b[constants.MPISSHPublicKey]) > 0) ||
|
||||
(len(a[corev1.SSHAuthPrivateKey]) > 0) != (len(b[corev1.SSHAuthPrivateKey]) > 0)) {
|
||||
isKeysEqual = false
|
||||
}
|
||||
return isKeysEqual && cmp.Equal(a, b, cmpopts.IgnoreMapEntries(func(k string, _ []byte) bool {
|
||||
return k == constants.MPISSHPublicKey || k == corev1.SSHAuthPrivateKey
|
||||
}))
|
||||
}
|
||||
+264
-71
@@ -24,12 +24,13 @@ import (
|
||||
"k8s.io/apimachinery/pkg/types"
|
||||
"k8s.io/apimachinery/pkg/util/intstr"
|
||||
"k8s.io/utils/ptr"
|
||||
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
||||
schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"
|
||||
"slices"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
jobsetplugin "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset"
|
||||
jobsetplgconsts "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset/constants"
|
||||
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
||||
schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"
|
||||
)
|
||||
|
||||
type JobSetWrapper struct {
|
||||
@@ -58,20 +59,27 @@ func MakeJobSetWrapper(namespace, name string) *JobSetWrapper {
|
||||
Containers: []corev1.Container{
|
||||
{
|
||||
Name: constants.ContainerDatasetInitializer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountDatasetInitializer,
|
||||
},
|
||||
VolumeMounts: []corev1.VolumeMount{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.DatasetMountPath,
|
||||
}},
|
||||
},
|
||||
{
|
||||
Name: constants.ContainerModelInitializer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountModelInitializer,
|
||||
},
|
||||
VolumeMounts: []corev1.VolumeMount{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.ModelMountPath,
|
||||
}},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1.Volume{
|
||||
jobsetplugin.VolumeInitializer,
|
||||
},
|
||||
Volumes: []corev1.Volume{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
|
||||
ClaimName: jobsetplgconsts.VolumeNameInitializer,
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -79,12 +87,6 @@ func MakeJobSetWrapper(namespace, name string) *JobSetWrapper {
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
DependsOn: []jobsetv1alpha2.DependsOn{
|
||||
{
|
||||
Name: constants.JobInitializer,
|
||||
Status: jobsetv1alpha2.DependencyComplete,
|
||||
},
|
||||
},
|
||||
Template: batchv1.JobTemplateSpec{
|
||||
Spec: batchv1.JobSpec{
|
||||
Template: corev1.PodTemplateSpec{
|
||||
@@ -93,14 +95,25 @@ func MakeJobSetWrapper(namespace, name string) *JobSetWrapper {
|
||||
{
|
||||
Name: constants.ContainerTrainer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountDatasetInitializer,
|
||||
jobsetplugin.VolumeMountModelInitializer,
|
||||
{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.DatasetMountPath,
|
||||
},
|
||||
{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.ModelMountPath,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1.Volume{
|
||||
jobsetplugin.VolumeInitializer,
|
||||
},
|
||||
Volumes: []corev1.Volume{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
|
||||
ClaimName: jobsetplgconsts.VolumeNameInitializer,
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -112,9 +125,11 @@ func MakeJobSetWrapper(namespace, name string) *JobSetWrapper {
|
||||
}
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) Replicas(replicas int32) *JobSetWrapper {
|
||||
for idx := range j.Spec.ReplicatedJobs {
|
||||
j.Spec.ReplicatedJobs[idx].Replicas = replicas
|
||||
func (j *JobSetWrapper) Replicas(replicas int32, rJobNames ...string) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if slices.Contains(rJobNames, rJob.Name) {
|
||||
j.Spec.ReplicatedJobs[i].Replicas = replicas
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
@@ -129,6 +144,50 @@ func (j *JobSetWrapper) NumNodes(numNodes int32) *JobSetWrapper {
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) Parallelism(rName string, p int32) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rName {
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = &p
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) Completions(rName string, c int32) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rName {
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Completions = &c
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) LauncherReplica() *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobTrainerNode {
|
||||
j.Spec.ReplicatedJobs = append(j.Spec.ReplicatedJobs, jobsetv1alpha2.ReplicatedJob{})
|
||||
copy(j.Spec.ReplicatedJobs[i+1:], j.Spec.ReplicatedJobs[i:])
|
||||
j.Spec.ReplicatedJobs[i] = jobsetv1alpha2.ReplicatedJob{
|
||||
Name: constants.JobLauncher,
|
||||
Template: batchv1.JobTemplateSpec{
|
||||
Spec: batchv1.JobSpec{
|
||||
Template: corev1.PodTemplateSpec{
|
||||
Spec: corev1.PodSpec{
|
||||
Containers: []corev1.Container{{
|
||||
Name: constants.ContainerLauncher,
|
||||
Command: []string{"mpirun"},
|
||||
Args: []string{"echo.sh"},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) ContainerTrainer(image string, command []string, args []string, res corev1.ResourceList) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobTrainerNode {
|
||||
@@ -158,6 +217,8 @@ func (j *JobSetWrapper) ContainerTrainerPorts(ports []corev1.ContainerPort) *Job
|
||||
return j
|
||||
}
|
||||
|
||||
// TODO: ContainerTrainerEnv should be migrated to generic wrapper, Env(rJobName, containerName string, envs ...corev1.EnvVar).
|
||||
|
||||
func (j *JobSetWrapper) ContainerTrainerEnv(env []corev1.EnvVar) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobTrainerNode {
|
||||
@@ -171,6 +232,50 @@ func (j *JobSetWrapper) ContainerTrainerEnv(env []corev1.EnvVar) *JobSetWrapper
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) Volumes(rJobName string, v ...corev1.Volume) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rJobName {
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Volumes = append(
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Volumes,
|
||||
v...,
|
||||
)
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) VolumeMounts(rJobName, containerName string, vms ...corev1.VolumeMount) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rJobName {
|
||||
for k, container := range j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers {
|
||||
if container.Name == containerName {
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[k].VolumeMounts = append(
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[k].VolumeMounts,
|
||||
vms...,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) Env(rJobName, containerName string, envs ...corev1.EnvVar) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rJobName {
|
||||
for k, container := range j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers {
|
||||
if container.Name == containerName {
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[k].Env = append(
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[k].Env,
|
||||
envs...,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) ContainerDatasetModelInitializer(image string, command []string, args []string, res corev1.ResourceList) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobInitializer {
|
||||
@@ -187,13 +292,14 @@ func (j *JobSetWrapper) ContainerDatasetModelInitializer(image string, command [
|
||||
return j
|
||||
}
|
||||
|
||||
// TODO: ContainerDatasetInitializerEnv should be migrated to generic wrapper, Env(rJobName, containerName string, envs ...corev1.EnvVar).
|
||||
|
||||
func (j *JobSetWrapper) ContainerDatasetInitializerEnv(env []corev1.EnvVar) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobInitializer {
|
||||
for k, container := range rJob.Template.Spec.Template.Spec.Containers {
|
||||
if container.Name == constants.ContainerDatasetInitializer {
|
||||
j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[k].Env = env
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -214,6 +320,8 @@ func (j *JobSetWrapper) ContainerDatasetInitializerEnvFrom(envFrom []corev1.EnvF
|
||||
return j
|
||||
}
|
||||
|
||||
// TODO: ContainerModelInitializerEnv should be migrated to generic wrapper, Env(rJobName, containerName string, envs ...corev1.EnvVar).
|
||||
|
||||
func (j *JobSetWrapper) ContainerModelInitializerEnv(env []corev1.EnvVar) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobInitializer {
|
||||
@@ -290,6 +398,18 @@ func (j *JobSetWrapper) Conditions(conditions ...metav1.Condition) *JobSetWrappe
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) DependsOn(depName string, depStatus jobsetv1alpha2.DependsOnStatus, rJobName string) *JobSetWrapper {
|
||||
for i, rJob := range j.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rJobName {
|
||||
j.Spec.ReplicatedJobs[i].DependsOn = append(j.Spec.ReplicatedJobs[i].DependsOn, jobsetv1alpha2.DependsOn{
|
||||
Name: depName,
|
||||
Status: depStatus,
|
||||
})
|
||||
}
|
||||
}
|
||||
return j
|
||||
}
|
||||
|
||||
func (j *JobSetWrapper) Obj() *jobsetv1alpha2.JobSet {
|
||||
return &j.JobSet
|
||||
}
|
||||
@@ -506,20 +626,27 @@ func MakeTrainingRuntimeWrapper(namespace, name string) *TrainingRuntimeWrapper
|
||||
Containers: []corev1.Container{
|
||||
{
|
||||
Name: constants.ContainerDatasetInitializer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountDatasetInitializer,
|
||||
},
|
||||
VolumeMounts: []corev1.VolumeMount{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.DatasetMountPath,
|
||||
}},
|
||||
},
|
||||
{
|
||||
Name: constants.ContainerModelInitializer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountModelInitializer,
|
||||
},
|
||||
VolumeMounts: []corev1.VolumeMount{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.ModelMountPath,
|
||||
}},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1.Volume{
|
||||
jobsetplugin.VolumeInitializer,
|
||||
},
|
||||
Volumes: []corev1.Volume{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
|
||||
ClaimName: jobsetplgconsts.VolumeNameInitializer,
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -527,12 +654,6 @@ func MakeTrainingRuntimeWrapper(namespace, name string) *TrainingRuntimeWrapper
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
DependsOn: []jobsetv1alpha2.DependsOn{
|
||||
{
|
||||
Name: constants.JobInitializer,
|
||||
Status: jobsetv1alpha2.DependencyComplete,
|
||||
},
|
||||
},
|
||||
Template: batchv1.JobTemplateSpec{
|
||||
Spec: batchv1.JobSpec{
|
||||
Template: corev1.PodTemplateSpec{
|
||||
@@ -541,14 +662,25 @@ func MakeTrainingRuntimeWrapper(namespace, name string) *TrainingRuntimeWrapper
|
||||
{
|
||||
Name: constants.ContainerTrainer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountDatasetInitializer,
|
||||
jobsetplugin.VolumeMountModelInitializer,
|
||||
{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.DatasetMountPath,
|
||||
},
|
||||
{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.ModelMountPath,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1.Volume{
|
||||
jobsetplugin.VolumeInitializer,
|
||||
},
|
||||
Volumes: []corev1.Volume{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
|
||||
ClaimName: jobsetplgconsts.VolumeNameInitializer,
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -614,20 +746,27 @@ func MakeClusterTrainingRuntimeWrapper(name string) *ClusterTrainingRuntimeWrapp
|
||||
Containers: []corev1.Container{
|
||||
{
|
||||
Name: constants.ContainerDatasetInitializer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountDatasetInitializer,
|
||||
},
|
||||
VolumeMounts: []corev1.VolumeMount{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.DatasetMountPath,
|
||||
}},
|
||||
},
|
||||
{
|
||||
Name: constants.ContainerModelInitializer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountModelInitializer,
|
||||
},
|
||||
VolumeMounts: []corev1.VolumeMount{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.ModelMountPath,
|
||||
}},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1.Volume{
|
||||
jobsetplugin.VolumeInitializer,
|
||||
},
|
||||
Volumes: []corev1.Volume{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
|
||||
ClaimName: jobsetplgconsts.VolumeNameInitializer,
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -635,12 +774,6 @@ func MakeClusterTrainingRuntimeWrapper(name string) *ClusterTrainingRuntimeWrapp
|
||||
},
|
||||
{
|
||||
Name: constants.JobTrainerNode,
|
||||
DependsOn: []jobsetv1alpha2.DependsOn{
|
||||
{
|
||||
Name: constants.JobInitializer,
|
||||
Status: jobsetv1alpha2.DependencyComplete,
|
||||
},
|
||||
},
|
||||
Template: batchv1.JobTemplateSpec{
|
||||
Spec: batchv1.JobSpec{
|
||||
Template: corev1.PodTemplateSpec{
|
||||
@@ -649,14 +782,25 @@ func MakeClusterTrainingRuntimeWrapper(name string) *ClusterTrainingRuntimeWrapp
|
||||
{
|
||||
Name: constants.ContainerTrainer,
|
||||
VolumeMounts: []corev1.VolumeMount{
|
||||
jobsetplugin.VolumeMountDatasetInitializer,
|
||||
jobsetplugin.VolumeMountModelInitializer,
|
||||
{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.DatasetMountPath,
|
||||
},
|
||||
{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
MountPath: constants.ModelMountPath,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Volumes: []corev1.Volume{
|
||||
jobsetplugin.VolumeInitializer,
|
||||
},
|
||||
Volumes: []corev1.Volume{{
|
||||
Name: jobsetplgconsts.VolumeNameInitializer,
|
||||
VolumeSource: corev1.VolumeSource{
|
||||
PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
|
||||
ClaimName: jobsetplgconsts.VolumeNameInitializer,
|
||||
},
|
||||
},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -699,6 +843,32 @@ func (s *TrainingRuntimeSpecWrapper) WithMLPolicy(mlPolicy *trainer.MLPolicy) *T
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *TrainingRuntimeSpecWrapper) LauncherReplica() *TrainingRuntimeSpecWrapper {
|
||||
for i, rJob := range s.Template.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobTrainerNode {
|
||||
s.Template.Spec.ReplicatedJobs = append(s.Template.Spec.ReplicatedJobs, jobsetv1alpha2.ReplicatedJob{})
|
||||
copy(s.Template.Spec.ReplicatedJobs[i+1:], s.Template.Spec.ReplicatedJobs[i:])
|
||||
s.Template.Spec.ReplicatedJobs[i] = jobsetv1alpha2.ReplicatedJob{
|
||||
Name: constants.JobLauncher,
|
||||
Template: batchv1.JobTemplateSpec{
|
||||
Spec: batchv1.JobSpec{
|
||||
Template: corev1.PodTemplateSpec{
|
||||
Spec: corev1.PodSpec{
|
||||
Containers: []corev1.Container{{
|
||||
Name: constants.ContainerLauncher,
|
||||
Command: []string{"mpirun"},
|
||||
Args: []string{"echo.sh"},
|
||||
}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *TrainingRuntimeSpecWrapper) ContainerTrainer(image string, command []string, args []string, res corev1.ResourceList) *TrainingRuntimeSpecWrapper {
|
||||
for i, rJob := range s.Template.Spec.ReplicatedJobs {
|
||||
if rJob.Name == constants.JobTrainerNode {
|
||||
@@ -763,6 +933,24 @@ func (s *TrainingRuntimeSpecWrapper) PodGroupPolicyCoschedulingSchedulingTimeout
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *TrainingRuntimeSpecWrapper) Parallelism(rJobName string, parallelism int32) *TrainingRuntimeSpecWrapper {
|
||||
for i, rJob := range s.Template.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rJobName {
|
||||
s.Template.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = ¶llelism
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *TrainingRuntimeSpecWrapper) Completions(rJobName string, completions int32) *TrainingRuntimeSpecWrapper {
|
||||
for i, rJob := range s.Template.Spec.ReplicatedJobs {
|
||||
if rJob.Name == rJobName {
|
||||
s.Template.Spec.ReplicatedJobs[i].Template.Spec.Completions = &completions
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *TrainingRuntimeSpecWrapper) Obj() trainer.TrainingRuntimeSpec {
|
||||
return s.TrainingRuntimeSpec
|
||||
}
|
||||
@@ -793,14 +981,14 @@ func (m *MLPolicyWrapper) TorchPolicy(numProcPerNode string, elasticPolicy *trai
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *MLPolicyWrapper) MPIPolicy(numProcPerNode *int32, MPImplementation *trainer.MPIImplementation, sshAuthMountPath *string, runLauncherAsWorker *bool) *MLPolicyWrapper {
|
||||
func (m *MLPolicyWrapper) MPIPolicy(numProcPerNode *int32, MPImplementation *trainer.MPIImplementation, sshAuthMountPath *string, runLauncherAsNode *bool) *MLPolicyWrapper {
|
||||
if m.MLPolicySource.MPI == nil {
|
||||
m.MLPolicySource.MPI = &trainer.MPIMLPolicySource{}
|
||||
}
|
||||
m.MLPolicySource.MPI.NumProcPerNode = numProcPerNode
|
||||
m.MLPolicySource.MPI.MPIImplementation = MPImplementation
|
||||
m.MLPolicySource.MPI.SSHAuthMountPath = sshAuthMountPath
|
||||
m.MLPolicySource.MPI.RunLauncherAsNode = runLauncherAsWorker
|
||||
m.MLPolicySource.MPI.RunLauncherAsNode = runLauncherAsNode
|
||||
return m
|
||||
}
|
||||
|
||||
@@ -887,7 +1075,7 @@ func (c *ConfigMapWrapper) WithData(data map[string]string) *ConfigMapWrapper {
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *ConfigMapWrapper) OwnerReference(gvk schema.GroupVersionKind, name, uid string) *ConfigMapWrapper {
|
||||
func (c *ConfigMapWrapper) ControllerReference(gvk schema.GroupVersionKind, name, uid string) *ConfigMapWrapper {
|
||||
c.OwnerReferences = append(c.OwnerReferences, metav1.OwnerReference{
|
||||
APIVersion: gvk.GroupVersion().String(),
|
||||
Kind: gvk.Kind,
|
||||
@@ -937,7 +1125,12 @@ func (s *SecretWrapper) WithData(data map[string][]byte) *SecretWrapper {
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *SecretWrapper) OwnerReference(gvk schema.GroupVersionKind, name, uid string) *SecretWrapper {
|
||||
func (s *SecretWrapper) WithImmutable(immutable bool) *SecretWrapper {
|
||||
s.Immutable = &immutable
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *SecretWrapper) ControllerReference(gvk schema.GroupVersionKind, name, uid string) *SecretWrapper {
|
||||
s.OwnerReferences = append(s.OwnerReferences, metav1.OwnerReference{
|
||||
APIVersion: gvk.GroupVersion().String(),
|
||||
Kind: gvk.Kind,
|
||||
|
||||
@@ -24,6 +24,7 @@ import (
|
||||
"k8s.io/apimachinery/pkg/util/validation/field"
|
||||
jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
|
||||
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
testingutil "github.com/kubeflow/trainer/pkg/util/testing"
|
||||
)
|
||||
|
||||
@@ -34,12 +35,12 @@ func TestValidateReplicatedJobs(t *testing.T) {
|
||||
}{
|
||||
"valid replicatedJobs": {
|
||||
rJobs: testingutil.MakeJobSetWrapper("ns", "valid").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().Spec.ReplicatedJobs,
|
||||
},
|
||||
"invalid replicas": {
|
||||
rJobs: testingutil.MakeJobSetWrapper("ns", "valid").
|
||||
Replicas(2).
|
||||
Replicas(2, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().Spec.ReplicatedJobs,
|
||||
wantError: field.ErrorList{
|
||||
field.Invalid(field.NewPath("spec").Child("template").Child("spec").Child("replicatedJobs").Index(0).Child("replicas"),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -26,6 +26,7 @@ import (
|
||||
"sigs.k8s.io/controller-runtime/pkg/client"
|
||||
|
||||
trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
|
||||
"github.com/kubeflow/trainer/pkg/constants"
|
||||
testingutil "github.com/kubeflow/trainer/pkg/util/testing"
|
||||
"github.com/kubeflow/trainer/test/integration/framework"
|
||||
"github.com/kubeflow/trainer/test/util"
|
||||
@@ -185,7 +186,9 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
},
|
||||
}
|
||||
runtime.Spec.Template.Spec = testingutil.MakeJobSetWrapper(ns.Name, "runtime").
|
||||
Replicas(1).Obj().Spec
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec
|
||||
return runtime
|
||||
}),
|
||||
ginkgo.Entry("Should succeed to default mpi.mpiImplementation=OpenMPI",
|
||||
@@ -200,7 +203,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
@@ -219,7 +222,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
@@ -239,7 +242,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
@@ -258,7 +261,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
@@ -278,7 +281,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
@@ -297,7 +300,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
@@ -317,7 +320,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
@@ -336,7 +339,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
|
||||
).
|
||||
JobSetSpec(
|
||||
testingutil.MakeJobSetWrapper(ns.Name, "jobset").
|
||||
Replicas(1).
|
||||
Replicas(1, constants.JobTrainerNode, constants.JobInitializer, constants.JobLauncher).
|
||||
Obj().
|
||||
Spec,
|
||||
).
|
||||
|
||||
Reference in New Issue
Block a user