hanzo-ml/trainer
1
0
Fork 0
mirror of https://github.com/hanzo-ml/trainer.git synced 2026-09-24 00:31:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Implement E2E for OpenMPI workload (#2500)

Browse Source 
Signed-off-by: Yuki Iwai <[email protected]>
This commit is contained in:
Yuki Iwai
2025-03-14 05:23:31 +00:00
committed by GitHub
parent f98108aa6b
commit a0b7d4d3ba
1 changed files with 37 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/e2e/e2e_test.go
+37 -2
View File
@@ -16,6 +16,7 @@ import (
const (
torchRuntime = "torch-distributed"
mpiRuntime = "mpi-distributed"
)
var _ = ginkgo.Describe("TrainJob e2e", func() {
@@ -44,11 +45,11 @@ var _ = ginkgo.Describe("TrainJob e2e", func() {
})
// These tests create TrainJob that reference supported runtime without any additional changes.
ginkgo.When("creating TrainJob", func() {
ginkgo.When("Creating TrainJob to perform the PyTorch workload", func() {
// Verify `torch-distributed` ClusterTrainingRuntime.
ginkgo.It("should create TrainJob with PyTorch runtime reference", func() {
// Create a TrainJob.
trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-test").
trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-test-torch").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), torchRuntime).
Obj()
@@ -79,4 +80,38 @@ var _ = ginkgo.Describe("TrainJob e2e", func() {
})
})
})
ginkgo.When("Creating TrainJob to perform OpenMPI workload", func() {
ginkgo.It("should create TrainJob with OpenMPI runtime reference", func() {
trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-test-openmpi").
RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), mpiRuntime).
Obj()
ginkgo.By("Create a TrainJob with mpi-distributed runtime reference", func() {
gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
})
// Wait for TrainJob to be in Succeeded status.
ginkgo.By("Wait for TrainJob to be in Succeeded status", func() {
gomega.Eventually(func(g gomega.Gomega) {
gotTrainJob := &trainer.TrainJob{}
g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
g.Expect(gotTrainJob.Status.Conditions).Should(gomega.BeComparableTo([]metav1.Condition{
{
Type: trainer.TrainJobCreated,
Status: metav1.ConditionTrue,
Reason: trainer.TrainJobJobsCreationSucceededReason,
Message: constants.TrainJobJobsCreationSucceededMessage,
},
{
Type: trainer.TrainJobComplete,
Status: metav1.ConditionTrue,
Reason: jobsetconsts.AllJobsCompletedReason,
Message: jobsetconsts.AllJobsCompletedMessage,
},
}, util.IgnoreConditions))
}, util.TimeoutE2E, util.Interval).Should(gomega.Succeed())
})
})
})
})