AWS Batch in Go

 

In a previous post we've used AWS batch using boto3.

In this post we will wrap usage of AWS batch using golang.

First we'll create the batch wrapper class.

package awsbatch

import (
 "fmt"
 "github.com/aws/aws-sdk-go/aws"
 "github.com/aws/aws-sdk-go/aws/session"
 "github.com/aws/aws-sdk-go/service/batch"
 "time"
)

type BatchWrapper struct {
 batch *batch.Batch
 jobs  []*string
}

func ProduceBatchWrapper() *BatchWrapper {
 awsSession, err := session.NewSession()
 if err != nil {
  panic(err)
 }

 awsBatch := batch.New(awsSession)
 return &BatchWrapper{
  batch: awsBatch,
  jobs:  []*string{},
 }
}

Next, we submit batches, and let them run in the background.

func (b *BatchWrapper) SubmitBatch(
 jobName string,
 environment map[string]string,
) {

 overrides := batch.ContainerOverrides{
  Command:     []*string{aws.String("/simulatorbackend")},
  Environment: []*batch.KeyValuePair{},
 }

 for key, value := range environment {
  overrides.Environment = append(
   overrides.Environment,
   &batch.KeyValuePair{
    Name:  aws.String(key),
    Value: aws.String(value),
   },
  )
 }

 input := batch.SubmitJobInput{
  JobName:            aws.String(jobName),
  JobQueue:           aws.String("my-batch-queue"),
  JobDefinition:      aws.String("my-batch-jobdef"),
  ContainerOverrides: &overrides,
 }

 output, err := b.batch.SubmitJob(&input)
 if err != nil {
  panic(err)
 }

 b.jobs = append(b.jobs, output.JobId)
}

And finally, we wait for all the batches to complete.

func (b *BatchWrapper) WaitForBatches() {
 input := batch.DescribeJobsInput{
  Jobs: b.jobs,
 }

 b.jobs = []*string{}

 for {
  output, err := b.batch.DescribeJobs(&input)
  if err != nil {
   panic(err)
  }

  allDone := true
  for _, job := range output.Jobs {
   if *job.Status == "FAILED" {
    panic("job failed")
   } else if *job.Status == "RUNNING" || *job.Status == "STARTING" || *job.Status == "SUBMITTED" {
    fmt.Printf("job %v id %v status %v\n",
     *job.JobName,
     *job.JobId,
     *job.Status,
    )

    allDone = false
   }
  }
  if allDone {
   return
  }
  time.Sleep(time.Second * 10)
 }
}