-
Notifications
You must be signed in to change notification settings - Fork 2k
feat: tpu_queued_resources_startup_script/create_network/time_bound #3907
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
78c9912
50cc2e8
c992903
a9ee1f7
b8b5fa5
954e4ec
5fc0b13
a59ec1a
a89ba88
c7c0e5b
012857d
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,123 @@ | ||||||||||||||||||
/* | ||||||||||||||||||
* 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. | ||||||||||||||||||
*/ | ||||||||||||||||||
|
||||||||||||||||||
'use strict'; | ||||||||||||||||||
|
||||||||||||||||||
async function main(tpuClient) { | ||||||||||||||||||
// [START tpu_queued_resources_network] | ||||||||||||||||||
// Import the TPUClient | ||||||||||||||||||
// TODO(developer): Uncomment below line before running the sample. | ||||||||||||||||||
// const {TpuClient} = require('@google-cloud/tpu').v2alpha1; | ||||||||||||||||||
const {Node, NetworkConfig, QueuedResource} = | ||||||||||||||||||
require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1; | ||||||||||||||||||
|
||||||||||||||||||
// Instantiate a tpuClient | ||||||||||||||||||
// TODO(developer): Uncomment below line before running the sample. | ||||||||||||||||||
// tpuClient = new TpuClient(); | ||||||||||||||||||
Comment on lines
+28
to
+29
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Consider instantiating the
Suggested change
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It is commented, becuase mocked tpuClient is used for testing and I pass mock to this function. If user wants to use real TPUClient, they will need to uncomment this line There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Understood. For clarity, perhaps a comment indicating that the
Suggested change
|
||||||||||||||||||
|
||||||||||||||||||
/** | ||||||||||||||||||
* TODO(developer): Update/uncomment these variables before running the sample. | ||||||||||||||||||
*/ | ||||||||||||||||||
iennae marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||||
// Project ID or project number of the Google Cloud project, where you want to create queued resource. | ||||||||||||||||||
const projectId = await tpuClient.getProjectId(); | ||||||||||||||||||
|
||||||||||||||||||
// The name of the network you want the node to connect to. The network should be assigned to your project. | ||||||||||||||||||
const networkName = 'compute-tpu-network'; | ||||||||||||||||||
|
||||||||||||||||||
// The region of the network, that you want the node to connect to. | ||||||||||||||||||
const region = 'us-central1'; | ||||||||||||||||||
|
||||||||||||||||||
// The name for your queued resource. | ||||||||||||||||||
const queuedResourceName = 'queued-resource-1'; | ||||||||||||||||||
|
||||||||||||||||||
// The name for your node. | ||||||||||||||||||
const nodeName = 'node-name-1'; | ||||||||||||||||||
|
||||||||||||||||||
// The zone in which to create the node. | ||||||||||||||||||
// For more information about supported TPU types for specific zones, | ||||||||||||||||||
// see https://cloud.google.com/tpu/docs/regions-zones | ||||||||||||||||||
const zone = `${region}-a`; | ||||||||||||||||||
|
||||||||||||||||||
// The accelerator type that specifies the version and size of the node you want to create. | ||||||||||||||||||
// For more information about supported accelerator types for each TPU version, | ||||||||||||||||||
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions. | ||||||||||||||||||
const tpuType = 'v2-8'; | ||||||||||||||||||
|
||||||||||||||||||
// Software version that specifies the version of the node runtime to install. For more information, | ||||||||||||||||||
// see https://cloud.google.com/tpu/docs/runtimes | ||||||||||||||||||
const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; | ||||||||||||||||||
|
||||||||||||||||||
async function callCreateQueuedResourceNetwork() { | ||||||||||||||||||
// Specify the network and subnetwork that you want to connect your TPU to. | ||||||||||||||||||
const networkConfig = new NetworkConfig({ | ||||||||||||||||||
enableExternalIps: true, | ||||||||||||||||||
network: `projects/${projectId}/global/networks/${networkName}`, | ||||||||||||||||||
subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`, | ||||||||||||||||||
}); | ||||||||||||||||||
|
||||||||||||||||||
// Create a node | ||||||||||||||||||
const node = new Node({ | ||||||||||||||||||
name: nodeName, | ||||||||||||||||||
zone, | ||||||||||||||||||
acceleratorType: tpuType, | ||||||||||||||||||
runtimeVersion: tpuSoftwareVersion, | ||||||||||||||||||
networkConfig, | ||||||||||||||||||
queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`, | ||||||||||||||||||
}); | ||||||||||||||||||
|
||||||||||||||||||
// Define parent for requests | ||||||||||||||||||
const parent = `projects/${projectId}/locations/${zone}`; | ||||||||||||||||||
|
||||||||||||||||||
// Create queued resource | ||||||||||||||||||
const queuedResource = new QueuedResource({ | ||||||||||||||||||
name: queuedResourceName, | ||||||||||||||||||
tpu: { | ||||||||||||||||||
nodeSpec: [ | ||||||||||||||||||
{ | ||||||||||||||||||
parent, | ||||||||||||||||||
node, | ||||||||||||||||||
nodeId: nodeName, | ||||||||||||||||||
}, | ||||||||||||||||||
], | ||||||||||||||||||
}, | ||||||||||||||||||
}); | ||||||||||||||||||
|
||||||||||||||||||
const request = { | ||||||||||||||||||
parent: `projects/${projectId}/locations/${zone}`, | ||||||||||||||||||
queuedResource, | ||||||||||||||||||
queuedResourceId: queuedResourceName, | ||||||||||||||||||
}; | ||||||||||||||||||
|
||||||||||||||||||
const [operation] = await tpuClient.createQueuedResource(request); | ||||||||||||||||||
|
||||||||||||||||||
// Wait for the create operation to complete. | ||||||||||||||||||
const [response] = await operation.promise(); | ||||||||||||||||||
|
||||||||||||||||||
// You can wait until TPU Node is READY, | ||||||||||||||||||
// and check its status using getTpuVm() from `tpu_vm_get` sample. | ||||||||||||||||||
return response; | ||||||||||||||||||
} | ||||||||||||||||||
return await callCreateQueuedResourceNetwork(); | ||||||||||||||||||
// [END tpu_queued_resources_network] | ||||||||||||||||||
} | ||||||||||||||||||
|
||||||||||||||||||
module.exports = main; | ||||||||||||||||||
|
||||||||||||||||||
// TODO(developer): Uncomment below lines before running the sample. | ||||||||||||||||||
// main(...process.argv.slice(2)).catch(err => { | ||||||||||||||||||
// console.error(err); | ||||||||||||||||||
// process.exitCode = 1; | ||||||||||||||||||
// }); |
Original file line number | Diff line number | Diff line change | ||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,127 @@ | ||||||||||||||||||||
/* | ||||||||||||||||||||
* 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. | ||||||||||||||||||||
*/ | ||||||||||||||||||||
|
||||||||||||||||||||
'use strict'; | ||||||||||||||||||||
|
||||||||||||||||||||
async function main(tpuClient) { | ||||||||||||||||||||
// [START tpu_queued_resources_startup_script] | ||||||||||||||||||||
// Import the TPUClient | ||||||||||||||||||||
// TODO(developer): Uncomment below line before running the sample. | ||||||||||||||||||||
// const {TpuClient} = require('@google-cloud/tpu').v2alpha1; | ||||||||||||||||||||
Comment on lines
+22
to
+23
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Since
Suggested change
|
||||||||||||||||||||
const {Node, NetworkConfig, QueuedResource} = | ||||||||||||||||||||
require('@google-cloud/tpu').protos.google.cloud.tpu.v2alpha1; | ||||||||||||||||||||
|
||||||||||||||||||||
// Instantiate a tpuClient | ||||||||||||||||||||
// TODO(developer): Uncomment below line before running the sample. | ||||||||||||||||||||
// tpuClient = new TpuClient(); | ||||||||||||||||||||
Comment on lines
+28
to
+29
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Consider instantiating the
Suggested change
|
||||||||||||||||||||
|
||||||||||||||||||||
/** | ||||||||||||||||||||
* TODO(developer): Update/uncomment these variables before running the sample. | ||||||||||||||||||||
*/ | ||||||||||||||||||||
Comment on lines
+31
to
+33
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Update the TODO with instructions on how to set up the required environment variables or configuration.
Suggested change
|
||||||||||||||||||||
// Project ID or project number of the Google Cloud project, where you want to create queued resource. | ||||||||||||||||||||
const projectId = await tpuClient.getProjectId(); | ||||||||||||||||||||
|
||||||||||||||||||||
// The name of the network you want the node to connect to. The network should be assigned to your project. | ||||||||||||||||||||
const networkName = 'compute-tpu-network'; | ||||||||||||||||||||
|
||||||||||||||||||||
// The region of the network, that you want the node to connect to. | ||||||||||||||||||||
const region = 'us-central1'; | ||||||||||||||||||||
|
||||||||||||||||||||
// The name for your queued resource. | ||||||||||||||||||||
const queuedResourceName = 'queued-resource-1'; | ||||||||||||||||||||
|
||||||||||||||||||||
// The name for your node. | ||||||||||||||||||||
const nodeName = 'node-name-1'; | ||||||||||||||||||||
|
||||||||||||||||||||
// The zone in which to create the node. | ||||||||||||||||||||
// For more information about supported TPU types for specific zones, | ||||||||||||||||||||
// see https://cloud.google.com/tpu/docs/regions-zones | ||||||||||||||||||||
const zone = `${region}-a`; | ||||||||||||||||||||
|
||||||||||||||||||||
// The accelerator type that specifies the version and size of the node you want to create. | ||||||||||||||||||||
// For more information about supported accelerator types for each TPU version, | ||||||||||||||||||||
// see https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#versions. | ||||||||||||||||||||
const tpuType = 'v2-8'; | ||||||||||||||||||||
|
||||||||||||||||||||
// Software version that specifies the version of the node runtime to install. For more information, | ||||||||||||||||||||
// see https://cloud.google.com/tpu/docs/runtimes | ||||||||||||||||||||
const tpuSoftwareVersion = 'tpu-vm-tf-2.14.1'; | ||||||||||||||||||||
|
||||||||||||||||||||
async function callCreateQueuedResourceStartupScript() { | ||||||||||||||||||||
// Create a node | ||||||||||||||||||||
const node = new Node({ | ||||||||||||||||||||
name: nodeName, | ||||||||||||||||||||
zone, | ||||||||||||||||||||
acceleratorType: tpuType, | ||||||||||||||||||||
runtimeVersion: tpuSoftwareVersion, | ||||||||||||||||||||
// Define network | ||||||||||||||||||||
networkConfig: new NetworkConfig({ | ||||||||||||||||||||
enableExternalIps: true, | ||||||||||||||||||||
network: `projects/${projectId}/global/networks/${networkName}`, | ||||||||||||||||||||
subnetwork: `projects/${projectId}/regions/${region}/subnetworks/${networkName}`, | ||||||||||||||||||||
}), | ||||||||||||||||||||
queuedResource: `projects/${projectId}/locations/${zone}/queuedResources/${queuedResourceName}`, | ||||||||||||||||||||
metadata: { | ||||||||||||||||||||
// The script updates numpy to the latest version and logs the output to a file. | ||||||||||||||||||||
'startup-script': `#!/bin/bash | ||||||||||||||||||||
echo "Hello World" > /var/log/hello.log | ||||||||||||||||||||
sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1`, | ||||||||||||||||||||
}, | ||||||||||||||||||||
}); | ||||||||||||||||||||
|
||||||||||||||||||||
// Define parent for requests | ||||||||||||||||||||
const parent = `projects/${projectId}/locations/${zone}`; | ||||||||||||||||||||
|
||||||||||||||||||||
// Create queued resource | ||||||||||||||||||||
const queuedResource = new QueuedResource({ | ||||||||||||||||||||
name: queuedResourceName, | ||||||||||||||||||||
tpu: { | ||||||||||||||||||||
nodeSpec: [ | ||||||||||||||||||||
{ | ||||||||||||||||||||
parent, | ||||||||||||||||||||
node, | ||||||||||||||||||||
nodeId: nodeName, | ||||||||||||||||||||
}, | ||||||||||||||||||||
], | ||||||||||||||||||||
}, | ||||||||||||||||||||
}); | ||||||||||||||||||||
|
||||||||||||||||||||
const request = { | ||||||||||||||||||||
parent: `projects/${projectId}/locations/${zone}`, | ||||||||||||||||||||
queuedResource, | ||||||||||||||||||||
queuedResourceId: queuedResourceName, | ||||||||||||||||||||
}; | ||||||||||||||||||||
|
||||||||||||||||||||
const [operation] = await tpuClient.createQueuedResource(request); | ||||||||||||||||||||
|
||||||||||||||||||||
// Wait for the create operation to complete. | ||||||||||||||||||||
const [response] = await operation.promise(); | ||||||||||||||||||||
|
||||||||||||||||||||
// You can wait until TPU Node is READY, | ||||||||||||||||||||
// and check its status using getTpuVm() from `tpu_vm_get` sample. | ||||||||||||||||||||
return response; | ||||||||||||||||||||
} | ||||||||||||||||||||
return await callCreateQueuedResourceStartupScript(); | ||||||||||||||||||||
// [END tpu_queued_resources_startup_script] | ||||||||||||||||||||
} | ||||||||||||||||||||
|
||||||||||||||||||||
module.exports = main; | ||||||||||||||||||||
|
||||||||||||||||||||
// TODO(developer): Uncomment below lines before running the sample. | ||||||||||||||||||||
// main(...process.argv.slice(2)).catch(err => { | ||||||||||||||||||||
// console.error(err); | ||||||||||||||||||||
// process.exitCode = 1; | ||||||||||||||||||||
// }); |
Uh oh!
There was an error while loading. Please reload this page.