volcano.sh/volcano@v1.9.0/test/e2e/jobseq/pytorch_plugin.go (about)

     1  package jobseq
     2  
     3  import (
     4  	. "github.com/onsi/ginkgo/v2"
     5  	. "github.com/onsi/gomega"
     6  	vcbatch "volcano.sh/apis/pkg/apis/batch/v1alpha1"
     7  	vcbus "volcano.sh/apis/pkg/apis/bus/v1alpha1"
     8  	e2eutil "volcano.sh/volcano/test/e2e/util"
     9  )
    10  
    11  var _ = Describe("Pytorch Plugin E2E Test", func() {
    12  	It("will run and complete finally", func() {
    13  		// Community CI can skip this use case, and enable this use case verification when releasing the version.
    14  		Skip("Pytorch's test image download fails probabilistically, causing the current use case to fail. ")
    15  		context := e2eutil.InitTestContext(e2eutil.Options{})
    16  		defer e2eutil.CleanupTestContext(context)
    17  
    18  		slot := e2eutil.OneCPU
    19  
    20  		spec := &e2eutil.JobSpec{
    21  			Name: "pytorch-job",
    22  			Min:  1,
    23  			Policies: []vcbatch.LifecyclePolicy{
    24  				{
    25  					Action: vcbus.CompleteJobAction,
    26  					Event:  vcbus.TaskCompletedEvent,
    27  				},
    28  			},
    29  			Plugins: map[string][]string{
    30  				"pytorch": {"--master=master", "--worker=worker", "--port=23456"},
    31  			},
    32  			Tasks: []e2eutil.TaskSpec{
    33  				{
    34  					Name:       "master",
    35  					Img:        e2eutil.DefaultPytorchImage,
    36  					Req:        slot,
    37  					Min:        1,
    38  					Rep:        1,
    39  					WorkingDir: "/home",
    40  					// Need sometime waiting for worker node ready
    41  					Command: `python3 /opt/pytorch-mnist/mnist.py --epochs=1`,
    42  				},
    43  				{
    44  					Name:       "worker",
    45  					Img:        e2eutil.DefaultPytorchImage,
    46  					Req:        slot,
    47  					Min:        2,
    48  					Rep:        2,
    49  					WorkingDir: "/home",
    50  					Command:    "python3 /opt/pytorch-mnist/mnist.py --epochs=1",
    51  				},
    52  			},
    53  		}
    54  
    55  		job := e2eutil.CreateJob(context, spec)
    56  		err := e2eutil.WaitJobPhases(context, job, []vcbatch.JobPhase{
    57  			vcbatch.Pending, vcbatch.Running, vcbatch.Completed})
    58  		Expect(err).NotTo(HaveOccurred())
    59  	})
    60  })