|
| 1 | +/* |
| 2 | + * Copyright 2021 Google LLC |
| 3 | + * |
| 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + * you may not use this file except in compliance with the License. |
| 6 | + * You may obtain a copy of the License at |
| 7 | + * |
| 8 | + * https://www.apache.org/licenses/LICENSE-2.0 |
| 9 | + * |
| 10 | + * Unless required by applicable law or agreed to in writing, software |
| 11 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | + * See the License for the specific language governing permissions and |
| 14 | + * limitations under the License. |
| 15 | + */ |
| 16 | + |
| 17 | +'use strict'; |
| 18 | + |
| 19 | +async function main() { |
| 20 | + // [START batch_create_gpu_job] |
| 21 | + // Imports the Batch library |
| 22 | + const batchLib = require('@google-cloud/batch'); |
| 23 | + const batch = batchLib.protos.google.cloud.batch.v1; |
| 24 | + |
| 25 | + // Instantiates a client |
| 26 | + const batchClient = new batchLib.v1.BatchServiceClient(); |
| 27 | + |
| 28 | + /** |
| 29 | + * TODO(developer): Update these variables before running the sample. |
| 30 | + */ |
| 31 | + const projectId = await batchClient.getProjectId(); |
| 32 | + // Name of the region you want to use to run the job. Regions that are |
| 33 | + // available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations |
| 34 | + const region = 'europe-central2'; |
| 35 | + // The name of the job that will be created. |
| 36 | + // It needs to be unique for each project and region pair. |
| 37 | + const jobName = 'batch-gpu-job'; |
| 38 | + // The GPU type. You can view a list of the available GPU types |
| 39 | + // by using the `gcloud compute accelerator-types list` command. |
| 40 | + const gpuType = 'nvidia-l4'; |
| 41 | + // The number of GPUs of the specified type. |
| 42 | + const gpuCount = 1; |
| 43 | + // Optional. When set to true, Batch fetches the drivers required for the GPU type |
| 44 | + // that you specify in the policy field from a third-party location, |
| 45 | + // and Batch installs them on your behalf. If you set this field to false (default), |
| 46 | + // you need to install GPU drivers manually to use any GPUs for this job. |
| 47 | + const installGpuDrivers = false; |
| 48 | + // Accelerator-optimized machine types are available to Batch jobs. See the list |
| 49 | + // of available types on: https://cloud.google.com/compute/docs/accelerator-optimized-machines |
| 50 | + const machineType = 'g2-standard-4'; |
| 51 | + |
| 52 | + // Define what will be done as part of the job. |
| 53 | + const runnable = new batch.Runnable({ |
| 54 | + script: new batch.Runnable.Script({ |
| 55 | + commands: ['-c', 'echo Hello world! This is task ${BATCH_TASK_INDEX}.'], |
| 56 | + }), |
| 57 | + }); |
| 58 | + |
| 59 | + const task = new batch.TaskSpec({ |
| 60 | + runnables: [runnable], |
| 61 | + maxRetryCount: 2, |
| 62 | + maxRunDuration: {seconds: 3600}, |
| 63 | + }); |
| 64 | + |
| 65 | + // Tasks are grouped inside a job using TaskGroups. |
| 66 | + const group = new batch.TaskGroup({ |
| 67 | + taskCount: 3, |
| 68 | + taskSpec: task, |
| 69 | + }); |
| 70 | + |
| 71 | + // Policies are used to define on what kind of virtual machines the tasks will run on. |
| 72 | + // In this case, we tell the system to use "e2-standard-4" machine type. |
| 73 | + // Read more about machine types here: https://cloud.google.com/compute/docs/machine-types |
| 74 | + const instancePolicy = new batch.AllocationPolicy.InstancePolicy({ |
| 75 | + machineType, |
| 76 | + // Accelerator describes Compute Engine accelerators to be attached to the VM |
| 77 | + accelerators: [ |
| 78 | + new batch.AllocationPolicy.Accelerator({ |
| 79 | + type: gpuType, |
| 80 | + count: gpuCount, |
| 81 | + installGpuDrivers, |
| 82 | + }), |
| 83 | + ], |
| 84 | + }); |
| 85 | + |
| 86 | + const allocationPolicy = new batch.AllocationPolicy.InstancePolicyOrTemplate({ |
| 87 | + instances: [{installGpuDrivers, policy: instancePolicy}], |
| 88 | + }); |
| 89 | + |
| 90 | + const job = new batch.Job({ |
| 91 | + name: jobName, |
| 92 | + taskGroups: [group], |
| 93 | + labels: {env: 'testing', type: 'script'}, |
| 94 | + allocationPolicy, |
| 95 | + // We use Cloud Logging as it's an option available out of the box |
| 96 | + logsPolicy: new batch.LogsPolicy({ |
| 97 | + destination: batch.LogsPolicy.Destination.CLOUD_LOGGING, |
| 98 | + }), |
| 99 | + }); |
| 100 | + // The job's parent is the project and region in which the job will run |
| 101 | + const parent = `projects/${projectId}/locations/${region}`; |
| 102 | + |
| 103 | + async function callCreateBatchGPUJob() { |
| 104 | + // Construct request |
| 105 | + const request = { |
| 106 | + parent, |
| 107 | + jobId: jobName, |
| 108 | + job, |
| 109 | + }; |
| 110 | + |
| 111 | + // Run request |
| 112 | + const [response] = await batchClient.createJob(request); |
| 113 | + console.log(JSON.stringify(response)); |
| 114 | + } |
| 115 | + |
| 116 | + callCreateBatchGPUJob(); |
| 117 | + // [END batch_create_gpu_job] |
| 118 | +} |
| 119 | + |
| 120 | +process.on('unhandledRejection', err => { |
| 121 | + console.error(err.message); |
| 122 | + process.exitCode = 1; |
| 123 | +}); |
| 124 | + |
| 125 | +main(); |
0 commit comments