# Copyright 2022 PerfKitBenchmarker Authors. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Module containing class for AWS's EMR services. Clusters can be created and deleted. """ import collections import dataclasses import gzip import json import logging import os from typing import Any, Dict from absl import flags from perfkitbenchmarker import disk from perfkitbenchmarker import dpb_constants from perfkitbenchmarker import dpb_service from perfkitbenchmarker import errors from perfkitbenchmarker import provider_info from perfkitbenchmarker import temp_dir from perfkitbenchmarker import vm_util from perfkitbenchmarker.providers.aws import aws_disk from perfkitbenchmarker.providers.aws import aws_dpb_emr_serverless_prices from perfkitbenchmarker.providers.aws import aws_network from perfkitbenchmarker.providers.aws import aws_virtual_machine from perfkitbenchmarker.providers.aws import s3 from perfkitbenchmarker.providers.aws import util FLAGS = flags.FLAGS flags.DEFINE_string( 'dpb_emr_release_label', None, 'DEPRECATED use dpb_service.version.' ) INVALID_STATES = ['TERMINATED_WITH_ERRORS', 'TERMINATED'] READY_CHECK_SLEEP = 30 READY_CHECK_TRIES = 60 READY_STATE = 'WAITING' JOB_WAIT_SLEEP = 30 EMR_TIMEOUT = 14400 disk_to_hdfs_map = { aws_disk.ST1: 'HDD (ST1)', aws_disk.GP2: 'SSD (GP2)', disk.LOCAL: 'Local SSD', } DATAPROC_TO_EMR_CONF_FILES = { # https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-configure-apps.html 'core': 'core-site', 'hdfs': 'hdfs-site', # https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-spark-configure.html 'spark': 'spark-defaults', } def _GetClusterConfiguration(cluster_properties: list[str]) -> str: """Return a JSON string containing dpb_cluster_properties.""" properties = collections.defaultdict(dict) for entry in cluster_properties: file, kv = entry.split(':') key, value = kv.split('=') if file not in DATAPROC_TO_EMR_CONF_FILES: raise errors.Config.InvalidValue( 'Unsupported EMR configuration file "{}". '.format(file) + 'Please add it to aws_dpb_emr.DATAPROC_TO_EMR_CONF_FILES.' ) properties[DATAPROC_TO_EMR_CONF_FILES[file]][key] = value json_conf = [] for file, props in properties.items(): json_conf.append({ # https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-configure-apps.html 'Classification': file, 'Properties': props, }) return json.dumps(json_conf) class EMRRetryableException(Exception): pass @dataclasses.dataclass(frozen=True) class _AwsDpbEmrServerlessJobRun: """Holds one EMR Serverless job run info, such as IDs and usage stats. Attributes: application_id: The application ID the job run belongs to. job_run_id: The job run's ID. region: Job run's region. memory_gb_hour: RAM GB * hour used. storage_gb_hour: Shuffle storage GB * hour used. vcpu_hour: vCPUs * hour used. """ application_id: str | None = None job_run_id: str | None = None region: str | None = None memory_gb_hour: float | None = None storage_gb_hour: float | None = None vcpu_hour: float | None = None def __bool__(self): """Returns if this represents a job run or is just a dummy placeholder.""" return None not in (self.application_id, self.job_run_id, self.region) def HasStats(self): """Returns whether there are stats collected for the job run.""" return None not in ( self.memory_gb_hour, self.storage_gb_hour, self.vcpu_hour, ) def ComputeJobRunCost(self) -> dpb_service.JobCosts: """Computes the cost of a run for region given this usage.""" if not self.HasStats(): return dpb_service.JobCosts() region_prices = aws_dpb_emr_serverless_prices.EMR_SERVERLESS_PRICES.get( self.region, {} ) memory_gb_hour_price = region_prices.get('memory_gb_hours') storage_gb_hour_price = region_prices.get('storage_gb_hours') vcpu_hour_price = region_prices.get('vcpu_hours') if ( memory_gb_hour_price is None or storage_gb_hour_price is None or vcpu_hour_price is None ): return dpb_service.JobCosts() vcpu_cost = self.vcpu_hour * vcpu_hour_price memory_cost = self.memory_gb_hour * memory_gb_hour_price storage_cost = self.storage_gb_hour * storage_gb_hour_price return dpb_service.JobCosts( total_cost=vcpu_cost + memory_cost + storage_cost, compute_cost=vcpu_cost, memory_cost=memory_cost, storage_cost=storage_cost, compute_units_used=self.vcpu_hour, memory_units_used=self.memory_gb_hour, storage_units_used=self.storage_gb_hour, compute_unit_cost=vcpu_hour_price, memory_unit_cost=memory_gb_hour_price, storage_unit_cost=storage_gb_hour_price, compute_unit_name='vCPU*hr', memory_unit_name='GB*hr', storage_unit_name='GB*hr', ) class AwsDpbEmr(dpb_service.BaseDpbService): """Object representing a AWS EMR cluster. Attributes: cluster_id: ID of the cluster. project: ID of the project in which the cluster is being launched. dpb_service_type: Set to 'emr'. cmd_prefix: Setting default prefix for the emr commands (region optional). network: Dedicated network for the EMR cluster storage_service: Region specific instance of S3 for bucket management. bucket_to_delete: Cluster associated bucket to be cleaned up. dpb_version: EMR version to use. """ CLOUD = provider_info.AWS SERVICE_TYPE = 'emr' SUPPORTS_NO_DYNALLOC = True def __init__(self, dpb_service_spec): super().__init__(dpb_service_spec) self.project = None self.cmd_prefix = list(util.AWS_PREFIX) if self.dpb_service_zone: self.region = util.GetRegionFromZone(self.dpb_service_zone) else: raise errors.Setup.InvalidSetupError( 'dpb_service_zone must be provided, for provisioning.' ) self.cmd_prefix += ['--region', self.region] self.network = aws_network.AwsNetwork.GetNetworkFromNetworkSpec( aws_network.AwsNetworkSpec(zone=self.dpb_service_zone) ) self.storage_service = s3.S3Service() self.storage_service.PrepareService(self.region) self.persistent_fs_prefix = 's3://' self.bucket_to_delete = None self._cluster_create_time: float | None = None self._cluster_ready_time: float | None = None self._cluster_delete_time: float | None = None if not self.GetDpbVersion(): raise errors.Setup.InvalidSetupError( 'dpb_service.version must be provided.' ) def GetDpbVersion(self) -> str | None: return FLAGS.dpb_emr_release_label or super().GetDpbVersion() def GetClusterCreateTime(self) -> float | None: """Returns the cluster creation time. On this implementation, the time returned is based on the timestamps reported by the EMR API (which is stored in the _cluster_create_time attribute). Returns: A float representing the creation time in seconds or None. """ return self._cluster_ready_time - self._cluster_create_time @property def security_group_id(self): """Returns the security group ID of this Cluster.""" return self.network.regional_network.vpc.default_security_group_id def _CreateDependencies(self): """Set up the ssh key.""" super()._CreateDependencies() aws_virtual_machine.AwsKeyFileManager.ImportKeyfile(self.region) def Create(self, restore: bool = False) -> None: """Overrides parent to register creation timeout as KNOWN_INTERMITTENT.""" try: super().Create() except vm_util.RetryError as e: raise errors.Resource.ProvisionTimeoutError from e def _Create(self): """Creates the cluster.""" name = 'pkb_' + FLAGS.run_uri # Set up ebs details if disk_spec is present in the config ebs_configuration = None if self.spec.worker_group.disk_spec: # Make sure nothing we are ignoring is included in the disk spec assert self.spec.worker_group.disk_spec.device_path is None assert self.spec.worker_group.disk_spec.disk_number is None assert self.spec.worker_group.disk_spec.provisioned_iops is None if self.spec.worker_group.disk_spec.disk_type != disk.LOCAL: ebs_configuration = { 'EbsBlockDeviceConfigs': [{ 'VolumeSpecification': { 'SizeInGB': self.spec.worker_group.disk_spec.disk_size, 'VolumeType': self.spec.worker_group.disk_spec.disk_type, }, 'VolumesPerInstance': self.spec.worker_group.disk_count, }] } # Create the specification for the master and the worker nodes instance_groups = [] core_instances = { 'InstanceCount': self.spec.worker_count, 'InstanceGroupType': 'CORE', 'InstanceType': self.spec.worker_group.vm_spec.machine_type, } if ebs_configuration: core_instances.update({'EbsConfiguration': ebs_configuration}) master_instance = { 'InstanceCount': 1, 'InstanceGroupType': 'MASTER', 'InstanceType': self.spec.worker_group.vm_spec.machine_type, } if ebs_configuration: master_instance.update({'EbsConfiguration': ebs_configuration}) instance_groups.append(core_instances) instance_groups.append(master_instance) # Spark SQL needs to access Hive cmd = self.cmd_prefix + [ 'emr', 'create-cluster', '--name', name, '--release-label', self.GetDpbVersion(), '--use-default-roles', '--instance-groups', json.dumps(instance_groups), '--application', 'Name=Spark', 'Name=Hadoop', 'Name=Hive', '--log-uri', self.base_dir, ] ec2_attributes = [ 'KeyName=' + aws_virtual_machine.AwsKeyFileManager.GetKeyNameForRun(), 'SubnetId=' + self.network.subnet.id, # Place all VMs in default security group for simplicity and speed of # provisioning 'EmrManagedMasterSecurityGroup=' + self.security_group_id, 'EmrManagedSlaveSecurityGroup=' + self.security_group_id, ] cmd += ['--ec2-attributes', ','.join(ec2_attributes)] if self.GetClusterProperties(): cmd += [ '--configurations', _GetClusterConfiguration(self.GetClusterProperties()), ] stdout, _, _ = vm_util.IssueCommand(cmd) result = json.loads(stdout) self.cluster_id = result['ClusterId'] logging.info('Cluster created with id %s', self.cluster_id) self._AddTags(util.MakeDefaultTags()) def _AddTags(self, tags_dict: dict[str, str]): tag_args = [f'{key}={value}' for key, value in tags_dict.items()] cmd = ( self.cmd_prefix + ['emr', 'add-tags', '--resource-id', self.cluster_id, '--tags'] + tag_args ) try: vm_util.IssueCommand(cmd) except errors.VmUtil.IssueCommandError as e: error_message = str(e) if 'ThrottlingException' in error_message: raise errors.Benchmarks.QuotaFailure.RateLimitExceededError( error_message ) from e raise def _Delete(self): if self.cluster_id: delete_cmd = self.cmd_prefix + [ 'emr', 'terminate-clusters', '--cluster-ids', self.cluster_id, ] vm_util.IssueCommand(delete_cmd, raise_on_failure=False) def _DeleteDependencies(self): super()._DeleteDependencies() aws_virtual_machine.AwsKeyFileManager.DeleteKeyfile(self.region) def _Exists(self): """Check to see whether the cluster exists.""" if not self.cluster_id: return False cmd = self.cmd_prefix + [ 'emr', 'describe-cluster', '--cluster-id', self.cluster_id, ] stdout, _, retcode = vm_util.IssueCommand(cmd, raise_on_failure=False) if retcode != 0: return False result = json.loads(stdout) end_datetime = ( result.get('Cluster', {}) .get('Status', {}) .get('Timeline', {}) .get('EndDateTime') ) if end_datetime is not None: self._cluster_delete_time = end_datetime if result['Cluster']['Status']['State'] in INVALID_STATES: return False else: return True def _IsReady(self): """Check to see if the cluster is ready.""" logging.info('Checking _Ready cluster: %s', self.cluster_id) cmd = self.cmd_prefix + [ 'emr', 'describe-cluster', '--cluster-id', self.cluster_id, ] stdout, _, _ = vm_util.IssueCommand(cmd) result = json.loads(stdout) # TODO(saksena): Handle error outcomees when spinning up emr clusters is_ready = result['Cluster']['Status']['State'] == READY_STATE if is_ready: self._cluster_create_time, self._cluster_ready_time = ( self._ParseClusterCreateTime(result) ) return is_ready @classmethod def _ParseClusterCreateTime( cls, data: dict[str, Any] ) -> tuple[float | None, float | None]: """Parses the cluster create & ready time from an API response dict.""" try: creation_ts = data['Cluster']['Status']['Timeline']['CreationDateTime'] ready_ts = data['Cluster']['Status']['Timeline']['ReadyDateTime'] return creation_ts, ready_ts except (LookupError, TypeError): return None, None def _GetCompletedJob(self, job_id): """See base class.""" cmd = self.cmd_prefix + [ 'emr', 'describe-step', '--cluster-id', self.cluster_id, '--step-id', job_id, ] stdout, stderr, retcode = vm_util.IssueCommand(cmd, raise_on_failure=False) if retcode: if 'ThrottlingException' in stderr: logging.warning( 'Rate limited while polling EMR step:\n%s\nRetrying.', stderr ) return None else: raise errors.VmUtil.IssueCommandError( f'Getting step status failed:\n{stderr}' ) result = json.loads(stdout) state = result['Step']['Status']['State'] if state == 'FAILED': raise dpb_service.JobSubmissionError( result['Step']['Status']['FailureDetails'] ) if state == 'COMPLETED': pending_time = result['Step']['Status']['Timeline']['CreationDateTime'] start_time = result['Step']['Status']['Timeline']['StartDateTime'] end_time = result['Step']['Status']['Timeline']['EndDateTime'] return dpb_service.JobResult( run_time=end_time - start_time, pending_time=start_time - pending_time ) def SubmitJob( self, jarfile=None, classname=None, pyspark_file=None, query_file=None, job_poll_interval=None, job_arguments=None, job_files=None, job_jars=None, job_py_files=None, job_stdout_file=None, job_type=None, properties=None, ): """See base class.""" if job_arguments: # Escape commas in arguments job_arguments = (arg.replace(',', '\\,') for arg in job_arguments) all_properties = self.GetJobProperties() all_properties.update(properties or {}) if job_type == 'hadoop': if not (jarfile or classname): raise ValueError('You must specify jarfile or classname.') if jarfile and classname: raise ValueError('You cannot specify both jarfile and classname.') arg_list = [] # Order is important if classname: # EMR does not support passing classnames as jobs. Instead manually # invoke `hadoop CLASSNAME` using command-runner.jar jarfile = 'command-runner.jar' arg_list = ['hadoop', classname] # Order is important arg_list += ['-D{}={}'.format(k, v) for k, v in all_properties.items()] if job_arguments: arg_list += job_arguments arg_spec = 'Args=[' + ','.join(arg_list) + ']' step_list = ['Jar=' + jarfile, arg_spec] elif job_type == dpb_constants.SPARK_JOB_TYPE: arg_list = [] if job_files: arg_list += ['--files', ','.join(job_files)] if job_py_files: arg_list += ['--py-files', ','.join(job_py_files)] if job_jars: arg_list += ['--jars', ','.join(job_jars)] for k, v in all_properties.items(): arg_list += ['--conf', '{}={}'.format(k, v)] # jarfile must be last before args arg_list += ['--class', classname, jarfile] if job_arguments: arg_list += job_arguments arg_spec = '[' + ','.join(arg_list) + ']' step_type_spec = 'Type=Spark' step_list = [step_type_spec, 'Args=' + arg_spec] elif job_type == dpb_constants.PYSPARK_JOB_TYPE: arg_list = [] if job_files: arg_list += ['--files', ','.join(job_files)] if job_jars: arg_list += ['--jars', ','.join(job_jars)] for k, v in all_properties.items(): arg_list += ['--conf', '{}={}'.format(k, v)] # pyspark_file must be last before args arg_list += [pyspark_file] if job_arguments: arg_list += job_arguments arg_spec = 'Args=[{}]'.format(','.join(arg_list)) step_list = ['Type=Spark', arg_spec] elif job_type == dpb_constants.SPARKSQL_JOB_TYPE: assert not job_arguments arg_list = [query_file] jar_spec = 'Jar="command-runner.jar"' for k, v in all_properties.items(): arg_list += ['--conf', '{}={}'.format(k, v)] arg_spec = 'Args=[spark-sql,-f,{}]'.format(','.join(arg_list)) step_list = [jar_spec, arg_spec] step_string = ','.join(step_list) step_cmd = self.cmd_prefix + [ 'emr', 'add-steps', '--cluster-id', self.cluster_id, '--steps', step_string, ] stdout, _, _ = vm_util.IssueCommand(step_cmd) result = json.loads(stdout) step_id = result['StepIds'][0] return self._WaitForJob(step_id, EMR_TIMEOUT, job_poll_interval) def DistributedCopy(self, source, destination, properties=None): """Method to copy data using a distributed job on the cluster.""" job_arguments = ['s3-dist-cp'] job_arguments.append('--src={}'.format(source)) job_arguments.append('--dest={}'.format(destination)) return self.SubmitJob( 'command-runner.jar', job_arguments=job_arguments, job_type=dpb_constants.HADOOP_JOB_TYPE, ) def GetHdfsType(self) -> str | None: """Gets human friendly disk type for metric metadata.""" try: return disk_to_hdfs_map[self.spec.worker_group.disk_spec.disk_type] except KeyError: raise errors.Setup.InvalidSetupError( f'Invalid disk_type={self.spec.worker_group.disk_spec.disk_type!r} in' ' spec.' ) from None def _FetchLogs(self, step_id: str) -> str | None: local_stdout_path = os.path.join( temp_dir.GetRunDirPath(), f'emr_{self.cluster_id}_{step_id}.stdout.gz', ) get_stdout_cmd = self.cmd_prefix + [ 's3', 'cp', os.path.join( self.base_dir, f'{self.cluster_id}/steps/{step_id}/stderr.gz' ), local_stdout_path, ] _, _, _ = vm_util.IssueCommand(get_stdout_cmd) with gzip.open(local_stdout_path, 'rt') as f: stdout = f.read() return stdout class AwsDpbEmrServerless( dpb_service.DpbServiceServerlessMixin, dpb_service.BaseDpbService ): """Resource that allows spawning EMR Serverless Jobs. Pre-initialization capacity is not supported yet. Docs: https://docs.aws.amazon.com/emr/latest/EMR-Serverless-UserGuide/emr-serverless.html """ CLOUD = provider_info.AWS SERVICE_TYPE = 'emr_serverless' def __init__(self, dpb_service_spec): # TODO(odiego): Refactor the AwsDpbEmr and AwsDpbEmrServerless into a # hierarchy or move common code to a parent class. super().__init__(dpb_service_spec) self.project = None self.cmd_prefix = list(util.AWS_PREFIX) if self.dpb_service_zone: self.region = util.GetRegionFromZone(self.dpb_service_zone) else: raise errors.Setup.InvalidSetupError( 'dpb_service_zone must be provided, for provisioning.' ) self.cmd_prefix += ['--region', self.region] self.storage_service = s3.S3Service() self.storage_service.PrepareService(self.region) self.persistent_fs_prefix = 's3://' self._cluster_create_time = None if not self.GetDpbVersion(): raise errors.Setup.InvalidSetupError( 'dpb_service.version must be provided. Versions follow the format: ' '"emr-x.y.z" and are listed at ' 'https://docs.aws.amazon.com/emr/latest/EMR-Serverless-UserGuide/' 'release-versions.html' ) self.role = FLAGS.aws_emr_serverless_role # Last job usage info self._job_run = _AwsDpbEmrServerlessJobRun() self._FillMetadata() def SubmitJob( self, jarfile=None, classname=None, pyspark_file=None, query_file=None, job_poll_interval=None, job_arguments=None, job_files=None, job_jars=None, job_py_files=None, job_stdout_file=None, job_type=None, properties=None, ): """See base class.""" assert job_type # Set vars according to job type. if job_type == dpb_constants.PYSPARK_JOB_TYPE: application_type = 'SPARK' spark_props = self.GetJobProperties() if job_py_files: spark_props['spark.submit.pyFiles'] = ','.join(job_py_files) job_driver_dict = { 'sparkSubmit': { 'entryPoint': pyspark_file, 'entryPointArguments': job_arguments, 'sparkSubmitParameters': ' '.join( f'--conf {prop}={val}' for prop, val in spark_props.items() ), } } else: raise NotImplementedError( f'Unsupported job type {job_type} for AWS EMR Serverless.' ) s3_monitoring_config = { 's3MonitoringConfiguration': { 'logUri': os.path.join(self.base_dir, 'logs') } } # Create the application. stdout, _, _ = vm_util.IssueCommand( self.cmd_prefix + [ 'emr-serverless', 'create-application', '--release-label', self.GetDpbVersion(), '--name', self.cluster_id, '--type', application_type, '--tags', json.dumps(util.MakeDefaultTags()), '--monitoring-configuration', json.dumps(s3_monitoring_config), ] ) result = json.loads(stdout) application_id = result['applicationId'] @vm_util.Retry( poll_interval=job_poll_interval, fuzz=0, retryable_exceptions=(EMRRetryableException,), ) def WaitTilApplicationReady(): result = self._GetApplication(application_id) if result['application']['state'] not in ('CREATED', 'STARTED'): raise EMRRetryableException( f'Application {application_id} not ready yet.' ) return result WaitTilApplicationReady() # Run the job. stdout, _, _ = vm_util.IssueCommand( self.cmd_prefix + [ 'emr-serverless', 'start-job-run', '--application-id', application_id, '--execution-role-arn', self.role, '--job-driver', json.dumps(job_driver_dict), ] ) result = json.loads(stdout) self._job_run = _AwsDpbEmrServerlessJobRun( application_id=result['applicationId'], job_run_id=result['jobRunId'], region=self.region, ) return self._WaitForJob(self._job_run, EMR_TIMEOUT, job_poll_interval) def CalculateLastJobCosts(self) -> dpb_service.JobCosts: @vm_util.Retry( fuzz=0, retryable_exceptions=(EMRRetryableException,), ) def WaitTilUsageMetricsAvailable(): self._CallGetJobRunApi(self._job_run) if not self._job_run.HasStats(): raise EMRRetryableException( 'Usage metrics not ready yet for EMR Serverless ' f'application_id={self._job_run.application_id!r} ' f'job_run_id={self._job_run.job_run_id!r}' ) if not self._job_run: return _AwsDpbEmrServerlessJobRun().ComputeJobRunCost() if not self._job_run.HasStats(): try: WaitTilUsageMetricsAvailable() except vm_util.TimeoutExceededRetryError: logging.warning('Timeout exceeded for retrieving usage metrics.') return self._job_run.ComputeJobRunCost() def GetJobProperties(self) -> Dict[str, str]: result = {'spark.dynamicAllocation.enabled': 'FALSE'} if self.spec.emr_serverless_core_count: result['spark.executor.cores'] = self.spec.emr_serverless_core_count result['spark.driver.cores'] = self.spec.emr_serverless_core_count if self.spec.emr_serverless_memory: result['spark.executor.memory'] = f'{self.spec.emr_serverless_memory}G' if self.spec.emr_serverless_executor_count: result['spark.executor.instances'] = ( self.spec.emr_serverless_executor_count ) if self.spec.worker_group.disk_spec.disk_size: result['spark.emr-serverless.driver.disk'] = ( f'{self.spec.worker_group.disk_spec.disk_size}G' ) result['spark.emr-serverless.executor.disk'] = ( f'{self.spec.worker_group.disk_spec.disk_size}G' ) result.update(super().GetJobProperties()) return result def _GetApplication(self, application_id): stdout, _, _ = vm_util.IssueCommand( self.cmd_prefix + [ 'emr-serverless', 'get-application', '--application-id', application_id, ] ) result = json.loads(stdout) return result def _ComputeJobRunCost( self, memory_gb_hour: float, storage_gb_hour: float, vcpu_hour: float ) -> dpb_service.JobCosts: region_prices = aws_dpb_emr_serverless_prices.EMR_SERVERLESS_PRICES.get( self.region, {} ) memory_gb_hour_price = region_prices.get('memory_gb_hours') storage_gb_hour_price = region_prices.get('storage_gb_hours') vcpu_hour_price = region_prices.get('vcpu_hours') if ( memory_gb_hour_price is None or storage_gb_hour_price is None or vcpu_hour_price is None ): return dpb_service.JobCosts() vcpu_cost = vcpu_hour * vcpu_hour_price memory_cost = memory_gb_hour * memory_gb_hour_price storage_cost = storage_gb_hour * storage_gb_hour_price return dpb_service.JobCosts( total_cost=vcpu_cost + memory_cost + storage_cost, compute_cost=vcpu_cost, memory_cost=memory_cost, storage_cost=storage_cost, compute_units_used=vcpu_hour, memory_units_used=memory_gb_hour, storage_units_used=storage_gb_hour, compute_unit_cost=vcpu_hour_price, memory_unit_cost=memory_gb_hour_price, storage_unit_cost=storage_gb_hour_price, compute_unit_name='vCPU*hr', memory_unit_name='GB*hr', storage_unit_name='GB*hr', ) def _GetCompletedJob(self, job_run): """See base class.""" return self._CallGetJobRunApi(job_run) def _CallGetJobRunApi( self, job_run: _AwsDpbEmrServerlessJobRun ) -> dpb_service.JobResult | None: """Performs EMR Serverless GetJobRun API call.""" cmd = self.cmd_prefix + [ 'emr-serverless', 'get-job-run', '--application-id', job_run.application_id, '--job-run-id', job_run.job_run_id, ] stdout, stderr, retcode = vm_util.IssueCommand(cmd, raise_on_failure=False) if retcode: if 'ThrottlingException' in stderr: logging.warning( 'Rate limited while polling EMR JobRun:\n%s\nRetrying.', stderr ) return None raise errors.VmUtil.IssueCommandError( f'Getting JobRun status failed:\n{stderr}' ) result = json.loads(stdout) state = result['jobRun']['state'] if state in ('FAILED', 'CANCELLED'): raise dpb_service.JobSubmissionError(result['jobRun'].get('stateDetails')) if state == 'SUCCESS': start_time = result['jobRun']['createdAt'] end_time = result['jobRun']['updatedAt'] self._job_run = self._ParseCostMetrics(result) return dpb_service.JobResult( run_time=end_time - start_time, fetch_output_fn=lambda: ( self._FetchLogs(job_run.application_id, job_run.job_run_id), None, ), ) def _FetchLogs(self, application_id: str, job_run_id: str) -> str | None: local_stdout_path = os.path.join( temp_dir.GetRunDirPath(), f'emrs8s_{application_id}_{job_run_id}.stdout.gz', ) get_stdout_cmd = self.cmd_prefix + [ 's3', 'cp', os.path.join( self.base_dir, f'logs/applications/{application_id}/jobs/{job_run_id}/' 'SPARK_DRIVER/stdout.gz', ), local_stdout_path, ] _, _, _ = vm_util.IssueCommand(get_stdout_cmd) with gzip.open(local_stdout_path, 'rt') as f: stdout = f.read() return stdout def _FillMetadata(self) -> None: """Gets a dict to initialize this DPB service instance's metadata.""" basic_data = self.metadata dpb_disk_size = self.spec.worker_group.disk_spec.disk_size or 'default' core_count = str(self.spec.emr_serverless_core_count) or 'default' cluster_shape = f'emr-serverless-{core_count}' cluster_size = str(self.spec.emr_serverless_executor_count) or 'default' self.metadata = { 'dpb_service': basic_data['dpb_service'], 'dpb_version': basic_data['dpb_version'], 'dpb_service_version': basic_data['dpb_service_version'], 'dpb_cluster_shape': cluster_shape, 'dpb_cluster_size': cluster_size, 'dpb_hdfs_type': basic_data['dpb_hdfs_type'], 'dpb_memory_per_node': self.spec.emr_serverless_memory or 'default', 'dpb_disk_size': dpb_disk_size, 'dpb_service_zone': basic_data['dpb_service_zone'], 'dpb_job_properties': basic_data['dpb_job_properties'], } def GetHdfsType(self) -> str | None: """Gets human friendly disk type for metric metadata.""" return 'default-disk' def _ParseCostMetrics( self, get_job_run_result: dict[Any, Any] ) -> _AwsDpbEmrServerlessJobRun: """Parses usage metrics from an EMR s8s GetJobRun API response.""" resource_utilization = get_job_run_result.get('jobRun', {}).get( 'totalResourceUtilization', {} ) return dataclasses.replace( self._job_run, memory_gb_hour=resource_utilization.get('memoryGBHour'), storage_gb_hour=resource_utilization.get('storageGBHour'), vcpu_hour=resource_utilization.get('vCPUHour'), )