Add wr.emr.submit_spark_step

Author: igorborgest

awswrangler/emr.py

@@ -1032,3 +1032,61 @@ def build_spark_step(
         region=region,
         boto3_session=boto3_session,
     )
+
+
+def submit_spark_step(
+    cluster_id: str,
+    path: str,
+    deploy_mode: str = "cluster",
+    docker_image: Optional[str] = None,
+    name: str = "my-step",
+    action_on_failure: str = "CONTINUE",
+    region: Optional[str] = None,
+    boto3_session: Optional[boto3.Session] = None,
+) -> str:
+    """Submit Spark Step.
+
+    Parameters
+    ----------
+    cluster_id : str
+        Cluster ID.
+    path : str
+        Script path. (e.g. s3://bucket/app.py)
+    deploy_mode : str
+        "cluster" | "client"
+    docker_image : str, optional
+        e.g. "{ACCOUNT_ID}.dkr.ecr.{REGION}.amazonaws.com/{IMAGE_NAME}:{TAG}"
+    name : str, optional
+        Step name.
+    action_on_failure : str
+        'TERMINATE_JOB_FLOW', 'TERMINATE_CLUSTER', 'CANCEL_AND_WAIT', 'CONTINUE'
+    region: str, optional
+        Region name to not get it from boto3.Session. (e.g. `us-east-1`)
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    str
+        Step ID.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> step_id = wr.emr.submit_spark_step(
+    >>>     cluster_id="cluster-id",
+    >>>     path="s3://bucket/emr/app.py"
+    >>> )
+
+    """
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    step = build_spark_step(
+        path=path,
+        deploy_mode=deploy_mode,
+        docker_image=docker_image,
+        name=name,
+        action_on_failure=action_on_failure,
+        region=region,
+        boto3_session=session,
+    )
+    return submit_steps(cluster_id=cluster_id, steps=[step], boto3_session=session)[0]

docs/source/api.rst

@@ -113,10 +113,12 @@ EMR
     get_cluster_state
     terminate_cluster
     submit_step
+    submit_spark_step
+    submit_ecr_credentials_refresh
     submit_steps
     build_step
+    build_spark_step
     get_step_state
-    update_ecr_credentials
 
 CloudWatch Logs
 ---------------

testing/test_awswrangler/test_emr.py

@@ -182,4 +182,5 @@ def test_docker(bucket, cloudformation_outputs):
             )
         ],
     )
+    wr.emr.submit_spark_step(cluster_id=cluster_id, path=f"s3://{bucket}/emr/test_docker.py")
     wr.emr.terminate_cluster(cluster_id=cluster_id)

tutorials/16 - EMR & Docker.ipynb

@@ -11,12 +11,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import awswrangler as wr\n",
-    "import boto3"
+    "import boto3\n",
+    "import getpass"
    ]
   },
   {
@@ -40,7 +41,6 @@
     }
    ],
    "source": [
-    "import getpass\n",
     "bucket = getpass.getpass()"
    ]
   },
@@ -164,7 +164,7 @@
     {
      "data": {
       "text/plain": [
-       "'s-3OPMPDCYGEGOT'"
+       "'s-1B0O45RWJL8CL'"
       ]
      },
      "execution_count": 5,
@@ -173,7 +173,7 @@
     }
    ],
    "source": [
-    "wr.emr.submit_ecr_credentials_refresh(cluster_id, path=f\"s3://{bucket}/emr/\")"
+    "wr.emr.submit_ecr_credentials_refresh(cluster_id, path=f\"s3://{bucket}/\")"
    ]
   },
   {
@@ -185,7 +185,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -201,11 +201,7 @@
     "print(f\"Wrangler version: {wr.__version__}\")\n",
     "\"\"\"\n",
     "\n",
-    "_ = boto3.client(\"s3\").put_object(\n",
-    "    Body=script,\n",
-    "    Bucket=bucket,\n",
-    "    Key=\"emr/test_docker.py\"\n",
-    ")"
+    "boto3.client(\"s3\").put_object(Body=script, Bucket=bucket, Key=\"test_docker.py\");"
    ]
   },
   {
@@ -217,15 +213,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
     "DOCKER_IMAGE = f\"{wr.get_account_id()}.dkr.ecr.us-east-1.amazonaws.com/emr-wrangler:emr-wrangler\"\n",
     "\n",
-    "step = wr.emr.build_spark_step(f\"s3://{bucket}/emr/test_docker.py\", docker_image=DOCKER_IMAGE)\n",
-    "\n",
-    "steps_ids = wr.emr.submit_steps(cluster_id, steps=[step])"
+    "step_id = wr.emr.submit_spark_step(\n",
+    "    cluster_id,\n",
+    "    f\"s3://{bucket}/test_docker.py\",\n",
+    "    docker_image=DOCKER_IMAGE\n",
+    ")"
    ]
   },
   {
@@ -237,11 +235,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "while wr.emr.get_step_state(cluster_id, steps_ids[0]) != \"COMPLETED\":\n",
+    "while wr.emr.get_step_state(cluster_id, step_id) != \"COMPLETED\":\n",
     "    pass"
    ]
   },
@@ -254,7 +252,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -270,7 +268,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -313,11 +311,10 @@
     "\n",
     "DOCKER_IMAGE = f\"{wr.get_account_id()}.dkr.ecr.us-east-1.amazonaws.com/emr-wrangler:emr-wrangler\"\n",
     "\n",
-    "steps_ids = wr.emr.submit_steps(\n",
-    "    cluster_id=cluster_id,\n",
-    "    steps=[\n",
-    "        wr.emr.build_spark_step(f\"s3://{bucket}/emr/test_docker.py\", docker_image=DOCKER_IMAGE)\n",
-    "    ]\n",
+    "step_id = wr.emr.submit_spark_step(\n",
+    "    cluster_id,\n",
+    "    f\"s3://{bucket}/test_docker.py\",\n",
+    "    docker_image=DOCKER_IMAGE\n",
     ")"
    ]
   },
@@ -327,7 +324,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "while wr.emr.get_step_state(cluster_id, steps_ids[0]) != \"COMPLETED\":\n",
+    "while wr.emr.get_step_state(cluster_id, step_id) != \"COMPLETED\":\n",
     "    pass\n",
     "\n",
     "wr.emr.terminate_cluster(cluster_id)"