-
Notifications
You must be signed in to change notification settings - Fork 170
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'master' into add-try-catch-protection-for-initialize-wo…
…rkers
- Loading branch information
Showing
25 changed files
with
958 additions
and
43 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,26 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 main_test | ||
|
||
import ( | ||
"testing" | ||
|
||
. "github.com/onsi/ginkgo/v2" | ||
. "github.com/onsi/gomega" | ||
) | ||
|
||
func TestMaster(t *testing.T) { | ||
RegisterFailHandler(Fail) | ||
RunSpecs(t, "Master Suite") | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,26 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 batchscheduler_test | ||
|
||
import ( | ||
"testing" | ||
|
||
. "github.com/onsi/ginkgo/v2" | ||
. "github.com/onsi/gomega" | ||
) | ||
|
||
func TestBatchscheduler(t *testing.T) { | ||
RegisterFailHandler(Fail) | ||
RunSpecs(t, "Batchscheduler Suite") | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,62 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 batchscheduler | ||
|
||
import ( | ||
"context" | ||
|
||
"github.com/intelligent-machine-learning/dlrover/go/master/pkg/common" | ||
"github.com/intelligent-machine-learning/dlrover/go/master/pkg/kubeutils" | ||
) | ||
|
||
// ElasticScheduler launches pods without waiting for all resouces of pod are ready | ||
type ElasticScheduler struct { | ||
KubeScheduler | ||
SchedulerName string | ||
} | ||
|
||
// NewElasticScheduler creates an elastic scheduler. | ||
func NewElasticScheduler() *ElasticScheduler { | ||
return &ElasticScheduler{ | ||
KubeScheduler: KubeScheduler{ | ||
toCreatePods: common.NewQueue(), | ||
}, | ||
SchedulerName: "elastic", | ||
} | ||
} | ||
|
||
// Start starts a routine to launch Pods. | ||
func (scheduler *ElasticScheduler) Start(ctx context.Context, jobContext *common.JobContext) { | ||
go scheduler.LoopToLaunchPods(ctx) | ||
} | ||
|
||
// DoScheduling creates/updates/deletes pods | ||
func (scheduler *ElasticScheduler) DoScheduling(jobContext *common.JobContext, plan *SchedulingPlan) { | ||
for replicaType, spec := range plan.ReplicaSpecs { | ||
for i := int32(0); i < spec.Replicas; i++ { | ||
replicaConfig := &kubeutils.ReplicaConfig{ | ||
Type: string(replicaType), | ||
ID: i, | ||
Number: spec.Replicas, | ||
Rank: i, | ||
} | ||
podConfig := &kubeutils.PodConfig{ | ||
Replica: replicaConfig, | ||
TemplateSpec: spec.Template.DeepCopy(), | ||
} | ||
pod := kubeutils.BuildPod(jobContext, podConfig) | ||
scheduler.toCreatePods.PushBack(pod) | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,62 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 batchscheduler | ||
|
||
import ( | ||
"fmt" | ||
|
||
commonv1 "github.com/intelligent-machine-learning/dlrover/go/elasticjob/pkg/common/api/v1" | ||
"github.com/intelligent-machine-learning/dlrover/go/master/pkg/common" | ||
|
||
. "github.com/onsi/ginkgo/v2" | ||
. "github.com/onsi/gomega" | ||
corev1 "k8s.io/api/core/v1" | ||
) | ||
|
||
var _ = Describe("Elastic", func() { | ||
It("Do scheduling to launch pods.", func() { | ||
jobContext := &common.JobContext{ | ||
NameSpace: "dlrover", | ||
Name: "train-demo", | ||
MasterHost: "127.0.0.1", | ||
MasterPort: 12345, | ||
} | ||
|
||
container := corev1.Container{ | ||
Name: "main", | ||
Image: "python:3.12.8", | ||
ImagePullPolicy: corev1.PullIfNotPresent, | ||
Command: []string{"/bin/bash", "-c", "echo 0"}, | ||
} | ||
replicas := make(map[commonv1.ReplicaType]*commonv1.ReplicaSpec) | ||
replicas["worker"] = &commonv1.ReplicaSpec{ | ||
Replicas: 3, | ||
Template: corev1.PodTemplateSpec{ | ||
Spec: corev1.PodSpec{ | ||
Containers: []corev1.Container{container}, | ||
RestartPolicy: corev1.RestartPolicyNever, | ||
}, | ||
}, | ||
} | ||
schedulingPlan := &SchedulingPlan{ReplicaSpecs: replicas} | ||
scheduler := NewElasticScheduler() | ||
scheduler.DoScheduling(jobContext, schedulingPlan) | ||
Expect(scheduler.toCreatePods.Len()).To(Equal(3)) | ||
for i := 0; i < 3; i++ { | ||
pod := scheduler.toCreatePods.PopFront().(*corev1.Pod) | ||
expectPodName := fmt.Sprintf("train-demo-worker-%d", i) | ||
Expect(pod.ObjectMeta.Name).To(Equal(expectPodName)) | ||
} | ||
}) | ||
}) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,14 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 batchscheduler |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,89 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 batchscheduler | ||
|
||
import ( | ||
"context" | ||
"time" | ||
|
||
elasticjob "github.com/intelligent-machine-learning/dlrover/go/elasticjob/api/v1alpha1" | ||
commonv1 "github.com/intelligent-machine-learning/dlrover/go/elasticjob/pkg/common/api/v1" | ||
"github.com/intelligent-machine-learning/dlrover/go/master/pkg/common" | ||
"github.com/intelligent-machine-learning/dlrover/go/master/pkg/kubeutils" | ||
logger "github.com/sirupsen/logrus" | ||
corev1 "k8s.io/api/core/v1" | ||
"k8s.io/apimachinery/pkg/api/errors" | ||
) | ||
|
||
// BatchScheduler creates/updates/deletes the batch pods of an elastic job. | ||
type BatchScheduler interface { | ||
Start(ctx context.Context, jobContext *common.JobContext) | ||
DoScheduling(jobContext *common.JobContext, plan *SchedulingPlan) | ||
} | ||
|
||
// SchedulingPlan is the scheduling plan to notify the scheduler CURD pods. | ||
type SchedulingPlan struct { | ||
// ReplicaSpecs is a map which contains the replica specification to create Pods. | ||
ReplicaSpecs map[commonv1.ReplicaType]*commonv1.ReplicaSpec | ||
|
||
// CreatedPods are Pods to be created. | ||
CreatedPods []*kubeutils.PodConfig | ||
|
||
// RemovedPods are Pods to be removed | ||
RemovedPods []*kubeutils.PodConfig | ||
|
||
// OwnerJob specifies a job to scale. | ||
OwnerJob *elasticjob.ElasticJob | ||
} | ||
|
||
// NewBatchScheduler creates a batch scheduler according to the scheduler name. | ||
func NewBatchScheduler(schedulerName string) BatchScheduler { | ||
if schedulerName == "elastic" || schedulerName == "" { | ||
scheduler := NewElasticScheduler() | ||
return scheduler | ||
} | ||
return nil | ||
} | ||
|
||
// KubeScheduler is the base scheduler to create/update/remove pods. | ||
type KubeScheduler struct { | ||
toCreatePods *common.Queue | ||
} | ||
|
||
// LoopToLaunchPods launches pods from the pod queue. | ||
func (scheduler *KubeScheduler) LoopToLaunchPods(ctx context.Context) { | ||
for { | ||
select { | ||
case <-ctx.Done(): | ||
logger.Infof("The loop to launch Pod exists.") | ||
default: | ||
for scheduler.toCreatePods.Len() > 0 { | ||
pod := scheduler.toCreatePods.PopFront().(*corev1.Pod) | ||
err := kubeutils.GlobalK8sClient.CreatePod(ctx, pod) | ||
if errors.IsAlreadyExists(err) { | ||
logger.Warnf("The pod %s already exists.", pod.ObjectMeta.Name) | ||
} else if errors.IsTooManyRequests(err) || errors.IsTimeout(err) || errors.IsServerTimeout(err) { | ||
logger.Warnf("Fail to create pod %s with err: %v", pod.ObjectMeta.Name, err) | ||
// Retry to create pod due to timeout. | ||
scheduler.toCreatePods.PushFront(pod) | ||
time.Sleep(5 * time.Second) | ||
} else { | ||
logger.Warnf("Fail to create pod %s with err: %v", pod.ObjectMeta.Name, err) | ||
panic(err.Error()) | ||
} | ||
} | ||
} | ||
time.Sleep(1 * time.Second) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,14 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 batchscheduler |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,26 @@ | ||
// Copyright 2025 The DLRover Authors. All rights reserved. | ||
// 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 common_test | ||
|
||
import ( | ||
"testing" | ||
|
||
. "github.com/onsi/ginkgo/v2" | ||
. "github.com/onsi/gomega" | ||
) | ||
|
||
func TestCommon(t *testing.T) { | ||
RegisterFailHandler(Fail) | ||
RunSpecs(t, "Common Suite") | ||
} |
Oops, something went wrong.