sdk/python/jobs/spark/submit_spark_pipeline_jobs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Submit pipeline jobs with Spark component (preview)\n",
    "Please see documentation page: [Submit Spark jobs in Azure Machine Learning (preview)](https://learn.microsoft.com/azure/machine-learning/how-to-submit-spark-jobs#spark-component-in-a-pipeline-job) for more details related to the code samples in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use an attached Synapse Spark pool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should have an attached Synapse Spark pool available in your workspace. Please see documentation page: [Attach and manage a Synapse Spark pool in Azure Machine Learning (preview)](https://learn.microsoft.com/azure/machine-learning/how-to-manage-synapse-spark-pool) for more details.\n",
    "\n",
    "**Note** - To ensure successful execution of spark job, the identity being used for the Spark job should be assigned **Contributor** and **Storage Blob Data Contributor** roles on the Azure storage account used for data input and output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using managed identity\n",
    "For an attached Synpase Spark pool, managed identity is compute identity of the attached Synapse Spark pool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient, dsl, spark, Input, Output\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.ml.entities import ManagedIdentityConfiguration\n",
    "from azure.ai.ml.constants import InputOutputModes\n",
    "\n",
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\"\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")\n",
    "\n",
    "spark_component = spark(\n",
    "    name=\"Spark Component 1\",\n",
    "    inputs={\n",
    "        \"titanic_data\": Input(type=\"uri_file\", mode=\"direct\"),\n",
    "    },\n",
    "    outputs={\n",
    "        \"wrangled_data\": Output(type=\"uri_folder\", mode=\"direct\"),\n",
    "    },\n",
    "    # The source folder of the component\n",
    "    code=\"./src\",\n",
    "    entry={\"file\": \"titanic.py\"},\n",
    "    driver_cores=1,\n",
    "    driver_memory=\"2g\",\n",
    "    executor_cores=2,\n",
    "    executor_memory=\"2g\",\n",
    "    executor_instances=2,\n",
    "    args=\"--titanic_data ${{inputs.titanic_data}} --wrangled_data ${{outputs.wrangled_data}}\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    description=\"Sample Pipeline with Spark component\",\n",
    ")\n",
    "def spark_pipeline(spark_input_data):\n",
    "    spark_step = spark_component(titanic_data=spark_input_data)\n",
    "    spark_step.inputs.titanic_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.outputs.wrangled_data = Output(\n",
    "        type=\"uri_folder\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/wrangled/\",\n",
    "    )\n",
    "    spark_step.outputs.wrangled_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.identity = ManagedIdentityConfiguration()\n",
    "    spark_step.compute = \"<ATTACHED_SPARK_POOL_NAME>\"\n",
    "\n",
    "\n",
    "pipeline = spark_pipeline(\n",
    "    spark_input_data=Input(\n",
    "        type=\"uri_file\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/titanic.csv\",\n",
    "    )\n",
    ")\n",
    "\n",
    "pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline,\n",
    "    experiment_name=\"Titanic-Spark-Pipeline-SDK-1\",\n",
    ")\n",
    "\n",
    "# Wait until the job completes\n",
    "ml_client.jobs.stream(pipeline_job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using user identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient, dsl, spark, Input, Output\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.ml.entities import UserIdentityConfiguration\n",
    "from azure.ai.ml.constants import InputOutputModes\n",
    "\n",
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\"\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")\n",
    "\n",
    "spark_component = spark(\n",
    "    name=\"Spark Component 2\",\n",
    "    inputs={\n",
    "        \"titanic_data\": Input(type=\"uri_file\", mode=\"direct\"),\n",
    "    },\n",
    "    outputs={\n",
    "        \"wrangled_data\": Output(type=\"uri_folder\", mode=\"direct\"),\n",
    "    },\n",
    "    # The source folder of the component\n",
    "    code=\"./src\",\n",
    "    entry={\"file\": \"titanic.py\"},\n",
    "    driver_cores=1,\n",
    "    driver_memory=\"2g\",\n",
    "    executor_cores=2,\n",
    "    executor_memory=\"2g\",\n",
    "    executor_instances=2,\n",
    "    args=\"--titanic_data ${{inputs.titanic_data}} --wrangled_data ${{outputs.wrangled_data}}\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    description=\"Sample Pipeline with Spark component\",\n",
    ")\n",
    "def spark_pipeline(spark_input_data):\n",
    "    spark_step = spark_component(titanic_data=spark_input_data)\n",
    "    spark_step.inputs.titanic_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.outputs.wrangled_data = Output(\n",
    "        type=\"uri_folder\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/wrangled/\",\n",
    "    )\n",
    "    spark_step.outputs.wrangled_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.identity = UserIdentityConfiguration()\n",
    "    spark_step.compute = \"<ATTACHED_SPARK_POOL_NAME>\"\n",
    "\n",
    "\n",
    "pipeline = spark_pipeline(\n",
    "    spark_input_data=Input(\n",
    "        type=\"uri_file\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/titanic.csv\",\n",
    "    )\n",
    ")\n",
    "\n",
    "pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline,\n",
    "    experiment_name=\"Titanic-Spark-Pipeline-SDK-2\",\n",
    ")\n",
    "\n",
    "# Wait until the job completes\n",
    "ml_client.jobs.stream(pipeline_job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using default identity\n",
    "Default identity for an attached Synpase Spark pool is managed identity (compute identity of the attached Synapse Spark pool)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient, dsl, spark, Input, Output\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.ml.constants import InputOutputModes\n",
    "\n",
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\"\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")\n",
    "\n",
    "spark_component = spark(\n",
    "    name=\"Spark Component 3\",\n",
    "    inputs={\n",
    "        \"titanic_data\": Input(type=\"uri_file\", mode=\"direct\"),\n",
    "    },\n",
    "    outputs={\n",
    "        \"wrangled_data\": Output(type=\"uri_folder\", mode=\"direct\"),\n",
    "    },\n",
    "    # The source folder of the component\n",
    "    code=\"./src\",\n",
    "    entry={\"file\": \"titanic.py\"},\n",
    "    driver_cores=1,\n",
    "    driver_memory=\"2g\",\n",
    "    executor_cores=2,\n",
    "    executor_memory=\"2g\",\n",
    "    executor_instances=2,\n",
    "    args=\"--titanic_data ${{inputs.titanic_data}} --wrangled_data ${{outputs.wrangled_data}}\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(description=\"Sample Pipeline with Spark component\")\n",
    "def spark_pipeline(spark_input_data):\n",
    "    spark_step = spark_component(titanic_data=spark_input_data)\n",
    "    spark_step.inputs.titanic_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.outputs.wrangled_data = Output(\n",
    "        type=\"uri_folder\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/wrangled/\",\n",
    "    )\n",
    "    spark_step.outputs.wrangled_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.compute = \"<ATTACHED_SPARK_POOL_NAME>\"\n",
    "\n",
    "\n",
    "pipeline = spark_pipeline(\n",
    "    spark_input_data=Input(\n",
    "        type=\"uri_file\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/titanic.csv\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline, experiment_name=\"Titanic-Spark-Pipeline-SDK-3\"\n",
    ")\n",
    "\n",
    "# Wait until the job completes\n",
    "ml_client.jobs.stream(pipeline_job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use a Managed (Automatic) Spark compute"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should have a Managed (Automatic) Spark compute available in your workspace. Please see [this documentation page](https://learn.microsoft.com/azure/machine-learning/interactive-data-wrangling-with-apache-spark-azure-ml#create-and-configure-managed-automatic-spark-compute-in-azure-machine-learning-notebooks) for more information about creating Managed (Automatic) Spark compute.\n",
    "\n",
    "**Note** - To ensure successful execution of pipeline job with Spark component, the identity being used for the job should be assigned **Contributor** and **Storage Blob Data Contributor** roles on the Azure storage account used for data input and output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using managed identity\n",
    "For a Managed (Automatic) Spark compute, managed identity is the user assigned managed identity attached to the Azure Machine Learning workspace. Please see [this documentation page to learn how to attach the user assigned managed identity](https://learn.microsoft.com/azure/machine-learning/how-to-submit-spark-jobs#ensuring-resource-access-for-spark-jobs) to the Azure Machine Learning workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient, dsl, spark, Input, Output\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.ml.entities import ManagedIdentityConfiguration\n",
    "from azure.ai.ml.constants import InputOutputModes\n",
    "\n",
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\"\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")\n",
    "\n",
    "spark_component = spark(\n",
    "    name=\"Spark Component 4\",\n",
    "    inputs={\n",
    "        \"titanic_data\": Input(type=\"uri_file\", mode=\"direct\"),\n",
    "    },\n",
    "    outputs={\n",
    "        \"wrangled_data\": Output(type=\"uri_folder\", mode=\"direct\"),\n",
    "    },\n",
    "    # The source folder of the component\n",
    "    code=\"./src\",\n",
    "    entry={\"file\": \"titanic.py\"},\n",
    "    driver_cores=1,\n",
    "    driver_memory=\"2g\",\n",
    "    executor_cores=2,\n",
    "    executor_memory=\"2g\",\n",
    "    executor_instances=2,\n",
    "    args=\"--titanic_data ${{inputs.titanic_data}} --wrangled_data ${{outputs.wrangled_data}}\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    description=\"Sample Pipeline with Spark component\",\n",
    ")\n",
    "def spark_pipeline(spark_input_data):\n",
    "    spark_step = spark_component(titanic_data=spark_input_data)\n",
    "    spark_step.inputs.titanic_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.outputs.wrangled_data = Output(\n",
    "        type=\"uri_folder\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/wrangled/\",\n",
    "    )\n",
    "    spark_step.outputs.wrangled_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.identity = ManagedIdentityConfiguration()\n",
    "    spark_step.resources = {\n",
    "        \"instance_type\": \"Standard_E8S_V3\",\n",
    "        \"runtime_version\": \"3.2.0\",\n",
    "    }\n",
    "\n",
    "\n",
    "pipeline = spark_pipeline(\n",
    "    spark_input_data=Input(\n",
    "        type=\"uri_file\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/titanic.csv\",\n",
    "    )\n",
    ")\n",
    "\n",
    "pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline,\n",
    "    experiment_name=\"Titanic-Spark-Pipeline-SDK-4\",\n",
    ")\n",
    "\n",
    "# Wait until the job completes\n",
    "ml_client.jobs.stream(pipeline_job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using user identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient, dsl, spark, Input, Output\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.ml.entities import UserIdentityConfiguration\n",
    "from azure.ai.ml.constants import InputOutputModes\n",
    "\n",
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\"\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")\n",
    "\n",
    "spark_component = spark(\n",
    "    name=\"Spark Component 5\",\n",
    "    inputs={\n",
    "        \"titanic_data\": Input(type=\"uri_file\", mode=\"direct\"),\n",
    "    },\n",
    "    outputs={\n",
    "        \"wrangled_data\": Output(type=\"uri_folder\", mode=\"direct\"),\n",
    "    },\n",
    "    # The source folder of the component\n",
    "    code=\"./src\",\n",
    "    entry={\"file\": \"titanic.py\"},\n",
    "    driver_cores=1,\n",
    "    driver_memory=\"2g\",\n",
    "    executor_cores=2,\n",
    "    executor_memory=\"2g\",\n",
    "    executor_instances=2,\n",
    "    args=\"--titanic_data ${{inputs.titanic_data}} --wrangled_data ${{outputs.wrangled_data}}\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    description=\"Sample Pipeline with Spark component\",\n",
    ")\n",
    "def spark_pipeline(spark_input_data):\n",
    "    spark_step = spark_component(titanic_data=spark_input_data)\n",
    "    spark_step.inputs.titanic_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.outputs.wrangled_data = Output(\n",
    "        type=\"uri_folder\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/wrangled/\",\n",
    "    )\n",
    "    spark_step.outputs.wrangled_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.identity = UserIdentityConfiguration()\n",
    "    spark_step.resources = {\n",
    "        \"instance_type\": \"Standard_E8S_V3\",\n",
    "        \"runtime_version\": \"3.2.0\",\n",
    "    }\n",
    "\n",
    "\n",
    "pipeline = spark_pipeline(\n",
    "    spark_input_data=Input(\n",
    "        type=\"uri_file\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/titanic.csv\",\n",
    "    )\n",
    ")\n",
    "\n",
    "pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline,\n",
    "    experiment_name=\"Titanic-Spark-Pipeline-SDK-5\",\n",
    ")\n",
    "\n",
    "# Wait until the job completes\n",
    "ml_client.jobs.stream(pipeline_job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using default identity\n",
    "Default identity for the Managed (Automatic) Spark compute is user identity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient, dsl, spark, Input, Output\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.ai.ml.constants import InputOutputModes\n",
    "\n",
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\"\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")\n",
    "\n",
    "spark_component = spark(\n",
    "    name=\"Spark Component 6\",\n",
    "    inputs={\n",
    "        \"titanic_data\": Input(type=\"uri_file\", mode=\"direct\"),\n",
    "    },\n",
    "    outputs={\n",
    "        \"wrangled_data\": Output(type=\"uri_folder\", mode=\"direct\"),\n",
    "    },\n",
    "    # The source folder of the component\n",
    "    code=\"./src\",\n",
    "    entry={\"file\": \"titanic.py\"},\n",
    "    driver_cores=1,\n",
    "    driver_memory=\"2g\",\n",
    "    executor_cores=2,\n",
    "    executor_memory=\"2g\",\n",
    "    executor_instances=2,\n",
    "    args=\"--titanic_data ${{inputs.titanic_data}} --wrangled_data ${{outputs.wrangled_data}}\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    description=\"Sample Pipeline with Spark component\",\n",
    ")\n",
    "def spark_pipeline(spark_input_data):\n",
    "    spark_step = spark_component(titanic_data=spark_input_data)\n",
    "    spark_step.inputs.titanic_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.outputs.wrangled_data = Output(\n",
    "        type=\"uri_folder\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/wrangled/\",\n",
    "    )\n",
    "    spark_step.outputs.wrangled_data.mode = InputOutputModes.DIRECT\n",
    "    spark_step.resources = {\n",
    "        \"instance_type\": \"Standard_E8S_V3\",\n",
    "        \"runtime_version\": \"3.2.0\",\n",
    "    }\n",
    "\n",
    "\n",
    "pipeline = spark_pipeline(\n",
    "    spark_input_data=Input(\n",
    "        type=\"uri_file\",\n",
    "        path=\"azureml://datastores/workspaceblobstore/paths/data/titanic.csv\",\n",
    "    )\n",
    ")\n",
    "\n",
    "pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline,\n",
    "    experiment_name=\"Titanic-Spark-Pipeline-SDK-6\",\n",
    ")\n",
    "\n",
    "# Wait until the job completes\n",
    "ml_client.jobs.stream(pipeline_job.name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK V2",
   "language": "python",
   "name": "python310-sdkv2"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "6aeff17a1aa7735c2f7cb3a6d691fe1b4d4c3b8d2d650f644ad0f24e1b8e3f3f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}