File: //snap/google-cloud-cli/current/lib/googlecloudsdk/api_lib/workbench/instances.py
# -*- coding: utf-8 -*- #
# Copyright 2023 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.
"""workbench instances api helper."""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
import enum
from googlecloudsdk.api_lib.workbench import util
from googlecloudsdk.command_lib.util.apis import arg_utils
from googlecloudsdk.command_lib.workbench import flags
from googlecloudsdk.core import exceptions
from googlecloudsdk.core import log
from googlecloudsdk.core import resources
_RESERVATION_AFFINITY_KEY = 'compute.googleapis.com/reservation-name'
def GetNetworkRelativeName(args):
if args.IsSpecified('network'):
return args.CONCEPTS.network.Parse().RelativeName()
def GetSubnetRelativeName(args):
if args.IsSpecified('subnet'):
try:
return args.CONCEPTS.subnet.Parse().RelativeName()
except AttributeError:
raise exceptions.Error(
'Subnet is not formatted correctly. Expected format:'
' projects/{project-id}/regions/{region}/subnetworks/{subnet-name}'
)
def CreateNetworkConfigMessage(args, messages):
"""Creates the network config for the instance.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
A list of NetworkInterface messages for the instance.
"""
if not (
args.IsSpecified('network')
or args.IsSpecified('subnet')
or args.IsSpecified('nic_type')
):
return []
network_name = None
subnet_name = None
nic_type = None
if args.IsSpecified('network'):
network_name = GetNetworkRelativeName(args)
if args.IsSpecified('subnet'):
subnet_name = GetSubnetRelativeName(args)
if args.IsSpecified('nic_type'):
nic_type = arg_utils.ChoiceEnumMapper(
arg_name='nic-type',
message_enum=messages.NetworkInterface.NicTypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x,
).GetEnumForChoice(arg_utils.EnumNameToChoice(args.nic_type))
return [messages.NetworkInterface(
network=network_name,
subnet=subnet_name,
nicType=nic_type,
)]
def CreateAcceleratorConfigMessage(args, messages):
"""Creates the accelerator config for the instance.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
A list of AcceleratorConfig messages for the instance.
"""
if not (
args.IsSpecified('accelerator_type')
or args.IsSpecified('accelerator_core_count')
):
return []
type_enum = None
if args.IsSpecified('accelerator_type'):
type_enum = arg_utils.ChoiceEnumMapper(
arg_name='accelerator-type',
message_enum=messages.AcceleratorConfig.TypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x,
).GetEnumForChoice(arg_utils.EnumNameToChoice(args.accelerator_type))
return [messages.AcceleratorConfig(
type=type_enum, coreCount=args.accelerator_core_count
)]
def CreateServiceAccountConfigMessage(args, messages):
if not args.IsSpecified('service_account_email'):
return []
return [messages.ServiceAccount(email=args.service_account_email)]
def CreateGPUDriverConfigMessage(args, messages):
if not (
args.IsSpecified('custom_gpu_driver_path')
or args.IsSpecified('install_gpu_driver')
):
return None
return messages.GPUDriverConfig(
customGpuDriverPath=args.custom_gpu_driver_path,
enableGpuDriver=args.install_gpu_driver,
)
def CreateReservationConfigMessage(args, messages):
"""Creates the reservation config for the instance.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
Reservation config for the instance.
"""
if not (
args.IsSpecified('reservation_type')
or args.IsSpecified('reservation_key')
or args.IsSpecified('reservation_values')
):
return None
reservation_type_enum = None
if args.IsSpecified('reservation_type'):
reservation_type_enum = flags.GetReservationTypeMapper(
messages
).GetEnumForChoice(args.reservation_type)
reservation_values = []
if args.IsSpecified('reservation_values'):
reservation_values = args.reservation_values
return messages.ReservationAffinity(
consumeReservationType=reservation_type_enum,
key=args.reservation_key,
values=reservation_values,
)
def GetBootDisk(args, messages):
"""Creates the boot disk config for the instance.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
Boot disk config for the instance.
"""
if not (
args.IsSpecified('boot_disk_size')
or args.IsSpecified('boot_disk_type')
or args.IsSpecified('boot_disk_encryption')
or args.IsSpecified('boot_disk_kms_key')
):
return None
boot_disk_message = messages.BootDisk
boot_disk_encryption_enum = None
boot_disk_type_enum = None
kms_key = None
if args.IsSpecified('boot_disk_type'):
boot_disk_type_enum = arg_utils.ChoiceEnumMapper(
arg_name='boot-disk-type',
message_enum=boot_disk_message.DiskTypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x,
).GetEnumForChoice(arg_utils.EnumNameToChoice(args.boot_disk_type))
if args.IsSpecified('boot_disk_encryption'):
boot_disk_encryption_enum = arg_utils.ChoiceEnumMapper(
arg_name='boot-disk-encryption',
message_enum=boot_disk_message.DiskEncryptionValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x,
).GetEnumForChoice(arg_utils.EnumNameToChoice(args.boot_disk_encryption))
if args.IsSpecified('boot_disk_kms_key'):
kms_key = args.CONCEPTS.boot_disk_kms_key.Parse().RelativeName()
return boot_disk_message(
diskType=boot_disk_type_enum,
diskEncryption=boot_disk_encryption_enum,
diskSizeGb=args.boot_disk_size,
kmsKey=kms_key,
)
def GetDataDisk(args, messages):
"""Creates the data disk config for the instance.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
Data disk config for the instance.
"""
if not (
args.IsSpecified('data_disk_size')
or args.IsSpecified('data_disk_type')
or args.IsSpecified('data_disk_encryption')
or args.IsSpecified('data_disk_kms_key')
or args.IsSpecified('data_disk_resource_policies')
):
return []
data_disk_message = messages.DataDisk
data_disk_encryption_enum = None
data_disk_type_enum = None
kms_key = None
resource_policies = []
if args.IsSpecified('data_disk_type'):
data_disk_type_enum = arg_utils.ChoiceEnumMapper(
arg_name='data-disk-type',
message_enum=data_disk_message.DiskTypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x,
).GetEnumForChoice(arg_utils.EnumNameToChoice(args.data_disk_type))
if args.IsSpecified('data_disk_encryption'):
data_disk_encryption_enum = arg_utils.ChoiceEnumMapper(
arg_name='data-disk-encryption',
message_enum=data_disk_message.DiskEncryptionValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x,
).GetEnumForChoice(arg_utils.EnumNameToChoice(args.data_disk_encryption))
if args.IsSpecified('data_disk_kms_key'):
kms_key = args.CONCEPTS.data_disk_kms_key.Parse().RelativeName()
if args.IsSpecified('data_disk_resource_policies'):
resource_policies = args.data_disk_resource_policies
return [
data_disk_message(
diskType=data_disk_type_enum,
diskEncryption=data_disk_encryption_enum,
diskSizeGb=args.data_disk_size,
kmsKey=kms_key,
resourcePolicies=resource_policies,
)
]
def CreateContainerImageFromArgs(args, messages):
if not (
args.IsSpecified('container_repository')
or args.IsSpecified('container_tag')
):
return None
container_image = messages.ContainerImage(
repository=args.container_repository, tag=args.container_tag
)
return container_image
def CreateVmImageFromArgs(args, messages):
"""Create VmImage Message from an environment or from args."""
if not (
args.IsSpecified('vm_image_project')
or args.IsSpecified('vm_image_family')
or args.IsSpecified('vm_image_name')
):
return None
vm_image = messages.VmImage()
vm_image.project = args.vm_image_project
if args.IsSpecified('vm_image_family'):
vm_image.family = args.vm_image_family
else:
vm_image.name = args.vm_image_name
return vm_image
def GetInstanceOwnersFromArgs(args):
if not args.IsSpecified('instance_owners'):
return []
return [args.instance_owners]
def GetLabelsFromArgs(args, messages):
if not args.IsSpecified('labels'):
return None
labels_message = messages.Instance.LabelsValue
return labels_message(additionalProperties=[
labels_message.AdditionalProperty(key=key, value=value)
for key, value in args.labels.items()
])
def GetTagsFromArgs(args):
if not args.IsSpecified('tags'):
return []
return args.tags
def GetMetadataFromArgs(args, messages):
if not args.IsSpecified('metadata'):
return None
metadata_message = messages.GceSetup.MetadataValue
return metadata_message(
additionalProperties=[
metadata_message.AdditionalProperty(key=key, value=value)
for key, value in args.metadata.items()
]
)
def GetShieldedInstanceConfigFromArgs(args, messages):
"""Creates the Shielded Instance Config message for the create/update request.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
Shielded Instance Config of the Instance message.
"""
if not (
args.IsSpecified('shielded_secure_boot')
or args.IsSpecified('shielded_vtpm')
or args.IsSpecified('shielded_integrity_monitoring')
):
return None
true_values = ['1', 'true', 'on', 'yes', 'y']
if args.IsSpecified('shielded_secure_boot'):
shielded_secure_boot = args.shielded_secure_boot.lower() in true_values
else:
shielded_secure_boot = None
if args.IsSpecified('shielded_vtpm'):
shielded_vtpm = args.shielded_vtpm.lower() in true_values
else:
shielded_vtpm = None
if args.IsSpecified('shielded_integrity_monitoring'):
shielded_integrity_monitoring = (
args.shielded_integrity_monitoring.lower() in true_values
)
else:
shielded_integrity_monitoring = None
shielded_instance_config_message = messages.ShieldedInstanceConfig
return shielded_instance_config_message(
enableIntegrityMonitoring=shielded_integrity_monitoring,
enableSecureBoot=shielded_secure_boot,
enableVtpm=shielded_vtpm,
)
def GetConfidentialInstanceConfigFromArgs(args, messages):
"""Creates the Confidential Instance Config message for the create request.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
Confidential Instance Config of the Instance message.
"""
if not args.IsSpecified('confidential_compute_type'):
return None
confidential_instance_type = arg_utils.ChoiceEnumMapper(
arg_name='confidential-compute-type',
message_enum=messages.ConfidentialInstanceConfig.ConfidentialInstanceTypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x,
).GetEnumForChoice(
arg_utils.EnumNameToChoice(args.confidential_compute_type)
)
return messages.ConfidentialInstanceConfig(
confidentialInstanceType=confidential_instance_type,
)
def GetMachineTypeFromArgs(args):
if not args.IsSpecified('machine_type'):
return None
return args.machine_type
def GetMinCpuPlatformFromArgs(args):
if not args.IsSpecified('min_cpu_platform'):
return None
return args.min_cpu_platform
def GetDisablePublicIpFromArgs(args):
if not args.IsSpecified('disable_public_ip'):
return None
return args.disable_public_ip
def GetEnableIpForwardingFromArgs(args):
if not args.IsSpecified('enable_ip_forwarding'):
return None
return args.enable_ip_forwarding
def GetEnableThirdPartyIdentityFromArgs(args):
if not args.IsSpecified('enable_third_party_identity'):
return None
return args.enable_third_party_identity
def GetEnableManagedEucFromArgs(args):
if not args.IsSpecified('enable_managed_euc'):
return None
return args.enable_managed_euc
def GetDisableProxyAccessFromArgs(args):
if not args.IsSpecified('disable_proxy_access'):
return None
return args.disable_proxy_access
def CreateInstance(args, messages, support_managed_euc):
"""Creates the Instance message for the create request.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
support_managed_euc: Whether the managed EUC feature is supported.
Returns:
Instance of the Instance message.
"""
gce_setup = messages.GceSetup(
acceleratorConfigs=CreateAcceleratorConfigMessage(args, messages),
bootDisk=GetBootDisk(args, messages),
containerImage=CreateContainerImageFromArgs(args, messages),
dataDisks=GetDataDisk(args, messages),
disablePublicIp=GetDisablePublicIpFromArgs(args),
enableIpForwarding=GetEnableIpForwardingFromArgs(args),
gpuDriverConfig=CreateGPUDriverConfigMessage(args, messages),
machineType=GetMachineTypeFromArgs(args),
minCpuPlatform=GetMinCpuPlatformFromArgs(args),
metadata=GetMetadataFromArgs(args, messages),
networkInterfaces=CreateNetworkConfigMessage(args, messages),
reservationAffinity=CreateReservationConfigMessage(args, messages),
serviceAccounts=CreateServiceAccountConfigMessage(args, messages),
shieldedInstanceConfig=GetShieldedInstanceConfigFromArgs(args, messages),
confidentialInstanceConfig=GetConfidentialInstanceConfigFromArgs(
args, messages
),
tags=GetTagsFromArgs(args),
vmImage=CreateVmImageFromArgs(args, messages),
)
instance = messages.Instance(
name=args.instance,
disableProxyAccess=GetDisableProxyAccessFromArgs(args),
gceSetup=gce_setup if gce_setup != messages.GceSetup() else None,
instanceOwners=GetInstanceOwnersFromArgs(args),
labels=GetLabelsFromArgs(args, messages),
enableThirdPartyIdentity=GetEnableThirdPartyIdentityFromArgs(args),
enableManagedEuc=GetEnableManagedEucFromArgs(args)
if support_managed_euc
else None,
)
return instance
def CreateInstanceCreateRequest(args, messages, support_managed_euc):
"""Creates the update mask for update Instance request.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the API.
support_managed_euc: Whether the managed EUC feature is supported.
Returns:
Update mask of the Instance message.
"""
parent = util.GetParentForInstance(args)
instance = CreateInstance(args, messages, support_managed_euc)
return messages.NotebooksProjectsLocationsInstancesCreateRequest(
parent=parent, instance=instance, instanceId=args.instance)
def CreateUpdateMask(args):
"""Creates the update mask for update Instance request.
Args:
args: Argparse object from Command.Run
Returns:
Update mask of the Instance message.
"""
mask_array = []
field_mask_dict = {
'accelerator_type': 'gce_setup.accelerator_configs.type',
'accelerator_core_count': 'gce_setup.accelerator_configs.core_count',
'install_gpu_driver': 'gce_setup.gpu_driver_config.enable_gpu_driver',
'custom_gpu_driver_path': (
'gce_setup.gpu_driver_config.custom_gpu_driver_path'
),
'shielded_secure_boot': (
'gce_setup.shielded_instance_config.enable_secure_boot'
),
'shielded_vtpm': 'gce_setup.shielded_instance_config.enable_vtpm',
'shielded_integrity_monitoring': (
'gce_setup.shielded_instance_config.enable_integrity_monitoring'
),
'labels': 'labels',
'metadata': 'gce_setup.metadata',
'machine_type': 'gce_setup.machine_type',
'tags': 'gce_setup.tags',
'container_repository': 'gce_setup.container_image.repository',
'container_tag': 'gce_setup.container_image.tag',
}
for key, value in sorted(field_mask_dict.items()):
if args.IsSpecified(key):
mask_array.append(value)
return ','.join(map(str, mask_array))
def GetInstanceResource(args):
return args.CONCEPTS.instance.Parse()
def CreateInstanceListRequest(args, messages):
parent = util.GetParentFromArgs(args)
return messages.NotebooksProjectsLocationsInstancesListRequest(parent=parent)
def CreateInstanceDeleteRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
return messages.NotebooksProjectsLocationsInstancesDeleteRequest(
name=instance)
def CreateInstanceDescribeRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
return messages.NotebooksProjectsLocationsInstancesGetRequest(name=instance)
def CreateInstanceResetRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
reset_request = messages.ResetInstanceRequest()
return messages.NotebooksProjectsLocationsInstancesResetRequest(
name=instance, resetInstanceRequest=reset_request)
def CreateInstanceRollbackRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
rollback_request = messages.RollbackInstanceRequest(
targetSnapshot=args.target_snapshot)
return messages.NotebooksProjectsLocationsInstancesRollbackRequest(
name=instance, rollbackInstanceRequest=rollback_request)
def CreateInstanceStartRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
start_request = messages.StartInstanceRequest()
return messages.NotebooksProjectsLocationsInstancesStartRequest(
name=instance, startInstanceRequest=start_request)
def CreateInstanceStopRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
stop_request = messages.StopInstanceRequest()
return messages.NotebooksProjectsLocationsInstancesStopRequest(
name=instance, stopInstanceRequest=stop_request)
def CreateGetConfigRequest(args, messages):
name = util.GetParentFromArgs(args)
return messages.NotebooksProjectsLocationsInstancesGetConfigRequest(name=name)
def CreateInstanceRestoreRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
restore_request = messages.RestoreInstanceRequest(
snapshot=messages.Snapshot(
projectId=args.snapshot_project, snapshotId=args.snapshot
)
)
return messages.NotebooksProjectsLocationsInstancesRestoreRequest(
name=instance, restoreInstanceRequest=restore_request
)
def UpdateInstance(args, messages):
"""Creates the Instance message for the update request.
Args:
args: Argparse object from Command.Run
messages: Module containing messages definition for the specified API.
Returns:
Instance of the Instance message.
"""
gce_setup = messages.GceSetup(
acceleratorConfigs=CreateAcceleratorConfigMessage(args, messages),
gpuDriverConfig=CreateGPUDriverConfigMessage(args, messages),
machineType=GetMachineTypeFromArgs(args),
metadata=GetMetadataFromArgs(args, messages),
shieldedInstanceConfig=GetShieldedInstanceConfigFromArgs(args, messages),
tags=GetTagsFromArgs(args),
containerImage=CreateContainerImageFromArgs(args, messages),
)
instance = messages.Instance(
name=args.instance,
gceSetup=gce_setup if gce_setup != messages.GceSetup() else None,
labels=GetLabelsFromArgs(args, messages),
)
return instance
def CreateInstanceUpdateRequest(args, messages):
instance = UpdateInstance(args, messages)
name = GetInstanceResource(args).RelativeName()
update_mask = CreateUpdateMask(args)
return messages.NotebooksProjectsLocationsInstancesPatchRequest(
instance=instance,
name=name,
updateMask=update_mask)
def CreateInstanceDiagnoseRequest(args, messages):
""""Create and return Diagnose request."""
instance = GetInstanceResource(args).RelativeName()
diagnostic_config = messages.DiagnosticConfig(
gcsBucket=args.gcs_bucket,
)
if args.IsSpecified('relative_path'):
diagnostic_config.relativePath = args.relative_path
if args.IsSpecified('enable_repair'):
diagnostic_config.enableRepairFlag = True
if args.IsSpecified('enable_packet_capture'):
diagnostic_config.enablePacketCaptureFlag = True
if args.IsSpecified('enable_copy_home_files'):
diagnostic_config.enableCopyHomeFilesFlag = True
return messages.NotebooksProjectsLocationsInstancesDiagnoseRequest(
name=instance, diagnoseInstanceRequest=messages.DiagnoseInstanceRequest(
diagnosticConfig=diagnostic_config))
def CreateInstanceResizeDisk(args, messages):
"""Create and return ResizeDisk request."""
instance = GetInstanceResource(args).RelativeName()
request = None
if args.IsSpecified('boot_disk_size'):
request = messages.ResizeDiskRequest(
bootDisk=messages.BootDisk(
diskSizeGb=args.boot_disk_size,
)
)
elif args.IsSpecified('data_disk_size'):
request = messages.ResizeDiskRequest(
dataDisk=messages.DataDisk(
diskSizeGb=args.data_disk_size,
)
)
return messages.NotebooksProjectsLocationsInstancesResizeDiskRequest(
notebookInstance=instance, resizeDiskRequest=request)
def CreateInstanceUpgradeRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
upgrade_request = messages.UpgradeInstanceRequest()
return messages.NotebooksProjectsLocationsInstancesUpgradeRequest(
name=instance, upgradeInstanceRequest=upgrade_request)
def CreateInstanceCheckUpgradabilityRequest(args, messages):
instance = GetInstanceResource(args).RelativeName()
return messages.NotebooksProjectsLocationsInstancesCheckUpgradabilityRequest(
notebookInstance=instance)
def GetInstanceURI(resource):
instance = resources.REGISTRY.ParseRelativeName(
resource.name, collection='notebooks.projects.locations.instances')
return instance.SelfLink()
class OperationType(enum.Enum):
CREATE = (log.CreatedResource, 'created')
UPDATE = (log.UpdatedResource, 'updated')
RESTORE = (log.RestoredResource, 'restored')
UPGRADE = (log.UpdatedResource, 'upgraded')
ROLLBACK = (log.UpdatedResource, 'rolled back')
DELETE = (log.DeletedResource, 'deleted')
RESET = (log.ResetResource, 'reset')
def HandleLRO(operation,
args,
instance_service,
release_track,
operation_type=OperationType.UPDATE):
"""Handles Long-running Operations for both cases of async.
Args:
operation: The operation to poll.
args: ArgParse instance containing user entered arguments.
instance_service: The service to get the resource after the long-running
operation completes.
release_track: base.ReleaseTrack object.
operation_type: Enum value of type OperationType indicating the kind of
operation to wait for.
Raises:
apitools.base.py.HttpError: if the request returns an HTTP error
Returns:
The Instance resource if synchronous, else the Operation Resource.
"""
logging_method = operation_type.value[0]
if args.async_:
logging_method(
util.GetOperationResource(operation.name, release_track),
kind='workbench instance {0}'.format(args.instance),
is_async=True)
return operation
else:
response = util.WaitForOperation(
operation,
'Waiting for operation on Instance [{}] to be {} with [{}]'.format(
args.instance, operation_type.value[1], operation.name),
service=instance_service,
release_track=release_track,
is_delete=(operation_type.value[1] == 'deleted'))
logging_method(
util.GetOperationResource(operation.name, release_track),
kind='workbench instance {0}'.format(args.instance),
is_async=False)
return response