sdk/python/jobs/spark/submit_spark_pipeline_jobs.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Submit pipeline jobs with Spark component (preview)\n",
    "Please see documentation page: [Submit Spark jobs in Azure Machine Learning](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.\n",
    "\n",
    "You can submit a Spark job from:\n",
    "- an Azure Machine Learning Notebook connected to an Azure Machine Learning compute instance. \n",
    "- [Visual Studio Code connected to an Azure Machine Learning compute instance](https://learn.microsoft.com/azure/machine-learning/how-to-set-up-vs-code-remote?tabs=studio).\n",
    "- your local computer that has [the Azure Machine Learning SDK for Python](https://learn.microsoft.com/python/api/overview/azure/ai-ml-readme) installed."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use an attached Synapse Spark pool"
   ]
  },
  {
   "attachments": {},
   "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](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."
   ]
  },
  {
   "attachments": {},
   "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",
    "attached_spark_pool_name = \"<ATTACHED_SPARK_POOL_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)"
   ]
  },
  {
   "attachments": {},
   "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",
    "attached_spark_pool_name_uai = \"<ATTACHED_SPARK_POOL_NAME_UAI>\"\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_uai\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)"
   ]
  },
  {
   "attachments": {},
   "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",
    "attached_spark_pool_name = \"<ATTACHED_SPARK_POOL_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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use a serverless Spark compute"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A serverless Spark compute is available in your workspace by default. Please see [this documentation page](https://learn.microsoft.com/azure/machine-learning/interactive-data-wrangling-with-apache-spark-azure-ml##serverless-spark-compute-in-azure-machine-learning-notebooks) for more information about serverless 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."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using managed identity\n",
    "For a serverless 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#attach-user-assigned-managed-identity-using-cli-v2) 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.3.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)"
   ]
  },
  {
   "attachments": {},
   "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.3.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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline job using default identity\n",
    "Default identity for the serverless 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.3.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
}