phi/aws/resource/emr/cluster.py

from typing import Optional, Any, Dict, List
from typing_extensions import Literal

from phi.aws.api_client import AwsApiClient
from phi.aws.resource.base import AwsResource
from phi.cli.console import print_info
from phi.utils.log import logger


class EmrCluster(AwsResource):
    """
    Reference:
    - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/emr.html
    """

    resource_type: Optional[str] = "EmrCluster"
    service_name: str = "emr"

    # Name of the cluster.
    name: str
    # The location in Amazon S3 to write the log files of the job flow.
    # If a value is not provided, logs are not created.
    log_uri: Optional[str] = None
    # The KMS key used for encrypting log files. If a value is not provided, the logs remain encrypted by AES-256.
    # This attribute is only available with Amazon EMR version 5.30.0 and later, excluding Amazon EMR 6.0.0.
    log_encryption_kms_key_id: Optional[str] = None
    # A JSON string for selecting additional features.
    additional_info: Optional[str] = None
    # The Amazon EMR release label, which determines the version of open-source application packages installed on the
    # cluster. Release labels are in the form emr-x.x.x,
    # where x.x.x is an Amazon EMR release version such as emr-5.14.0 .
    release_label: Optional[str] = None
    # A specification of the number and type of Amazon EC2 instances.
    instances: Optional[Dict[str, Any]] = None
    # A list of steps to run.
    steps: Optional[List[Dict[str, Any]]] = None
    # A list of bootstrap actions to run before Hadoop starts on the cluster nodes.
    bootstrap_actions: Optional[List[Dict[str, Any]]] = None
    # For Amazon EMR releases 3.x and 2.x. For Amazon EMR releases 4.x and later, use Applications.
    # A list of strings that indicates third-party software to use.
    supported_products: Optional[List[str]]
    new_supported_products: Optional[List[Dict[str, Any]]] = None
    # Applies to Amazon EMR releases 4.0 and later.
    # A case-insensitive list of applications for Amazon EMR to install and configure when launching the cluster.
    applications: Optional[List[Dict[str, Any]]] = None
    # For Amazon EMR releases 4.0 and later. The list of configurations supplied for the EMR cluster you are creating.
    configurations: Optional[List[Dict[str, Any]]] = None
    # Also called instance profile and EC2 role. An IAM role for an EMR cluster.
    # The EC2 instances of the cluster assume this role. The default role is EMR_EC2_DefaultRole.
    # In order to use the default role, you must have already created it using the CLI or console.
    job_flow_role: Optional[str] = None
    # he IAM role that Amazon EMR assumes in order to access Amazon Web Services resources on your behalf.
    service_role: Optional[str] = None
    # A list of tags to associate with a cluster and propagate to Amazon EC2 instances.
    tags: Optional[List[Dict[str, str]]] = None
    # The name of a security configuration to apply to the cluster.
    security_configuration: Optional[str] = None
    # An IAM role for automatic scaling policies. The default role is EMR_AutoScaling_DefaultRole.
    # The IAM role provides permissions that the automatic scaling feature requires to launch and terminate EC2
    # instances in an instance group.
    auto_scaling_role: Optional[str] = None
    scale_down_behavior: Optional[Literal["TERMINATE_AT_INSTANCE_HOUR", "TERMINATE_AT_TASK_COMPLETION"]] = None
    custom_ami_id: Optional[str] = None
    # The size, in GiB, of the Amazon EBS root device volume of the Linux AMI that is used for each EC2 instance.
    ebs_root_volume_size: Optional[int] = None
    repo_upgrade_on_boot: Optional[Literal["SECURITY", "NONE"]] = None
    # Attributes for Kerberos configuration when Kerberos authentication is enabled using a security configuration.
    kerberos_attributes: Optional[Dict[str, str]] = None
    # Specifies the number of steps that can be executed concurrently.
    # The default value is 1 . The maximum value is 256 .
    step_concurrency_level: Optional[int] = None
    # The specified managed scaling policy for an Amazon EMR cluster.
    managed_scaling_policy: Optional[Dict[str, Any]] = None
    placement_group_configs: Optional[List[Dict[str, Any]]] = None
    # The auto-termination policy defines the amount of idle time in seconds after which a cluster terminates.
    auto_termination_policy: Optional[Dict[str, int]] = None

    # provided by api on create
    # A unique identifier for the job flow.
    job_flow_id: Optional[str] = None
    # The Amazon Resource Name (ARN) of the cluster.
    cluster_arn: Optional[str] = None
    # ClusterSummary returned on read
    cluster_summary: Optional[Dict] = None

    def _create(self, aws_client: AwsApiClient) -> bool:
        """Creates the EmrCluster

        Args:
            aws_client: The AwsApiClient for the current cluster
        """

        print_info(f"Creating {self.get_resource_type()}: {self.get_resource_name()}")
        try:
            # create a dict of args which are not null, otherwise aws type validation fails
            not_null_args: Dict[str, Any] = {}

            if self.log_uri:
                not_null_args["LogUri"] = self.log_uri
            if self.log_encryption_kms_key_id:
                not_null_args["LogEncryptionKmsKeyId"] = self.log_encryption_kms_key_id
            if self.additional_info:
                not_null_args["AdditionalInfo"] = self.additional_info
            if self.release_label:
                not_null_args["ReleaseLabel"] = self.release_label
            if self.instances:
                not_null_args["Instances"] = self.instances
            if self.steps:
                not_null_args["Steps"] = self.steps
            if self.bootstrap_actions:
                not_null_args["BootstrapActions"] = self.bootstrap_actions
            if self.supported_products:
                not_null_args["SupportedProducts"] = self.supported_products
            if self.new_supported_products:
                not_null_args["NewSupportedProducts"] = self.new_supported_products
            if self.applications:
                not_null_args["Applications"] = self.applications
            if self.configurations:
                not_null_args["Configurations"] = self.configurations
            if self.job_flow_role:
                not_null_args["JobFlowRole"] = self.job_flow_role
            if self.service_role:
                not_null_args["ServiceRole"] = self.service_role
            if self.tags:
                not_null_args["Tags"] = self.tags
            if self.security_configuration:
                not_null_args["SecurityConfiguration"] = self.security_configuration
            if self.auto_scaling_role:
                not_null_args["AutoScalingRole"] = self.auto_scaling_role
            if self.scale_down_behavior:
                not_null_args["ScaleDownBehavior"] = self.scale_down_behavior
            if self.custom_ami_id:
                not_null_args["CustomAmiId"] = self.custom_ami_id
            if self.ebs_root_volume_size:
                not_null_args["EbsRootVolumeSize"] = self.ebs_root_volume_size
            if self.repo_upgrade_on_boot:
                not_null_args["RepoUpgradeOnBoot"] = self.repo_upgrade_on_boot
            if self.kerberos_attributes:
                not_null_args["KerberosAttributes"] = self.kerberos_attributes
            if self.step_concurrency_level:
                not_null_args["StepConcurrencyLevel"] = self.step_concurrency_level
            if self.managed_scaling_policy:
                not_null_args["ManagedScalingPolicy"] = self.managed_scaling_policy
            if self.placement_group_configs:
                not_null_args["PlacementGroupConfigs"] = self.placement_group_configs
            if self.auto_termination_policy:
                not_null_args["AutoTerminationPolicy"] = self.auto_termination_policy

            # Get the service_client
            service_client = self.get_service_client(aws_client)

            # Create EmrCluster
            create_response = service_client.run_job_flow(
                Name=self.name,
                **not_null_args,
            )
            logger.debug(f"create_response type: {type(create_response)}")
            logger.debug(f"create_response: {create_response}")

            self.job_flow_id = create_response.get("JobFlowId", None)
            self.cluster_arn = create_response.get("ClusterArn", None)
            self.active_resource = create_response
            if self.active_resource is not None:
                print_info(f"{self.get_resource_type()}: {self.get_resource_name()} created")
                logger.debug(f"JobFlowId: {self.job_flow_id}")
                logger.debug(f"ClusterArn: {self.cluster_arn}")
                return True
        except Exception as e:
            logger.error(f"{self.get_resource_type()} could not be created.")
            logger.error(e)
        return False

    def post_create(self, aws_client: AwsApiClient) -> bool:
        ## Wait for Cluster to be created
        if self.wait_for_create:
            try:
                print_info("Waiting for EmrCluster to be active.")
                if self.job_flow_id is not None:
                    waiter = self.get_service_client(aws_client).get_waiter("cluster_running")
                    waiter.wait(
                        ClusterId=self.job_flow_id,
                        WaiterConfig={
                            "Delay": self.waiter_delay,
                            "MaxAttempts": self.waiter_max_attempts,
                        },
                    )
                else:
                    logger.warning("Skipping waiter, No ClusterId found")
            except Exception as e:
                logger.error("Waiter failed.")
                logger.error(e)
        return True

    def _read(self, aws_client: AwsApiClient) -> Optional[Any]:
        """Returns the EmrCluster

        Args:
            aws_client: The AwsApiClient for the current cluster
        """
        from botocore.exceptions import ClientError

        logger.debug(f"Reading {self.get_resource_type()}: {self.get_resource_name()}")
        try:
            service_client = self.get_service_client(aws_client)
            list_response = service_client.list_clusters()
            # logger.debug(f"list_response type: {type(list_response)}")
            # logger.debug(f"list_response: {list_response}")

            cluster_summary_list = list_response.get("Clusters", None)
            if cluster_summary_list is not None and isinstance(cluster_summary_list, list):
                for _cluster_summary in cluster_summary_list:
                    cluster_name = _cluster_summary.get("Name", None)
                    if cluster_name == self.name:
                        self.active_resource = _cluster_summary
                        break

            if self.active_resource is None:
                logger.debug(f"No {self.get_resource_type()} found")
                return None

            # logger.debug(f"EmrCluster: {self.active_resource}")
            self.job_flow_id = self.active_resource.get("Id", None)
            self.cluster_arn = self.active_resource.get("ClusterArn", None)
        except ClientError as ce:
            logger.debug(f"ClientError: {ce}")
        except Exception as e:
            logger.error(f"Error reading {self.get_resource_type()}.")
            logger.error(e)
        return self.active_resource

    def _delete(self, aws_client: AwsApiClient) -> bool:
        """Deletes the EmrCluster

        Args:
            aws_client: The AwsApiClient for the current cluster
        """

        print_info(f"Deleting {self.get_resource_type()}: {self.get_resource_name()}")
        try:
            # populate self.job_flow_id
            self._read(aws_client)

            service_client = self.get_service_client(aws_client)
            self.active_resource = None

            if self.job_flow_id:
                service_client.terminate_job_flows(JobFlowIds=[self.job_flow_id])
                print_info(f"{self.get_resource_type()}: {self.get_resource_name()} deleted")
            else:
                logger.error("Could not find cluster id")
            return True
        except Exception as e:
            logger.error(f"{self.get_resource_type()} could not be deleted.")
            logger.error("Please try again or delete resources manually.")
            logger.error(e)
        return False