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

feat: Add a public function to create runtime info objects (#2837)

Browse Source 
Signed-off-by: Tomas Tormo <[email protected]>
This commit is contained in:
Tomas Tormo
2025-09-22 15:47:02 +00:00
committed by GitHub
parent cbb653e620
commit 71ac41935d
3 changed files with 26 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/runtime/core/clustertrainingruntime.go
+12 -1
View File
@@ -57,7 +57,18 @@ func (r *ClusterTrainingRuntime) NewObjects(ctx context.Context, trainJob *train
if err := r.client.Get(ctx, client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}, &clTrainingRuntime); err != nil {
return nil, fmt.Errorf("%w: %w", errorNotFoundSpecifiedClusterTrainingRuntime, err)
}
return r.buildObjects(ctx, trainJob, clTrainingRuntime.Spec.Template, clTrainingRuntime.Spec.MLPolicy, clTrainingRuntime.Spec.PodGroupPolicy)
info, err := r.RuntimeInfo(trainJob, clTrainingRuntime.Spec.Template, clTrainingRuntime.Spec.MLPolicy, clTrainingRuntime.Spec.PodGroupPolicy)
if err != nil {
return nil, err
}
return r.framework.RunComponentBuilderPlugins(ctx, info, trainJob)
}
func (r *ClusterTrainingRuntime) RuntimeInfo(
trainJob *trainer.TrainJob, runtimeTemplateSpec any, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy,
) (*runtime.Info, error) {
return r.TrainingRuntime.RuntimeInfo(trainJob, runtimeTemplateSpec, mlPolicy, podGroupPolicy)
}
func (r *ClusterTrainingRuntime) TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error) {
pkg/runtime/core/trainingruntime.go
+13 -5
View File
@@ -85,13 +85,21 @@ func (r *TrainingRuntime) NewObjects(ctx context.Context, trainJob *trainer.Trai
if err != nil {
return nil, fmt.Errorf("%w: %w", errorNotFoundSpecifiedTrainingRuntime, err)
}
return r.buildObjects(ctx, trainJob, trainingRuntime.Spec.Template, trainingRuntime.Spec.MLPolicy, trainingRuntime.Spec.PodGroupPolicy)
info, err := r.RuntimeInfo(trainJob, trainingRuntime.Spec.Template, trainingRuntime.Spec.MLPolicy, trainingRuntime.Spec.PodGroupPolicy)
if err != nil {
return nil, err
}
return r.framework.RunComponentBuilderPlugins(ctx, info, trainJob)
}
func (r *TrainingRuntime) buildObjects(
ctx context.Context, trainJob *trainer.TrainJob, jobSetTemplateSpec trainer.JobSetTemplateSpec, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy,
) ([]any, error) {
func (r *TrainingRuntime) RuntimeInfo(
trainJob *trainer.TrainJob, runtimeTemplateSpec any, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy,
) (*runtime.Info, error) {
jobSetTemplateSpec, ok := runtimeTemplateSpec.(trainer.JobSetTemplateSpec)
if !ok {
return nil, fmt.Errorf("unsupported runtimeTemplateSpec")
}
info, err := r.newRuntimeInfo(trainJob, jobSetTemplateSpec, mlPolicy, podGroupPolicy)
if err != nil {
return nil, err
@@ -108,7 +116,7 @@ func (r *TrainingRuntime) buildObjects(
return nil, err
}
return r.framework.RunComponentBuilderPlugins(ctx, info, trainJob)
return info, nil
}
func (r *TrainingRuntime) newRuntimeInfo(
pkg/runtime/interface.go
+1
View File
@@ -37,6 +37,7 @@ type Runtime interface {
// https://github.com/kubernetes/kubernetes/pull/129313 becomes available
NewObjects(ctx context.Context, trainJob *trainer.TrainJob) ([]any, error)
RuntimeInfo(trainJob *trainer.TrainJob, runtimeTemplateSpec any, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy) (*Info, error)
TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error)
EventHandlerRegistrars() []ReconcilerBuilder
ValidateObjects(ctx context.Context, old, new *trainer.TrainJob) (admission.Warnings, field.ErrorList)