{
  "cells": [
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Notebook as a Step using Notebooks Executor\n",
     "The following sample shows how to use the notebook executor as part of a Vertex AI\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Libraries and Variables"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "!which pip"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "!pip install kfp --upgrade -q\n",
     "!pip install --upgrade google-cloud-aiplatform -q\n",
     "!pip install --upgrade google-cloud-pipeline-components -q"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import kfp\n",
     "import os\n",
     "\n",
     "from datetime import datetime\n",
     "from google.cloud import aiplatform\n",
     "from kfp.v2 import compiler\n",
     "import google.cloud.aiplatform as aip"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "kfp.__version__"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Variables\n",
     "PROJECT_ID = '<YOUR_PROJECT_ID>'\n",
     "REGION = 'us-central1' # This is where the KMS and Metastore resides from the project configuration\n",
     "\n",
     "ROOT_PATH = \"..\"\n",
     "PIPELINE_ROOT_PATH = f'gs://{PROJECT_ID}-vertex-root'\n",
     "PACKAGE_PATH = 'notebook-as-a-step-sample-pipeline.json'\n",
     "\n",
     "RUNNABLES_PATH = './runnables'\n",
     "COMPONENT_YAML_PATH = os.path.join(ROOT_PATH, 'component.yaml')\n",
     "if not os.path.isfile(COMPONENT_YAML_PATH):\n",
     "    print(f'COMPONENT_YAML_PATH does not exist')\n",
     "\n",
     "WORKING_BUCKET_NAME = f'{PROJECT_ID}-naas'\n",
     "INPUT_NOTEBOOK_FILE = f'gs://{WORKING_BUCKET_NAME}/runnables/run_create_bucket.ipynb'\n",
     "OUTPUT_NOTEBOOK_FOLDER = f'gs://{WORKING_BUCKET_NAME}/outputs'\n",
     "\n",
     "INPUT_NOTEBOOK_FILE"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "!gcloud config set project \"{PROJECT_ID}\"\n",
     "!gcloud config list"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# The service account used by the Pipelines must have access to this bucket.\n",
     "!gsutil ls \"{PIPELINE_ROOT_PATH}\" || gsutil mb -l \"{REGION}\" \"{PIPELINE_ROOT_PATH}\"\n",
     "!gsutil cp -r \"{RUNNABLES_PATH}/*\" \"gs://{WORKING_BUCKET_NAME}/runnables\"\n",
     "!gsutil ls \"gs://{WORKING_BUCKET_NAME}/runnables\""
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Prerequisites\n",
     "You need to configure your project as detailed in https://cloud.google.com/vertex-ai/docs/pipelines/configure-project"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "@kfp.dsl.pipeline(\n",
     "    name=\"notebook-as-a-step-sample\",\n",
     "    pipeline_root=PIPELINE_ROOT_PATH)\n",
     "def pipeline(\n",
     "        project: str, \n",
     "        execution_id: str, \n",
     "        input_notebook_file:str, \n",
     "        output_notebook_folder:str,\n",
     "        location:str,\n",
     "        master_type:str,\n",
     "        container_image_uri:str):\n",
     "    \n",
     "    execute_notebook_component = kfp.components.load_component_from_file(COMPONENT_YAML_PATH)\n",
     "    execute_notebook_op = execute_notebook_component(\n",
     "        project=project,\n",
     "        execution_id=execution_id,\n",
     "        input_notebook_file=input_notebook_file,\n",
     "        output_notebook_folder=output_notebook_folder,\n",
     "        location=location,\n",
     "        master_type=master_type,\n",
     "        container_image_uri=container_image_uri,\n",
     "        parameters=f'PROJECT_ID={project},EXECUTION_ID={execution_id}'\n",
     "    )"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "For additional details about:\n",
     "- Google Cloud Notebook Executor template, see the [ExecutionTemplate API documentation](https://cloud.google.com/notebooks/docs/reference/rest/v1/ExecutionTemplate).\n",
     "- Pipeline types, see [kfp.dsl.types](https://kubeflow-pipelines.readthedocs.io/en/latest/_modules/kfp/dsl/types.html).\n",
     "- Notebook parameters, see [Papermill parameters](https://papermill.readthedocs.io/en/latest/usage-parameterize.html)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Run pipeline"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "NOW = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
     "JOB_ID=f'naas-{NOW}'\n",
     "EXECUTION_ID=f'naas_{NOW}'\n",
     "\n",
     "compiler.Compiler().compile(\n",
     "    pipeline_func=pipeline, package_path=PACKAGE_PATH)\n",
     "\n",
     "job = aip.PipelineJob(\n",
     "    display_name='notebook-executor-pipeline',\n",
     "    template_path=PACKAGE_PATH,\n",
     "    job_id=JOB_ID,\n",
     "    parameter_values={\n",
     "        'project': PROJECT_ID,\n",
     "        'execution_id': EXECUTION_ID,\n",
     "        'input_notebook_file': INPUT_NOTEBOOK_FILE,\n",
     "        'output_notebook_folder': OUTPUT_NOTEBOOK_FOLDER,\n",
     "        'location': 'us-central1',\n",
     "        'master_type': 'n1-standard-4',\n",
     "        'container_image_uri': 'gcr.io/deeplearning-platform-release/base-cpu'\n",
     "    },\n",
     ")\n",
     "\n",
     "job.submit()"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": []
   }
  ],
  "metadata": {
   "environment": {
    "name": "managed-notebooks.m78",
    "type": "gcloud",
    "uri": "gcr.io/deeplearning-platform-release/managed-notebooks:m78"
   },
   "interpreter": {
    "hash": "7b5f95d16b1bb0de859c4a3a5a9d93554fa0e5fd056930a88b2762108d6d7c14"
   },
   "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.8.8"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 4
 }
 