{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "483b2f3f",
      "metadata": {
        "id": "483b2f3f"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "81db39f3",
      "metadata": {
        "id": "81db39f3"
      },
      "source": [
        "# Vertex Pipelines: Dataflow Python Job OP\n",
        "\n",
        "## Overview\n",
        "This notebook shows how to use the `DataflowPythonJobOp` to create a Python Dataflow Job component. `DataflowPythonJobOp` creates a pipeline component that prepares data by submitting an Apache Beam job (authored in Python) to Cloud Dataflow for execution. The Python Beam code is run with Cloud Dataflow Runner. learn more about [Google Cloud Dataflow Runner](https://beam.apache.org/documentation/runners/dataflow/) here.\n",
        "\n",
        "\n",
        "For more details on `DataflowPythonJobOp` interface please see the [API doc](https://google-cloud-pipeline-components.readthedocs.io/)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "94160504",
      "metadata": {
        "id": "94160504"
      },
      "source": [
        "### Install required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a5db0dbb",
      "metadata": {
        "id": "a5db0dbb"
      },
      "outputs": [],
      "source": [
        "!pip3 install  -U google-cloud-pipeline-components -q"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8e2d800b",
      "metadata": {
        "id": "8e2d800b"
      },
      "source": [
        "## Before you begin\n",
        "Set your Project ID, Location, Pipeline Root, and a few parameters required for the Dataflow sample."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f8da5659",
      "metadata": {
        "id": "f8da5659"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = 'YOUR_PROJECT_ID'\n",
        "LOCATION = \"us-central1\"\n",
        "PIPELINE_ROOT = 'gs://YOUR_BUCKET_NAME' # No ending slash\n",
        "\n",
        "# Dataflow sample parameters\n",
        "PIPELINE_NAME = 'dataflow-pipeline-sample'\n",
        "OUTPUT_FILE = '{}/wc/wordcount.out'.format(PIPELINE_ROOT)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6bb598bc",
      "metadata": {
        "id": "6bb598bc"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "501938b8",
      "metadata": {
        "id": "501938b8"
      },
      "outputs": [],
      "source": [
        "from google_cloud_pipeline_components.experimental.dataflow import DataflowPythonJobOp\n",
        "from google_cloud_pipeline_components.experimental.wait_gcp_resources import WaitGcpResourcesOp"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "847962bd",
      "metadata": {
        "id": "847962bd"
      },
      "source": [
        "## Create a pipeline using DataflowPythonJobOp and WaitGcpResourcesOp\n",
        "In this section we create a pipeline using the `DataflowPythonJobOp` and the [Apache Beam WordCount Examples](https://beam.apache.org/get-started/wordcount-example/). Then we use the 'WaitGcpResourcesOp' to poll the resource status and wait for it to finish.\n",
        "To use the 'WaitGcpResourcesOp' component, first create the `DataflowPythonJobOp` component which outputs a JSON formatted [gcp_resources proto](https://github.com/kubeflow/pipelines/tree/master/components/google-cloud/google_cloud_pipeline_components/experimental/proto), then pass it to the wait component."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ed9621e0",
      "metadata": {
        "id": "ed9621e0"
      },
      "outputs": [],
      "source": [
        "import kfp.dsl as dsl\n",
        "import json\n",
        "\n",
        "@dsl.pipeline(\n",
        "    name=PIPELINE_NAME,\n",
        "    description='Dataflow launch python pipeline'\n",
        ")\n",
        "def pipeline(\n",
        "    python_file_path:str = 'gs://ml-pipeline-playground/samples/dataflow/wc/wc.py',\n",
        "    project_id:str = PROJECT_ID,\n",
        "    location:str = LOCATION,\n",
        "    staging_dir:str = PIPELINE_ROOT,\n",
        "    requirements_file_path:str = 'gs://ml-pipeline-playground/samples/dataflow/wc/requirements.txt',\n",
        "):\n",
        "    dataflow_python_op = DataflowPythonJobOp(\n",
        "        project=project_id,\n",
        "        location=location,\n",
        "        python_module_path=python_file_path,\n",
        "        temp_location = staging_dir,\n",
        "        requirements_file_path = requirements_file_path,\n",
        "        args = ['--output', OUTPUT_FILE],\n",
        "    )\n",
        "    dataflow_wait_op = WaitGcpResourcesOp(\n",
        "      gcp_resources = dataflow_python_op.outputs[\"gcp_resources\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d90f0434",
      "metadata": {
        "id": "d90f0434"
      },
      "source": [
        "You can proceed to complie and run the pipeline from here as usual."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "DataflowPythonJobOp_sample.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "/piper/depot/google3/third_party/py/google_cloud_pipeline_components/google_cloud_pipeline_components/experimental/dataflow/python_job/DataflowPythonJobOp_sample.ipynb?workspaceId=chavoshi:dataflow_component::citc",
          "timestamp": 1634795827499
        },
        {
          "file_id": "1QWlRC8HvyvZuFek3kaUQ-8DuQQIUyY6_",
          "timestamp": 1634795779033
        }
      ]
    },
    "environment": {
      "name": "common-cpu.m73",
      "type": "gcloud",
      "uri": "gcr.io/deeplearning-platform-release/base-cpu:m73"
    },
    "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.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
