github.com/kubeflow/training-operator@v1.7.0/test_job/test_util/v1/test_job_util.go (about)

     1  // Copyright 2018 The Kubeflow Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package v1
    16  
    17  import (
    18  	"time"
    19  
    20  	apiv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
    21  	v1 "k8s.io/api/core/v1"
    22  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    23  
    24  	testjobv1 "github.com/kubeflow/training-operator/test_job/apis/test_job/v1"
    25  )
    26  
    27  func NewTestJob(worker int) *testjobv1.TestJob {
    28  	testJob := &testjobv1.TestJob{
    29  		TypeMeta: metav1.TypeMeta{
    30  			Kind: testjobv1.Kind,
    31  		},
    32  		ObjectMeta: metav1.ObjectMeta{
    33  			Name:      TestJobName,
    34  			Namespace: metav1.NamespaceDefault,
    35  		},
    36  		Spec: testjobv1.TestJobSpec{
    37  			TestReplicaSpecs: make(map[testjobv1.TestReplicaType]*apiv1.ReplicaSpec),
    38  		},
    39  	}
    40  
    41  	if worker > 0 {
    42  		worker := int32(worker)
    43  		workerReplicaSpec := &apiv1.ReplicaSpec{
    44  			Replicas: &worker,
    45  			Template: NewTestReplicaSpecTemplate(),
    46  		}
    47  		testJob.Spec.TestReplicaSpecs[testjobv1.TestReplicaTypeWorker] = workerReplicaSpec
    48  	}
    49  
    50  	return testJob
    51  }
    52  
    53  func NewTestReplicaSpecTemplate() v1.PodTemplateSpec {
    54  	return v1.PodTemplateSpec{
    55  		Spec: v1.PodSpec{
    56  			Containers: []v1.Container{
    57  				v1.Container{
    58  					Name:  testjobv1.DefaultContainerName,
    59  					Image: TestImageName,
    60  					Args:  []string{"Fake", "Fake"},
    61  					Ports: []v1.ContainerPort{
    62  						v1.ContainerPort{
    63  							Name:          testjobv1.DefaultPortName,
    64  							ContainerPort: testjobv1.DefaultPort,
    65  						},
    66  					},
    67  				},
    68  			},
    69  		},
    70  	}
    71  }
    72  
    73  func SetTestJobCompletionTime(testJob *testjobv1.TestJob) {
    74  	now := metav1.Time{Time: time.Now()}
    75  	testJob.Status.CompletionTime = &now
    76  }