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
