Skip to content

Commit

Permalink
feat: create job with custom labels sample
Browse files Browse the repository at this point in the history
  • Loading branch information
BigBlackWolf committed Dec 17, 2024
1 parent e567a2b commit 8954812
Show file tree
Hide file tree
Showing 2 changed files with 168 additions and 0 deletions.
114 changes: 114 additions & 0 deletions batch/create_with_job_labels.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
// Copyright 2024 Google LLC
//
// 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
//
// https://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 snippets

// [START batch_labels_job]
import (
"context"
"fmt"
"io"

batch "cloud.google.com/go/batch/apiv1"
"cloud.google.com/go/batch/apiv1/batchpb"
durationpb "google.golang.org/protobuf/types/known/durationpb"
)

// Creates and runs a job with custom labels for runnable.
func createJobWithCustomJobLabels(w io.Writer, projectID, region, jobName string) (*batchpb.Job, error) {
ctx := context.Background()
batchClient, err := batch.NewClient(ctx)
if err != nil {
return nil, fmt.Errorf("batchClient error: %w", err)
}
defer batchClient.Close()

runn := &batchpb.Runnable{
Executable: &batchpb.Runnable_Script_{
Script: &batchpb.Runnable_Script{
Command: &batchpb.Runnable_Script_Text{
Text: "echo Hello world from script 1 for task ${BATCH_TASK_INDEX}",
},
},
},
}

taskSpec := &batchpb.TaskSpec{
ComputeResource: &batchpb.ComputeResource{
// CpuMilli is milliseconds per cpu-second. This means the task requires 2 whole CPUs.
CpuMilli: 2000,
MemoryMib: 16,
},
MaxRunDuration: &durationpb.Duration{
Seconds: 3600,
},
MaxRetryCount: 2,
Runnables: []*batchpb.Runnable{runn},
}

taskGroups := []*batchpb.TaskGroup{
{
TaskCount: 4,
TaskSpec: taskSpec,
},
}

// Policies are used to define on what kind of virtual machines the tasks will run on.
// In this case, we tell the system to use "e2-standard-4" machine type.
// Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
allocationPolicy := &batchpb.AllocationPolicy{
Instances: []*batchpb.AllocationPolicy_InstancePolicyOrTemplate{{
PolicyTemplate: &batchpb.AllocationPolicy_InstancePolicyOrTemplate_Policy{
Policy: &batchpb.AllocationPolicy_InstancePolicy{
MachineType: "e2-standard-4",
},
},
}},
}

// We use Cloud Logging as it's an out of the box available option
logsPolicy := &batchpb.LogsPolicy{
Destination: batchpb.LogsPolicy_CLOUD_LOGGING,
}

// Setting some labels for job
labels := map[string]string{
"env": "dev",
"type": "single command",
}

job := &batchpb.Job{
Name: jobName,
TaskGroups: taskGroups,
AllocationPolicy: allocationPolicy,
LogsPolicy: logsPolicy,
Labels: labels,
}

request := &batchpb.CreateJobRequest{
Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, region),
JobId: jobName,
Job: job,
}

created_job, err := batchClient.CreateJob(ctx, request)
if err != nil {
return nil, fmt.Errorf("unable to create job: %w", err)
}

fmt.Fprintf(w, "Job created: %v\n", created_job)
return created_job, nil
}

// [END batch_labels_job]
54 changes: 54 additions & 0 deletions batch/create_with_job_labels_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
// Copyright 2024 Google LLC
//
// 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
//
// https://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 snippets

import (
"bytes"
"fmt"
"math/rand"
"strings"
"testing"
"time"

"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
)

func TestCreateJobWithLabelsRunnable(t *testing.T) {
var r *rand.Rand = rand.New(
rand.NewSource(time.Now().UnixNano()))
tc := testutil.SystemTest(t)
jobName := fmt.Sprintf("test-job-go-%v-%v", time.Now().Format("2006-01-02"), r.Int())
region := "us-central1"

buf := &bytes.Buffer{}

job, err := createJobWithCustomJobLabels(buf, tc.ProjectID, region, jobName)

if err != nil {
t.Errorf("createJobWithCustomJobLabels got err: %v", err)
}
if got := buf.String(); !strings.Contains(got, "Job created") {
t.Errorf("createJobWithCustomJobLabels got %q, expected %q", got, "Job created")
}

labels := job.GetLabels()
if labels["env"] != "dev" || labels["type"] != "single command" {
t.Errorf("labels weren't set")
}

if err := deleteJob(buf, tc.ProjectID, region, jobName); err != nil {
t.Errorf("deleteJob got err: %v", err)
}
}

0 comments on commit 8954812

Please sign in to comment.