From 54f78ad46cfec6d5f4379b59bebf1aabd392b0eb Mon Sep 17 00:00:00 2001 From: Stepan Rasputny Date: Tue, 17 Dec 2024 18:25:35 +0100 Subject: [PATCH] feat: batch create with custom network sample (#4872) * feat: batch create with custom network sample * added buffer check --- batch/create_with_custom_network.go | 117 +++++++++++++++++++++++ batch/create_with_custom_network_test.go | 58 +++++++++++ 2 files changed, 175 insertions(+) create mode 100644 batch/create_with_custom_network.go create mode 100644 batch/create_with_custom_network_test.go diff --git a/batch/create_with_custom_network.go b/batch/create_with_custom_network.go new file mode 100644 index 0000000000..93aefd5758 --- /dev/null +++ b/batch/create_with_custom_network.go @@ -0,0 +1,117 @@ +// 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_create_custom_network] +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 network +func createJobWithCustomNetwork(w io.Writer, projectID, region, jobName, networkName, subnetworkName 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", + }, + }, + }}, + Network: &batchpb.AllocationPolicy_NetworkPolicy{ + NetworkInterfaces: []*batchpb.AllocationPolicy_NetworkInterface{ + { + // Set the network to the specified network name within the project + Network: fmt.Sprintf("projects/%s/global/networks/%s", projectID, networkName), + // Set the subnetwork to the specified subnetwork within the region + Subnetwork: fmt.Sprintf("projects/%s/regions/%s/subnetworks/%s", projectID, region, subnetworkName), + }, + }, + }, + } + + // We use Cloud Logging as it's an out of the box available option + logsPolicy := &batchpb.LogsPolicy{ + Destination: batchpb.LogsPolicy_CLOUD_LOGGING, + } + + job := &batchpb.Job{ + Name: jobName, + TaskGroups: taskGroups, + AllocationPolicy: allocationPolicy, + LogsPolicy: logsPolicy, + } + + 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_create_custom_network] diff --git a/batch/create_with_custom_network_test.go b/batch/create_with_custom_network_test.go new file mode 100644 index 0000000000..15624e2b08 --- /dev/null +++ b/batch/create_with_custom_network_test.go @@ -0,0 +1,58 @@ +// 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 TestCreateJobWithCustomNetwork(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" + networkName, subnetworkName := "default", "default" + + buf := &bytes.Buffer{} + + job, err := createJobWithCustomNetwork(buf, tc.ProjectID, region, jobName, networkName, subnetworkName) + + if err != nil { + t.Errorf("createJobWithCustomNetwork got err: %v", err) + } + if got := buf.String(); !strings.Contains(got, "Job created") { + t.Errorf("createJobWithCustomNetwork got %q, expected %q", got, "Job created") + } + + expectedNetwork := fmt.Sprintf("projects/%s/global/networks/%s", tc.ProjectID, networkName) + expectedSubnetwork := fmt.Sprintf("projects/%s/regions/%s/subnetworks/%s", tc.ProjectID, region, subnetworkName) + + interfaces := job.GetAllocationPolicy().GetNetwork().GetNetworkInterfaces() + if interfaces[0].GetNetwork() != expectedNetwork || interfaces[0].GetSubnetwork() != expectedSubnetwork { + t.Errorf("Network wasn't set") + } + + if err := deleteJob(buf, tc.ProjectID, region, jobName); err != nil { + t.Errorf("deleteJob got err: %v", err) + } +}