def GetClusterConfig(args, dataproc, project_id, compute_resources, beta=False): """Get dataproc cluster configuration. Args: args: Arguments parsed from argparse.ArgParser. dataproc: Dataproc object that contains client, messages, and resources project_id: Dataproc project ID compute_resources: compute resource for cluster beta: use BETA only features Returns: cluster_config: Dataproc cluster configuration """ master_accelerator_type = None worker_accelerator_type = None master_accelerator_count = None worker_accelerator_count = None if beta: if args.master_accelerator: master_accelerator_type = args.master_accelerator['type'] master_accelerator_count = args.master_accelerator.get('count', 1) if args.worker_accelerator: worker_accelerator_type = args.worker_accelerator['type'] worker_accelerator_count = args.worker_accelerator.get('count', 1) # Resolve non-zonal GCE resources # We will let the server resolve short names of zonal resources because # if auto zone is requested, we will not know the zone before sending the # request image_ref = args.image and compute_resources.Parse( args.image, params={'project': project_id}, collection='compute.images') network_ref = args.network and compute_resources.Parse( args.network, params={'project': project_id}, collection='compute.networks') subnetwork_ref = args.subnet and compute_resources.Parse( args.subnet, params={ 'project': project_id, 'region': properties.VALUES.compute.region.GetOrFail, }, collection='compute.subnetworks') timeout_str = str(args.initialization_action_timeout) + 's' init_actions = [ dataproc.messages.NodeInitializationAction( executableFile=exe, executionTimeout=timeout_str) for exe in (args.initialization_actions or []) ] # Increase the client timeout for each initialization action. args.timeout += args.initialization_action_timeout * len(init_actions) expanded_scopes = compute_helpers.ExpandScopeAliases(args.scopes) software_config = dataproc.messages.SoftwareConfig( imageVersion=args.image_version) master_boot_disk_size_gb = args.master_boot_disk_size_gb if args.master_boot_disk_size: master_boot_disk_size_gb = (api_utils.BytesToGb( args.master_boot_disk_size)) worker_boot_disk_size_gb = args.worker_boot_disk_size_gb if args.worker_boot_disk_size: worker_boot_disk_size_gb = (api_utils.BytesToGb( args.worker_boot_disk_size)) preemptible_worker_boot_disk_size_gb = (api_utils.BytesToGb( args.preemptible_worker_boot_disk_size)) if args.single_node or args.num_workers == 0: # Explicitly specifying --num-workers=0 gives you a single node cluster, # but if --num-workers is omitted, args.num_workers is None (not 0), and # this property will not be set args.properties[constants.ALLOW_ZERO_WORKERS_PROPERTY] = 'true' if args.properties: software_config.properties = encoding.DictToMessage( args.properties, dataproc.messages.SoftwareConfig.PropertiesValue) if beta: if args.components: software_config_cls = dataproc.messages.SoftwareConfig software_config.optionalComponents.extend( list( map( software_config_cls. OptionalComponentsValueListEntryValuesEnum, args.components))) gce_cluster_config = dataproc.messages.GceClusterConfig( networkUri=network_ref and network_ref.SelfLink(), subnetworkUri=subnetwork_ref and subnetwork_ref.SelfLink(), internalIpOnly=args.no_address, serviceAccount=args.service_account, serviceAccountScopes=expanded_scopes, zoneUri=properties.VALUES.compute.zone.GetOrFail()) if args.tags: gce_cluster_config.tags = args.tags if args.metadata: flat_metadata = dict( (k, v) for d in args.metadata for k, v in d.items()) gce_cluster_config.metadata = encoding.DictToMessage( flat_metadata, dataproc.messages.GceClusterConfig.MetadataValue) master_accelerators = [] if master_accelerator_type: master_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=master_accelerator_type, acceleratorCount=master_accelerator_count)) worker_accelerators = [] if worker_accelerator_type: worker_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=worker_accelerator_type, acceleratorCount=worker_accelerator_count)) cluster_config = dataproc.messages.ClusterConfig( configBucket=args.bucket, gceClusterConfig=gce_cluster_config, masterConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_masters, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.master_machine_type, accelerators=master_accelerators, diskConfig=GetDiskConfig(dataproc, args.master_boot_disk_type, master_boot_disk_size_gb, args.num_master_local_ssds)), workerConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_workers, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.worker_machine_type, accelerators=worker_accelerators, diskConfig=GetDiskConfig( dataproc, args.worker_boot_disk_type, worker_boot_disk_size_gb, args.num_worker_local_ssds, )), initializationActions=init_actions, softwareConfig=software_config, ) if beta: cluster_config.masterConfig.minCpuPlatform = args.master_min_cpu_platform cluster_config.workerConfig.minCpuPlatform = args.worker_min_cpu_platform if beta: lifecycle_config = dataproc.messages.LifecycleConfig() changed_config = False if args.max_age is not None: lifecycle_config.autoDeleteTtl = str(args.max_age) + 's' changed_config = True if args.expiration_time is not None: lifecycle_config.autoDeleteTime = times.FormatDateTime( args.expiration_time) changed_config = True if args.max_idle is not None: lifecycle_config.idleDeleteTtl = str(args.max_idle) + 's' changed_config = True if changed_config: cluster_config.lifecycleConfig = lifecycle_config if beta and hasattr(args.CONCEPTS, 'kms_key'): kms_ref = args.CONCEPTS.kms_key.Parse() if kms_ref: encryption_config = dataproc.messages.EncryptionConfig() encryption_config.gcePdKmsKeyName = kms_ref.RelativeName() cluster_config.encryptionConfig = encryption_config else: # Did user use any gce-pd-kms-key flags? for keyword in [ 'gce-pd-kms-key', 'gce-pd-kms-key-project', 'gce-pd-kms-key-location', 'gce-pd-kms-key-keyring' ]: if getattr(args, keyword.replace('-', '_'), None): raise exceptions.ArgumentError( '--gce-pd-kms-key was not fully specified.') # Secondary worker group is optional. However, users may specify # future pVMs configuration at creation time. if (args.num_preemptible_workers is not None or preemptible_worker_boot_disk_size_gb is not None or args.preemptible_worker_boot_disk_type is not None or (beta and args.worker_min_cpu_platform is not None)): cluster_config.secondaryWorkerConfig = ( dataproc.messages.InstanceGroupConfig( numInstances=args.num_preemptible_workers, diskConfig=GetDiskConfig( dataproc, args.preemptible_worker_boot_disk_type, preemptible_worker_boot_disk_size_gb, None, ))) if beta and args.worker_min_cpu_platform: cluster_config.secondaryWorkerConfig.minCpuPlatform = ( args.worker_min_cpu_platform) return cluster_config
def GetClusterConfig(args, dataproc, project_id, compute_resources, beta=False, include_deprecated=True, include_ttl_config=False, include_gke_platform_args=False): """Get dataproc cluster configuration. Args: args: Arguments parsed from argparse.ArgParser. dataproc: Dataproc object that contains client, messages, and resources project_id: Dataproc project ID compute_resources: compute resource for cluster beta: use BETA only features include_deprecated: whether to include deprecated args include_ttl_config: whether to include Scheduled Delete(TTL) args include_gke_platform_args: whether to include GKE-based cluster args Returns: cluster_config: Dataproc cluster configuration """ master_accelerator_type = None worker_accelerator_type = None secondary_worker_accelerator_type = None if args.master_accelerator: master_accelerator_type = args.master_accelerator['type'] master_accelerator_count = args.master_accelerator.get('count', 1) if args.worker_accelerator: worker_accelerator_type = args.worker_accelerator['type'] worker_accelerator_count = args.worker_accelerator.get('count', 1) secondary_worker_accelerator = _FirstNonNone( args.secondary_worker_accelerator, args.preemptible_worker_accelerator) if secondary_worker_accelerator: secondary_worker_accelerator_type = secondary_worker_accelerator[ 'type'] secondary_worker_accelerator_count = secondary_worker_accelerator.get( 'count', 1) # Resolve non-zonal GCE resources # We will let the server resolve short names of zonal resources because # if auto zone is requested, we will not know the zone before sending the # request image_ref = args.image and compute_resources.Parse( args.image, params={'project': project_id}, collection='compute.images') network_ref = args.network and compute_resources.Parse( args.network, params={'project': project_id}, collection='compute.networks') subnetwork_ref = args.subnet and compute_resources.Parse( args.subnet, params={ 'project': project_id, 'region': properties.VALUES.compute.region.GetOrFail, }, collection='compute.subnetworks') timeout_str = six.text_type(args.initialization_action_timeout) + 's' init_actions = [ dataproc.messages.NodeInitializationAction( executableFile=exe, executionTimeout=timeout_str) for exe in (args.initialization_actions or []) ] # Increase the client timeout for each initialization action. args.timeout += args.initialization_action_timeout * len(init_actions) expanded_scopes = compute_helpers.ExpandScopeAliases(args.scopes) software_config = dataproc.messages.SoftwareConfig( imageVersion=args.image_version) if include_deprecated: master_boot_disk_size_gb = args.master_boot_disk_size_gb else: master_boot_disk_size_gb = None if args.master_boot_disk_size: master_boot_disk_size_gb = (api_utils.BytesToGb( args.master_boot_disk_size)) if include_deprecated: worker_boot_disk_size_gb = args.worker_boot_disk_size_gb else: worker_boot_disk_size_gb = None if args.worker_boot_disk_size: worker_boot_disk_size_gb = (api_utils.BytesToGb( args.worker_boot_disk_size)) secondary_worker_boot_disk_size_gb = (api_utils.BytesToGb( _FirstNonNone(args.secondary_worker_boot_disk_size, args.preemptible_worker_boot_disk_size))) if args.single_node or args.num_workers == 0: # Explicitly specifying --num-workers=0 gives you a single node cluster, # but if --num-workers is omitted, args.num_workers is None (not 0), and # this property will not be set args.properties[constants.ALLOW_ZERO_WORKERS_PROPERTY] = 'true' if args.properties: software_config.properties = encoding.DictToAdditionalPropertyMessage( args.properties, dataproc.messages.SoftwareConfig.PropertiesValue, sort_items=True) if args.components: software_config_cls = dataproc.messages.SoftwareConfig software_config.optionalComponents.extend( list( map( software_config_cls. OptionalComponentsValueListEntryValuesEnum, args.components))) gce_cluster_config = dataproc.messages.GceClusterConfig( networkUri=network_ref and network_ref.SelfLink(), subnetworkUri=subnetwork_ref and subnetwork_ref.SelfLink(), internalIpOnly=args.no_address, serviceAccount=args.service_account, serviceAccountScopes=expanded_scopes, zoneUri=properties.VALUES.compute.zone.GetOrFail()) reservation_affinity = GetReservationAffinity(args, dataproc) gce_cluster_config.reservationAffinity = reservation_affinity if args.tags: gce_cluster_config.tags = args.tags if args.metadata: flat_metadata = collections.OrderedDict([(k, v) for d in args.metadata for k, v in d.items()]) gce_cluster_config.metadata = encoding.DictToAdditionalPropertyMessage( flat_metadata, dataproc.messages.GceClusterConfig.MetadataValue) master_accelerators = [] if master_accelerator_type: master_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=master_accelerator_type, acceleratorCount=master_accelerator_count)) worker_accelerators = [] if worker_accelerator_type: worker_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=worker_accelerator_type, acceleratorCount=worker_accelerator_count)) secondary_worker_accelerators = [] if secondary_worker_accelerator_type: secondary_worker_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=secondary_worker_accelerator_type, acceleratorCount=secondary_worker_accelerator_count)) cluster_config = dataproc.messages.ClusterConfig( configBucket=args.bucket, gceClusterConfig=gce_cluster_config, masterConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_masters, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.master_machine_type, accelerators=master_accelerators, diskConfig=GetDiskConfig(dataproc, args.master_boot_disk_type, master_boot_disk_size_gb, args.num_master_local_ssds), minCpuPlatform=args.master_min_cpu_platform), workerConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_workers, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.worker_machine_type, accelerators=worker_accelerators, diskConfig=GetDiskConfig( dataproc, args.worker_boot_disk_type, worker_boot_disk_size_gb, args.num_worker_local_ssds, ), minCpuPlatform=args.worker_min_cpu_platform), initializationActions=init_actions, softwareConfig=software_config, ) if args.kerberos_config_file or args.kerberos_root_principal_password_uri: cluster_config.securityConfig = dataproc.messages.SecurityConfig() if args.kerberos_config_file: cluster_config.securityConfig.kerberosConfig = ParseKerberosConfigFile( dataproc, args.kerberos_config_file) else: kerberos_config = dataproc.messages.KerberosConfig() kerberos_config.enableKerberos = True if args.kerberos_root_principal_password_uri: kerberos_config.rootPrincipalPasswordUri = \ args.kerberos_root_principal_password_uri kerberos_kms_ref = args.CONCEPTS.kerberos_kms_key.Parse() kerberos_config.kmsKeyUri = kerberos_kms_ref.RelativeName() cluster_config.securityConfig.kerberosConfig = kerberos_config if args.autoscaling_policy: cluster_config.autoscalingConfig = dataproc.messages.AutoscalingConfig( policyUri=args.CONCEPTS.autoscaling_policy.Parse().RelativeName()) if include_ttl_config: lifecycle_config = dataproc.messages.LifecycleConfig() changed_config = False if args.max_age is not None: lifecycle_config.autoDeleteTtl = six.text_type(args.max_age) + 's' changed_config = True if args.expiration_time is not None: lifecycle_config.autoDeleteTime = times.FormatDateTime( args.expiration_time) changed_config = True if args.max_idle is not None: lifecycle_config.idleDeleteTtl = six.text_type(args.max_idle) + 's' changed_config = True if changed_config: cluster_config.lifecycleConfig = lifecycle_config if hasattr(args.CONCEPTS, 'kms_key'): kms_ref = args.CONCEPTS.kms_key.Parse() if kms_ref: encryption_config = dataproc.messages.EncryptionConfig() encryption_config.gcePdKmsKeyName = kms_ref.RelativeName() cluster_config.encryptionConfig = encryption_config else: # Did user use any gce-pd-kms-key flags? for keyword in [ 'gce-pd-kms-key', 'gce-pd-kms-key-project', 'gce-pd-kms-key-location', 'gce-pd-kms-key-keyring' ]: if getattr(args, keyword.replace('-', '_'), None): raise exceptions.ArgumentError( '--gce-pd-kms-key was not fully specified.') # Secondary worker group is optional. However, users may specify # future pVMs configuration at creation time. num_secondary_workers = _FirstNonNone(args.num_secondary_workers, args.num_preemptible_workers) secondary_worker_boot_disk_type = _FirstNonNone( args.secondary_worker_boot_disk_type, args.preemptible_worker_boot_disk_type) num_secondary_worker_local_ssds = _FirstNonNone( args.num_secondary_worker_local_ssds, args.num_preemptible_worker_local_ssds) if (num_secondary_workers is not None or secondary_worker_boot_disk_size_gb is not None or secondary_worker_boot_disk_type is not None or num_secondary_worker_local_ssds is not None or args.worker_min_cpu_platform is not None or args.secondary_worker_type != 'unspecified'): cluster_config.secondaryWorkerConfig = ( dataproc.messages.InstanceGroupConfig( numInstances=num_secondary_workers, accelerators=secondary_worker_accelerators, diskConfig=GetDiskConfig( dataproc, secondary_worker_boot_disk_type, secondary_worker_boot_disk_size_gb, num_secondary_worker_local_ssds, ), minCpuPlatform=args.worker_min_cpu_platform, preemptibility=_GetType(dataproc, args.secondary_worker_type))) if include_gke_platform_args: if args.enable_component_gateway: cluster_config.endpointConfig = dataproc.messages.EndpointConfig( enableHttpPortAccess=args.enable_component_gateway) if args.gke_cluster is not None: location = args.zone or args.region target_gke_cluster = 'projects/{0}/locations/{1}/clusters/{2}'.format( project_id, location, args.gke_cluster) cluster_config.gkeClusterConfig = dataproc.messages.GkeClusterConfig( namespacedGkeDeploymentTarget=dataproc.messages. NamespacedGkeDeploymentTarget( targetGkeCluster=target_gke_cluster, clusterNamespace=args.gke_cluster_namespace)) cluster_config.gceClusterConfig = None cluster_config.masterConfig = None cluster_config.workerConfig = None cluster_config.secondaryWorkerConfig = None return cluster_config
def Run(self, args): self.ValidateArgs(args) client = self.context['dataproc_client'] messages = self.context['dataproc_messages'] cluster_ref = util.ParseCluster(args.name, self.context) compute_resources = compute_helpers.GetComputeResources( self.ReleaseTrack(), args.name) main_accelerator_type = None worker_accelerator_type = None main_accelerator_count = None worker_accelerator_count = None if self.ReleaseTrack() == base.ReleaseTrack.BETA: if args.main_accelerator: main_accelerator_type = args.main_accelerator['type'] main_accelerator_count = args.main_accelerator.get('count', 1) if args.worker_accelerator: worker_accelerator_type = args.worker_accelerator['type'] worker_accelerator_count = args.worker_accelerator.get( 'count', 1) # Resolve GCE resources zone_ref = compute_resources.Parse(None, collection='compute.zones') image_ref = args.image and compute_resources.Parse( args.image, collection='compute.images') main_machine_type_ref = (args.main_machine_type and compute_resources.Parse( args.main_machine_type, collection='compute.machineTypes')) worker_machine_type_ref = (args.worker_machine_type and compute_resources.Parse( args.worker_machine_type, collection='compute.machineTypes')) network_ref = args.network and compute_resources.Parse( args.network, collection='compute.networks') subnetwork_ref = args.subnet and compute_resources.Parse( args.subnet, collection='compute.subnetworks') main_accelerator_type_ref = ( main_accelerator_type and compute_resources.Parse( main_accelerator_type, collection='compute.acceleratorTypes')) worker_accelerator_type_ref = ( worker_accelerator_type and compute_resources.Parse(worker_accelerator_type, collection='compute.acceleratorTypes')) init_actions = [] timeout_str = str(args.initialization_action_timeout) + 's' if args.initialization_actions: init_actions = [ messages.NodeInitializationAction(executableFile=exe, executionTimeout=timeout_str) for exe in args.initialization_actions ] expanded_scopes = compute_helpers.ExpandScopeAliases(args.scopes) software_config = messages.SoftwareConfig( imageVersion=args.image_version) main_boot_disk_size_gb = args.main_boot_disk_size_gb if args.main_boot_disk_size: main_boot_disk_size_gb = (api_utils.BytesToGb( args.main_boot_disk_size)) worker_boot_disk_size_gb = args.worker_boot_disk_size_gb if args.worker_boot_disk_size: worker_boot_disk_size_gb = (api_utils.BytesToGb( args.worker_boot_disk_size)) preemptible_worker_boot_disk_size_gb = (api_utils.BytesToGb( args.preemptible_worker_boot_disk_size)) if args.single_node: args.properties[constants.ALLOW_ZERO_WORKERS_PROPERTY] = 'true' if args.properties: software_config.properties = encoding.DictToMessage( args.properties, messages.SoftwareConfig.PropertiesValue) gce_cluster_config = messages.GceClusterConfig( networkUri=network_ref and network_ref.SelfLink(), subnetworkUri=subnetwork_ref and subnetwork_ref.SelfLink(), serviceAccount=args.service_account, serviceAccountScopes=expanded_scopes, zoneUri=zone_ref and zone_ref.SelfLink()) if args.tags: gce_cluster_config.tags = args.tags if args.metadata: flat_metadata = dict( (k, v) for d in args.metadata for k, v in d.items()) gce_cluster_config.metadata = encoding.DictToMessage( flat_metadata, messages.GceClusterConfig.MetadataValue) main_accelerators = [] if main_accelerator_type: main_accelerators.append( messages.AcceleratorConfig( acceleratorTypeUri=main_accelerator_type_ref and main_accelerator_type_ref.SelfLink(), acceleratorCount=main_accelerator_count)) worker_accelerators = [] if worker_accelerator_type: worker_accelerators.append( messages.AcceleratorConfig( acceleratorTypeUri=worker_accelerator_type_ref and worker_accelerator_type_ref.SelfLink(), acceleratorCount=worker_accelerator_count)) cluster_config = messages.ClusterConfig( configBucket=args.bucket, gceClusterConfig=gce_cluster_config, mainConfig=messages.InstanceGroupConfig( numInstances=args.num_mains, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=main_machine_type_ref and main_machine_type_ref.SelfLink(), accelerators=main_accelerators, diskConfig=messages.DiskConfig( bootDiskSizeGb=main_boot_disk_size_gb, numLocalSsds=args.num_main_local_ssds, ), ), workerConfig=messages.InstanceGroupConfig( numInstances=args.num_workers, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=worker_machine_type_ref and worker_machine_type_ref.SelfLink(), accelerators=worker_accelerators, diskConfig=messages.DiskConfig( bootDiskSizeGb=worker_boot_disk_size_gb, numLocalSsds=args.num_worker_local_ssds, ), ), initializationActions=init_actions, softwareConfig=software_config, ) # Secondary worker group is optional. However, users may specify # future pVM disk size at creation time. if (args.num_preemptible_workers is not None or preemptible_worker_boot_disk_size_gb is not None): cluster_config.secondaryWorkerConfig = ( messages.InstanceGroupConfig( numInstances=args.num_preemptible_workers, diskConfig=messages.DiskConfig( bootDiskSizeGb=preemptible_worker_boot_disk_size_gb, )) ) cluster = messages.Cluster(config=cluster_config, clusterName=cluster_ref.clusterName, projectId=cluster_ref.projectId) self.ConfigureCluster(messages, args, cluster) operation = client.projects_regions_clusters.Create( messages.DataprocProjectsRegionsClustersCreateRequest( projectId=cluster_ref.projectId, region=cluster_ref.region, cluster=cluster)) if args. async: log.status.write('Creating [{0}] with operation [{1}].'.format( cluster_ref, operation.name)) return operation = util.WaitForOperation( operation, self.context, 'Waiting for cluster creation operation') get_request = messages.DataprocProjectsRegionsClustersGetRequest( projectId=cluster_ref.projectId, region=cluster_ref.region, clusterName=cluster_ref.clusterName) cluster = client.projects_regions_clusters.Get(get_request) if cluster.status.state == ( messages.ClusterStatus.StateValueValuesEnum.RUNNING): log.CreatedResource(cluster_ref) else: log.error('Create cluster failed!') if operation.details: log.error('Details:\n' + operation.details) return cluster
def Run(self, args): self.ValidateArgs(args) client = self.context['dataproc_client'] messages = self.context['dataproc_messages'] cluster_ref = util.ParseCluster(args.name, self.context) config_helper = compute_helpers.ConfigurationHelper.FromContext( self.context) compute_uris = config_helper.ResolveGceUris(args.name, args.image, args.master_machine_type, args.worker_machine_type, args.network, args.subnet) init_actions = [] timeout_str = str(args.initialization_action_timeout) + 's' if args.initialization_actions: init_actions = [ messages.NodeInitializationAction(executableFile=exe, executionTimeout=timeout_str) for exe in args.initialization_actions ] expanded_scopes = compute_helpers.ExpandScopeAliases(args.scopes) software_config = messages.SoftwareConfig( imageVersion=args.image_version) master_boot_disk_size_gb = args.master_boot_disk_size_gb if args.master_boot_disk_size: master_boot_disk_size_gb = (api_utils.BytesToGb( args.master_boot_disk_size)) worker_boot_disk_size_gb = args.worker_boot_disk_size_gb if args.worker_boot_disk_size: worker_boot_disk_size_gb = (api_utils.BytesToGb( args.worker_boot_disk_size)) preemptible_worker_boot_disk_size_gb = (api_utils.BytesToGb( args.preemptible_worker_boot_disk_size)) if args.properties: software_config.properties = encoding.DictToMessage( args.properties, messages.SoftwareConfig.PropertiesValue) gce_cluster_config = messages.GceClusterConfig( networkUri=compute_uris['network'], subnetworkUri=compute_uris['subnetwork'], serviceAccountScopes=expanded_scopes, zoneUri=compute_uris['zone']) if args.tags: gce_cluster_config.tags = args.tags if args.metadata: flat_metadata = dict( (k, v) for d in args.metadata for k, v in d.items()) gce_cluster_config.metadata = encoding.DictToMessage( flat_metadata, messages.GceClusterConfig.MetadataValue) cluster_config = messages.ClusterConfig( configBucket=args.bucket, gceClusterConfig=gce_cluster_config, masterConfig=messages.InstanceGroupConfig( numInstances=args.num_masters, imageUri=compute_uris['image'], machineTypeUri=compute_uris['master_machine_type'], diskConfig=messages.DiskConfig( bootDiskSizeGb=master_boot_disk_size_gb, numLocalSsds=args.num_master_local_ssds, ), ), workerConfig=messages.InstanceGroupConfig( numInstances=args.num_workers, imageUri=compute_uris['image'], machineTypeUri=compute_uris['worker_machine_type'], diskConfig=messages.DiskConfig( bootDiskSizeGb=worker_boot_disk_size_gb, numLocalSsds=args.num_worker_local_ssds, ), ), initializationActions=init_actions, softwareConfig=software_config, ) # Secondary worker group is optional. However, users may specify # future pVM disk size at creation time. if (args.num_preemptible_workers is not None or preemptible_worker_boot_disk_size_gb is not None): cluster_config.secondaryWorkerConfig = ( messages.InstanceGroupConfig( numInstances=args.num_preemptible_workers, diskConfig=messages.DiskConfig( bootDiskSizeGb=preemptible_worker_boot_disk_size_gb, )) ) cluster = messages.Cluster(config=cluster_config, clusterName=cluster_ref.clusterName, projectId=cluster_ref.projectId) operation = client.projects_regions_clusters.Create( messages.DataprocProjectsRegionsClustersCreateRequest( projectId=cluster_ref.projectId, region=cluster_ref.region, cluster=cluster)) if args. async: log.status.write('Creating [{0}] with operation [{1}].'.format( cluster_ref, operation.name)) return operation = util.WaitForOperation( operation, self.context, 'Waiting for cluster creation operation') cluster = client.projects_regions_clusters.Get(cluster_ref.Request()) if cluster.status.state == ( messages.ClusterStatus.StateValueValuesEnum.RUNNING): log.CreatedResource(cluster_ref) else: log.error('Create cluster failed!') if operation.details: log.error('Details:\n' + operation.details) return cluster
def GetClusterConfig(args, dataproc, project_id, compute_resources, use_accelerators=False, use_auto_delete_ttl=False, use_min_cpu_platform=False): """Get dataproc cluster configuration. Args: args: Arguments parsed from argparse.ArgParser. dataproc: Dataproc object that contains client, messages, and resources project_id: Dataproc project ID compute_resources: compute resource for cluster use_accelerators: use accelerators in BETA only. use_auto_delete_ttl: use to configure auto-delete/TTL in BETA only. use_min_cpu_platform: use to configure minimum CPU platform for instances. Returns: cluster_config: Dataproc cluster configuration """ master_accelerator_type = None worker_accelerator_type = None master_accelerator_count = None worker_accelerator_count = None if use_accelerators: if args.master_accelerator: master_accelerator_type = args.master_accelerator['type'] master_accelerator_count = args.master_accelerator.get('count', 1) if args.worker_accelerator: worker_accelerator_type = args.worker_accelerator['type'] worker_accelerator_count = args.worker_accelerator.get('count', 1) # Resolve non-zonal GCE resources # We will let the server resolve short names of zonal resources because # if auto zone is requested, we will not know the zone before sending the # request image_ref = args.image and compute_resources.Parse( args.image, params={'project': project_id}, collection='compute.images') network_ref = args.network and compute_resources.Parse( args.network, params={'project': project_id}, collection='compute.networks') subnetwork_ref = args.subnet and compute_resources.Parse( args.subnet, params={ 'project': project_id, 'region': properties.VALUES.compute.region.GetOrFail, }, collection='compute.subnetworks') timeout_str = str(args.initialization_action_timeout) + 's' init_actions = [ dataproc.messages.NodeInitializationAction( executableFile=exe, executionTimeout=timeout_str) for exe in (args.initialization_actions or []) ] # Increase the client timeout for each initialization action. args.timeout += args.initialization_action_timeout * len(init_actions) expanded_scopes = compute_helpers.ExpandScopeAliases(args.scopes) software_config = dataproc.messages.SoftwareConfig( imageVersion=args.image_version) master_boot_disk_size_gb = args.master_boot_disk_size_gb if args.master_boot_disk_size: master_boot_disk_size_gb = (api_utils.BytesToGb( args.master_boot_disk_size)) worker_boot_disk_size_gb = args.worker_boot_disk_size_gb if args.worker_boot_disk_size: worker_boot_disk_size_gb = (api_utils.BytesToGb( args.worker_boot_disk_size)) preemptible_worker_boot_disk_size_gb = (api_utils.BytesToGb( args.preemptible_worker_boot_disk_size)) if args.single_node or args.num_workers == 0: # Explicitly specifying --num-workers=0 gives you a single node cluster, # but if --num-workers is omitted, args.num_workers is None (not 0), and # this property will not be set args.properties[constants.ALLOW_ZERO_WORKERS_PROPERTY] = 'true' if args.properties: software_config.properties = encoding.DictToMessage( args.properties, dataproc.messages.SoftwareConfig.PropertiesValue) gce_cluster_config = dataproc.messages.GceClusterConfig( networkUri=network_ref and network_ref.SelfLink(), subnetworkUri=subnetwork_ref and subnetwork_ref.SelfLink(), internalIpOnly=args.no_address, serviceAccount=args.service_account, serviceAccountScopes=expanded_scopes, zoneUri=properties.VALUES.compute.zone.GetOrFail()) if args.tags: gce_cluster_config.tags = args.tags if args.metadata: flat_metadata = dict( (k, v) for d in args.metadata for k, v in d.items()) gce_cluster_config.metadata = encoding.DictToMessage( flat_metadata, dataproc.messages.GceClusterConfig.MetadataValue) master_accelerators = [] if master_accelerator_type: master_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=master_accelerator_type, acceleratorCount=master_accelerator_count)) worker_accelerators = [] if worker_accelerator_type: worker_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=worker_accelerator_type, acceleratorCount=worker_accelerator_count)) cluster_config = dataproc.messages.ClusterConfig( configBucket=args.bucket, gceClusterConfig=gce_cluster_config, masterConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_masters, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.master_machine_type, accelerators=master_accelerators, diskConfig=dataproc.messages.DiskConfig( bootDiskSizeGb=master_boot_disk_size_gb, numLocalSsds=args.num_master_local_ssds, ), ), workerConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_workers, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.worker_machine_type, accelerators=worker_accelerators, diskConfig=dataproc.messages.DiskConfig( bootDiskSizeGb=worker_boot_disk_size_gb, numLocalSsds=args.num_worker_local_ssds, ), ), initializationActions=init_actions, softwareConfig=software_config, ) if use_min_cpu_platform: cluster_config.masterConfig.minCpuPlatform = args.master_min_cpu_platform cluster_config.workerConfig.minCpuPlatform = args.worker_min_cpu_platform if use_auto_delete_ttl: lifecycle_config = dataproc.messages.LifecycleConfig() changed_config = False if args.max_age is not None: lifecycle_config.autoDeleteTtl = str(args.max_age) + 's' changed_config = True if args.expiration_time is not None: lifecycle_config.autoDeleteTime = times.FormatDateTime( args.expiration_time) changed_config = True if args.max_idle is not None: lifecycle_config.idleDeleteTtl = str(args.max_idle) + 's' changed_config = True if changed_config: cluster_config.lifecycleConfig = lifecycle_config # Secondary worker group is optional. However, users may specify # future pVM disk size at creation time. if (args.num_preemptible_workers is not None or preemptible_worker_boot_disk_size_gb is not None): cluster_config.secondaryWorkerConfig = ( dataproc.messages.InstanceGroupConfig( numInstances=args.num_preemptible_workers, diskConfig=dataproc.messages.DiskConfig( bootDiskSizeGb=preemptible_worker_boot_disk_size_gb, ))) if use_min_cpu_platform and args.worker_min_cpu_platform: secondary_worker_config = cluster_config.secondaryWorkerConfig secondary_worker_config.minCpuPlatform = args.worker_min_cpu_platform return cluster_config
def Run(self, args): self.ValidateArgs(args) dataproc = dp.Dataproc() cluster_ref = dataproc.ParseCluster(args.name) compute_resources = compute_helpers.GetComputeResources( self.ReleaseTrack(), args.name) master_accelerator_type = None worker_accelerator_type = None master_accelerator_count = None worker_accelerator_count = None if self.ReleaseTrack() == base.ReleaseTrack.BETA: if args.master_accelerator: master_accelerator_type = args.master_accelerator['type'] master_accelerator_count = args.master_accelerator.get('count', 1) if args.worker_accelerator: worker_accelerator_type = args.worker_accelerator['type'] worker_accelerator_count = args.worker_accelerator.get('count', 1) # Resolve non-zonal GCE resources # We will let the server resolve short names of zonal resources because # if auto zone is requested, we will not know the zone before sending the # request image_ref = args.image and compute_resources.Parse( args.image, params={'project': cluster_ref.projectId}, collection='compute.images') network_ref = args.network and compute_resources.Parse( args.network, params={'project': cluster_ref.projectId}, collection='compute.networks') subnetwork_ref = args.subnet and compute_resources.Parse( args.subnet, params={ 'project': cluster_ref.projectId, 'region': properties.VALUES.compute.region.GetOrFail, }, collection='compute.subnetworks') timeout_str = str(args.initialization_action_timeout) + 's' init_actions = [ dataproc.messages.NodeInitializationAction( executableFile=exe, executionTimeout=timeout_str) for exe in (args.initialization_actions or [])] # Increase the client timeout for each initialization action. args.timeout += args.initialization_action_timeout * len(init_actions) expanded_scopes = compute_helpers.ExpandScopeAliases(args.scopes) software_config = dataproc.messages.SoftwareConfig( imageVersion=args.image_version) master_boot_disk_size_gb = args.master_boot_disk_size_gb if args.master_boot_disk_size: master_boot_disk_size_gb = ( api_utils.BytesToGb(args.master_boot_disk_size)) worker_boot_disk_size_gb = args.worker_boot_disk_size_gb if args.worker_boot_disk_size: worker_boot_disk_size_gb = ( api_utils.BytesToGb(args.worker_boot_disk_size)) preemptible_worker_boot_disk_size_gb = ( api_utils.BytesToGb(args.preemptible_worker_boot_disk_size)) if args.single_node: args.properties[constants.ALLOW_ZERO_WORKERS_PROPERTY] = 'true' if args.properties: software_config.properties = encoding.DictToMessage( args.properties, dataproc.messages.SoftwareConfig.PropertiesValue) gce_cluster_config = dataproc.messages.GceClusterConfig( networkUri=network_ref and network_ref.SelfLink(), subnetworkUri=subnetwork_ref and subnetwork_ref.SelfLink(), internalIpOnly=args.no_address, serviceAccount=args.service_account, serviceAccountScopes=expanded_scopes, zoneUri=properties.VALUES.compute.zone.GetOrFail()) if args.tags: gce_cluster_config.tags = args.tags if args.metadata: flat_metadata = dict((k, v) for d in args.metadata for k, v in d.items()) gce_cluster_config.metadata = encoding.DictToMessage( flat_metadata, dataproc.messages.GceClusterConfig.MetadataValue) master_accelerators = [] if master_accelerator_type: master_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=master_accelerator_type, acceleratorCount=master_accelerator_count)) worker_accelerators = [] if worker_accelerator_type: worker_accelerators.append( dataproc.messages.AcceleratorConfig( acceleratorTypeUri=worker_accelerator_type, acceleratorCount=worker_accelerator_count)) cluster_config = dataproc.messages.ClusterConfig( configBucket=args.bucket, gceClusterConfig=gce_cluster_config, masterConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_masters, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.master_machine_type, accelerators=master_accelerators, diskConfig=dataproc.messages.DiskConfig( bootDiskSizeGb=master_boot_disk_size_gb, numLocalSsds=args.num_master_local_ssds,),), workerConfig=dataproc.messages.InstanceGroupConfig( numInstances=args.num_workers, imageUri=image_ref and image_ref.SelfLink(), machineTypeUri=args.worker_machine_type, accelerators=worker_accelerators, diskConfig=dataproc.messages.DiskConfig( bootDiskSizeGb=worker_boot_disk_size_gb, numLocalSsds=args.num_worker_local_ssds,),), initializationActions=init_actions, softwareConfig=software_config,) # Secondary worker group is optional. However, users may specify # future pVM disk size at creation time. if (args.num_preemptible_workers is not None or preemptible_worker_boot_disk_size_gb is not None): cluster_config.secondaryWorkerConfig = ( dataproc.messages.InstanceGroupConfig( numInstances=args.num_preemptible_workers, diskConfig=dataproc.messages.DiskConfig( bootDiskSizeGb=preemptible_worker_boot_disk_size_gb, ))) cluster = dataproc.messages.Cluster( config=cluster_config, clusterName=cluster_ref.clusterName, projectId=cluster_ref.projectId) self.ConfigureCluster(dataproc.messages, args, cluster) operation = dataproc.client.projects_regions_clusters.Create( dataproc.messages.DataprocProjectsRegionsClustersCreateRequest( projectId=cluster_ref.projectId, region=cluster_ref.region, cluster=cluster)) if args.async: log.status.write( 'Creating [{0}] with operation [{1}].'.format( cluster_ref, operation.name)) return operation = dataproc.WaitForOperation( operation, message='Waiting for cluster creation operation', timeout_s=args.timeout) get_request = dataproc.messages.DataprocProjectsRegionsClustersGetRequest( projectId=cluster_ref.projectId, region=cluster_ref.region, clusterName=cluster_ref.clusterName) cluster = dataproc.client.projects_regions_clusters.Get(get_request) if cluster.status.state == ( dataproc.messages.ClusterStatus.StateValueValuesEnum.RUNNING): zone_uri = cluster.config.gceClusterConfig.zoneUri zone_short_name = zone_uri.split('/')[-1] # Log the URL of the cluster log.CreatedResource( cluster_ref, # Also indicate which zone the cluster was placed in. This is helpful # if the server picked a zone (auto zone) details='Cluster placed in zone [{0}]'.format(zone_short_name)) else: log.error('Create cluster failed!') if operation.details: log.error('Details:\n' + operation.details) return cluster
def Run(self, args): client = self.context['dataproc_client'] messages = self.context['dataproc_messages'] cluster_ref = util.ParseCluster(args.name, self.context) config_helper = compute_helpers.ConfigurationHelper.FromContext( self.context) compute_uris = config_helper.ResolveGceUris(args.name, args.image, args.master_machine_type, args.worker_machine_type, args.network) init_actions = [] timeout_str = str(args.initialization_action_timeout) + 's' if args.initialization_actions: init_actions = [ messages.NodeInitializationAction(executableFile=exe, executionTimeout=timeout_str) for exe in args.initialization_actions ] expanded_scopes = compute_helpers.ExpandScopeAliases(args.scopes) cluster_config = messages.ClusterConfiguration( configurationBucket=args.bucket, gceClusterConfiguration=messages.GceClusterConfiguration( networkUri=compute_uris['network'], serviceAccountScopes=expanded_scopes, zoneUri=compute_uris['zone'], ), masterConfiguration=messages.InstanceGroupConfiguration( imageUri=compute_uris['image'], machineTypeUri=compute_uris['master_machine_type'], diskConfiguration=messages.DiskConfiguration( bootDiskSizeGb=args.master_boot_disk_size_gb, numLocalSsds=args.num_master_local_ssds, ), ), workerConfiguration=messages.InstanceGroupConfiguration( numInstances=args.num_workers, imageUri=compute_uris['image'], machineTypeUri=compute_uris['worker_machine_type'], diskConfiguration=messages.DiskConfiguration( bootDiskSizeGb=args.worker_boot_disk_size_gb, numLocalSsds=args.num_worker_local_ssds, ), ), initializationActions=init_actions, softwareConfiguration=messages.SoftwareConfiguration( imageVersion=args.image_version), ) # Secondary worker group is optional. if args.num_preemptible_workers is not None: cluster_config.secondaryWorkerConfiguration = ( messages.InstanceGroupConfiguration( numInstances=args.num_preemptible_workers)) cluster = messages.Cluster(configuration=cluster_config, clusterName=cluster_ref.clusterName, projectId=cluster_ref.projectId) operation = client.projects_clusters.Create(cluster) operation = util.WaitForOperation( operation, self.context, 'Waiting for cluster creation operation') cluster = client.projects_clusters.Get(cluster_ref.Request()) if cluster.status.state == ( messages.ClusterStatus.StateValueValuesEnum.RUNNING): log.CreatedResource(cluster_ref) else: log.error('Create cluster failed!') if operation.details: log.error('Details:\n' + operation.details) return cluster