         sagemaker-core/sagemaker-core-pyspark-processing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Processing using Apache Spark and SageMaker Processing\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Apache Spark is a unified analytics engine for large-scale data processing. The Spark framework is often used within the context of machine learning workflows to run data transformation or feature engineering workloads at scale. Amazon SageMaker provides a set of prebuilt Docker images that include Apache Spark and other dependencies needed to run distributed data processing jobs on Amazon SageMaker. This example notebook demonstrates how to use the prebuilt Spark images on SageMaker Processing using the SageMaker Core Python SDK."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Runtime\n",
    "\n",
    "This notebook takes approximately 22 minutes to run."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "1. [Setup](#Setup)\n",
    "1. [Example 1: Running a basic PySpark application](#Example-1:-Running-a-basic-PySpark-application)\n",
    "1. [Example 2: Specify additional Python and jar file dependencies](#Example-2:-Specify-additional-Python-and-jar-file-dependencies)\n",
    "1. [Example 3: Run a Java/Scala Spark application](#Example-3:-Run-a-Java/Scala-Spark-application)\n",
    "1. [Example 4: Specifying additional Spark configuration](#Example-4:-Specifying-additional-Spark-configuration)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install the latest SageMaker Core Python SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip uninstall sagemaker-core -y\n",
    "!pip install pip --upgrade --quiet\n",
    "!pip install sagemaker-core --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Restart your notebook kernel after upgrading the SDK*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Running a basic PySpark application\n",
    "\n",
    "The first example is a basic Spark MLlib data processing script. This script will take a raw data set and do some transformations on it such as string indexing and one hot encoding.\n",
    "\n",
    "### Setup S3 bucket locations and roles\n",
    "\n",
    "First, setup some locations in the default SageMaker bucket to store the raw input datasets and the Spark job output. Here, you'll also define the role that will be used to run all SageMaker Processing jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sagemaker\n",
    "from time import gmtime, strftime\n",
    "from sagemaker_core.helper.session_helper import get_execution_role, Session\n",
    "\n",
    "sagemaker_logger = logging.getLogger(\"sagemaker\")\n",
    "sagemaker_logger.setLevel(logging.INFO)\n",
    "sagemaker_logger.addHandler(logging.StreamHandler())\n",
    "\n",
    "sagemaker_session = Session()\n",
    "REGION_NAME = sagemaker_session._region_name\n",
    "role = get_execution_role()\n",
    "s3_bucket_name = sagemaker.Session().default_bucket()\n",
    "default_bucket_prefix = sagemaker.Session().default_bucket_prefix\n",
    "default_bucket_prefix_path = \"\"\n",
    "\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    default_bucket_prefix_path = f\"/{default_bucket_prefix}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, you'll download the example dataset from a SageMaker staging bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch the dataset from the SageMaker bucket\n",
    "import boto3\n",
    "\n",
    "s3 = boto3.client(\"s3\")\n",
    "s3.download_file(\n",
    "    f\"sagemaker-sample-files\", \"datasets/tabular/uci_abalone/abalone.csv\", \"./data/abalone.csv\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write the PySpark script\n",
    "\n",
    "The source for a preprocessing script is in the cell below. The cell uses the `%%writefile` directive to save this file locally. This script does some basic feature engineering on a raw input dataset. In this example, the dataset is the [Abalone Data Set](https://archive.ics.uci.edu/ml/datasets/abalone) and the code below performs string indexing, one hot encoding, vector assembly, and combines them into a pipeline to perform these transformations in order. The script then does an 80-20 split to produce training and validation datasets as output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./code/preprocess.py\n",
    "from __future__ import print_function\n",
    "from __future__ import unicode_literals\n",
    "\n",
    "import argparse\n",
    "import csv\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "import time\n",
    "\n",
    "import pyspark\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.feature import (\n",
    "    OneHotEncoder,\n",
    "    StringIndexer,\n",
    "    VectorAssembler,\n",
    "    VectorIndexer,\n",
    ")\n",
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import (\n",
    "    DoubleType,\n",
    "    StringType,\n",
    "    StructField,\n",
    "    StructType,\n",
    ")\n",
    "\n",
    "\n",
    "def csv_line(data):\n",
    "    r = \",\".join(str(d) for d in data[1])\n",
    "    return str(data[0]) + \",\" + r\n",
    "\n",
    "\n",
    "def main():\n",
    "    parser = argparse.ArgumentParser(description=\"app inputs and outputs\")\n",
    "    parser.add_argument(\"--s3_input_bucket\", type=str, help=\"s3 input bucket\")\n",
    "    parser.add_argument(\"--s3_input_key_prefix\", type=str, help=\"s3 input key prefix\")\n",
    "    parser.add_argument(\"--s3_output_bucket\", type=str, help=\"s3 output bucket\")\n",
    "    parser.add_argument(\"--s3_output_key_prefix\", type=str, help=\"s3 output key prefix\")\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    spark = SparkSession.builder.appName(\"PySparkApp\").getOrCreate()\n",
    "\n",
    "    # This is needed to save RDDs which is the only way to write nested Dataframes into CSV format\n",
    "    spark.sparkContext._jsc.hadoopConfiguration().set(\n",
    "        \"mapred.output.committer.class\", \"org.apache.hadoop.mapred.FileOutputCommitter\"\n",
    "    )\n",
    "\n",
    "    # Defining the schema corresponding to the input data. The input data does not contain the headers\n",
    "    schema = StructType(\n",
    "        [\n",
    "            StructField(\"sex\", StringType(), True),\n",
    "            StructField(\"length\", DoubleType(), True),\n",
    "            StructField(\"diameter\", DoubleType(), True),\n",
    "            StructField(\"height\", DoubleType(), True),\n",
    "            StructField(\"whole_weight\", DoubleType(), True),\n",
    "            StructField(\"shucked_weight\", DoubleType(), True),\n",
    "            StructField(\"viscera_weight\", DoubleType(), True),\n",
    "            StructField(\"shell_weight\", DoubleType(), True),\n",
    "            StructField(\"rings\", DoubleType(), True),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    # Downloading the data from S3 into a Dataframe\n",
    "    total_df = spark.read.csv(\n",
    "        (\"s3://\" + os.path.join(args.s3_input_bucket, args.s3_input_key_prefix, \"abalone.csv\")),\n",
    "        header=False,\n",
    "        schema=schema,\n",
    "    )\n",
    "\n",
    "    # StringIndexer on the sex column which has categorical value\n",
    "    sex_indexer = StringIndexer(inputCol=\"sex\", outputCol=\"indexed_sex\")\n",
    "\n",
    "    # one-hot-encoding is being performed on the string-indexed sex column (indexed_sex)\n",
    "    sex_encoder = OneHotEncoder(inputCol=\"indexed_sex\", outputCol=\"sex_vec\")\n",
    "\n",
    "    # vector-assembler will bring all the features to a 1D vector for us to save easily into CSV format\n",
    "    assembler = VectorAssembler(\n",
    "        inputCols=[\n",
    "            \"sex_vec\",\n",
    "            \"length\",\n",
    "            \"diameter\",\n",
    "            \"height\",\n",
    "            \"whole_weight\",\n",
    "            \"shucked_weight\",\n",
    "            \"viscera_weight\",\n",
    "            \"shell_weight\",\n",
    "        ],\n",
    "        outputCol=\"features\",\n",
    "    )\n",
    "\n",
    "    # The pipeline is comprised of the steps added above\n",
    "    pipeline = Pipeline(stages=[sex_indexer, sex_encoder, assembler])\n",
    "\n",
    "    # This step trains the feature transformers\n",
    "    model = pipeline.fit(total_df)\n",
    "\n",
    "    # This step transforms the dataset with information obtained from the previous fit\n",
    "    transformed_total_df = model.transform(total_df)\n",
    "\n",
    "    # Split the overall dataset into 80-20 training and validation\n",
    "    (train_df, validation_df) = transformed_total_df.randomSplit([0.8, 0.2])\n",
    "\n",
    "    # Convert the train dataframe to RDD to save in CSV format and upload to S3\n",
    "    train_rdd = train_df.rdd.map(lambda x: (x.rings, x.features))\n",
    "    train_lines = train_rdd.map(csv_line)\n",
    "    train_lines.saveAsTextFile(\n",
    "        \"s3://\" + os.path.join(args.s3_output_bucket, args.s3_output_key_prefix, \"train\")\n",
    "    )\n",
    "\n",
    "    # Convert the validation dataframe to RDD to save in CSV format and upload to S3\n",
    "    validation_rdd = validation_df.rdd.map(lambda x: (x.rings, x.features))\n",
    "    validation_lines = validation_rdd.map(csv_line)\n",
    "    validation_lines.saveAsTextFile(\n",
    "        \"s3://\" + os.path.join(args.s3_output_bucket, args.s3_output_key_prefix, \"validation\")\n",
    "    )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the SageMaker Processing Job\n",
    "\n",
    "Next, you'll use the `ProcessingJob` class to define a Spark job and run it using SageMaker Processing. A few things to note in the definition of the `ProcessingJob`:\n",
    "\n",
    "* This is a multi-node job with two m5.xlarge instances (which is specified via the `instance_count` and `instance_type` parameters)\n",
    "* Spark framework version 3.1 image is specified via the `image_uri` parameter\n",
    "* The PySpark script defined above is passed via via the `ProcessingInput` class\n",
    "* Command-line arguments to the PySpark script (such as the S3 input and output locations) are passed via the `arguments` parameter\n",
    "* Spark event logs will be offloaded to the S3 location in `spark_event_logs` folder.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from sagemaker_core.shapes import (\n",
    "    ProcessingInput,\n",
    "    ProcessingResources,\n",
    "    AppSpecification,\n",
    "    ProcessingS3Input,\n",
    "    ProcessingOutputConfig,\n",
    ")\n",
    "from sagemaker_core.shapes import (\n",
    "    ProcessingResources,\n",
    "    ProcessingClusterConfig,\n",
    "    ProcessingOutput,\n",
    "    ProcessingS3Output,\n",
    ")\n",
    "from sagemaker_core.resources import ProcessingJob\n",
    "\n",
    "# Upload the raw input dataset to a unique S3 location\n",
    "timestamp_prefix = strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "prefix = \"sagemaker/spark-preprocess-demo/{}\".format(timestamp_prefix)\n",
    "\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    prefix = f\"{default_bucket_prefix}/{prefix}\"\n",
    "\n",
    "input_prefix_abalone = \"{}/input/raw/abalone\".format(prefix)\n",
    "input_preprocessed_prefix_abalone = \"{}/input/preprocessed/abalone\".format(prefix)\n",
    "\n",
    "base_job_name = \"sm-spark\"\n",
    "final_job_name = base_job_name + \"-\" + timestamp_prefix\n",
    "\n",
    "input_script_prefix = \"{}/input/code\".format(final_job_name)\n",
    "\n",
    "# uploading required data to S3 for reference\n",
    "# uploading abolone.csv to S3 bucket\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./data/abalone.csv\", bucket=s3_bucket_name, key_prefix=input_prefix_abalone\n",
    ")\n",
    "\n",
    "# uploading preprocess.py to S3 bucket\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./code/preprocess.py\", bucket=s3_bucket_name, key_prefix=input_script_prefix\n",
    ")\n",
    "\n",
    "# initializing ProcessingInputs,ProcessingResources,ProcessingOutputConfig and AppSpecification configurations\n",
    "processing_input = ProcessingInput(\n",
    "    input_name=\"code\",\n",
    "    s3_input=ProcessingS3Input(\n",
    "        s3_uri=f\"s3://{s3_bucket_name}{default_bucket_prefix_path}/{final_job_name}/input/code/preprocess.py\",\n",
    "        # s3_uri=\"s3://sagemaker-us-east-1-774297356213/sm-spark-2024-08-30-05-25-18-294/input/code/preprocess.py\",\n",
    "        s3_data_type=\"S3Prefix\",\n",
    "        local_path=\"/opt/ml/processing/input/code\",\n",
    "        s3_input_mode=\"File\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "processing_output_config = ProcessingOutputConfig(\n",
    "    outputs=[\n",
    "        ProcessingOutput(\n",
    "            output_name=\"output-1\",\n",
    "            s3_output=ProcessingS3Output(\n",
    "                s3_uri=f\"s3://{s3_bucket_name}/{prefix}/spark_event_logs\",\n",
    "                local_path=\"/opt/ml/processing/spark-events/\",\n",
    "                s3_upload_mode=\"Continuous\",\n",
    "            ),\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "processing_resources = ProcessingResources(\n",
    "    cluster_config=ProcessingClusterConfig(\n",
    "        instance_count=2, instance_type=\"ml.m5.xlarge\", volume_size_in_gb=30\n",
    "    )\n",
    ")\n",
    "\n",
    "app_specification = AppSpecification(\n",
    "    image_uri=\"173754725891.dkr.ecr.us-east-1.amazonaws.com/sagemaker-spark-processing:3.1-cpu\",\n",
    "    container_entrypoint=[\n",
    "        \"smspark-submit\",\n",
    "        \"--local-spark-event-logs-dir\",\n",
    "        \"/opt/ml/processing/spark-events/\",\n",
    "        \"/opt/ml/processing/input/code/preprocess.py\",\n",
    "    ],\n",
    "    container_arguments=[\n",
    "        \"--s3_input_bucket\",\n",
    "        f\"{s3_bucket_name}\",\n",
    "        \"--s3_input_key_prefix\",\n",
    "        f\"{input_prefix_abalone}\",\n",
    "        \"--s3_output_bucket\",\n",
    "        f\"{s3_bucket_name}\",\n",
    "        \"--s3_output_key_prefix\",\n",
    "        f\"{input_preprocessed_prefix_abalone}\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Run the processing job\n",
    "processing_job_obj = ProcessingJob.create(\n",
    "    processing_job_name=final_job_name,\n",
    "    processing_resources=processing_resources,\n",
    "    app_specification=app_specification,\n",
    "    role_arn=role,\n",
    "    processing_inputs=[processing_input],\n",
    "    processing_output_config=processing_output_config,\n",
    ")\n",
    "\n",
    "processing_job_obj.wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validate Data Processing Results\n",
    "\n",
    "Next, validate the output of our data preprocessing job by looking at the first 5 rows of the output dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Top 5 rows from s3://{}/{}/train/\".format(s3_bucket_name, input_preprocessed_prefix_abalone))\n",
    "!aws s3 cp --quiet s3://$s3_bucket_name/$input_preprocessed_prefix_abalone/train/part-00000 - | head -n5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Specify additional Python and jar file dependencies\n",
    "\n",
    "The next example demonstrates a scenario where additional Python file dependencies are required by the PySpark script. You'll use a sample PySpark script that requires additional user-defined functions (UDFs) defined in a local module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./code/hello_py_spark_app.py\n",
    "import argparse\n",
    "import time\n",
    "\n",
    "# Import local module to test spark-submit--py-files dependencies\n",
    "import hello_py_spark_udfs as udfs\n",
    "from pyspark.sql import SparkSession, SQLContext\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import IntegerType\n",
    "import time\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    print(\"Hello World, this is PySpark!\")\n",
    "\n",
    "    parser = argparse.ArgumentParser(description=\"inputs and outputs\")\n",
    "    parser.add_argument(\"--input\", type=str, help=\"path to input data\")\n",
    "    parser.add_argument(\"--output\", required=False, type=str, help=\"path to output data\")\n",
    "    args = parser.parse_args()\n",
    "    spark = SparkSession.builder.appName(\"SparkTestApp\").getOrCreate()\n",
    "    sqlContext = SQLContext(spark.sparkContext)\n",
    "\n",
    "    # Load test data set\n",
    "    inputPath = args.input\n",
    "    outputPath = args.output\n",
    "    salesDF = spark.read.json(inputPath)\n",
    "    salesDF.printSchema()\n",
    "\n",
    "    salesDF.createOrReplaceTempView(\"sales\")\n",
    "\n",
    "    # Define a UDF that doubles an integer column\n",
    "    # The UDF function is imported from local module to test spark-submit--py-files dependencies\n",
    "    double_udf_int = udf(udfs.double_x, IntegerType())\n",
    "\n",
    "    # Save transformed data set to disk\n",
    "    salesDF.select(\"date\", \"sale\", double_udf_int(\"sale\").alias(\"sale_double\")).write.json(\n",
    "        outputPath\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating `hello_py_spark_udfs.py` inside `code` folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./code/hello_py_spark_udfs.py\n",
    "def double_x(x):\n",
    "    return x + x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a processing job with Python file dependencies\n",
    "\n",
    "Then, you'll create a processing job where the additional Python file dependencies are specified via the `py-files` input name in the `ProcessingInput` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker_core.shapes import (\n",
    "    ProcessingInput,\n",
    "    ProcessingResources,\n",
    "    AppSpecification,\n",
    "    ProcessingS3Input,\n",
    "    ProcessingOutputConfig,\n",
    ")\n",
    "from sagemaker_core.shapes import (\n",
    "    ProcessingResources,\n",
    "    ProcessingClusterConfig,\n",
    "    ProcessingOutput,\n",
    "    ProcessingS3Output,\n",
    ")\n",
    "from sagemaker_core.resources import ProcessingJob\n",
    "\n",
    "# Upload the raw input dataset to a unique S3 location\n",
    "timestamp_prefix = strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "prefix = \"sagemaker/spark-preprocess-demo/{}\".format(timestamp_prefix)\n",
    "\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    prefix = f\"{default_bucket_prefix}/{prefix}\"\n",
    "\n",
    "input_prefix_sales = \"{}/input/sales\".format(prefix)\n",
    "output_prefix_sales = \"{}/output/sales\".format(prefix)\n",
    "input_s3_uri = \"s3://{}/{}\".format(s3_bucket_name, input_prefix_sales)\n",
    "output_s3_uri = \"s3://{}/{}\".format(s3_bucket_name, output_prefix_sales)\n",
    "\n",
    "base_job_name = \"sm-spark-udfs\"\n",
    "final_job_name = base_job_name + \"-\" + timestamp_prefix\n",
    "\n",
    "input_script_prefix = \"{}/input/code\".format(final_job_name)\n",
    "input_pyfiles_prefix = \"{}/input/py-files\".format(final_job_name)\n",
    "\n",
    "# uploading required data to S3 for reference\n",
    "# uploading data.jsonl to S3\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./data/data.jsonl\", bucket=s3_bucket_name, key_prefix=input_prefix_sales\n",
    ")\n",
    "\n",
    "# uploading hello_py_spark_app.py to S3\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./code/hello_py_spark_app.py\", bucket=s3_bucket_name, key_prefix=input_script_prefix\n",
    ")\n",
    "\n",
    "# uploading hello_py_spark_udfs.py to S3\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./code/hello_py_spark_udfs.py\", bucket=s3_bucket_name, key_prefix=input_pyfiles_prefix\n",
    ")\n",
    "\n",
    "# initializing ProcessingInputs,ProcessingResources and AppSpecification configurations\n",
    "# providing processing script\n",
    "processing_input_code = ProcessingInput(\n",
    "    input_name=\"code\",\n",
    "    s3_input=ProcessingS3Input(\n",
    "        s3_uri=f\"s3://{s3_bucket_name}{default_bucket_prefix_path}/{final_job_name}/input/code/hello_py_spark_app.py\",\n",
    "        s3_data_type=\"S3Prefix\",\n",
    "        local_path=\"/opt/ml/processing/input/code\",\n",
    "        s3_input_mode=\"File\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "# providing py files\n",
    "processing_input_pyfiles = ProcessingInput(\n",
    "    input_name=\"py-files\",\n",
    "    s3_input=ProcessingS3Input(\n",
    "        s3_uri=f\"s3://{s3_bucket_name}{default_bucket_prefix_path}/{final_job_name}/input/py-files\",\n",
    "        s3_data_type=\"S3Prefix\",\n",
    "        local_path=\"/opt/ml/processing/input/py-files\",\n",
    "        s3_input_mode=\"File\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "# providing processing resources\n",
    "processing_resources = ProcessingResources(\n",
    "    cluster_config=ProcessingClusterConfig(\n",
    "        instance_count=2, instance_type=\"ml.m5.xlarge\", volume_size_in_gb=30\n",
    "    )\n",
    ")\n",
    "\n",
    "# providing app specification\n",
    "app_specification = AppSpecification(\n",
    "    image_uri=\"173754725891.dkr.ecr.us-east-1.amazonaws.com/sagemaker-spark-processing:3.1-cpu\",\n",
    "    container_entrypoint=[\n",
    "        \"smspark-submit\",\n",
    "        \"--py-files\",\n",
    "        \"/opt/ml/processing/input/py-files\",\n",
    "        \"/opt/ml/processing/input/code/hello_py_spark_app.py\",\n",
    "    ],\n",
    "    container_arguments=[\n",
    "        \"--input\",\n",
    "        f\"s3://{s3_bucket_name}/{input_prefix_sales}\",\n",
    "        \"--output\",\n",
    "        f\"s3://{s3_bucket_name}/{output_prefix_sales}\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Run the processing job\n",
    "processing_job_obj = ProcessingJob.create(\n",
    "    processing_job_name=final_job_name,\n",
    "    processing_resources=processing_resources,\n",
    "    app_specification=app_specification,\n",
    "    role_arn=role,\n",
    "    processing_inputs=[processing_input_code, processing_input_pyfiles],\n",
    ")\n",
    "\n",
    "processing_job_obj.wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validate Data Processing Results\n",
    "\n",
    "Next, validate the output of the Spark job by ensuring that the output URI contains the Spark `_SUCCESS` file along with the output json lines file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Output files in {}\".format(output_s3_uri))\n",
    "!aws s3 ls $output_s3_uri/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 3: Run a Java/Scala Spark application\n",
    "\n",
    "In the next example, you'll take a Spark application jar (located in `./code/spark-test-app.jar`) that is already built and run it using SageMaker Processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker_core.shapes import (\n",
    "    ProcessingInput,\n",
    "    ProcessingResources,\n",
    "    AppSpecification,\n",
    "    ProcessingS3Input,\n",
    "    ProcessingOutputConfig,\n",
    ")\n",
    "from sagemaker_core.shapes import (\n",
    "    ProcessingResources,\n",
    "    ProcessingClusterConfig,\n",
    "    ProcessingOutput,\n",
    "    ProcessingS3Output,\n",
    ")\n",
    "from sagemaker_core.resources import ProcessingJob\n",
    "\n",
    "# Upload the raw input dataset to S3\n",
    "timestamp_prefix = strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "prefix = \"sagemaker/spark-preprocess-demo/{}\".format(timestamp_prefix)\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    prefix = f\"{default_bucket_prefix}/{prefix}\"\n",
    "\n",
    "input_prefix_sales = \"{}/input/sales\".format(prefix)\n",
    "output_prefix_sales = \"{}/output/sales\".format(prefix)\n",
    "input_s3_uri = \"s3://{}/{}\".format(s3_bucket_name, input_prefix_sales)\n",
    "output_s3_uri = \"s3://{}/{}\".format(s3_bucket_name, output_prefix_sales)\n",
    "\n",
    "\n",
    "base_job_name = \"sm-spark-java\"\n",
    "final_job_name = base_job_name + \"-\" + timestamp_prefix\n",
    "\n",
    "input_script_prefix = \"{}/input/code\".format(final_job_name)\n",
    "input_pyfiles_prefix = \"{}/input/py-files\".format(final_job_name)\n",
    "\n",
    "# uploading required data to S3 for reference\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./data/data.jsonl\", bucket=s3_bucket_name, key_prefix=input_prefix_sales\n",
    ")\n",
    "\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./code/spark-test-app.jar\", bucket=s3_bucket_name, key_prefix=input_script_prefix\n",
    ")\n",
    "\n",
    "# initializing ProcessingInputs,ProcessingResources and AppSpecification configurations\n",
    "processing_input_code = ProcessingInput(\n",
    "    input_name=\"code\",\n",
    "    s3_input=ProcessingS3Input(\n",
    "        s3_uri=f\"s3://{s3_bucket_name}{default_bucket_prefix_path}/{final_job_name}/input/code/spark-test-app.jar\",\n",
    "        # s3_uri=\"s3://sagemaker-us-east-1-774297356213/sm-spark-2024-08-30-05-25-18-294/input/code/preprocess.py\",\n",
    "        s3_data_type=\"S3Prefix\",\n",
    "        local_path=\"/opt/ml/processing/input/code\",\n",
    "        s3_input_mode=\"File\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "processing_resources = ProcessingResources(\n",
    "    cluster_config=ProcessingClusterConfig(\n",
    "        instance_count=2, instance_type=\"ml.m5.xlarge\", volume_size_in_gb=30\n",
    "    )\n",
    ")\n",
    "\n",
    "\n",
    "app_specification = AppSpecification(\n",
    "    image_uri=\"173754725891.dkr.ecr.us-east-1.amazonaws.com/sagemaker-spark-processing:3.1-cpu\",\n",
    "    container_entrypoint=[\n",
    "        \"smspark-submit\",\n",
    "        \"--class\",\n",
    "        \"com.amazonaws.sagemaker.spark.test.HelloJavaSparkApp\",\n",
    "        \"/opt/ml/processing/input/code/spark-test-app.jar\",\n",
    "    ],\n",
    "    container_arguments=[\n",
    "        \"--input\",\n",
    "        f\"s3://{s3_bucket_name}/{input_prefix_sales}\",\n",
    "        \"--output\",\n",
    "        f\"s3://{s3_bucket_name}/{output_prefix_sales}\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Run the processing job\n",
    "processing_job_obj = ProcessingJob.create(\n",
    "    processing_job_name=final_job_name,\n",
    "    processing_resources=processing_resources,\n",
    "    app_specification=app_specification,\n",
    "    role_arn=role,\n",
    "    processing_inputs=[processing_input_code, processing_input_pyfiles],\n",
    ")\n",
    "\n",
    "# waiting for the processing job to be completed\n",
    "processing_job_obj.wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 4: Specifying additional Spark configuration\n",
    "\n",
    "Overriding Spark configuration is crucial for a number of tasks such as tuning your Spark application or configuring the Hive metastore. Using the SageMaker Python SDK, you can easily override Spark/Hive/Hadoop configuration.\n",
    "\n",
    "The next example demonstrates this by overriding Spark executor memory/cores.\n",
    "\n",
    "For more information on configuring your Spark application, see the EMR documentation on [Configuring Applications](https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-configure-apps.html)\n",
    "\n",
    "#### Creating configuration.json file for overriding Spark executor memory/cores "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./code/configuration.json\n",
    "[{\"Classification\": \"spark-defaults\", \"Properties\": {\"spark.executor.memory\": \"2g\", \"spark.executor.cores\": \"1\"}}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker_core.shapes import (\n",
    "    ProcessingInput,\n",
    "    ProcessingResources,\n",
    "    AppSpecification,\n",
    "    ProcessingS3Input,\n",
    "    ProcessingOutputConfig,\n",
    ")\n",
    "from sagemaker_core.shapes import (\n",
    "    ProcessingResources,\n",
    "    ProcessingClusterConfig,\n",
    "    ProcessingOutput,\n",
    "    ProcessingS3Output,\n",
    ")\n",
    "from sagemaker_core.resources import ProcessingJob\n",
    "\n",
    "# Upload the raw input dataset to a unique S3 location\n",
    "timestamp_prefix = strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "prefix = \"sagemaker/spark-preprocess-demo/{}\".format(timestamp_prefix)\n",
    "\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    prefix = f\"{default_bucket_prefix}/{prefix}\"\n",
    "\n",
    "input_prefix_abalone = \"{}/input/raw/abalone\".format(prefix)\n",
    "input_preprocessed_prefix_abalone = \"{}/input/preprocessed/abalone\".format(prefix)\n",
    "\n",
    "# base job name\n",
    "base_job_name = \"sm-spark\"\n",
    "final_job_name = base_job_name + \"-\" + timestamp_prefix\n",
    "\n",
    "input_script_prefix = \"{}/input/code\".format(final_job_name)\n",
    "input_conf_prefix = \"{}/input/conf\".format(final_job_name)\n",
    "\n",
    "# uploading required data to S3 for reference\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./data/abalone.csv\", bucket=s3_bucket_name, key_prefix=input_prefix_abalone\n",
    ")\n",
    "\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./code/preprocess.py\", bucket=s3_bucket_name, key_prefix=input_script_prefix\n",
    ")\n",
    "\n",
    "sagemaker_session.upload_data(\n",
    "    path=\"./code/configuration.json\", bucket=s3_bucket_name, key_prefix=input_conf_prefix\n",
    ")\n",
    "\n",
    "# initializing ProcessingInputs,ProcessingResources and AppSpecification configurations\n",
    "processing_input_code = ProcessingInput(\n",
    "    input_name=\"code\",\n",
    "    s3_input=ProcessingS3Input(\n",
    "        s3_uri=f\"s3://{s3_bucket_name}{default_bucket_prefix_path}/{final_job_name}/input/code/preprocess.py\",\n",
    "        # s3_uri=\"s3://sagemaker-us-east-1-774297356213/sm-spark-2024-08-30-05-25-18-294/input/code/preprocess.py\",\n",
    "        s3_data_type=\"S3Prefix\",\n",
    "        local_path=\"/opt/ml/processing/input/code\",\n",
    "        s3_input_mode=\"File\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "processing_input_conf = ProcessingInput(\n",
    "    input_name=\"conf\",\n",
    "    s3_input=ProcessingS3Input(\n",
    "        s3_uri=f\"s3://{s3_bucket_name}{default_bucket_prefix_path}/{final_job_name}/input/conf/configuration.json\",\n",
    "        # s3_uri=\"s3://sagemaker-us-east-1-774297356213/sm-spark-2024-08-30-05-25-18-294/input/code/preprocess.py\",\n",
    "        s3_data_type=\"S3Prefix\",\n",
    "        local_path=\"/opt/ml/processing/input/conf\",\n",
    "        s3_input_mode=\"File\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "\n",
    "processing_resources = ProcessingResources(\n",
    "    cluster_config=ProcessingClusterConfig(\n",
    "        instance_count=2, instance_type=\"ml.m5.xlarge\", volume_size_in_gb=30\n",
    "    )\n",
    ")\n",
    "\n",
    "app_specification = AppSpecification(\n",
    "    image_uri=\"173754725891.dkr.ecr.us-east-1.amazonaws.com/sagemaker-spark-processing:3.1-cpu\",\n",
    "    container_entrypoint=[\"smspark-submit\", \"/opt/ml/processing/input/code/preprocess.py\"],\n",
    "    container_arguments=[\n",
    "        \"--s3_input_bucket\",\n",
    "        s3_bucket_name,\n",
    "        \"--s3_input_key_prefix\",\n",
    "        input_prefix_abalone,\n",
    "        \"--s3_output_bucket\",\n",
    "        s3_bucket_name,\n",
    "        \"--s3_output_key_prefix\",\n",
    "        input_preprocessed_prefix_abalone,\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Run the processing job\n",
    "processing_job_obj = ProcessingJob.create(\n",
    "    processing_job_name=final_job_name,\n",
    "    processing_resources=processing_resources,\n",
    "    app_specification=app_specification,\n",
    "    role_arn=role,\n",
    "    processing_inputs=[processing_input_code, processing_input_conf],\n",
    ")\n",
    "\n",
    "# waiting for the processing job to be completed\n",
    "processing_job_obj.wait()"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}