591 lines
19 KiB
Python
591 lines
19 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2024 Google LLC. All Rights Reserved.
|
|
#
|
|
# 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
|
|
#
|
|
# http://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.
|
|
"""Utilities for flags for `gcloud colab-enterprise` commands."""
|
|
|
|
from googlecloudsdk.api_lib.colab_enterprise import runtime_templates as runtime_templates_util
|
|
from googlecloudsdk.api_lib.colab_enterprise import runtimes as runtimes_util
|
|
from googlecloudsdk.calliope import arg_parsers
|
|
from googlecloudsdk.calliope import base
|
|
from googlecloudsdk.calliope.concepts import concepts
|
|
from googlecloudsdk.calliope.concepts import deps
|
|
from googlecloudsdk.command_lib.colab_enterprise import completers
|
|
from googlecloudsdk.command_lib.compute.networks import flags as compute_network_flags
|
|
from googlecloudsdk.command_lib.compute.networks.subnets import flags as compute_subnet_flags
|
|
from googlecloudsdk.command_lib.iam import iam_util
|
|
from googlecloudsdk.command_lib.util.concepts import concept_parsers
|
|
from googlecloudsdk.core import properties
|
|
|
|
_accelerator_choices = [
|
|
'NVIDIA_TESLA_V100',
|
|
'NVIDIA_TESLA_T4',
|
|
'NVIDIA_TESLA_A100',
|
|
'NVIDIA_A100_80GB',
|
|
'NVIDIA_L4',
|
|
]
|
|
_disk_choices = ['PD_STANDARD', 'PD_SSD', 'PD_BALANCED', 'PD_EXTREME']
|
|
_post_startup_script_behavior_choices = [
|
|
'POST_STARTUP_SCRIPT_BEHAVIOR_UNSPECIFIED',
|
|
'RUN_ONCE',
|
|
'RUN_EVERY_START',
|
|
'DOWNLOAD_AND_RUN_EVERY_START',
|
|
]
|
|
|
|
|
|
def GetRegionAttributeConfig():
|
|
"""Get the attribute config for the region resource.
|
|
|
|
Returns:
|
|
The resource attribute for the region.
|
|
"""
|
|
return concepts.ResourceParameterAttributeConfig(
|
|
name='region',
|
|
help_text='Cloud region for the {resource}.',
|
|
fallthroughs=[deps.PropertyFallthrough(properties.VALUES.colab.region)],
|
|
)
|
|
|
|
|
|
def AddRuntimeTemplateResourceArg(parser, verb, is_positional):
|
|
"""Add a resource argument for a runtime template to the parser.
|
|
|
|
Args:
|
|
parser: argparse parser for the command.
|
|
verb: str, the verb to describe the resource, such as 'to update'.
|
|
is_positional: bool, True for if arg is positional; False if flag.
|
|
"""
|
|
|
|
def GetRuntimeTemplateResourceSpec(resource_name='runtime template'):
|
|
return concepts.ResourceSpec(
|
|
'aiplatform.projects.locations.notebookRuntimeTemplates',
|
|
resource_name=resource_name,
|
|
projectsId=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
|
|
locationsId=GetRegionAttributeConfig(),
|
|
disable_auto_completers=True,
|
|
)
|
|
|
|
if is_positional:
|
|
arg_name = 'runtime_template'
|
|
prefixes = False
|
|
flag_name_overrides = None
|
|
fallthroughs = None
|
|
else:
|
|
arg_name = '--runtime-template'
|
|
prefixes = True
|
|
flag_name_overrides = {'region': ''}
|
|
fallthroughs = {'region': ['--region']}
|
|
concept_parsers.ConceptParser.ForResource(
|
|
arg_name,
|
|
GetRuntimeTemplateResourceSpec(),
|
|
'Unique name of the runtime template {}. This was optionally provided by'
|
|
' setting --runtime-template-id in the create runtime-template command,'
|
|
' or was system-generated if unspecified.'.format(verb),
|
|
required=True,
|
|
prefixes=prefixes,
|
|
flag_name_overrides=flag_name_overrides,
|
|
command_level_fallthroughs=fallthroughs,
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddRuntimeResourceArg(parser, verb):
|
|
"""Add a resource argument for a runtime template to the parser.
|
|
|
|
Args:
|
|
parser: argparse parser for the command.
|
|
verb: str, the verb to describe the resource, such as 'to update'.
|
|
"""
|
|
|
|
def GetRuntimeResourceSpec(resource_name='runtime'):
|
|
return concepts.ResourceSpec(
|
|
'aiplatform.projects.locations.notebookRuntimes',
|
|
resource_name=resource_name,
|
|
projectsId=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
|
|
locationsId=GetRegionAttributeConfig(),
|
|
disable_auto_completers=True,
|
|
)
|
|
|
|
concept_parsers.ConceptParser.ForResource(
|
|
'runtime',
|
|
GetRuntimeResourceSpec(),
|
|
'Unique name of the runtime {}. This was optionally provided by setting'
|
|
' --runtime-id in the create runtime command, or was system-generated if'
|
|
' unspecified.'.format(verb),
|
|
required=True,
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddRegionResourceArg(parser, verb):
|
|
"""Add a resource argument for a Vertex AI region to the parser.
|
|
|
|
Args:
|
|
parser: argparse parser for the command.
|
|
verb: str, the verb to describe the resource, such as 'to update'.
|
|
"""
|
|
|
|
region_resource_spec = concepts.ResourceSpec(
|
|
'aiplatform.projects.locations',
|
|
resource_name='region',
|
|
locationsId=GetRegionAttributeConfig(),
|
|
projectsId=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
|
|
disable_auto_completers=True,
|
|
)
|
|
|
|
concept_parsers.ConceptParser.ForResource(
|
|
'--region',
|
|
region_resource_spec,
|
|
'Cloud region {}. Please see '
|
|
' https://cloud.google.com/colab/docs/locations for a list of supported'
|
|
' regions.'.format(verb),
|
|
required=True,
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddNetworkResourceArg(help_text, parser):
|
|
"""Adds Resource arg for network to the parser.
|
|
|
|
Args:
|
|
help_text: str, the help text for the flag.
|
|
parser: argparse parser for the command.
|
|
"""
|
|
|
|
def GetNetworkResourceSpec():
|
|
"""Constructs and returns the Resource specification for network."""
|
|
|
|
def NetworkAttributeConfig():
|
|
return concepts.ResourceParameterAttributeConfig(
|
|
name='network',
|
|
help_text=help_text,
|
|
completer=compute_network_flags.NetworksCompleter,
|
|
)
|
|
|
|
return concepts.ResourceSpec(
|
|
'compute.networks',
|
|
resource_name='network',
|
|
network=NetworkAttributeConfig(),
|
|
project=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
|
|
disable_auto_completers=False,
|
|
)
|
|
|
|
concept_parsers.ConceptParser.ForResource(
|
|
'--network', GetNetworkResourceSpec(), help_text
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddSubnetworkResourceArg(help_text, parser):
|
|
"""Adds Resource arg for subnetwork to the parser.
|
|
|
|
Args:
|
|
help_text: str, the help text for the flag.
|
|
parser: argparse parser for the command.
|
|
"""
|
|
|
|
def GetSubnetResourceSpec():
|
|
"""Constructs and returns the Resource specification for Subnet."""
|
|
|
|
def SubnetAttributeConfig():
|
|
return concepts.ResourceParameterAttributeConfig(
|
|
name='subnetwork',
|
|
help_text=help_text,
|
|
completer=compute_subnet_flags.SubnetworksCompleter,
|
|
)
|
|
|
|
def RegionAttributeConfig():
|
|
return concepts.ResourceParameterAttributeConfig(
|
|
name='subnetwork-region',
|
|
help_text=(
|
|
'Google Cloud region of this subnetwork '
|
|
'https://cloud.google.com/compute/docs/regions-zones/#locations.'
|
|
),
|
|
completer=completers.RegionCompleter,
|
|
)
|
|
|
|
return concepts.ResourceSpec(
|
|
'compute.subnetworks',
|
|
resource_name='subnetwork',
|
|
subnetwork=SubnetAttributeConfig(),
|
|
region=RegionAttributeConfig(),
|
|
project=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
|
|
)
|
|
|
|
concept_parsers.ConceptParser.ForResource(
|
|
'--subnetwork', GetSubnetResourceSpec(), help_text
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddKmsKeyResourceArg(parser, help_text):
|
|
"""Adds Resource arg for KMS key to the parser.
|
|
|
|
Args:
|
|
parser: argparse parser for the command.
|
|
help_text: str, the help text for the flag.
|
|
"""
|
|
|
|
def GetKmsKeyResourceSpec():
|
|
|
|
def KmsKeyAttributeConfig():
|
|
# For anchor attribute, help text is generated automatically.
|
|
return concepts.ResourceParameterAttributeConfig(name='kms-key')
|
|
|
|
def KmsKeyringAttributeConfig():
|
|
return concepts.ResourceParameterAttributeConfig(
|
|
name='kms-keyring', help_text='KMS keyring id of the {resource}.'
|
|
)
|
|
|
|
def KmsLocationAttributeConfig():
|
|
return concepts.ResourceParameterAttributeConfig(
|
|
name='kms-location', help_text='Cloud location for the {resource}.'
|
|
)
|
|
|
|
def KmsProjectAttributeConfig():
|
|
return concepts.ResourceParameterAttributeConfig(
|
|
name='kms-project', help_text='Cloud project id for the {resource}.'
|
|
)
|
|
|
|
return concepts.ResourceSpec(
|
|
'cloudkms.projects.locations.keyRings.cryptoKeys',
|
|
resource_name='key',
|
|
cryptoKeysId=KmsKeyAttributeConfig(),
|
|
keyRingsId=KmsKeyringAttributeConfig(),
|
|
locationsId=KmsLocationAttributeConfig(),
|
|
projectsId=KmsProjectAttributeConfig(),
|
|
)
|
|
|
|
concept_parsers.ConceptParser.ForResource(
|
|
'--kms-key',
|
|
GetKmsKeyResourceSpec(),
|
|
help_text,
|
|
required=False,
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddAsyncFlag(parser):
|
|
"""Adds the --async flags to the given parser."""
|
|
base.ASYNC_FLAG.AddToParser(parser)
|
|
|
|
|
|
# This can be used for both to providing compute configuration for Workbench
|
|
# executions and creating Colab runtime templates since CustomEnvironmentSpec is
|
|
# a subset of RuntimeTemplate.
|
|
def AddCustomEnvSpecFlags(parser, workbench_execution=False):
|
|
"""Construct args to provide a custom compute spec for notebook execution.
|
|
|
|
Args:
|
|
parser: argparse parser for the command.
|
|
workbench_execution: bool, true if these flags are for creating a Workbench
|
|
execution.
|
|
"""
|
|
# `Runtime` is decoupled from Workbench for clarity to end users.
|
|
vm_name = (
|
|
'execution environment' if workbench_execution else 'runtime'
|
|
)
|
|
machine_spec_group = parser.add_group(
|
|
help=f'The machine configuration of the {vm_name}.',
|
|
)
|
|
machine_spec_group.add_argument(
|
|
'--machine-type',
|
|
required=False,
|
|
help=f'The Compute Engine machine type selected for the {vm_name}.',
|
|
default='e2-standard-4',
|
|
)
|
|
machine_spec_group.add_argument(
|
|
'--accelerator-type',
|
|
help=f'The type of hardware accelerator used by the {vm_name}. If'
|
|
' specified, --accelerator-count must also be specified.',
|
|
choices=_accelerator_choices,
|
|
default=None,
|
|
)
|
|
machine_spec_group.add_argument(
|
|
'--accelerator-count',
|
|
type=int,
|
|
help='The number of accelerators used by the runtime.',
|
|
)
|
|
disk_spec_group = parser.add_group(
|
|
help=(
|
|
f'The configuration for the data disk of the {vm_name}.'
|
|
),
|
|
)
|
|
disk_spec_group.add_argument(
|
|
'--disk-type',
|
|
help='The type of the persistent disk.',
|
|
choices=_disk_choices,
|
|
default='PD_STANDARD',
|
|
)
|
|
disk_spec_group.add_argument(
|
|
'--disk-size-gb',
|
|
help=f'The disk size of the {vm_name} in GB. If specified, the'
|
|
' --disk-type must also be specified. The minimum size is 10GB and the'
|
|
' maximum is 65536GB.',
|
|
type=int,
|
|
default=100,
|
|
)
|
|
network_spec_group = parser.add_group(
|
|
help=f'The network configuration for the {vm_name}.',
|
|
)
|
|
AddNetworkResourceArg(
|
|
help_text=f'The name of the VPC that this {vm_name} is in.',
|
|
parser=network_spec_group,
|
|
)
|
|
AddSubnetworkResourceArg(
|
|
help_text=(
|
|
f'The name of the subnetwork that this {vm_name} is in.'
|
|
),
|
|
parser=network_spec_group,
|
|
)
|
|
# Since default for this flag is True, the documentation will only show
|
|
# --no-enable-internet-access. Even though now true/false values will always
|
|
# be specified in API request instead of defaulting to none, this gives a
|
|
# better UX than providing both the flag and its negative when default is
|
|
# True.
|
|
network_spec_group.add_argument(
|
|
'--enable-internet-access',
|
|
action='store_true',
|
|
dest='enable_internet_access',
|
|
default=True,
|
|
help=f'Enable public internet access for the {vm_name}.',
|
|
)
|
|
|
|
|
|
def AddRuntimeTemplateSoftwareConfigFlags(parser):
|
|
"""Adds Resource arg for runtime template to the parser."""
|
|
|
|
software_config_group = parser.add_group(
|
|
hidden=True,
|
|
help='The software configuration of the runtime template.',
|
|
)
|
|
software_config_group.add_argument(
|
|
'--post-startup-script',
|
|
required=False,
|
|
help='Post startup script in raw string to execute on the runtime.',
|
|
)
|
|
software_config_group.add_argument(
|
|
'--post-startup-script-url',
|
|
required=False,
|
|
help=(
|
|
'Post startup script URL to execute on the runtime. This can be a'
|
|
' public or private Google Cloud Storage object. In the form of'
|
|
' gs://bucket/object or'
|
|
' https://storage.googleapis.com/bucket/object.'
|
|
),
|
|
)
|
|
software_config_group.add_argument(
|
|
'--post-startup-script-behavior',
|
|
required=False,
|
|
help=(
|
|
'The behavior of the post startup script. The default if passing a'
|
|
' post-startup-script-url is RUN_ONCE.'
|
|
),
|
|
choices=_post_startup_script_behavior_choices,
|
|
)
|
|
software_config_group.add_argument(
|
|
'--set-env-vars',
|
|
required=False,
|
|
help='Set environment variables used by the runtime.',
|
|
type=arg_parsers.ArgDict(),
|
|
metavar='KEY=VALUE',
|
|
)
|
|
|
|
|
|
def AddCreateRuntimeTemplateFlags(parser):
|
|
"""Construct groups and arguments specific to runtime template creation."""
|
|
|
|
AddRegionResourceArg(parser, 'to create runtime template')
|
|
parser.add_argument(
|
|
'--runtime-template-id',
|
|
required=False,
|
|
help='The id of the runtime template. If not specified, a random id will'
|
|
' be generated.',
|
|
)
|
|
parser.add_argument(
|
|
'--display-name',
|
|
required=True,
|
|
help='The display name of the runtime template.',
|
|
)
|
|
runtime_template_group = parser.add_group(
|
|
help='Configuration of the runtime template',
|
|
)
|
|
runtime_template_group.add_argument(
|
|
'--description',
|
|
required=False,
|
|
help='The description of the runtime template.',
|
|
)
|
|
AddCustomEnvSpecFlags(runtime_template_group)
|
|
AddRuntimeTemplateSoftwareConfigFlags(runtime_template_group)
|
|
runtime_template_group.add_argument(
|
|
'--labels',
|
|
help='Add labels to identify and group the runtime template.',
|
|
type=arg_parsers.ArgDict(),
|
|
metavar='KEY=VALUE',
|
|
)
|
|
runtime_template_group.add_argument(
|
|
'--idle-shutdown-timeout',
|
|
help=(
|
|
'The duration after which the runtime is automatically shut down. An'
|
|
' input of 0s disables the idle shutdown feature, and a valid range'
|
|
" is [10m, 24h]. See '$ gcloud topic datetimes' for details on"
|
|
' formatting the input duration.'
|
|
),
|
|
type=arg_parsers.Duration(),
|
|
default='3h',
|
|
)
|
|
runtime_template_group.add_argument(
|
|
'--enable-euc',
|
|
action='store_true',
|
|
dest='enable_euc',
|
|
help='Enable end user credential access for the runtime.',
|
|
default=True,
|
|
)
|
|
runtime_template_group.add_argument(
|
|
'--enable-secure-boot',
|
|
action='store_true',
|
|
dest='enable_secure_boot',
|
|
help='Enables secure boot for the runtime. Disabled by default.',
|
|
default=False,
|
|
)
|
|
runtime_template_group.add_argument(
|
|
'--network-tags',
|
|
type=arg_parsers.ArgList(),
|
|
metavar='TAGS',
|
|
help='Applies the given Compute Engine tags to the runtime.',
|
|
)
|
|
AddKmsKeyResourceArg(
|
|
runtime_template_group,
|
|
'The Cloud KMS encryption key (customer-managed encryption key) used to'
|
|
' protect the runtime. The key must be in the same region as the runtime.'
|
|
' If not specified, Google-managed encryption keys will be used.',
|
|
)
|
|
AddAsyncFlag(parser)
|
|
|
|
|
|
def AddDeleteRuntimeTemplateFlags(parser):
|
|
"""Construct groups and arguments specific to runtime template deletion."""
|
|
AddRuntimeTemplateResourceArg(parser, 'to delete', is_positional=True)
|
|
AddAsyncFlag(parser)
|
|
|
|
|
|
def AddDescribeRuntimeTemplateFlags(parser):
|
|
"""Construct groups and arguments specific to describing a runtime template."""
|
|
AddRuntimeTemplateResourceArg(parser, 'to describe', is_positional=True)
|
|
|
|
|
|
def AddListRuntimeTemplatesFlags(parser):
|
|
"""Construct groups and arguments specific to listing runtime templates."""
|
|
AddRegionResourceArg(parser, 'for which to list all runtime templates')
|
|
parser.display_info.AddUriFunc(runtime_templates_util.GetRuntimeTemplateUri)
|
|
|
|
|
|
def AddFlagsToAddIamPolicyBinding(parser):
|
|
"""Construct arguments for adding an IAM policy binding to a runtime template."""
|
|
AddRuntimeTemplateResourceArg(
|
|
parser, 'to add IAM policy binding to', is_positional=True
|
|
)
|
|
iam_util.AddArgsForAddIamPolicyBinding(parser)
|
|
|
|
|
|
def AddGetIamPolicyFlags(parser):
|
|
"""Construct arguments for getting the IAM policy for a runtime template."""
|
|
AddRuntimeTemplateResourceArg(
|
|
parser, 'to get IAM policy for', is_positional=True
|
|
)
|
|
|
|
|
|
def AddSetIamPolicyBindingFlags(parser):
|
|
"""Construct arguments for setting the IAM policy for a runtime template."""
|
|
AddRuntimeTemplateResourceArg(
|
|
parser, 'to set IAM policy for', is_positional=True
|
|
)
|
|
iam_util.AddArgForPolicyFile(parser)
|
|
|
|
|
|
def AddRemoveIamPolicyBindingFlags(parser):
|
|
"""Construct arguments for removing an IAM policy binding from a runtime template."""
|
|
AddRuntimeTemplateResourceArg(
|
|
parser, 'to remove IAM policy from', is_positional=True
|
|
)
|
|
iam_util.AddArgsForRemoveIamPolicyBinding(parser)
|
|
|
|
|
|
def AddCreateRuntimeFlags(parser):
|
|
"""Construct arguments for creating a runtime."""
|
|
|
|
AddRegionResourceArg(parser, verb='to create runtime')
|
|
|
|
AddRuntimeTemplateResourceArg(
|
|
parser, 'to configure the runtime with', is_positional=False
|
|
)
|
|
parser.add_argument(
|
|
'--runtime-id',
|
|
required=False,
|
|
help=(
|
|
'The id of the runtime to create. If not specified, a random id will'
|
|
' be generated.'
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
'--display-name',
|
|
required=True,
|
|
help='The display name of the runtime to create.',
|
|
)
|
|
parser.add_argument('--description', required=False, help='The description')
|
|
parser.add_argument(
|
|
'--runtime-user',
|
|
required=False,
|
|
help=(
|
|
'User email for the runtime owner. Runtimes can only be used by the'
|
|
' owner. If a user is not provided, the gcloud user will be assumed'
|
|
' to be the owner. The user cannot be a service account.'
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
'--labels',
|
|
help='Add labels to identify and group the runtime template.',
|
|
type=arg_parsers.ArgDict(),
|
|
metavar='KEY=VALUE',
|
|
)
|
|
|
|
AddAsyncFlag(parser)
|
|
|
|
|
|
def AddDescribeRuntimeFlags(parser):
|
|
"""Construct arguments specific to describing a runtime."""
|
|
AddRuntimeResourceArg(parser, verb='to describe')
|
|
|
|
|
|
def AddListRuntimeFlags(parser):
|
|
"""Construct arguments specific to listing runtimes."""
|
|
AddRegionResourceArg(parser, verb='for which to list all runtimes')
|
|
parser.display_info.AddUriFunc(runtimes_util.GetRuntimeUri)
|
|
|
|
|
|
def AddDeleteRuntimeFlags(parser):
|
|
"""Construct arguments specific to deleting a runtime."""
|
|
AddRuntimeResourceArg(parser, verb='to delete')
|
|
AddAsyncFlag(parser)
|
|
|
|
|
|
def AddUpgradeRuntimeFlags(parser):
|
|
"""Construct arguments specific to upgrading a runtime."""
|
|
AddRuntimeResourceArg(parser, verb='to upgrade')
|
|
AddAsyncFlag(parser)
|
|
|
|
|
|
def AddStartRuntimeFlags(parser):
|
|
"""Construct arguments specific to starting a stopped runtime."""
|
|
AddRuntimeResourceArg(parser, verb='to start')
|
|
AddAsyncFlag(parser)
|
|
|
|
|
|
def AddStopRuntimeFlags(parser):
|
|
"""Construct arguments specific to stopping a runtime."""
|
|
AddRuntimeResourceArg(parser, verb='to stop')
|
|
AddAsyncFlag(parser)
|