ai-pipeline/03-DataBricksComputePi.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure DataBricks integration with Kubeflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Azure Databricks Package provides a set of Kubeflow Pipeline Tasks (Ops) which let us manipulate Databricks resources using the [Azure Databricks Operator for Kubernetes](https://github.com/microsoft/azure-databricks-operator). This makes the user experience much nicer, and less error prone, than using the ResourceOp to manipulate these Databricks resources.\n",
    "\n",
    "## Supported Ops\n",
    "\n",
    "* CreateClusterOp, to create a cluster in Databricks.\n",
    "* DeleteClusterOp, to delete a cluster created with CreateClusterOp.\n",
    "* CreateJobOp, to create a Spark job in Databricks.\n",
    "* DeleteJobOp, to delete a job created with CreateJobOp.\n",
    "* SubmitRunOp, to submit a job run in Databricks.\n",
    "* DeleteRunOp, to delete a run submitted with SubmitRunOp.\n",
    "* CreateSecretScopeOp, to create a secret scope in Databricks.\n",
    "* DeleteSecretScopeOp, to delete a secret scope created with CreateSecretScopeOp.\n",
    "* ImportWorkspaceItemOp, to import an item into a Databricks Workspace.\n",
    "* DeleteWorkspaceItemOp, to delete an item imported with ImportWorkspaceItemOp.\n",
    "* CreateDbfsBlockOp, to create Dbfs Block in Databricks.\n",
    "* DeleteDbfsBlockOp, to delete Dbfs Block created with CreateDbfsBlockOp.\n",
    "\n",
    "For each of these there are two ways a Kubeflow user can create the Ops:\n",
    "\n",
    "* By passing the complete Databricks spec for the Op within a Python Dictionary.\n",
    "* By using named parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings; warnings.simplefilter('ignore')\n",
    "\n",
    "import os\n",
    "\n",
    "x_auth_token = os.environ.get('X_AUTH_TOKEN')\n",
    "\n",
    "print('X-Auth-Token for DataBricks: ', x_auth_token)\n",
    "\n",
    "os.environ['X_AUTH_TOKEN'] = x_auth_token"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define a pipeline function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kfp.dsl as dsl\n",
    "import kfp.compiler as compiler\n",
    "import databricks\n",
    "\n",
    "def create_cluster(cluster_name):\n",
    "    return databricks.CreateClusterOp(\n",
    "        name=\"createcluster\",\n",
    "        cluster_name=cluster_name,\n",
    "        spark_version=\"5.3.x-scala2.11\",\n",
    "        node_type_id=\"Standard_D3_v2\",\n",
    "        spark_conf={\n",
    "            \"spark.speculation\": \"true\"\n",
    "        },\n",
    "        num_workers=2\n",
    "    )\n",
    "\n",
    "def submit_run(run_name, cluster_id, parameter):\n",
    "    return databricks.SubmitRunOp(\n",
    "        name=\"submitrun\",\n",
    "        run_name=run_name,\n",
    "        existing_cluster_id=cluster_id,\n",
    "        libraries=[{\"jar\": \"dbfs:/docs/sparkpi.jar\"}],\n",
    "        spark_jar_task={\n",
    "            \"main_class_name\": \"org.apache.spark.examples.SparkPi\",\n",
    "            \"parameters\": [parameter]\n",
    "        }\n",
    "    )\n",
    "\n",
    "def delete_run(run_name):\n",
    "    return databricks.DeleteRunOp(\n",
    "        name=\"deleterun\",\n",
    "        run_name=run_name\n",
    "    )\n",
    "\n",
    "def delete_cluster(cluster_name):\n",
    "    return databricks.DeleteClusterOp(\n",
    "        name=\"deletecluster\",\n",
    "        cluster_name=cluster_name\n",
    "    )\n",
    "\n",
    "@dsl.pipeline(\n",
    "    name=\"DatabricksCluster\",\n",
    "    description=\"A toy pipeline that computes an approximation to pi with Azure Databricks.\"\n",
    ")\n",
    "def pipeline_calc(cluster_name=\"test-cluster\", run_name=\"test-run\", parameter=\"10\"):\n",
    "    create_cluster_task = create_cluster(cluster_name)\n",
    "    submit_run_task = submit_run(run_name, create_cluster_task.outputs[\"cluster_id\"], parameter)\n",
    "    delete_run_task = delete_run(run_name)\n",
    "    delete_run_task.after(submit_run_task)\n",
    "    delete_cluster_task = delete_cluster(cluster_name)\n",
    "    delete_cluster_task.after(delete_run_task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compile the pipeline\n",
    "\n",
    "Compile the pipeline into a tar package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    compiler.Compiler()._create_and_write_workflow(\n",
    "        pipeline_func=pipeline_calc,\n",
    "        package_path=\"pipeline_calc.tar.gz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit and run the pipeline with parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run = client.run_pipeline(exp.id, 'pipeline-databricks-' + time.strftime(\"%Y%m%d-%H%M%S\"), 'pipeline_databricks.zip',\n",
    "                          params={'cluster_name': 'test-cluster',\n",
    "                                  'run_name': 'test-run',\n",
    "                                  'parameter': '10'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Additional examples\n",
    "\n",
    "More sample pipelines can be found in:\n",
    "\n",
    "* [samples/contrib/azure-samples/databricks-pipelines](https://github.com/kubeflow/pipelines/blob/master/samples/contrib/azure-samples/databricks-pipelines) \n",
    "* [samples/contrib/azure-samples/kfp-azure-databricks/tests](https://github.com/kubeflow/pipelines/blob/master/samples/contrib/azure-samples/kfp-azure-databricks/tests)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Additional information\n",
    "\n",
    "* [Kubeflow Pipelines](https://www.kubeflow.org/docs/pipelines/)\n",
    "* [Azure Databricks documentation](https://docs.microsoft.com/azure/azure-databricks/)\n",
    "* [Azure Databricks Operator for Kubernetes](https://github.com/microsoft/azure-databricks-operator)\n",
    "* [Golang SDK for DataBricks REST API 2.0 and Azure DataBricks REST API 2.0, used by Azure Databricks Operator](https://github.com/xinsnake/databricks-sdk-golang)\n",
    "* [Databricks REST API 2.0](https://docs.databricks.com/dev-tools/api/latest/index.html)\n",
    "* [Azure Databricks REST API 2.0](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/)\n",
    "\n",
    "The following articles provide information on the supported spec fields for the supported Databricks Ops:\n",
    "\n",
    "* **Cluster Ops**: [Azure Databricks Cluster API](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/clusters)\n",
    "* **Job Ops**: [Azure Databricks Jobs API](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/jobs)\n",
    "* **Run Ops**: [Azure Databricks Jobs API - Runs Submit](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/jobs#--runs-submit)\n",
    "* **Secret Scope Ops**: [Azure Databricks Secrets API](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/secrets)\n",
    "* **Workspace Item Ops**: [Azure Databricks Workspace API](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/workspace)\n",
    "* **DbfsBlock Ops**: [Azure Databricks DBFS API](https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/dbfs)"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}