1.emrspark_lib.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `emrspark_lib` plugin creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from botocore.exceptions import ClientError\n",
    "import subprocess\n",
    "import json\n",
    "from pprint import pprint\n",
    "import requests\n",
    "import configparser\n",
    "import time\n",
    "\n",
    "import logging\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.ERROR)\n",
    "\n",
    "\n",
    "config = configparser.ConfigParser()\n",
    "config.read('airflow/config.cfg')\n",
    "\n",
    "REGION_NAME = config['AWS']['REGION_NAME']\n",
    "CLUSTER_NAME = config['AWS']['CLUSTER_NAME']\n",
    "\n",
    "# When empty, use the first available VPC\n",
    "VPC_ID = config['AWS']['VPC_ID']\n",
    "# When empty, use the first available subnet\n",
    "# NOTE: Subnet must have an internet gateway within its routes.\n",
    "SUBNET_ID = config['AWS']['SUBNET_ID']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_boto_clients(region_name, config=None):\n",
    "    # If access and secret keys are empty, use the one stored by the OS.\n",
    "    if config != None and config['AWS']['AWS_ACCESS_KEY_ID'] != '' and config['AWS']['AWS_SECRET_ACCESS_KEY'] != '':    \n",
    "        ec2 = boto3.client('ec2', region_name=region_name,\n",
    "                           aws_access_key_id=config['AWS']['AWS_ACCESS_KEY_ID'],\n",
    "                           aws_secret_access_key=config['AWS']['AWS_SECRET_ACCESS_KEY']\n",
    "                          )\n",
    "        emr = boto3.client('emr', region_name=region_name,\n",
    "                           aws_access_key_id=config['AWS']['AWS_ACCESS_KEY_ID'],\n",
    "                           aws_secret_access_key=config['AWS']['AWS_SECRET_ACCESS_KEY']\n",
    "                          )\n",
    "        iam = boto3.client('iam', region_name=region_name,\n",
    "                           aws_access_key_id=config['AWS']['AWS_ACCESS_KEY_ID'],\n",
    "                           aws_secret_access_key=config['AWS']['AWS_SECRET_ACCESS_KEY']\n",
    "                          )\n",
    "    else:\n",
    "        ec2 = boto3.client('ec2', region_name=region_name)\n",
    "        emr = boto3.client('emr', region_name=region_name)\n",
    "        iam = boto3.client('iam', region_name=region_name)\n",
    "    return (ec2, emr, iam)\n",
    "\n",
    "\n",
    "def get_first_available_vpc(ec2_client):\n",
    "    return ec2_client.describe_vpcs().get('Vpcs', [{}])[0].get('VpcId', '')\n",
    "\n",
    "\n",
    "def get_first_available_subnet(ec2_client, vpc_id):\n",
    "    return ec2_client.describe_subnets(Filters=[{'Name': 'vpc-id', 'Values': [vpc_id]}, {'Name': 'state', 'Values': ['available']}])['Subnets'][0].get('SubnetId', '')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ec2, emr, iam = get_boto_clients(REGION_NAME, config=config)\n",
    "\n",
    "if VPC_ID == '':\n",
    "    VPC_ID = get_first_available_vpc(ec2)\n",
    "\n",
    "if SUBNET_ID == '':\n",
    "    SUBNET_ID = get_first_available_subnet(ec2, VPC_ID)\n",
    "    \n",
    "# emrlib.create_cluster()\n",
    "print('vpc:', VPC_ID)\n",
    "print('subnet:', SUBNET_ID)\n",
    "print('region:', ec2.meta.region_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing getting default ip address.\n",
    "ip = requests.get('https://api.ipify.org').text\n",
    "print('My public IP address is:', ip)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Security Group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_security_group(ec2_client, name, desc, vpc_id, ip=None):\n",
    "    \"\"\" Create a security group\n",
    "    Args:\n",
    "        - ec2_client (boto3.EC2.Client): EC2 client object.\n",
    "        - name (string): Name of Security Group\n",
    "        - desc (string): Description of Security Group\n",
    "        - vpc_id (string): Name of VPC. If empty, use the first available VPC\n",
    "        - ip (string): The IP address of this machine. Only this machine can connect to the cluster.\n",
    "                       If empty, use https://api.ipify.org service to get public IP address.\n",
    "    Return:\n",
    "    \n",
    "        dict: {\n",
    "            'KeyFingerprint': 'string',\n",
    "            'KeyMaterial': 'string',\n",
    "            'KeyName': 'string',\n",
    "            'KeyPairId': 'string'\n",
    "        }\n",
    "    \"\"\"\n",
    "    region = ec2_client.meta.region_name\n",
    "    security_group_id = None\n",
    "    \n",
    "    try:\n",
    "        # Do not create if we found an existing Security Group\n",
    "        response = ec2_client.describe_security_groups(\n",
    "            Filters=[\n",
    "                {'Name':'group-name', 'Values': [name]}\n",
    "            ]\n",
    "        )\n",
    "        groups = response['SecurityGroups']\n",
    "\n",
    "        if ip is None:\n",
    "            ip = requests.get('https://api.ipify.org').text\n",
    "\n",
    "        if len(groups) > 0:\n",
    "            # Update the rule to use the new IP address\n",
    "            \n",
    "            security_group_id = groups[0]['GroupId']\n",
    "            logging.info('Found Security Group: %s in vpc %s (%s).' % (security_group_id, vpc_id, region))\n",
    "\n",
    "            ip_permissions = groups[0]['IpPermissions']\n",
    "            for ip_permission in ip_permissions:\n",
    "                # Delete all rules that listens to TCP port 8998\n",
    "                if ip_permission[\"IpProtocol\"] == 'tcp' and ip_permission[\"FromPort\"] == 8998 and ip_permission[\"FromPort\"] == 8998:\n",
    "                    cidr_ip = ip_permission['IpRanges'][0]['CidrIp']\n",
    "                    revoke_status = ec2_client.revoke_security_group_ingress(\n",
    "                        GroupId=security_group_id,\n",
    "                        IpPermissions=[\n",
    "                            {'IpProtocol': 'tcp',\n",
    "                             'FromPort': 8998,\n",
    "                             'ToPort': 8998,\n",
    "                             'IpRanges': [{'CidrIp': cidr_ip}]\n",
    "                            }\n",
    "                        ])\n",
    "            \n",
    "            # Create a new inbound rule that listens to this machine's IP\n",
    "            data = ec2_client.authorize_security_group_ingress(\n",
    "                GroupId=security_group_id,\n",
    "                IpPermissions=[\n",
    "                    {'IpProtocol': 'tcp',\n",
    "                     'FromPort': 8998,\n",
    "                     'ToPort': 8998,\n",
    "                     'IpRanges': [{'CidrIp': '{}/32'.format(ip)}]}\n",
    "                ])\n",
    "            return groups[0]['GroupId']\n",
    "        else:\n",
    "            response = ec2_client.create_security_group(GroupName=name,\n",
    "                                                 Description=desc,\n",
    "                                                 VpcId=vpc_id)\n",
    "            security_group_id = response['GroupId']\n",
    "            logging.info('New Security Group created: %s in vpc %s (%s).' % (security_group_id, vpc_id, region))\n",
    "\n",
    "            data = ec2_client.authorize_security_group_ingress(\n",
    "                GroupId=security_group_id,\n",
    "                IpPermissions=[\n",
    "                    {'IpProtocol': 'tcp',\n",
    "                     'FromPort': 8998,\n",
    "                     'ToPort': 8998,\n",
    "                     'IpRanges': [{'CidrIp': '{}/32'.format(ip)}]}\n",
    "                ])\n",
    "            return security_group_id\n",
    "    except ClientError as e:\n",
    "        logging.error(e)\n",
    "    return security_group_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "master_sg_id = create_security_group(ec2, '{}SG'.format(CLUSTER_NAME), 'Master SG for {}'.format(CLUSTER_NAME), VPC_ID)\n",
    "slave_sg_id = create_security_group(ec2, '{}SlaveSG'.format(CLUSTER_NAME), 'Slave SG for {}'.format(CLUSTER_NAME), VPC_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Default Roles and Key Pair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_default_roles(iam_client):\n",
    "    try:\n",
    "        iam_client.remove_role_from_instance_profile(InstanceProfileName='EMR_EC2_DefaultRole', RoleName='EMR_EC2_DefaultRole')\n",
    "        iam_client.delete_instance_profile(InstanceProfileName='EMR_EC2_DefaultRole')\n",
    "        iam_client.detach_role_policy(RoleName='EMR_EC2_DefaultRole', PolicyArn='arn:aws:iam::aws:policy/service-role/AmazonElasticMapReduceforEC2Role')\n",
    "        iam_client.delete_role(RoleName='EMR_EC2_DefaultRole')\n",
    "        iam_client.detach_role_policy(RoleName='EMR_DefaultRole', PolicyArn='arn:aws:iam::aws:policy/service-role/AmazonElasticMapReduceRole')\n",
    "        iam_client.delete_role(RoleName='EMR_DefaultRole')\n",
    "    except iam_client.exceptions.NoSuchEntityException:\n",
    "        pass\n",
    "\n",
    "def create_default_roles(iam_client):\n",
    "    # Recreate default roles\n",
    "    try:\n",
    "        job_flow_role = iam_client.get_role(RoleName='EMR_EC2_DefaultRole')\n",
    "        service_role = iam_client.get_role(RoleName='EMR_DefaultRole')\n",
    "        instance_profile = iam_client.get_instance_profile(InstanceProfileName='EMR_EC2_DefaultRole')\n",
    "    except iam_client.exceptions.NoSuchEntityException:\n",
    "        logging.info(\"Output of create_default_roles:\\n{}\".format(\n",
    "            json.loads(subprocess.check_output(['aws', 'emr', 'create-default-roles']))))\n",
    "\n",
    "\n",
    "def create_key_pair(ec2_client, key_name):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        - ec2_client (boto3.EC2.Client): EC2 client object.\n",
    "        - key_name (string): Name of key, usually 'xxx_pem'\n",
    "    Return:\n",
    "    \n",
    "        dict: {\n",
    "            'KeyFingerprint': 'string',\n",
    "            'KeyMaterial': 'string',\n",
    "            'KeyName': 'string',\n",
    "            'KeyPairId': 'string'\n",
    "        }\n",
    "    \"\"\"\n",
    "    response = ec2_client.describe_key_pairs(Filters=[\n",
    "        {'Name': 'key-name',\n",
    "         'Values': [key_name]\n",
    "        }\n",
    "    ])\n",
    "    keypairs = response['KeyPairs']\n",
    "    if len(keypairs) == 0:\n",
    "        keypair = ec2_client.create_key_pair(KeyName=key_name)\n",
    "        logging.info(\"keypair {} created:\\n{}\".format(key_name, keypair))\n",
    "    else:\n",
    "        keypair = keypairs[0]\n",
    "    return keypair\n",
    "\n",
    "\n",
    "def wait_for_roles(iam_client, job_flow_role_name='EMR_EC2_DefaultRole', service_role_name='EMR_DefaultRole', instance_profile_name='EMR_EC2_DefaultRole'):\n",
    "    role_names = [job_flow_role_name, service_role_name]\n",
    "    ok = False\n",
    "    while ok == False:\n",
    "        ok = True\n",
    "        for role_name in role_names:\n",
    "            try:\n",
    "                role = iam_client.get_role(RoleName=role_name)\n",
    "                logging.info(\"Role {} is ready\".format(role_name))\n",
    "            except iam_client.exceptions.NoSuchEntityException:\n",
    "                logging.info(\"Role {} is not ready. Waiting...\".format(role_name))\n",
    "                ok = False\n",
    "        try:\n",
    "            instance_profile = iam_client.get_instance_profile(InstanceProfileName=instance_profile_name)\n",
    "            logging.info(\"Instance Profile {} is ready\".format(instance_profile_name))\n",
    "        except iam_client.exceptions.NoSuchEntityException:\n",
    "            logging.info(\"Instance Profile {} is not ready. Waiting...\".format(instance_profile_name))\n",
    "            ok = False\n",
    "            \n",
    "        if ok == False:\n",
    "            time.sleep(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "create_default_roles(iam)\n",
    "# Wait a bit until the roles are ready, otherwise we'd get Failed to authorize instance profile arn.../instance-profile/EMR_EC2_DefaultRole\n",
    "wait_for_roles(iam)\n",
    "\n",
    "keypair = create_key_pair(ec2, '{}_pem'.format(CLUSTER_NAME))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create EMR Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClusterError(Exception):\n",
    "    def __init__(self, last_guess):\n",
    "        self.last_guess = last_guess\n",
    "\n",
    "        \n",
    "def get_cluster_status(emr_client, cluster_id):\n",
    "    cluster = emr_client.describe_cluster(ClusterId=cluster_id)\n",
    "    return cluster['Cluster']['Status']['State']\n",
    "\n",
    "\n",
    "def is_cluster_terminated(emr_client, cluster_id):\n",
    "    cluster = emr_client.describe_cluster(ClusterId=cluster_id)\n",
    "    return 'TERMINATED' in cluster['Cluster']['Status']['State']\n",
    "\n",
    "\n",
    "def create_emr_cluster(emr_client, cluster_name, master_sg, slave_sg, keypair_name, subnet_id, job_flow_role='EMR_EC2_DefaultRole', service_role='EMR_DefaultRole', release_label='emr-5.9.0',\n",
    "                   master_instance_type='m3.xlarge', num_core_nodes=3, core_node_instance_type='m3.xlarge'):\n",
    "    \"\"\" Create an EMR cluster\n",
    "    Args:\n",
    "        - subnet_id (string): If empty, use first available VPC (VPC is inferred from Security Groups)\n",
    "    \"\"\"\n",
    "    # Avoid recreating cluster\n",
    "    clusters = emr_client.list_clusters(ClusterStates=['STARTING', 'RUNNING', 'WAITING', 'BOOTSTRAPPING'])\n",
    "    active_clusters = [i for i in clusters['Clusters'] if i['Name'] == cluster_name]\n",
    "    if len(active_clusters) > 0:\n",
    "        return active_clusters[0]['Id']\n",
    "    else:\n",
    "        # Create cluster\n",
    "\n",
    "        # To avoid error:\n",
    "        #    botocore.exceptions.ClientError: An error occurred (ValidationException) when calling the RunJobFlow operation: Invalid InstanceProfile: EMR_EC2_DefaultRole.\n",
    "        # We use do while in here\n",
    "        ok = False\n",
    "        while ok == False:\n",
    "            try:\n",
    "                cluster_response = emr_client.run_job_flow(\n",
    "                    Name=cluster_name,\n",
    "                    ReleaseLabel=release_label,\n",
    "                    Instances={\n",
    "                        'InstanceGroups': [\n",
    "                            {\n",
    "                                'Name': \"Master nodes\",\n",
    "                                'Market': 'ON_DEMAND',\n",
    "                                'InstanceRole': 'MASTER',\n",
    "                                'InstanceType': master_instance_type,\n",
    "                                'InstanceCount': 1\n",
    "                            },\n",
    "                            {\n",
    "                                'Name': \"Slave nodes\",\n",
    "                                'Market': 'ON_DEMAND',\n",
    "                                'InstanceRole': 'CORE',\n",
    "                                'InstanceType': core_node_instance_type,\n",
    "                                'InstanceCount': num_core_nodes\n",
    "                            }\n",
    "                        ],\n",
    "                        'KeepJobFlowAliveWhenNoSteps': True,\n",
    "                        'Ec2SubnetId': subnet_id,\n",
    "                        'Ec2KeyName' : keypair_name,\n",
    "                        'EmrManagedMasterSecurityGroup': master_sg,\n",
    "                        'EmrManagedSlaveSecurityGroup': slave_sg\n",
    "                    },\n",
    "                    VisibleToAllUsers=True,\n",
    "                    JobFlowRole=job_flow_role,\n",
    "                    ServiceRole=service_role,\n",
    "                    Applications=[\n",
    "                        { 'Name': 'hadoop' },\n",
    "                        { 'Name': 'spark' },\n",
    "                        { 'Name': 'hive' },\n",
    "                        { 'Name': 'livy' },\n",
    "                        { 'Name': 'zeppelin' }\n",
    "                    ]\n",
    "                )\n",
    "                ok = True\n",
    "            except ClientError as e:\n",
    "                logging.info(e)\n",
    "        cluster_id = cluster_response['JobFlowId']\n",
    "        cluster_state = get_cluster_status(emr_client, cluster_id)\n",
    "        if cluster_state != 'STARTING':\n",
    "            reason = emr_client.describe_cluster(ClusterId=cluster_id)['Cluster']['Status']['StateChangeReason']\n",
    "            raise Exception(\"Cluster error: {} - {}\".format(reason['Code'], reason['Message']))\n",
    "            \n",
    "        exit_loop = False\n",
    "        while exit_loop == False:\n",
    "            cluster_state = get_cluster_status(emr_client, cluster_id)\n",
    "            if cluster_state == 'WAITING':\n",
    "                exit_loop = True\n",
    "            elif 'TERMINATED' in cluster_state:\n",
    "                exit_loop = True\n",
    "                raise Exception(\"Cluser terminated:\\n{}\".format(emr_client.describe_cluster(ClusterId=cluster_id)))\n",
    "            else:\n",
    "                logging.info(\"Cluster is {}. Waiting for completion...\".format(cluster_state))\n",
    "                time.sleep(10)\n",
    "        logging.info(\"Cluser created:\\n{}\".format(emr_client.describe_cluster(ClusterId=cluster_id)))\n",
    "        return cluster_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "cluster_id = create_emr_cluster(emr, CLUSTER_NAME, master_sg_id, slave_sg_id, keypair['KeyName'], SUBNET_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emr.describe_cluster(ClusterId=cluster_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Spark Session\n",
    "\n",
    "Wait until the cluster is in WAITING state and then create a spark session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_cluster_ready(emr_client, cluster_id):\n",
    "    return get_cluster_status(emr_client, cluster_id) == 'WAITING'\n",
    "\n",
    "\n",
    "def get_cluster_dns(emr_client, cluster_id):\n",
    "    cluster = emr_client.describe_cluster(ClusterId=cluster_id)\n",
    "    return cluster['Cluster']['MasterPublicDnsName']\n",
    "\n",
    "\n",
    "def spark_url(master_dns, location='', port=8998):\n",
    "    \"\"\"Get spark session url.\"\"\"\n",
    "    return 'http://{}:{}{}'.format(master_dns, port, location)\n",
    "\n",
    "\n",
    "def kill_spark_session(master_dns, session_headers, port=8998):\n",
    "    session_url = spark_url(master_dns, location=session_headers['Location'], port=port)\n",
    "    requests.delete(session_url, headers={'Content-Type': 'application/json'})\n",
    "\n",
    "\n",
    "def kill_spark_session_by_id(master_dns, session_id, port=8998):\n",
    "    session_url = spark_url(master_dns, location='/sessions/{}'.format(session_id), port=port)\n",
    "    requests.delete(session_url, headers={'Content-Type': 'application/json'})\n",
    "    \n",
    "\n",
    "def kill_all_inactive_spark_sessions(master_dns):\n",
    "    response = requests.get(spark_url(master_dns, location='/sessions'))\n",
    "    spark_sessions = response.json()['sessions']\n",
    "    logging.info(\"Killing all inactive spark sessions\")\n",
    "    for spark_session in spark_sessions:\n",
    "        if spark_session['state'] in ['idle', 'dead'] :\n",
    "            kill_spark_session_by_id(master_dns, spark_session['id'])\n",
    "            logging.info(\"Killed {} spark session id {}\".format(spark_session['state'],\n",
    "                                                                spark_session['id']))\n",
    "    \n",
    "\n",
    "def kill_all_spark_sessions(master_dns):\n",
    "    response = requests.get(spark_url(master_dns, location='/sessions'))\n",
    "    spark_sessions = response.json()['sessions']\n",
    "    logging.info(\"Killing all spark sessions\")\n",
    "    for spark_session in spark_sessions:\n",
    "        kill_spark_session_by_id(master_dns, spark_session['id'])\n",
    "        logging.info(\"Killed {} spark session id {}\".format(spark_session['state'],\n",
    "                                                            spark_session['id']))\n",
    "        \n",
    "        \n",
    "def create_spark_session(master_dns, port=8998):\n",
    "    session_url = spark_url(master_dns, location='/sessions', port=port)\n",
    "    data = {'kind': 'pyspark', \n",
    "            \"conf\" : {\"spark.jars.packages\" : \"saurfang:spark-sas7bdat:2.0.0-s_2.11\",\n",
    "                      \"spark.driver.extraJavaOptions\" : \"-Dlog4jspark.root.logger=WARN,console\"\n",
    "                     }\n",
    "           }\n",
    "    headers = {'Content-Type': 'application/json'}\n",
    "    response = requests.post(session_url, data=json.dumps(data), headers=headers)\n",
    "    logging.info(\"Sent spark session creation command to {}\".format(session_url))\n",
    "    logging.info(\"Response headers: {}\".format(response.headers))\n",
    "    logging.info(response.json())\n",
    "    if 'Location' not in response.headers:\n",
    "        raise Exception(\"Spark session creation failed. This is usually due \" + \\\n",
    "                        \"to too many spark sessions on the server. \" + \\\n",
    "                        \"Please run kill_all_inactive_spark_sessions function.\")\n",
    "    return response.headers\n",
    "\n",
    "\n",
    "def wait_for_spark(master_dns, session_headers, port=8998):\n",
    "    \"\"\"Wait until status is idle\"\"\"\n",
    "    status = ''\n",
    "    logging.info(\"Session headers: {}\".format(session_headers))\n",
    "    session_url = spark_url(master_dns, location=session_headers['Location'], port=port)\n",
    "    while status not in ['idle', 'dead']:\n",
    "        response = requests.get(session_url, headers=session_headers)\n",
    "        status = response.json()['state']\n",
    "        logging.info(\"Spark session status: {}\".format(status))\n",
    "        if status == 'dead':\n",
    "            raise Exception(\"Spark session is dead:\\nResponse status code: {}\\nHeaders: {}\\nContent: {}\" \\\n",
    "                            .format(response.status_code, response.headers, json.dumps(response.content)))\n",
    "        else:\n",
    "            time.sleep(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_dns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kill_all_inactive_spark_sessions(cluster_dns)\n",
    "\n",
    "response = requests.get(spark_url(cluster_dns, location='/sessions'))\n",
    "print(response.status_code)\n",
    "print(response.headers)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_dns = get_cluster_dns(emr, cluster_id)\n",
    "ss_headers = create_spark_session(cluster_dns)\n",
    "print(ss_headers)\n",
    "wait_for_spark(cluster_dns, ss_headers)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Send Spark Jobs\n",
    "\n",
    "We will try pulling some stock market data from Quandl and QuoteMedia. Stock names are available here:\n",
    "\n",
    "- NASDAQ: https://old.nasdaq.com/screening/companies-by-name.aspx?letter=0&exchange=nasdaq&render=download\n",
    "- AMEX: https://old.nasdaq.com/screening/companies-by-name.aspx?letter=0&exchange=amex&render=download\n",
    "- NYSE: https://old.nasdaq.com/screening/companies-by-name.aspx?letter=0&exchange=nyse&render=download"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test Quandl request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv('https://old.nasdaq.com/screening/companies-by-name.aspx?letter=0&exchange=nasdaq&render=download')\n",
    "df.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exchange_map = {\n",
    "#     'nasdaq': 'FNSQ',\n",
    "#     'nyse': 'FNYX'\n",
    "# }\n",
    "\n",
    "def get_short_interests_pandas(exchange, ticker, api_key):\n",
    "    response = requests.get(\"https://www.quandl.com/api/v3/datasets/FINRA/{}_{}?api_key={}\".format(exchange, ticker, api_key))\n",
    "    if response.status_code == 200:\n",
    "        response_obj = response.json()\n",
    "        return pd.DataFrame(data=response_obj['dataset']['data'], columns=response_obj['dataset']['column_names'])\n",
    "    else:\n",
    "        raise Exception(\"Error when connecting to Quandl API.\")\n",
    "\n",
    "df = get_short_interests_pandas('FNYX', 'FB', config['Quandl']['API_KEY'])\n",
    "print(df.describe())\n",
    "df.head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quandl request through Spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "\n",
    "spark = spark = SparkSession \\\n",
    "        .builder \\\n",
    "        .config(\"spark.jars.packages\", \"org.apache.hadoop:hadoop-aws:2.7.0\") \\\n",
    "        .getOrCreate()\n",
    "\n",
    "def get_short_interests(spark, exchange, ticker, api_key):\n",
    "    url = \"https://www.quandl.com/api/v3/datasets/FINRA/{}_{}?api_key={}\".format(exchange, ticker, api_key)\n",
    "#     spark.sparkContext.addFile(url)\n",
    "#     response = spark.read.json(\"file://{}\".format(SparkFiles.get(\"{}_{}\".format(exchange, ticker))))\n",
    "#     print(response)\n",
    "    result = requests.get(url).json()\n",
    "    df = spark.createDataFrame(result['dataset']['data'], result['dataset']['column_names'])\n",
    "    df.createOrReplaceTempView('test')\n",
    "    table = spark.sql(\"SELECT * FROM test\")\n",
    "    table_path = \"test_data/test_table\"\n",
    "    table.write.mode('overwrite').parquet(table_path)\n",
    "\n",
    "df = get_short_interests(spark, 'FNYX', 'FB', config['Quandl']['API_KEY'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Send Spark job to download from Quandl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def push_args_into_code(code, args):\n",
    "    # Include arguments into the code (at the top of the file)\n",
    "    args_str = \"\"\n",
    "    for key, value in args.items():\n",
    "        if isinstance(value, str):\n",
    "            args_str += \"{}='{}'\\n\".format(key, value.replace(\"'\", \"\\'\"))\n",
    "        else:\n",
    "            args_str += \"{}={}\\n\".format(key, value)\n",
    "    code = args_str + code\n",
    "    return code\n",
    "\n",
    "\n",
    "def push_helpers_into_code(code, helpers):\n",
    "    # Include helpers into the code (at the top of the file)\n",
    "    code = helpers + \"\\n\" + code\n",
    "    return code\n",
    "\n",
    "    \n",
    "def submit_spark_job(master_dns, session_headers, code, args={}, helpers='', port=8998):\n",
    "    statements_url = spark_url(master_dns, location=session_headers['Location'] + \"/statements\", port=port)\n",
    "\n",
    "    code = push_helpers_into_code(code, helpers)\n",
    "    code = push_args_into_code(code, args)\n",
    "\n",
    "    job = {'code': code}\n",
    "    response = requests.post(statements_url, data=json.dumps(job),\n",
    "                             headers={'Content-Type': 'application/json'})\n",
    "    if response.status_code not in [200, 201]:\n",
    "        raise Exception(\"Spark job sending error:\\nResponse status code: {}\\nHeaders: {}\\nContent: {}\" \\\n",
    "                        .format(response.status_code, response.headers, response.content))\n",
    "    else:\n",
    "        logging.info(\"Spark job sending successful:\\nResponse status code: {}\\nHeaders: {}\\nContent: {}\" \\\n",
    "                     .format(response.status_code, response.headers, response.content))\n",
    "    return response.headers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(push_args_into_code(\"print(val1 + str(val2))\\nprint(val3)\", args={'val1': \"string\", 'val2': 123, 'val3': None}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ss_headers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "exchange = 'FNYX'\n",
    "ticker = 'FB'\n",
    "code = \"\"\"\n",
    "import requests\n",
    "url = \"https://www.quandl.com/api/v3/datasets/FINRA/{exchange}_{ticker}?api_key={quandl_api}\"\n",
    "result = requests.get(url).json()\n",
    "df = spark.createDataFrame(result['dataset']['data'], result['dataset']['column_names'])\n",
    "df.createOrReplaceTempView('test')\n",
    "table = spark.sql(\"SELECT * FROM test\")\n",
    "table_path = \"s3://short-interest-effect/data/test_table\"\n",
    "table.write.mode('overwrite').parquet(table_path)\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "code = code.format(exchange=exchange, ticker=ticker, quandl_api=config['Quandl']['API_KEY'])\n",
    "job_response_headers = submit_spark_job(cluster_dns, ss_headers, code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_response_headers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_response_headers['Location'].split('/statements', 1)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Track Spark job status\n",
    "\n",
    "The following code can be run several times to check the result of the above statement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def track_spark_job(master_dns, job_response_headers, port=8998):\n",
    "    job_status = ''\n",
    "    session_url = spark_url(master_dns, location=job_response_headers['Location'].split('/statements', 1)[0], port=port)\n",
    "    statement_url = spark_url(master_dns, location=job_response_headers['Location'], port=port)\n",
    "        \n",
    "    while job_status not in ['available']:\n",
    "        # If a statement takes longer than a few milliseconds to execute, Livy returns early and provides\n",
    "        # a statement URL that can be pooled until it is complete:\n",
    "        statement_response = requests.get(statement_url, headers={'Content-Type': 'application/json'})\n",
    "        job_status = statement_response.json()['state']\n",
    "        logging.info('Spark Job status: ' + job_status)\n",
    "        logging.info(\"Response: {}\".format(json.dumps(statement_response.json(), indent=4)))\n",
    "        if 'progress' in statement_response.json():\n",
    "            logging.info('Progress: ' + str(statement_response.json()['progress']))\n",
    "\n",
    "        if job_status == 'idle':\n",
    "            raise ValueError(\"track_spark_job error. Looks like you have passed spark session headers for the second parameter. \"+\n",
    "                             \"Pass in spark job response headers instead.\")\n",
    "\n",
    "        if job_status != 'available':\n",
    "            time.sleep(5)\n",
    "            \n",
    "    final_job_status = statement_response.json()['output']['status']\n",
    "\n",
    "    # Get the logs\n",
    "    log_lines = requests.get(session_url + '/log', \n",
    "                             headers={'Content-Type': 'application/json'}).json()['log']\n",
    "    logging.info(\"Log from the cluster:\\n{}\".format(\"\\n\".join(log_lines)))\n",
    "    logging.info('Final job Status: ' + final_job_status)\n",
    "\n",
    "    if final_job_status == 'error':\n",
    "        logging.info('Statement exception: ' + statement_response.json()['output']['evalue'])\n",
    "        for trace in statement_response.json()['output']['traceback']:\n",
    "            logging.info(trace)\n",
    "        raise ValueError('Stopped because the final job status was \"error\".')\n",
    "    \n",
    "    return (final_job_status, log_lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "track_spark_job(cluster_dns, job_response_headers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Kill Spark session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kill_spark_session(master_dns, session_headers, port=8998):\n",
    "    session_url = spark_url(master_dns, location=session_headers['Location'], port=port)\n",
    "    requests.delete(session_url, headers={'Content-Type': 'application/json'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kill_spark_session(cluster_dns, ss_headers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Delete Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_cluster(emr_client, cluster_id):\n",
    "    try:\n",
    "        response = emr_client.terminate_job_flows(JobFlowIds=[cluster_id])\n",
    "        \n",
    "        cluster_removed = False\n",
    "        while cluster_removed == False:\n",
    "            if is_cluster_terminated(emr_client, cluster_id):\n",
    "                cluster_removed = True\n",
    "            else:\n",
    "                state = get_cluster_status(emr_client, cluster_id)\n",
    "                logging.info(\"Cluster {} has not been terminated (Current cluster state: {}). waiting until the status is TERMINATED...\". \\\n",
    "                             format(cluster_id, state))\n",
    "                time.sleep(10)\n",
    "                \n",
    "        print('Cluster {} Deleted'.format(cluster_id))\n",
    "    except ClientError as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_cluster(emr, cluster_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Delete Key Pair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ec2.delete_key_pair(KeyName=keypair['KeyName'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Delete Security Group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_security_group(ec2_client, sgid):\n",
    "    # Delete security group\n",
    "    try:\n",
    "        ec2res = boto3.resource('ec2')\n",
    "        sg = ec2res.SecurityGroup(sgid)\n",
    "        if len(sg.ip_permissions) > 0:\n",
    "            for ip_permission in sg.ip_permissions:\n",
    "                for group_pair in ip_permission['UserIdGroupPairs']:\n",
    "                    if 'GroupName' in group_pair:\n",
    "                        del(group_pair['GroupName'])\n",
    "            sg.revoke_ingress(IpPermissions=sg.ip_permissions)\n",
    "        response = ec2_client.delete_security_group(GroupId=sgid)\n",
    "        logging.info('Security Group {} Deleted'.format(sgid))\n",
    "    except ClientError as e:\n",
    "        logging.error(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.setLevel(logging.INFO)\n",
    "delete_security_group(ec2, master_sg_id)\n",
    "time.sleep(2)\n",
    "delete_security_group(ec2, slave_sg_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "- EMR creation that works: https://github.com/dai-dao/udacity-data-engineering-capstone/blob/master/dags/lib/emr_lib.py\n",
    "- On Security Group Creation and Deletion: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/ec2-example-security-group.html\n",
    "- On how to recreate EMR_EC2_DefaultRole: https://aws.amazon.com/premiumsupport/knowledge-center/emr-default-role-invalid/\n",
    "- Using Apache Livy with Spark on EMR: https://aws.amazon.com/blogs/big-data/orchestrate-apache-spark-applications-using-aws-step-functions-and-apache-livy/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varnames = [None]\n",
    "status = True\n",
    "for var in varnames:\n",
    "    try:\n",
    "        if var != None:\n",
    "            status = status and False\n",
    "        else:\n",
    "            status = status and True\n",
    "    except KeyError as e:\n",
    "        status = status and True\n",
    "status"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}