def CreateSetMachineTypeRequest(args):
instance = GetInstanceResource(args).RelativeName()
set_machine_request = util.GetMessages().SetInstanceMachineTypeRequest(
machineType=args.machine_type)
return util.GetMessages(
).NotebooksProjectsLocationsInstancesSetMachineTypeRequest(
name=instance, setInstanceMachineTypeRequest=set_machine_request)
def CreateInstanceRegisterRequest(args):
instance = GetInstanceResource(args)
parent = util.GetLocationResource(instance.locationsId,
instance.projectsId).RelativeName()
register_request = util.GetMessages().RegisterInstanceRequest(
instanceId=instance.Name())
return util.GetMessages().NotebooksProjectsLocationsInstancesRegisterRequest(
parent=parent, registerInstanceRequest=register_request)
def CreateSetLabelsRequest(args):
instance = GetInstanceResource(args).RelativeName()
set_label_request = util.GetMessages().SetInstanceLabelsRequest()
labels_message = util.GetMessages().SetInstanceLabelsRequest.LabelsValue
set_label_request.labels = labels_message(additionalProperties=[
labels_message.AdditionalProperty(key=key, value=value)
for key, value in args.labels.items()
])
return util.GetMessages().NotebooksProjectsLocationsInstancesSetLabelsRequest(
name=instance, setInstanceLabelsRequest=set_label_request)
def Run(self, args):
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
instance_service = client.projects_locations_instances
if args.IsSpecified('accelerator_type') or args.IsSpecified(
'accelerator_core_count'):
operation = instance_service.SetAccelerator(
instance_util.CreateSetAcceleratorRequest(args, messages))
instance_util.HandleLRO(
operation,
args,
instance_service,
release_track,
operation_type=instance_util.OperationType.UPDATE)
if args.IsSpecified('labels'):
operation = instance_service.SetLabels(
instance_util.CreateSetLabelsRequest(args, messages))
instance_util.HandleLRO(
operation,
args,
instance_service,
release_track,
operation_type=instance_util.OperationType.UPDATE)
if args.IsSpecified('machine_type'):
operation = instance_service.SetMachineType(
instance_util.CreateSetMachineTypeRequest(args, messages))
instance_util.HandleLRO(
operation,
args,
instance_service,
release_track,
operation_type=instance_util.OperationType.UPDATE)
def CreateVmImageFromArgs(args):
vm_image = util.GetMessages().VmImage(project=args.vm_image_project)
if args.IsSpecified('vm_image_family'):
vm_image.imageFamily = args.vm_image_family
else:
vm_image.imageName = args.vm_image_name
return vm_image
def GetMetadataFromArgs():
if args.IsSpecified('metadata'):
metadata_message = util.GetMessages().Instance.MetadataValue
return metadata_message(additionalProperties=[
metadata_message.AdditionalProperty(key=key, value=value)
for key, value in args.metadata.items()
])
return None
def GetLabelsFromArgs():
if args.IsSpecified('labels'):
labels_message = util.GetMessages().Instance.LabelsValue
return labels_message(additionalProperties=[
labels_message.AdditionalProperty(key=key, value=value)
for key, value in args.labels.items()
])
return None
def Run(self, args):
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
instance_service = client.projects_locations_instances
result = instance_service.Get(
instance_util.CreateInstanceDescribeRequest(args, messages))
return result
def Run(self, args):
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
environment_service = client.projects_locations_environments
result = environment_service.Get(
env_util.CreateEnvironmentDescribeRequest(args, messages))
return result
def CreateSetAcceleratorRequest(args):
"""Create and return Accelerator update request."""
instance = GetInstanceResource(args).RelativeName()
set_acc_request = util.GetMessages().SetInstanceAcceleratorRequest()
accelerator_config = util.GetMessages().SetInstanceAcceleratorRequest
if args.IsSpecified('accelerator_core_count'):
set_acc_request.coreCount = args.accelerator_core_count
if args.IsSpecified('accelerator_type'):
type_enum = arg_utils.ChoiceEnumMapper(
arg_name='accelerator-type',
message_enum=accelerator_config.TypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
arg_utils.EnumNameToChoice(args.accelerator_type))
set_acc_request.type = type_enum
return util.GetMessages(
).NotebooksProjectsLocationsInstancesSetAcceleratorRequest(
name=instance, setInstanceAcceleratorRequest=set_acc_request)
def CreateContainerImageFromArgs():
if args.IsSpecified('environment'):
return GetContainerImageFromExistingEnvironment()
if args.IsSpecified('container_repository'):
container_image = util.GetMessages().ContainerImage(
repository=args.container_repository, tag=args.container_tag)
return container_image
return None
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
environment_service = client.projects_locations_environments
operation = environment_service.Create(
env_util.CreateEnvironmentCreateRequest(args, messages))
return env_util.HandleLRO(operation, args, environment_service,
release_track)
def GetDiskEncryption():
instance_message = util.GetMessages().Instance
type_enum = None
if args.IsSpecified('disk_encryption'):
type_enum = arg_utils.ChoiceEnumMapper(
arg_name='disk-encryption',
message_enum=instance_message.BootDiskTypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
arg_utils.EnumNameToChoice(args.disk_encryption))
return type_enum
def CreateAcceleratorConfigMessage():
accelerator_config = util.GetMessages().AcceleratorConfig
type_enum = None
if args.IsSpecified('accelerator_type'):
type_enum = arg_utils.ChoiceEnumMapper(
arg_name='accelerator-type',
message_enum=accelerator_config.TypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
arg_utils.EnumNameToChoice(args.accelerator_type))
return accelerator_config(
type=type_enum, coreCount=args.accelerator_core_count)
def CreateVmImageFromArgs():
"""Create VmImage Message from an environment or from args."""
if args.IsSpecified('environment'):
return GetVmImageFromExistingEnvironment()
if args.IsSpecified('vm_image_project'):
vm_image = util.GetMessages().VmImage(project=args.vm_image_project)
if args.IsSpecified('vm_image_family'):
vm_image.imageFamily = args.vm_image_family
else:
vm_image.imageName = args.vm_image_name
return vm_image
return None
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
location_service = client.projects_locations
return list_pager.YieldFromList(
location_service,
loc_util.CreateLocationListRequest(args, messages),
field='locations',
limit=args.limit,
batch_size_attribute='pageSize')
def Run(self, args):
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
instance_service = client.projects_locations_instances
operation = instance_service.Reset(
instance_util.CreateInstanceResetRequest(args, messages))
return instance_util.HandleLRO(
operation,
args,
instance_service,
release_track,
operation_type=instance_util.OperationType.RESET)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
instance_service = client.projects_locations_instances
operation = instance_service.Create(
instance_util.CreateInstanceCreateRequest(args, client, messages))
return instance_util.HandleLRO(
operation,
args,
instance_service,
release_track,
operation_type=instance_util.OperationType.CREATE)
def Run(self, args):
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
instance_service = client.projects_locations_instances
if args.IsSpecified('target_snapshot'):
operation = instance_service.Rollback(
instance_util.CreateInstanceRollbackRequest(args, messages))
return instance_util.HandleLRO(
operation,
args,
instance_service,
release_track,
operation_type=instance_util.OperationType.ROLLBACK)
def Run(self, args):
"""This is what gets called when the user runs this command."""
release_track = self.ReleaseTrack()
client = util.GetClient(release_track)
messages = util.GetMessages(release_track)
if (not args.IsSpecified('location')) and (
not properties.VALUES.notebooks.location.IsExplicitlySet()):
raise parser_errors.RequiredError(argument='--location')
instance_service = client.projects_locations_instances
return list_pager.YieldFromList(
instance_service,
instance_util.CreateInstanceListRequest(args, messages),
field='instances',
limit=args.limit,
batch_size_attribute='pageSize')
def CreateEnvironment(args):
"""Creates the Environment message for the create request.
Args:
args: Argparse object from Command.Run
Returns:
Instance of the Environment message.
"""
def CreateContainerImageFromArgs(args):
container_image = util.GetMessages().ContainerImage(
repository=args.container_repository, tag=args.container_tag)
return container_image
def CreateVmImageFromArgs(args):
vm_image = util.GetMessages().VmImage(project=args.vm_image_project)
if args.IsSpecified('vm_image_family'):
vm_image.imageFamily = args.vm_image_family
else:
vm_image.imageName = args.vm_image_name
return vm_image
if args.IsSpecified('vm_image_project'):
vm_image = CreateVmImageFromArgs(args)
else:
container_image = CreateContainerImageFromArgs(args)
environment = util.GetMessages().Environment(
name=args.environment,
description=args.description,
displayName=args.display_name,
postStartupScript=args.post_startup_script)
if args.IsSpecified('vm_image_project'):
environment.vmImage = vm_image
else:
environment.containerImage = container_image
return environment
def CreateInstanceDescribeRequest(args):
instance = GetInstanceResource(args).RelativeName()
return util.GetMessages().NotebooksProjectsLocationsInstancesGetRequest(
name=instance)
def CreateEnvironmentDescribeRequest(args):
environment = GetEnvironmentResource(args).RelativeName()
return util.GetMessages().NotebooksProjectsLocationsEnvironmentsGetRequest(
name=environment)
def CreateEnvironmentListRequest(args):
parent = util.GetParentFromArgs(args)
return util.GetMessages().NotebooksProjectsLocationsEnvironmentsListRequest(
parent=parent)
def CreateEnvironmentCreateRequest(args):
parent = util.GetParentForEnvironment(args)
environment = CreateEnvironment(args)
return util.GetMessages().NotebooksProjectsLocationsEnvironmentsCreateRequest(
parent=parent, environment=environment, environmentId=args.environment)
def CreateContainerImageFromArgs(args):
container_image = util.GetMessages().ContainerImage(
repository=args.container_repository, tag=args.container_tag)
return container_image
def CreateInstanceResetRequest(args):
instance = GetInstanceResource(args).RelativeName()
reset_request = util.GetMessages().ResetInstanceRequest()
return util.GetMessages().NotebooksProjectsLocationsInstancesResetRequest(
name=instance, resetInstanceRequest=reset_request)
def CreateInstanceStopRequest(args):
instance = GetInstanceResource(args).RelativeName()
stop_request = util.GetMessages().StopInstanceRequest()
return util.GetMessages().NotebooksProjectsLocationsInstancesStopRequest(
name=instance, stopInstanceRequest=stop_request)
def CreateLocationListRequest(args):
project_name = util.GetProjectResource(args.project).RelativeName()
return util.GetMessages().NotebooksProjectsLocationsListRequest(
name=project_name)
def CreateInstance(args):
"""Creates the Instance message for the create request.
Args:
args: Argparse object from Command.Run
Returns:
Instance of the Instance message.
"""
def GetContainerImageFromExistingEnvironment():
environment_service = util.GetClient().projects_locations_environments
result = environment_service.Get(
env_util.CreateEnvironmentDescribeRequest(args))
return result.containerImage
def GetVmImageFromExistingEnvironment():
environment_service = util.GetClient().projects_locations_environments
result = environment_service.Get(
env_util.CreateEnvironmentDescribeRequest(args))
return result.vmImage
def GetKmsRelativeName():
if args.IsSpecified('kms_key'):
return args.CONCEPTS.kms_key.Parse().RelativeName()
def GetNetworkRelativeName():
if args.IsSpecified('network'):
return args.CONCEPTS.network.Parse().RelativeName()
def GetSubnetRelativeName():
if args.IsSpecified('subnet'):
return args.CONCEPTS.subnet.Parse().RelativeName()
def CreateAcceleratorConfigMessage():
accelerator_config = util.GetMessages().AcceleratorConfig
type_enum = None
if args.IsSpecified('accelerator_type'):
type_enum = arg_utils.ChoiceEnumMapper(
arg_name='accelerator-type',
message_enum=accelerator_config.TypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
arg_utils.EnumNameToChoice(args.accelerator_type))
return accelerator_config(
type=type_enum, coreCount=args.accelerator_core_count)
def GetBootDisk():
type_enum = None
if args.IsSpecified('boot_disk_type'):
instance_message = util.GetMessages().Instance
type_enum = arg_utils.ChoiceEnumMapper(
arg_name='boot-disk-type',
message_enum=instance_message.BootDiskTypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
arg_utils.EnumNameToChoice(args.boot_disk_type))
return type_enum
def GetDiskEncryption():
instance_message = util.GetMessages().Instance
type_enum = None
if args.IsSpecified('disk_encryption'):
type_enum = arg_utils.ChoiceEnumMapper(
arg_name='disk-encryption',
message_enum=instance_message.BootDiskTypeValueValuesEnum,
include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
arg_utils.EnumNameToChoice(args.disk_encryption))
return type_enum
def CreateContainerImageFromArgs():
if args.IsSpecified('environment'):
return GetContainerImageFromExistingEnvironment()
if args.IsSpecified('container_repository'):
container_image = util.GetMessages().ContainerImage(
repository=args.container_repository, tag=args.container_tag)
return container_image
return None
def CreateVmImageFromArgs():
"""Create VmImage Message from an environment or from args."""
if args.IsSpecified('environment'):
return GetVmImageFromExistingEnvironment()
if args.IsSpecified('vm_image_project'):
vm_image = util.GetMessages().VmImage(project=args.vm_image_project)
if args.IsSpecified('vm_image_family'):
vm_image.imageFamily = args.vm_image_family
else:
vm_image.imageName = args.vm_image_name
return vm_image
return None
def GetInstanceOwnersFromArgs():
if args.IsSpecified('instance_owners'):
return [args.instance_owners]
return []
def GetLabelsFromArgs():
if args.IsSpecified('labels'):
labels_message = util.GetMessages().Instance.LabelsValue
return labels_message(additionalProperties=[
labels_message.AdditionalProperty(key=key, value=value)
for key, value in args.labels.items()
])
return None
def GetMetadataFromArgs():
if args.IsSpecified('metadata'):
metadata_message = util.GetMessages().Instance.MetadataValue
return metadata_message(additionalProperties=[
metadata_message.AdditionalProperty(key=key, value=value)
for key, value in args.metadata.items()
])
return None
instance = util.GetMessages().Instance(
name=args.instance,
postStartupScript=args.post_startup_script,
bootDiskSizeGb=args.boot_disk_size,
customGpuDriverPath=args.custom_gpu_driver_path,
instanceOwners=GetInstanceOwnersFromArgs(),
kmsKey=GetKmsRelativeName(),
machineType=args.machine_type,
network=GetNetworkRelativeName(),
noProxyAccess=args.no_proxy_access,
noPublicIp=args.no_public_ip,
serviceAccount=args.service_account,
subnet=GetSubnetRelativeName(),
vmImage=CreateVmImageFromArgs(),
acceleratorConfig=CreateAcceleratorConfigMessage(),
bootDiskType=GetBootDisk(),
containerImage=CreateContainerImageFromArgs(),
diskEncryption=GetDiskEncryption(),
labels=GetLabelsFromArgs(),
metadata=GetMetadataFromArgs(),
installGpuDriver=args.install_gpu_driver,
)
return instance