feat: Friday updates

notebooks/community/ml_ops/stage2/get_started_automl_training.ipynb

@@ -679,6 +679,142 @@
         "print(model_evaluation)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "deploy_model:mbsdk,automatic"
+      },
+      "source": [
+        "## Deploy the model\n",
+        "\n",
+        "Next, deploy your model for online prediction. To deploy the model, you invoke the `deploy` method."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "deploy_model:mbsdk,automatic"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint = model.deploy()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "make_prediction"
+      },
+      "source": [
+        "## Send a online prediction request\n",
+        "\n",
+        "Send a online prediction to your deployed model."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "get_test_item"
+      },
+      "source": [
+        "### Get test item\n",
+        "\n",
+        "You will use an arbitrary example out of the dataset as a test item. Don't be concerned that the example was likely used in training the model -- we just want to demonstrate how to make a prediction."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "get_test_item:automl,icn,csv"
+      },
+      "outputs": [],
+      "source": [
+        "test_item = !gsutil cat $IMPORT_FILE | head -n1\n",
+        "if len(str(test_item[0]).split(\",\")) == 3:\n",
+        "    _, test_item, test_label = str(test_item[0]).split(\",\")\n",
+        "else:\n",
+        "    test_item, test_label = str(test_item[0]).split(\",\")\n",
+        "\n",
+        "print(test_item, test_label)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "predict_request:mbsdk,icn"
+      },
+      "source": [
+        "### Make the prediction\n",
+        "\n",
+        "Now that your `Model` resource is deployed to an `Endpoint` resource, you can do online predictions by sending prediction requests to the Endpoint resource.\n",
+        "\n",
+        "#### Request\n",
+        "\n",
+        "Since in this example your test item is in a Cloud Storage bucket, you open and read the contents of the image using `tf.io.gfile.Gfile()`. To pass the test data to the prediction service, you encode the bytes into base64 -- which makes the content safe from modification while transmitting binary data over the network.\n",
+        "\n",
+        "The format of each instance is:\n",
+        "\n",
+        "    { 'content': { 'b64': base64_encoded_bytes } }\n",
+        "\n",
+        "Since the `predict()` method can take multiple items (instances), send your single test item as a list of one test item.\n",
+        "\n",
+        "#### Response\n",
+        "\n",
+        "The response from the `predict()` call is a Python dictionary with the following entries:\n",
+        "\n",
+        "- `ids`: The internal assigned unique identifiers for each prediction request.\n",
+        "- `displayNames`: The class names for each class label.\n",
+        "- `confidences`: The predicted confidence, between 0 and 1, per class label.\n",
+        "- `deployed_model_id`: The Vertex AI identifier for the deployed Model resource which did the predictions."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "predict_request:mbsdk,icn"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "\n",
+        "with tf.io.gfile.GFile(test_item, \"rb\") as f:\n",
+        "    content = f.read()\n",
+        "\n",
+        "# The format of each instance should conform to the deployed model's prediction input schema.\n",
+        "instances = [{\"content\": base64.b64encode(content).decode(\"utf-8\")}]\n",
+        "\n",
+        "prediction = endpoint.predict(instances=instances)\n",
+        "\n",
+        "print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "source": [
+        "## Undeploy the model\n",
+        "\n",
+        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.undeploy_all()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -759,6 +895,28 @@
         "dataset.delete()"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "endpoint_delete:mbsdk"
+      },
+      "source": [
+        "#### Delete the endpoint\n",
+        "\n",
+        "The method 'delete()' will delete the endpoint."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "endpoint_delete:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.delete()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -988,6 +1146,52 @@
         "print(model_evaluation)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "deploy_model:mbsdk,dedicated"
+      },
+      "source": [
+        "## Deploy the model\n",
+        "\n",
+        "Next, deploy your model for online prediction. To deploy the model, you invoke the `deploy` method, with the following parameters:\n",
+        "\n",
+        "- `machine_type`: The type of compute machine."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "deploy_model:mbsdk,dedicated"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint = model.deploy(machine_type=\"n1-standard-4\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "source": [
+        "## Undeploy the model\n",
+        "\n",
+        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.undeploy_all()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1063,6 +1267,28 @@
         "dataset.delete()"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "endpoint_delete:mbsdk"
+      },
+      "source": [
+        "#### Delete the endpoint\n",
+        "\n",
+        "The method 'delete()' will delete the endpoint."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "endpoint_delete:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.delete()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1287,6 +1513,50 @@
         "print(model_evaluation)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "deploy_model:mbsdk,automatic"
+      },
+      "source": [
+        "## Deploy the model\n",
+        "\n",
+        "Next, deploy your model for online prediction. To deploy the model, you invoke the `deploy` method."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "deploy_model:mbsdk,automatic"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint = model.deploy()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "source": [
+        "## Undeploy the model\n",
+        "\n",
+        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.undeploy_all()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1331,6 +1601,28 @@
         "dataset.delete()"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "endpoint_delete:mbsdk"
+      },
+      "source": [
+        "#### Delete the endpoint\n",
+        "\n",
+        "The method 'delete()' will delete the endpoint."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "endpoint_delete:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.delete()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/community/ml_ops/stage2/get_started_bqml_training.ipynb

@@ -84,6 +84,7 @@
         "\n",
         "- `BigQueryML Training`\n",
         "- `Vertex AI Model resource`\n",
+        "- `Vertex AI Vizier.\n",
         "\n",
         "The steps performed include:\n",
         "\n",
@@ -91,7 +92,8 @@
         "- Train a BQML model.\n",
         "- Evaluate the BQML model.\n",
         "- Export the BQML model as a cloud model.\n",
-        "- Upload the exported model as a Vertex AI Model resource."
+        "- Upload the exported model as a Vertex AI Model resource.\n",
+        "- Hyperparameter tune a BQML model with Vertex AI Vizier."
       ]
     },
     {
@@ -607,7 +609,10 @@
       "source": [
         "### Train BQML model\n",
         "\n",
-        "Next, you create and train a BQML tabular classification model from the public dataset penguins and store the model in your project.\n",
+        "Next, you create and train a BQML tabular classification model from the public dataset penguins and store the model in your project using the `CREATE MODEL` statement. The model configuration is specified in the `OPTIONS` statement as follows:\n",
+        "\n",
+        "- `model_type`: The type and archictecture of tabular model to train, e.g., DNN classification.\n",
+        "- `labels`: The column which are the labels.\n",
         "\n",
         "Learn more about [The CREATE MODEL statement](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-create)."
       ]
@@ -725,14 +730,251 @@
       },
       "outputs": [],
       "source": [
-        "bq_model = aip.Model.upload(\n",
+        "model = aip.Model.upload(\n",
         "    display_name=\"penguins_\" + TIMESTAMP,\n",
-        "    artifact_uri=SERVING_MODEL_DIR,\n",
+        "    artifact_uri=MODEL_DIR,\n",
         "    serving_container_image_uri=DEPLOY_IMAGE,\n",
         "    sync=True,\n",
         ")"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "deploy_model:mbsdk,all"
+      },
+      "source": [
+        "## Deploy the model\n",
+        "\n",
+        "Next, deploy your model for online prediction. To deploy the model, you invoke the `deploy` method, with the following parameters:\n",
+        "\n",
+        "- `deployed_model_display_name`: A human readable name for the deployed model.\n",
+        "- `traffic_split`: Percent of traffic at the endpoint that goes to this model, which is specified as a dictionary of one or more key/value pairs.\n",
+        "If only one model, then specify as { \"0\": 100 }, where \"0\" refers to this model being uploaded and 100 means 100% of the traffic.\n",
+        "If there are existing models on the endpoint, for which the traffic will be split, then use model_id to specify as { \"0\": percent, model_id: percent, ... }, where model_id is the model id of an existing model to the deployed endpoint. The percents must add up to 100.\n",
+        "- `machine_type`: The type of machine to use for training.\n",
+        "- `accelerator_type`: The hardware accelerator type.\n",
+        "- `accelerator_count`: The number of accelerators to attach to a worker replica.\n",
+        "- `starting_replica_count`: The number of compute instances to initially provision.\n",
+        "- `max_replica_count`: The maximum number of compute instances to scale to. In this tutorial, only one instance is provisioned."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "deploy_model:mbsdk,all"
+      },
+      "outputs": [],
+      "source": [
+        "DEPLOYED_NAME = \"penguins-\" + TIMESTAMP\n",
+        "\n",
+        "TRAFFIC_SPLIT = {\"0\": 100}\n",
+        "\n",
+        "MIN_NODES = 1\n",
+        "MAX_NODES = 1\n",
+        "\n",
+        "if DEPLOY_GPU:\n",
+        "    endpoint = model.deploy(\n",
+        "        deployed_model_display_name=DEPLOYED_NAME,\n",
+        "        traffic_split=TRAFFIC_SPLIT,\n",
+        "        machine_type=DEPLOY_COMPUTE,\n",
+        "        accelerator_type=DEPLOY_GPU.name,\n",
+        "        accelerator_count=DEPLOY_NGPU,\n",
+        "        min_replica_count=MIN_NODES,\n",
+        "        max_replica_count=MAX_NODES,\n",
+        "    )\n",
+        "else:\n",
+        "    endpoint = model.deploy(\n",
+        "        deployed_model_display_name=DEPLOYED_NAME,\n",
+        "        traffic_split=TRAFFIC_SPLIT,\n",
+        "        machine_type=DEPLOY_COMPUTE,\n",
+        "        min_replica_count=MIN_NODES,\n",
+        "        max_replica_count=MAX_NODES,\n",
+        "    )"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "source": [
+        "## Undeploy the model\n",
+        "\n",
+        "When you are done doing predictions, you undeploy the model from the `Endpoint` resouce. This deprovisions all compute resources and ends billing for the deployed model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "undeploy_model:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.undeploy_all()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "model_delete:mbsdk"
+      },
+      "source": [
+        "#### Delete the model\n",
+        "\n",
+        "The method 'delete()' will delete the model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "model_delete:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "model.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bqml_create_model:vizier"
+      },
+      "source": [
+        "### Hyperparameter Tune and train a BQML model\n",
+        "\n",
+        "Next, you train a BQML tabular classification model with hyperparameter tuning using the Vertex AI Vizier service. The hyperparameter settings are specified in the `OPTIONS` statement as follows:\n",
+        "\n",
+        "- `HPARAM_TUNING_ALGORITHM`: The algorithm for selecting the next trial parameters.\n",
+        "- `num_trials`: The number of trials.\n",
+        "- `max_parallel_trials`: The number of trials to do in parallel.\n",
+        "\n",
+        "Learn more about [Hyperparameter tuning for CREATE MODEL statements](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-hyperparameter-tuning)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bqml_create_model:vizier"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_NAME = \"penguins\"\n",
+        "MODEL_QUERY = f\"\"\"\n",
+        "CREATE OR REPLACE MODEL `{BQ_DATASET_NAME}.{MODEL_NAME}`\n",
+        "OPTIONS(\n",
+        "    model_type='DNN_CLASSIFIER',\n",
+        "    labels = ['species'],\n",
+        "    num_trials=10,\n",
+        "    max_parallel_trials=2,\n",
+        "    HPARAM_TUNING_ALGORITHM = 'VIZIER_DEFAULT'\n",
+        "    )\n",
+        "AS\n",
+        "SELECT *\n",
+        "FROM `{BQ_TABLE}`\n",
+        "\"\"\"\n",
+        "\n",
+        "job = bqclient.query(MODEL_QUERY)\n",
+        "print(job.errors, job.state)\n",
+        "\n",
+        "while job.running():\n",
+        "    from time import sleep\n",
+        "\n",
+        "    sleep(30)\n",
+        "    print(\"Running ...\")\n",
+        "print(job.errors, job.state)\n",
+        "\n",
+        "tblname = job.ddl_target_table\n",
+        "tblname = \"{}.{}\".format(tblname.dataset_id, tblname.table_id)\n",
+        "print(\"{} created in {}\".format(tblname, job.ended - job.started))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bqml_eval_model"
+      },
+      "source": [
+        "### Evaluate the BQML trained model\n",
+        "\n",
+        "Next, retrieve the model evaluation for the trained BQML model.\n",
+        "\n",
+        "Learn more about [The ML.EVALUATE function](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-evaluate)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bqml_eval_model"
+      },
+      "outputs": [],
+      "source": [
+        "EVAL_QUERY = f\"\"\"\n",
+        "SELECT *\n",
+        "FROM\n",
+        "  ML.EVALUATE(MODEL {BQ_DATASET_NAME}.{MODEL_NAME})\n",
+        "ORDER BY  roc_auc desc\n",
+        "LIMIT 1\"\"\"\n",
+        "\n",
+        "job = bqclient.query(EVAL_QUERY)\n",
+        "results = job.result().to_dataframe()\n",
+        "print(results)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bqml_create_model:xai"
+      },
+      "source": [
+        "### Train a BQML model with Explainability\n",
+        "\n",
+        "Next, you train the same BQML model, but this time you enable Vertex AI Explainability on the model predictions by adding the option:\n",
+        "\n",
+        "- `ENABLE_GLOBAL_EXPLAIN`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bqml_create_model:xai"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_NAME = \"penguins\"\n",
+        "MODEL_QUERY = f\"\"\"\n",
+        "CREATE OR REPLACE MODEL `{BQ_DATASET_NAME}.{MODEL_NAME}`\n",
+        "OPTIONS(\n",
+        "    model_type='DNN_CLASSIFIER',\n",
+        "    labels = ['species'],\n",
+        "    ENABLE_GLOBAL_EXPLAIN = True\n",
+        "    )\n",
+        "AS\n",
+        "SELECT *\n",
+        "FROM `{BQ_TABLE}`\n",
+        "\"\"\"\n",
+        "\n",
+        "job = bqclient.query(MODEL_QUERY)\n",
+        "print(job.errors, job.state)\n",
+        "\n",
+        "while job.running():\n",
+        "    from time import sleep\n",
+        "\n",
+        "    sleep(30)\n",
+        "    print(\"Running ...\")\n",
+        "print(job.errors, job.state)\n",
+        "\n",
+        "tblname = job.ddl_target_table\n",
+        "tblname = \"{}.{}\".format(tblname.dataset_id, tblname.table_id)\n",
+        "print(\"{} created in {}\".format(tblname, job.ended - job.started))"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/community/ml_ops/stage3/get_started_with_dataflow_pipeline_components.ipynb

@@ -715,7 +715,7 @@
         "    project_id: str = PROJECT_ID,\n",
         "    location: str = REGION,\n",
         "    staging_dir: str = PIPELINE_ROOT,\n",
-        "    args: list = [\"--output\", GCS_WC_OUT],\n",
+        "    args: list = [\"--output\", GCS_WC_OUT, \"--runner\", \"DataflowRunner\"],\n",
         "    requirements_file_path: str = GCS_REQUIREMENTS_TXT,\n",
         "):\n",
         "\n",
@@ -907,7 +907,44 @@
       "source": [
         "%%writefile requirements.txt\n",
         "apache-beam\n",
-        "tensorflow-transform"
+        "tensorflow-transform==1.2.0"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "writefile:setup,split"
+      },
+      "source": [
+        "### Write the setup.py (installs) for the Dataflow workers\n",
+        "\n",
+        "Next, create the `setup.py` file to specify Python modules that are required to be installed for executing the Dataflow workers -- in this case, `tensorflow-transform` is required."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "writefile:setup,split"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile setup.py\n",
+        "import setuptools\n",
+        "\n",
+        "REQUIRED_PACKAGES = [\n",
+        "    'tensorflow-transform==1.2.0',\n",
+        "]\n",
+        "PACKAGE_NAME = 'my_package'\n",
+        "PACKAGE_VERSION = '0.0.1'\n",
+        "setuptools.setup(\n",
+        "    name=PACKAGE_NAME,\n",
+        "    version=PACKAGE_VERSION,\n",
+        "    description='Demo for split transformation',\n",
+        "    install_requires=REQUIRED_PACKAGES,\n",
+        "    author=\"cdpe@google.com\",\n",
+        "    packages=setuptools.find_packages()\n",
+        ")"
       ]
     },
     {
@@ -918,7 +955,7 @@
       "source": [
         "### Copy python module and requirements file to Cloud Storage\n",
         "\n",
-        "Next, you copy the Python module and requirements file to your Cloud Storage bucket.\n",
+        "Next, you copy the Python module, requirements and setup file to your Cloud Storage bucket.\n",
         "\n",
         "Additional, you set the Cloud Storage location for the output of the Apache Beam dataset split pipeline."
       ]
@@ -934,7 +971,9 @@
         "GCS_SPLIT_PY = BUCKET_NAME + \"/split.py\"\n",
         "! gsutil cp split.py $GCS_SPLIT_PY\n",
         "GCS_REQUIREMENTS_TXT = BUCKET_NAME + \"/requirements.txt\"\n",
-        "! gsutil cp requirements.txt $GCS_REQUIREMENTS_TXT"
+        "! gsutil cp requirements.txt $GCS_REQUIREMENTS_TXT\n",
+        "GCS_SETUP_PY = BUCKET_NAME + \"/setup.py\"\n",
+        "! gsutil cp setup.py $GCS_SETUP_PY"
       ]
     },
     {
@@ -979,7 +1018,7 @@
         "\n",
         "Learn more about [Google Cloud Pipeline Component for Dataflow](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-0.2.0/google_cloud_pipeline_components.experimental.dataflow.html)\n",
         "\n",
-        "Additional, you add `--requirements_file` to the input args, such that the workers -- i.e., WaitGcpResourcesOp -- share the same pip installation requirements as the Dataflow component."
+        "Additional, you add `--runner=DataflowRunner` to the input args, to tell the component to use Dataflow instead of DirectRunner for the Apache Beam job."
       ]
     },
     {
@@ -1004,12 +1043,16 @@
         "        BUCKET_NAME,\n",
         "        \"--bq_table\",\n",
         "        BQ_TABLE,\n",
-        "        \"--requirements_file\",\n",
-        "        GCS_REQUIREMENTS_TXT,\n",
+        "        \"--runner\",\n",
+        "        \"DataflowRunner\",\n",
+        "        \"--setup_file\",\n",
+        "        GCS_SETUP_PY,\n",
         "    ],\n",
         "    requirements_file_path: str = GCS_REQUIREMENTS_TXT,\n",
         "):\n",
-        "    # DataflowPythonJobOp.component_spec.implementation.container.image = \"gcr.io/ml-pipeline/google-cloud-pipeline-components:v0.2.0_dataflow_logs_fix\"\n",
+        "    DataflowPythonJobOp.component_spec.implementation.container.image = (\n",
+        "        \"gcr.io/ml-pipeline/google-cloud-pipeline-components:v0.2.0_dataflow_logs_fix\"\n",
+        "    )\n",
         "    dataflow_python_op = DataflowPythonJobOp(\n",
         "        project=project_id,\n",
         "        location=location,\n",