diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb index fe963973b..a02e0bf32 100644 --- a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb @@ -38,20 +38,20 @@ "id": "JAPoU8Sm5E6e" }, "source": [ - "\u003ctable align=\"left\"\u003e\n", + "\n", "\n", - " \u003ctd\u003e\n", - " \u003ca href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb\"\u003e\n", - " \u003cimg src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"\u003e Run in Colab\n", - " \u003c/a\u003e\n", - " \u003c/td\u003e\n", - " \u003ctd\u003e\n", - " \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb\"\u003e\n", - " \u003cimg src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"\u003e\n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", " View on GitHub\n", - " \u003c/a\u003e\n", - " \u003c/td\u003e\n", - "\u003c/table\u003e" + " \n", + "
" ] }, { @@ -177,7 +177,7 @@ "# Google Cloud Notebook requires dependencies to be installed with '--user'\n", "USER_FLAG = \"\"\n", "if IS_GOOGLE_CLOUD_NOTEBOOK:\n", - " USER_FLAG = \"--user\"" + " USER_FLAG = \"--user\"" ] }, { @@ -218,11 +218,11 @@ "import os\n", "\n", "if not os.getenv(\"IS_TESTING\"):\n", - " # Automatically restart kernel after installs\n", - " import IPython\n", + " # Automatically restart kernel after installs\n", + " import IPython\n", "\n", - " app = IPython.Application.instance()\n", - " app.kernel.do_shutdown(True)" + " app = IPython.Application.instance()\n", + " app.kernel.do_shutdown(True)" ] }, { @@ -235,7 +235,7 @@ "\n", "### Select a GPU runtime\n", "\n", - "**Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select \"Runtime --\u003e Change runtime type \u003e GPU\"**" + "**Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select \"Runtime --> Change runtime type > GPU\"**" ] }, { @@ -287,9 +287,9 @@ "\n", "# Get your Google Cloud project ID from gcloud\n", "if not os.getenv(\"IS_TESTING\"):\n", - " shell_output=!gcloud config list --format 'value(core.project)' 2\u003e/dev/null\n", - " PROJECT_ID = shell_output[0]\n", - " print(\"Project ID: \", PROJECT_ID)" + " shell_output=!gcloud config list --format 'value(core.project)' 2>/dev/null\n", + " PROJECT_ID = shell_output[0]\n", + " print(\"Project ID: \", PROJECT_ID)" ] }, { @@ -310,7 +310,7 @@ "outputs": [], "source": [ "if PROJECT_ID == \"\" or PROJECT_ID is None:\n", - " PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}" + " PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}" ] }, { @@ -400,16 +400,16 @@ "\n", "# If on Google Cloud Notebooks, then don't execute this code\n", "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n", - " if \"google.colab\" in sys.modules:\n", - " from google.colab import auth as google_auth\n", + " if \"google.colab\" in sys.modules:\n", + " from google.colab import auth as google_auth\n", "\n", - " google_auth.authenticate_user()\n", + " google_auth.authenticate_user()\n", "\n", - " # If you are running this notebook locally, replace the string below with the\n", - " # path to your service account key and run this cell to authenticate your GCP\n", - " # account.\n", - " elif not os.getenv(\"IS_TESTING\"):\n", - " %env GOOGLE_APPLICATION_CREDENTIALS ''" + " # If you are running this notebook locally, replace the string below with the\n", + " # path to your service account key and run this cell to authenticate your GCP\n", + " # account.\n", + " elif not os.getenv(\"IS_TESTING\"):\n", + " %env GOOGLE_APPLICATION_CREDENTIALS ''" ] }, { @@ -455,7 +455,7 @@ "outputs": [], "source": [ "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n", - " BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP" + " BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP" ] }, { @@ -517,17 +517,10 @@ "source": [ "import os\n", "import sys\n", - "from typing import Optional\n", "\n", "from google_cloud_pipeline_components import aiplatform as gcc_aip\n", - "\n", - "from kfp.v2.components import create_component_from_func\n", - "from kfp.v2 import compiler\n", - "from kfp.v2 import dsl\n", - "from kfp.v2.google.client import AIPlatformClient\n", - "\n", - "import numpy as np\n", - "import tensorflow as tf" + "from kfp.v2 import compiler, dsl\n", + "from kfp.v2.google.client import AIPlatformClient" ] }, { @@ -583,11 +576,15 @@ "OUTPUT_COMPONENT_SPEC = \"output-component.yaml\" # Output component specification file.\n", "\n", "# BigQuery parameters (used for the Generator, Ingester, Logger)\n", - "BIGQUERY_TMP_FILE = \"tmp.json\" # Temporary file for storing data to be loaded into BigQuery.\n", + "BIGQUERY_TMP_FILE = (\n", + " \"tmp.json\" # Temporary file for storing data to be loaded into BigQuery.\n", + ")\n", "BIGQUERY_MAX_ROWS = 5 # Maximum number of rows of data in BigQuery to ingest.\n", "\n", "# Dataset parameters\n", - "TFRECORD_FILE = f\"{BUCKET_NAME}/trainer_input_path/*\" # TFRecord file to be used for training.\n", + "TFRECORD_FILE = (\n", + " f\"{BUCKET_NAME}/trainer_input_path/*\" # TFRecord file to be used for training.\n", + ")\n", "\n", "# Logger parameters (also used for the Logger hook in the prediction container)\n", "LOGGER_PUBSUB_TOPIC = \"logger-pubsub-topic\" # Pub/Sub topic name for the Logger.\n", @@ -611,8 +608,8 @@ "After pipeline construction, create the *Simulator* to send simulated MovieLens prediction requests, create the *Logger* to asynchronously log prediction inputs and results, and create the *Trigger* to trigger re-training.\n", "\n", "Here's the entire workflow:\n", - "1. The startup pipeline has the following components: Generator --\u003e Ingester --\u003e Trainer --\u003e Deployer. This pipeline only runs once.\n", - "2. Then, the Simulator generates prediction requests (e.g. every 5 mins), and the Logger gets invoked immediately at each prediction request and logs each prediction request asynchronously into BigQuery. The Trigger runs the re-training pipeline (e.g. every 30 mins) with the following components: Ingester --\u003e Trainer --\u003e Deploy.\n", + "1. The startup pipeline has the following components: Generator --> Ingester --> Trainer --> Deployer. This pipeline only runs once.\n", + "2. Then, the Simulator generates prediction requests (e.g. every 5 mins), and the Logger gets invoked immediately at each prediction request and logs each prediction request asynchronously into BigQuery. The Trigger runs the re-training pipeline (e.g. every 30 mins) with the following components: Ingester --> Trainer --> Deploy.\n", "\n", "You can find the KFP SDK documentation [here](https://www.kubeflow.org/docs/components/pipelines/sdk/sdk-overview/)." ] @@ -652,26 +649,6 @@ "! python3 -m unittest src.generator.test_generator_component" ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "pq8v8XL6Gilp" - }, - "source": [ - "#### Load the Generator component" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "jjA4aAYyGilp" - }, - "outputs": [], - "source": [ - "from src.generator import generator_component" - ] - }, { "cell_type": "markdown", "metadata": { @@ -707,26 +684,6 @@ "! python3 -m unittest src.ingester.test_ingester_component" ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "C_3qnZykOWQP" - }, - "source": [ - "#### Get the Ingester component" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "TVyfpoZIOWQP" - }, - "outputs": [], - "source": [ - "from src.ingester import ingester_component" - ] - }, { "cell_type": "markdown", "metadata": { @@ -753,9 +710,13 @@ "outputs": [], "source": [ "# Trainer parameters\n", - "TRAINING_ARTIFACTS_DIR = f\"{BUCKET_NAME}/artifacts\" # Root directory for training artifacts.\n", + "TRAINING_ARTIFACTS_DIR = (\n", + " f\"{BUCKET_NAME}/artifacts\" # Root directory for training artifacts.\n", + ")\n", "TRAINING_REPLICA_COUNT = \"1\" # Number of replica to run the custom training job.\n", - "TRAINING_MACHINE_TYPE = \"n1-standard-4\" # Type of machine to run the custom training job.\n", + "TRAINING_MACHINE_TYPE = (\n", + " \"n1-standard-4\" # Type of machine to run the custom training job.\n", + ")\n", "TRAINING_ACCELERATOR_TYPE = \"ACCELERATOR_TYPE_UNSPECIFIED\" # Type of accelerators to run the custom training job.\n", "TRAINING_ACCELERATOR_COUNT = \"0\" # Number of accelerators for the custom training job." ] @@ -780,26 +741,6 @@ "! python3 -m unittest src.trainer.test_trainer_component" ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "liJ8SnFNWkL9" - }, - "source": [ - "#### Get the Trainer component" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "qijhwmrGWkL-" - }, - "outputs": [], - "source": [ - "from src.trainer import trainer_component" - ] - }, { "cell_type": "markdown", "metadata": { @@ -825,7 +766,9 @@ "outputs": [], "source": [ "# Deployer parameters\n", - "TRAINED_POLICY_DISPLAY_NAME = \"movielens-trained-policy\" # Display name of the uploaded and deployed policy.\n", + "TRAINED_POLICY_DISPLAY_NAME = (\n", + " \"movielens-trained-policy\" # Display name of the uploaded and deployed policy.\n", + ")\n", "ENDPOINT_DISPLAY_NAME = \"movielens-endpoint\" # Display name of the prediction endpoint.\n", "ENDPOINT_MACHINE_TYPE = \"n1-standard-4\" # Type of machine of the prediction endpoint." ] @@ -893,7 +836,7 @@ ")\n", "\n", "with open(f\"{PREDICTION_CONTAINER_DIR}/cloudbuild.yaml\", \"w\") as fp:\n", - " fp.write(cloudbuild_yaml)" + " fp.write(cloudbuild_yaml)" ] }, { @@ -929,8 +872,7 @@ "cell_type": "code", "execution_count": null, "metadata": { - "id": "9uHbODeXibYd", - "scrolled": true + "id": "9uHbODeXibYd" }, "outputs": [], "source": [ @@ -958,21 +900,30 @@ }, "outputs": [], "source": [ - "@dsl.pipeline(\n", - " pipeline_root=PIPELINE_ROOT,\n", - " name=f\"{PIPELINE_NAME}-startup\")\n", + "from kfp.components import load_component_from_url\n", + "\n", + "generate_op = load_component_from_url(\n", + " \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/component.yaml\"\n", + ")\n", + "ingest_op = load_component_from_url(\n", + " \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml\"\n", + ")\n", + "train_op = load_component_from_url(\n", + " \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml\"\n", + ")\n", + "\n", + "\n", + "@dsl.pipeline(pipeline_root=PIPELINE_ROOT, name=f\"{PIPELINE_NAME}-startup\")\n", "def pipeline(\n", " # Pipeline configs\n", " project_id: str,\n", " raw_data_path: str,\n", " training_artifacts_dir: str,\n", - "\n", " # BigQuery configs\n", " bigquery_dataset_id: str,\n", " bigquery_location: str,\n", " bigquery_table_id: str,\n", " bigquery_max_rows: int = 10000,\n", - "\n", " # TF-Agents RL configs\n", " batch_size: int = 8,\n", " rank_k: int = 20,\n", @@ -980,134 +931,115 @@ " driver_steps: int = 3,\n", " num_epochs: int = 5,\n", " tikhonov_weight: float = 0.01,\n", - " agent_alpha: float = 10) -\u003e None:\n", - " \"\"\"Authors a RL pipeline for MovieLens movie recommendation system.\n", + " agent_alpha: float = 10,\n", + ") -> None:\n", + " \"\"\"Authors a RL pipeline for MovieLens movie recommendation system.\n", "\n", - " Integrates the Generator, Ingester, Trainer and Deployer components. This\n", - " pipeline generates initial training data with a random policy and runs once\n", - " as the initiation of the system.\n", + " Integrates the Generator, Ingester, Trainer and Deployer components. This\n", + " pipeline generates initial training data with a random policy and runs once\n", + " as the initiation of the system.\n", "\n", - " Args:\n", - " project_id: GCP project ID. This is required because otherwise the BigQuery\n", - " client will use the ID of the tenant GCP project created as a result of\n", - " KFP, which doesn't have proper access to BigQuery.\n", - " raw_data_path: Path to MovieLens 100K's \"u.data\" file.\n", - " training_artifacts_dir: Path to store the Trainer artifacts (trained policy).\n", + " Args:\n", + " project_id: GCP project ID. This is required because otherwise the BigQuery\n", + " client will use the ID of the tenant GCP project created as a result of\n", + " KFP, which doesn't have proper access to BigQuery.\n", + " raw_data_path: Path to MovieLens 100K's \"u.data\" file.\n", + " training_artifacts_dir: Path to store the Trainer artifacts (trained policy).\n", "\n", - " bigquery_dataset: A string of the BigQuery dataset ID in the format of\n", - " \"project.dataset\".\n", - " bigquery_location: A string of the BigQuery dataset location.\n", - " bigquery_table_id: A string of the BigQuery table ID in the format of\n", - " \"project.dataset.table\".\n", - " bigquery_max_rows: Optional; maximum number of rows to ingest.\n", + " bigquery_dataset: A string of the BigQuery dataset ID in the format of\n", + " \"project.dataset\".\n", + " bigquery_location: A string of the BigQuery dataset location.\n", + " bigquery_table_id: A string of the BigQuery table ID in the format of\n", + " \"project.dataset.table\".\n", + " bigquery_max_rows: Optional; maximum number of rows to ingest.\n", "\n", - " batch_size: Optional; batch size of environment generated quantities eg.\n", - " rewards.\n", - " rank_k: Optional; rank for matrix factorization in the MovieLens environment;\n", - " also the observation dimension.\n", - " num_actions: Optional; number of actions (movie items) to choose from.\n", - " driver_steps: Optional; number of steps to run per batch.\n", - " num_epochs: Optional; number of training epochs.\n", - " tikhonov_weight: Optional; LinUCB Tikhonov regularization weight of the\n", - " Trainer.\n", - " agent_alpha: Optional; LinUCB exploration parameter that multiplies the\n", - " confidence intervals of the Trainer.\n", - " \"\"\"\n", - " # Run the Generator component.\n", - " generate_op = create_component_from_func(\n", - " func=generator_component.generate_movielens_dataset_for_bigquery,\n", - " output_component_file=f\"generator-{OUTPUT_COMPONENT_SPEC}\",\n", - " packages_to_install=[\n", - " \"google-cloud-bigquery==2.20.0\",\n", - " \"tensorflow==2.5.0\",\n", - " \"tf-agents==0.8.0\",\n", - " ])\n", - " generate_task = generate_op(\n", - " project_id=project_id,\n", - " raw_data_path=raw_data_path,\n", - " batch_size=batch_size,\n", - " rank_k=rank_k,\n", - " num_actions=num_actions,\n", - " driver_steps=driver_steps,\n", - " bigquery_tmp_file=BIGQUERY_TMP_FILE,\n", - " bigquery_dataset_id=bigquery_dataset_id,\n", - " bigquery_location=bigquery_location,\n", - " bigquery_table_id=bigquery_table_id)\n", + " batch_size: Optional; batch size of environment generated quantities eg.\n", + " rewards.\n", + " rank_k: Optional; rank for matrix factorization in the MovieLens environment;\n", + " also the observation dimension.\n", + " num_actions: Optional; number of actions (movie items) to choose from.\n", + " driver_steps: Optional; number of steps to run per batch.\n", + " num_epochs: Optional; number of training epochs.\n", + " tikhonov_weight: Optional; LinUCB Tikhonov regularization weight of the\n", + " Trainer.\n", + " agent_alpha: Optional; LinUCB exploration parameter that multiplies the\n", + " confidence intervals of the Trainer.\n", + " \"\"\"\n", + " # Run the Generator component.\n", + " generate_task = generate_op(\n", + " project_id=project_id,\n", + " raw_data_path=raw_data_path,\n", + " batch_size=batch_size,\n", + " rank_k=rank_k,\n", + " num_actions=num_actions,\n", + " driver_steps=driver_steps,\n", + " bigquery_tmp_file=BIGQUERY_TMP_FILE,\n", + " bigquery_dataset_id=bigquery_dataset_id,\n", + " bigquery_location=bigquery_location,\n", + " bigquery_table_id=bigquery_table_id,\n", + " )\n", "\n", - " # Run the Ingester component.\n", - " ingest_op = create_component_from_func(\n", - " func=ingester_component.ingest_bigquery_dataset_into_tfrecord,\n", - " output_component_file=f\"ingester-{OUTPUT_COMPONENT_SPEC}\",\n", - " packages_to_install=[\n", - " \"google-cloud-bigquery==2.20.0\",\n", - " \"tensorflow==2.5.0\",\n", - " ])\n", - " ingest_task = ingest_op(\n", - " project_id=project_id,\n", - " bigquery_table_id=generate_task.outputs[\"bigquery_table_id\"],\n", - " bigquery_max_rows=bigquery_max_rows,\n", - " tfrecord_file=TFRECORD_FILE)\n", + " # Run the Ingester component.\n", + " ingest_task = ingest_op(\n", + " project_id=project_id,\n", + " bigquery_table_id=generate_task.outputs[\"bigquery_table_id\"],\n", + " bigquery_max_rows=bigquery_max_rows,\n", + " tfrecord_file=TFRECORD_FILE,\n", + " )\n", "\n", - " # Run the Trainer component and submit custom job to Vertex AI.\n", - " train_op = create_component_from_func(\n", - " func=trainer_component.training_op,\n", - " output_component_file=f\"trainer-{OUTPUT_COMPONENT_SPEC}\",\n", - " packages_to_install=[\n", - " \"tensorflow==2.5.0\",\n", - " \"tf-agents==0.8.0\",\n", - " ])\n", - " train_task = train_op(\n", - " training_artifacts_dir=training_artifacts_dir,\n", - " tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n", - " num_epochs=num_epochs,\n", - " rank_k=rank_k,\n", - " num_actions=num_actions,\n", - " tikhonov_weight=tikhonov_weight,\n", - " agent_alpha=agent_alpha)\n", + " # Run the Trainer component and submit custom job to Vertex AI.\n", + " train_task = train_op(\n", + " training_artifacts_dir=training_artifacts_dir,\n", + " tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n", + " num_epochs=num_epochs,\n", + " rank_k=rank_k,\n", + " num_actions=num_actions,\n", + " tikhonov_weight=tikhonov_weight,\n", + " agent_alpha=agent_alpha,\n", + " )\n", "\n", - " worker_pool_specs = [\n", - " {\n", - " \"containerSpec\": {\n", - " \"imageUri\":train_task.container.image,\n", - " },\n", - " \"replicaCount\": TRAINING_REPLICA_COUNT,\n", - " \"machineSpec\": {\n", - " \"machineType\": TRAINING_MACHINE_TYPE,\n", - " \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n", - " \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n", - " },\n", - " },\n", - " ]\n", - " train_task.custom_job_spec = {\n", - " \"displayName\": train_task.name,\n", - " \"jobSpec\": {\n", - " \"workerPoolSpecs\": worker_pool_specs,\n", - " }\n", - " }\n", + " worker_pool_specs = [\n", + " {\n", + " \"containerSpec\": {\n", + " \"imageUri\": train_task.container.image,\n", + " },\n", + " \"replicaCount\": TRAINING_REPLICA_COUNT,\n", + " \"machineSpec\": {\n", + " \"machineType\": TRAINING_MACHINE_TYPE,\n", + " \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n", + " \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n", + " },\n", + " },\n", + " ]\n", + " train_task.custom_job_spec = {\n", + " \"displayName\": train_task.name,\n", + " \"jobSpec\": {\n", + " \"workerPoolSpecs\": worker_pool_specs,\n", + " },\n", + " }\n", "\n", - " # Run the Deployer components.\n", - " # Upload the trained policy as a model.\n", - " model_upload_op = gcc_aip.ModelUploadOp(\n", - " project=project_id,\n", - " display_name=TRAINED_POLICY_DISPLAY_NAME,\n", - " artifact_uri=training_artifacts_dir,\n", - " serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n", - " )\n", - " # Model uploading has to occur after training completes.\n", - " model_upload_op.after(train_task)\n", - " # Create a Vertex AI endpoint. (This operation can occur in parallel with\n", - " # the Generator, Ingester, Trainer components.)\n", - " endpoint_create_op = gcc_aip.EndpointCreateOp(\n", - " project=project_id,\n", - " display_name=ENDPOINT_DISPLAY_NAME)\n", - " # Deploy the uploaded, trained policy to the created endpoint. (This operation\n", - " # has to occur after both model uploading and endpoint creation complete.)\n", - " model_deploy_op = gcc_aip.ModelDeployOp(\n", - " project=project_id,\n", - " endpoint=endpoint_create_op.outputs[\"endpoint\"],\n", - " model=model_upload_op.outputs[\"model\"],\n", - " deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n", - " machine_type=ENDPOINT_MACHINE_TYPE)" + " # Run the Deployer components.\n", + " # Upload the trained policy as a model.\n", + " model_upload_op = gcc_aip.ModelUploadOp(\n", + " project=project_id,\n", + " display_name=TRAINED_POLICY_DISPLAY_NAME,\n", + " artifact_uri=train_task.outputs[\"training_artifacts_dir\"],\n", + " serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n", + " )\n", + " # Create a Vertex AI endpoint. (This operation can occur in parallel with\n", + " # the Generator, Ingester, Trainer components.)\n", + " endpoint_create_op = gcc_aip.EndpointCreateOp(\n", + " project=project_id, display_name=ENDPOINT_DISPLAY_NAME\n", + " )\n", + " # Deploy the uploaded, trained policy to the created endpoint. (This operation\n", + " # has to occur after both model uploading and endpoint creation complete.)\n", + " gcc_aip.ModelDeployOp(\n", + " project=project_id,\n", + " endpoint=endpoint_create_op.outputs[\"endpoint\"],\n", + " model=model_upload_op.outputs[\"model\"],\n", + " deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n", + " machine_type=ENDPOINT_MACHINE_TYPE,\n", + " )" ] }, { @@ -1119,13 +1051,10 @@ "outputs": [], "source": [ "# Compile the authored pipeline.\n", - "compiler.Compiler().compile(pipeline_func=pipeline, \n", - " package_path=PIPELINE_SPEC_PATH)\n", + "compiler.Compiler().compile(pipeline_func=pipeline, package_path=PIPELINE_SPEC_PATH)\n", "\n", "# Createa Vertex AI client.\n", - "api_client = AIPlatformClient(\n", - " project_id=PROJECT_ID,\n", - " region=REGION)\n", + "api_client = AIPlatformClient(project_id=PROJECT_ID, region=REGION)\n", "\n", "# Create a pipeline run job.\n", "response = api_client.create_run_from_job_spec(\n", @@ -1135,13 +1064,13 @@ " \"project_id\": PROJECT_ID,\n", " \"raw_data_path\": RAW_DATA_PATH,\n", " \"training_artifacts_dir\": TRAINING_ARTIFACTS_DIR,\n", - "\n", " # BigQuery configs\n", " \"bigquery_dataset_id\": BIGQUERY_DATASET_ID,\n", " \"bigquery_location\": BIGQUERY_LOCATION,\n", " \"bigquery_table_id\": BIGQUERY_TABLE_ID,\n", " },\n", - " enable_caching=ENABLE_CACHING)" + " enable_caching=ENABLE_CACHING,\n", + ")" ] }, { @@ -1154,7 +1083,7 @@ "\n", "Create the Simulator to [obtain observations](https://github.com/tensorflow/agents/blob/v0.8.0/tf_agents/bandits/environments/movielens_py_environment.py#L118-L125) from the MovieLens simulation environment, formats them, and sends prediction requests to the Vertex AI endpoint.\n", "\n", - "The workflow is: Cloud Scheduler --\u003e Pub/Sub --\u003e Cloud Functions --\u003e Endpoint\n", + "The workflow is: Cloud Scheduler --> Pub/Sub --> Cloud Functions --> Endpoint\n", "\n", "In production, this Simulator logic can be modified to that of gathering real-world input features as observations, getting prediction results from the endpoint and communicating those results to real-world users.\n", "\n", @@ -1170,11 +1099,19 @@ "outputs": [], "source": [ "# Simulator parameters\n", - "SIMULATOR_PUBSUB_TOPIC = \"simulator-pubsub-topic\" # Pub/Sub topic name for the Simulator.\n", - "SIMULATOR_CLOUD_FUNCTION = \"simulator-cloud-function\" # Cloud Functions name for the Simulator.\n", - "SIMULATOR_SCHEDULER_JOB = \"simulator-scheduler-job\" # Cloud Scheduler cron job name for the Simulator.\n", + "SIMULATOR_PUBSUB_TOPIC = (\n", + " \"simulator-pubsub-topic\" # Pub/Sub topic name for the Simulator.\n", + ")\n", + "SIMULATOR_CLOUD_FUNCTION = (\n", + " \"simulator-cloud-function\" # Cloud Functions name for the Simulator.\n", + ")\n", + "SIMULATOR_SCHEDULER_JOB = (\n", + " \"simulator-scheduler-job\" # Cloud Scheduler cron job name for the Simulator.\n", + ")\n", "SIMULATOR_SCHEDULE = \"*/5 * * * *\" # Cloud Scheduler cron job schedule for the Simulator. Eg. \"*/5 * * * *\" means every 5 mins.\n", - "SIMULATOR_SCHEDULER_MESSAGE = \"simulator-message\" # Cloud Scheduler message for the Simulator." + "SIMULATOR_SCHEDULER_MESSAGE = (\n", + " \"simulator-message\" # Cloud Scheduler message for the Simulator.\n", + ")" ] }, { @@ -1239,12 +1176,14 @@ }, "outputs": [], "source": [ - "scheduler_job_args = \" \".join([\n", - " SIMULATOR_SCHEDULER_JOB,\n", - " f\"--schedule='{SIMULATOR_SCHEDULE}'\",\n", - " f\"--topic={SIMULATOR_PUBSUB_TOPIC}\",\n", - " f\"--message-body={SIMULATOR_SCHEDULER_MESSAGE}\",\n", - "])\n", + "scheduler_job_args = \" \".join(\n", + " [\n", + " SIMULATOR_SCHEDULER_JOB,\n", + " f\"--schedule='{SIMULATOR_SCHEDULE}'\",\n", + " f\"--topic={SIMULATOR_PUBSUB_TOPIC}\",\n", + " f\"--message-body={SIMULATOR_SCHEDULER_MESSAGE}\",\n", + " ]\n", + ")\n", "\n", "! echo $scheduler_job_args" ] @@ -1299,15 +1238,17 @@ }, "outputs": [], "source": [ - "ENV_VARS = \",\".join([\n", - " f\"PROJECT_ID={PROJECT_ID}\",\n", - " f\"REGION={REGION}\",\n", - " f\"ENDPOINT_ID={ENDPOINT_ID}\",\n", - " f\"RAW_DATA_PATH={RAW_DATA_PATH}\",\n", - " f\"BATCH_SIZE={BATCH_SIZE}\",\n", - " f\"RANK_K={RANK_K}\",\n", - " f\"NUM_ACTIONS={NUM_ACTIONS}\",\n", - "])\n", + "ENV_VARS = \",\".join(\n", + " [\n", + " f\"PROJECT_ID={PROJECT_ID}\",\n", + " f\"REGION={REGION}\",\n", + " f\"ENDPOINT_ID={ENDPOINT_ID}\",\n", + " f\"RAW_DATA_PATH={RAW_DATA_PATH}\",\n", + " f\"BATCH_SIZE={BATCH_SIZE}\",\n", + " f\"RANK_K={RANK_K}\",\n", + " f\"NUM_ACTIONS={NUM_ACTIONS}\",\n", + " ]\n", + ")\n", "\n", "! echo $ENV_VARS" ] @@ -1344,7 +1285,7 @@ "\n", "The Logger is triggered by a hook in the prediction code. At each prediction request, the prediction code messages a Pub/Sub topic, which triggers the Logger code.\n", "\n", - "The workflow is: prediction container code (at prediction request) --\u003e Pub/Sub --\u003e Cloud Functions (logging predictions back to BigQuery)\n", + "The workflow is: prediction container code (at prediction request) --> Pub/Sub --> Cloud Functions (logging predictions back to BigQuery)\n", "\n", "In production, this Logger logic can be modified to that of gathering real-world feedback (rewards) based on observations and predicted actions.\n", "\n", @@ -1414,17 +1355,19 @@ }, "outputs": [], "source": [ - "ENV_VARS = \",\".join([\n", - " f\"PROJECT_ID={PROJECT_ID}\",\n", - " f\"RAW_DATA_PATH={RAW_DATA_PATH}\",\n", - " f\"BATCH_SIZE={BATCH_SIZE}\",\n", - " f\"RANK_K={RANK_K}\",\n", - " f\"NUM_ACTIONS={NUM_ACTIONS}\",\n", - " f\"BIGQUERY_TMP_FILE={BIGQUERY_TMP_FILE}\",\n", - " f\"BIGQUERY_DATASET_ID={BIGQUERY_DATASET_ID}\",\n", - " f\"BIGQUERY_LOCATION={BIGQUERY_LOCATION}\",\n", - " f\"BIGQUERY_TABLE_ID={BIGQUERY_TABLE_ID}\",\n", - "])\n", + "ENV_VARS = \",\".join(\n", + " [\n", + " f\"PROJECT_ID={PROJECT_ID}\",\n", + " f\"RAW_DATA_PATH={RAW_DATA_PATH}\",\n", + " f\"BATCH_SIZE={BATCH_SIZE}\",\n", + " f\"RANK_K={RANK_K}\",\n", + " f\"NUM_ACTIONS={NUM_ACTIONS}\",\n", + " f\"BIGQUERY_TMP_FILE={BIGQUERY_TMP_FILE}\",\n", + " f\"BIGQUERY_DATASET_ID={BIGQUERY_DATASET_ID}\",\n", + " f\"BIGQUERY_LOCATION={BIGQUERY_LOCATION}\",\n", + " f\"BIGQUERY_TABLE_ID={BIGQUERY_TABLE_ID}\",\n", + " ]\n", + ")\n", "\n", "! echo $ENV_VARS" ] @@ -1481,121 +1424,115 @@ }, "outputs": [], "source": [ - "@dsl.pipeline(\n", - " pipeline_root=PIPELINE_ROOT,\n", - " name=f\"{PIPELINE_NAME}-retraining\")\n", + "from kfp.components import load_component_from_url\n", + "\n", + "ingest_op = load_component_from_url(\n", + " \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml\"\n", + ")\n", + "train_op = load_component_from_url(\n", + " \"https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/68d6cf46ee22a9b9295d62ea71996150baf8db94/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml\"\n", + ")\n", + "\n", + "\n", + "@dsl.pipeline(pipeline_root=PIPELINE_ROOT, name=f\"{PIPELINE_NAME}-retraining\")\n", "def pipeline(\n", " # Pipeline configs\n", " project_id: str,\n", " training_artifacts_dir: str,\n", - "\n", " # BigQuery configs\n", " bigquery_table_id: str,\n", " bigquery_max_rows: int = 10000,\n", - "\n", " # TF-Agents RL configs\n", " rank_k: int = 20,\n", " num_actions: int = 20,\n", " num_epochs: int = 5,\n", " tikhonov_weight: float = 0.01,\n", - " agent_alpha: float = 10) -\u003e None:\n", - " \"\"\"Authors a re-training pipeline for MovieLens movie recommendation system.\n", + " agent_alpha: float = 10,\n", + ") -> None:\n", + " \"\"\"Authors a re-training pipeline for MovieLens movie recommendation system.\n", "\n", - " Integrates the Ingester, Trainer and Deployer components.\n", + " Integrates the Ingester, Trainer and Deployer components.\n", "\n", - " Args:\n", - " project_id: GCP project ID. This is required because otherwise the BigQuery\n", - " client will use the ID of the tenant GCP project created as a result of\n", - " KFP, which doesn't have proper access to BigQuery.\n", - " training_artifacts_dir: Path to store the Trainer artifacts (trained policy).\n", + " Args:\n", + " project_id: GCP project ID. This is required because otherwise the BigQuery\n", + " client will use the ID of the tenant GCP project created as a result of\n", + " KFP, which doesn't have proper access to BigQuery.\n", + " training_artifacts_dir: Path to store the Trainer artifacts (trained policy).\n", "\n", - " bigquery_table_id: A string of the BigQuery table ID in the format of\n", - " \"project.dataset.table\".\n", - " bigquery_max_rows: Optional; maximum number of rows to ingest.\n", + " bigquery_table_id: A string of the BigQuery table ID in the format of\n", + " \"project.dataset.table\".\n", + " bigquery_max_rows: Optional; maximum number of rows to ingest.\n", "\n", - " rank_k: Optional; rank for matrix factorization in the MovieLens environment;\n", - " also the observation dimension.\n", - " num_actions: Optional; number of actions (movie items) to choose from.\n", - " num_epochs: Optional; number of training epochs.\n", - " tikhonov_weight: Optional; LinUCB Tikhonov regularization weight of the\n", - " Trainer.\n", - " agent_alpha: Optional; LinUCB exploration parameter that multiplies the\n", - " confidence intervals of the Trainer.\n", - " \"\"\"\n", - " # Run the Ingester component.\n", - " ingest_op = create_component_from_func(\n", - " func=ingester_component.ingest_bigquery_dataset_into_tfrecord,\n", - " output_component_file=f\"ingester-{OUTPUT_COMPONENT_SPEC}\",\n", - " packages_to_install=[\n", - " \"google-cloud-bigquery==2.20.0\",\n", - " \"tensorflow==2.5.0\",\n", - " ])\n", - " ingest_task = ingest_op(\n", - " project_id=project_id,\n", - " bigquery_table_id=bigquery_table_id,\n", - " bigquery_max_rows=bigquery_max_rows,\n", - " tfrecord_file=TFRECORD_FILE)\n", + " rank_k: Optional; rank for matrix factorization in the MovieLens environment;\n", + " also the observation dimension.\n", + " num_actions: Optional; number of actions (movie items) to choose from.\n", + " num_epochs: Optional; number of training epochs.\n", + " tikhonov_weight: Optional; LinUCB Tikhonov regularization weight of the\n", + " Trainer.\n", + " agent_alpha: Optional; LinUCB exploration parameter that multiplies the\n", + " confidence intervals of the Trainer.\n", + " \"\"\"\n", + " # Run the Ingester component.\n", + " ingest_task = ingest_op(\n", + " project_id=project_id,\n", + " bigquery_table_id=bigquery_table_id,\n", + " bigquery_max_rows=bigquery_max_rows,\n", + " tfrecord_file=TFRECORD_FILE,\n", + " )\n", "\n", - " # Run the Trainer component and submit custom job to Vertex AI.\n", - " train_op = create_component_from_func(\n", - " func=trainer_component.training_op,\n", - " output_component_file=f\"trainer-{OUTPUT_COMPONENT_SPEC}\",\n", - " packages_to_install=[\n", - " \"tensorflow==2.5.0\",\n", - " \"tf-agents==0.8.0\",\n", - " ])\n", - " train_task = train_op(\n", - " training_artifacts_dir=training_artifacts_dir,\n", - " tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n", - " num_epochs=num_epochs,\n", - " rank_k=rank_k,\n", - " num_actions=num_actions,\n", - " tikhonov_weight=tikhonov_weight,\n", - " agent_alpha=agent_alpha)\n", + " # Run the Trainer component and submit custom job to Vertex AI.\n", + " train_task = train_op(\n", + " training_artifacts_dir=training_artifacts_dir,\n", + " tfrecord_file=ingest_task.outputs[\"tfrecord_file\"],\n", + " num_epochs=num_epochs,\n", + " rank_k=rank_k,\n", + " num_actions=num_actions,\n", + " tikhonov_weight=tikhonov_weight,\n", + " agent_alpha=agent_alpha,\n", + " )\n", "\n", - " worker_pool_specs = [\n", - " {\n", - " \"containerSpec\": {\n", - " \"imageUri\":train_task.container.image,\n", - " },\n", - " \"replicaCount\": TRAINING_REPLICA_COUNT,\n", - " \"machineSpec\": {\n", - " \"machineType\": TRAINING_MACHINE_TYPE,\n", - " \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n", - " \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n", - " },\n", - " },\n", - " ]\n", - " train_task.custom_job_spec = {\n", - " \"displayName\": train_task.name,\n", - " \"jobSpec\": {\n", - " \"workerPoolSpecs\": worker_pool_specs,\n", - " }\n", - " }\n", + " worker_pool_specs = [\n", + " {\n", + " \"containerSpec\": {\n", + " \"imageUri\": train_task.container.image,\n", + " },\n", + " \"replicaCount\": TRAINING_REPLICA_COUNT,\n", + " \"machineSpec\": {\n", + " \"machineType\": TRAINING_MACHINE_TYPE,\n", + " \"acceleratorType\": TRAINING_ACCELERATOR_TYPE,\n", + " \"acceleratorCount\": TRAINING_ACCELERATOR_COUNT,\n", + " },\n", + " },\n", + " ]\n", + " train_task.custom_job_spec = {\n", + " \"displayName\": train_task.name,\n", + " \"jobSpec\": {\n", + " \"workerPoolSpecs\": worker_pool_specs,\n", + " },\n", + " }\n", "\n", - " # Run the Deployer components.\n", - " # Upload the trained policy as a model.\n", - " model_upload_op = gcc_aip.ModelUploadOp(\n", - " project=project_id,\n", - " display_name=TRAINED_POLICY_DISPLAY_NAME,\n", - " artifact_uri=training_artifacts_dir,\n", - " serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n", - " )\n", - " # Model uploading has to occur after training completes.\n", - " model_upload_op.after(train_task)\n", - " # Create a Vertex AI endpoint. (This operation can occur in parallel with\n", - " # the Generator, Ingester, Trainer components.)\n", - " endpoint_create_op = gcc_aip.EndpointCreateOp(\n", - " project=project_id,\n", - " display_name=ENDPOINT_DISPLAY_NAME)\n", - " # Deploy the uploaded, trained policy to the created endpoint. (This operation\n", - " # has to occur after both model uploading and endpoint creation complete.)\n", - " model_deploy_op = gcc_aip.ModelDeployOp(\n", - " project=project_id,\n", - " endpoint=endpoint_create_op.outputs[\"endpoint\"],\n", - " model=model_upload_op.outputs[\"model\"],\n", - " deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n", - " machine_type=ENDPOINT_MACHINE_TYPE)" + " # Run the Deployer components.\n", + " # Upload the trained policy as a model.\n", + " model_upload_op = gcc_aip.ModelUploadOp(\n", + " project=project_id,\n", + " display_name=TRAINED_POLICY_DISPLAY_NAME,\n", + " artifact_uri=train_task.outputs[\"training_artifacts_dir\"],\n", + " serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n", + " )\n", + " # Create a Vertex AI endpoint. (This operation can occur in parallel with\n", + " # the Generator, Ingester, Trainer components.)\n", + " endpoint_create_op = gcc_aip.EndpointCreateOp(\n", + " project=project_id, display_name=ENDPOINT_DISPLAY_NAME\n", + " )\n", + " # Deploy the uploaded, trained policy to the created endpoint. (This operation\n", + " # has to occur after both model uploading and endpoint creation complete.)\n", + " gcc_aip.ModelDeployOp(\n", + " project=project_id,\n", + " endpoint=endpoint_create_op.outputs[\"endpoint\"],\n", + " model=model_upload_op.outputs[\"model\"],\n", + " deployed_model_display_name=TRAINED_POLICY_DISPLAY_NAME,\n", + " machine_type=ENDPOINT_MACHINE_TYPE,\n", + " )" ] }, { @@ -1607,13 +1544,10 @@ "outputs": [], "source": [ "# Compile the authored pipeline.\n", - "compiler.Compiler().compile(pipeline_func=pipeline, \n", - " package_path=PIPELINE_SPEC_PATH)\n", + "compiler.Compiler().compile(pipeline_func=pipeline, package_path=PIPELINE_SPEC_PATH)\n", "\n", "# Createa Vertex AI client.\n", - "api_client = AIPlatformClient(\n", - " project_id=PROJECT_ID,\n", - " region=REGION)\n", + "api_client = AIPlatformClient(project_id=PROJECT_ID, region=REGION)\n", "\n", "# Schedule a recurring pipeline.\n", "response = api_client.create_schedule_from_job_spec(\n", @@ -1623,10 +1557,10 @@ " # Pipeline configs\n", " \"project_id\": PROJECT_ID,\n", " \"training_artifacts_dir\": TRAINING_ARTIFACTS_DIR,\n", - "\n", " # BigQuery config\n", " \"bigquery_table_id\": BIGQUERY_TABLE_ID,\n", - " })\n", + " },\n", + ")\n", "response[\"name\"]" ] }, @@ -1675,84 +1609,12 @@ "metadata": { "colab": { "collapsed_sections": [], - "name": "pipeline_reinforcement_learning_vertex_ai.ipynb", - "private_outputs": true, - "provenance": [ - { - "file_id": "/piper/depot/google3/cloud/ml/growth/experiments/reinforcement_learning/third_party/pipeline_reinforcement_learning_vertex_ai/pipeline_reinforcement_learning_vertex_ai.ipynb?workspaceId=feiyangyu:demo-links::citc", - "timestamp": 1628009146717 - }, - { - "file_id": "1jkxispXKEQcOKLXaw0mqrKe3mZraZiXg", - "timestamp": 1628008968400 - }, - { - "file_id": "1ljF4QzhcoEOrBfwmKo78Mqa3iWXYsmjO", - "timestamp": 1627697719794 - }, - { - "file_id": "1OAK4Kz0Naqm9dNYQxieLZ4TXP-N0mEAL", - "timestamp": 1627660909950 - }, - { - "file_id": "1nXXJGPmNSxvueyzRs7BLXoXseH67Dv7o", - "timestamp": 1627602309156 - }, - { - "file_id": "1XgFiNJ5eQOjuZ6ZaN6xNQ33xqsp416Ng", - "timestamp": 1627503527404 - }, - { - "file_id": "1w8p_X_kEop_Dhij_azn1Rx_Tv3Rfisy1", - "timestamp": 1627080324892 - }, - { - "file_id": "/piper/depot/google3/cloud/ml/growth/experiments/reinforcement_learning/demo/pipeline_reinforcement_learning_vertex_ai/pipeline_reinforcement_learning_vertex_ai.ipynb?workspaceId=feiyangyu:pipeline-construction-trainer::citc", - "timestamp": 1626910322761 - }, - { - "file_id": "1b0HtLvgNMsQQfHgpc5d2W1B_uJAzPqM9", - "timestamp": 1626909409390 - }, - { - "file_id": "1QYImZG95fKpseQw9XK46RqOJotcFomZJ", - "timestamp": 1626454662899 - }, - { - "file_id": "1UR0_lxXy6avD274M-EpfYUGUXeRjny17", - "timestamp": 1626375689136 - }, - { - "file_id": "1gbjpy_AzJ_z9nMWre1pTwMhuv_7jTixv", - "timestamp": 1626219338889 - }, - { - "file_id": "1tquvXjsY_vs9Z_1ZEJu1u9888KaArkP6", - "timestamp": 1626116646590 - } - ] - }, - "environment": { - "name": "common-cu110.m69", - "type": "gcloud", - "uri": "gcr.io/deeplearning-platform-release/base-cu110:m69" + "name": "mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb", + "toc_visible": true }, "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, diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/component.yaml b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/component.yaml new file mode 100644 index 000000000..e7575679f --- /dev/null +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/component.yaml @@ -0,0 +1,359 @@ +name: Generate movielens dataset for bigquery +description: Generates BigQuery training data using a MovieLens simulation environment. +inputs: +- name: project_id + type: String + description: |- + GCP project ID. This is required because otherwise the BigQuery + client will use the ID of the tenant GCP project created as a result of + KFP, which doesn't have proper access to BigQuery. +- {name: raw_data_path, type: String, description: Path to MovieLens 100K's "u.data" + file.} +- {name: batch_size, type: Integer, description: Batch size of environment generated + quantities eg. rewards.} +- name: rank_k + type: Integer + description: |- + Rank for matrix factorization in the MovieLens environment; also + the observation dimension. +- {name: num_actions, type: Integer, description: Number of actions (movie items) + to choose from.} +- {name: driver_steps, type: Integer, description: Number of steps to run per batch.} +- {name: bigquery_tmp_file, type: String, description: Path to a JSON file containing + the training dataset.} +- name: bigquery_dataset_id + type: String + description: |- + A string of the BigQuery dataset ID in the format of + "project.dataset". +- {name: bigquery_location, type: String, description: A string of the BigQuery dataset + location.} +- name: bigquery_table_id + type: String + description: |- + A string of the BigQuery table ID in the format of + "project.dataset.table". +outputs: +- {name: bigquery_dataset_id, type: String} +- {name: bigquery_location, type: String} +- {name: bigquery_table_id, type: String} +implementation: + container: + image: tensorflow/tensorflow:2.5.0 + command: + - sh + - -c + - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location + 'google-cloud-bigquery==2.20.0' 'tensorflow==2.5.0' 'tf-agents==0.8.0' || PIP_DISABLE_PIP_VERSION_CHECK=1 + python3 -m pip install --quiet --no-warn-script-location 'google-cloud-bigquery==2.20.0' + 'tensorflow==2.5.0' 'tf-agents==0.8.0' --user) && "$0" "$@" + - sh + - -ec + - | + program_path=$(mktemp) + printf "%s" "$0" > "$program_path" + python3 -u "$program_path" "$@" + - | + def generate_movielens_dataset_for_bigquery( + project_id, + raw_data_path, + batch_size, + rank_k, + num_actions, + driver_steps, + bigquery_tmp_file, + bigquery_dataset_id, + bigquery_location, + bigquery_table_id + ): + """Generates BigQuery training data using a MovieLens simulation environment. + + Serves as the Generator pipeline component: + 1. Generates `trajectories.Trajectory` data by applying a random policy on + MovieLens simulation environment. + 2. Converts `trajectories.Trajectory` data to JSON format. + 3. Loads JSON-formatted data into BigQuery. + + This function is to be built into a Kubeflow Pipelines (KFP) component. As a + result, this function must be entirely self-contained. This means that the + import statements and helper functions must reside within itself. + + Args: + project_id: GCP project ID. This is required because otherwise the BigQuery + client will use the ID of the tenant GCP project created as a result of + KFP, which doesn't have proper access to BigQuery. + raw_data_path: Path to MovieLens 100K's "u.data" file. + batch_size: Batch size of environment generated quantities eg. rewards. + rank_k: Rank for matrix factorization in the MovieLens environment; also + the observation dimension. + num_actions: Number of actions (movie items) to choose from. + driver_steps: Number of steps to run per batch. + bigquery_tmp_file: Path to a JSON file containing the training dataset. + bigquery_dataset_id: A string of the BigQuery dataset ID in the format of + "project.dataset". + bigquery_location: A string of the BigQuery dataset location. + bigquery_table_id: A string of the BigQuery table ID in the format of + "project.dataset.table". + + Returns: + A NamedTuple of (`bigquery_dataset_id`, `bigquery_location`, + `bigquery_table_id`). + """ + # pylint: disable=g-import-not-at-top + import collections + import json + from typing import Any, Dict + + from google.cloud import bigquery + + from tf_agents import replay_buffers + from tf_agents import trajectories + from tf_agents.bandits.agents.examples.v2 import trainer + from tf_agents.bandits.environments import movielens_py_environment + from tf_agents.drivers import dynamic_step_driver + from tf_agents.environments import tf_py_environment + from tf_agents.policies import random_tf_policy + + def generate_simulation_data( + raw_data_path, + batch_size, + rank_k, + num_actions, + driver_steps): + """Generates `trajectories.Trajectory` data from the simulation environment. + + Constructs a MovieLens simulation environment, and generates a set of + `trajectories.Trajectory` data using a random policy. + + Args: + raw_data_path: Path to MovieLens 100K's "u.data" file. + batch_size: Batch size of environment generated quantities eg. rewards. + rank_k: Rank for matrix factorization in the MovieLens environment; also + the observation dimension. + num_actions: Number of actions (movie items) to choose from. + driver_steps: Number of steps to run per batch. + + Returns: + A replay buffer holding randomly generated`trajectories.Trajectory` data. + """ + # Create MovieLens simulation environment. + env = movielens_py_environment.MovieLensPyEnvironment( + raw_data_path, + rank_k, + batch_size, + num_movies=num_actions, + csv_delimiter="\t") + environment = tf_py_environment.TFPyEnvironment(env) + + # Define random policy for collecting data. + random_policy = random_tf_policy.RandomTFPolicy( + action_spec=environment.action_spec(), + time_step_spec=environment.time_step_spec()) + + # Use replay buffer and observers to keep track of Trajectory data. + data_spec = random_policy.trajectory_spec + replay_buffer = trainer.get_replay_buffer(data_spec, environment.batch_size, + driver_steps) + observers = [replay_buffer.add_batch] + + # Run driver to apply the random policy in the simulation environment. + driver = dynamic_step_driver.DynamicStepDriver( + env=environment, + policy=random_policy, + num_steps=driver_steps * environment.batch_size, + observers=observers) + driver.run() + + return replay_buffer + + def build_dict_from_trajectory( + trajectory): + """Builds a dict from `trajectory` data. + + Args: + trajectory: A `trajectories.Trajectory` object. + + Returns: + A dict holding the same data as `trajectory`. + """ + trajectory_dict = { + "step_type": trajectory.step_type.numpy().tolist(), + "observation": [{ + "observation_batch": batch + } for batch in trajectory.observation.numpy().tolist()], + "action": trajectory.action.numpy().tolist(), + "policy_info": trajectory.policy_info, + "next_step_type": trajectory.next_step_type.numpy().tolist(), + "reward": trajectory.reward.numpy().tolist(), + "discount": trajectory.discount.numpy().tolist(), + } + return trajectory_dict + + def write_replay_buffer_to_file( + replay_buffer, + batch_size, + dataset_file): + """Writes replay buffer data to a file, each JSON in one line. + + Each `trajectories.Trajectory` object in `replay_buffer` will be written as + one line to the `dataset_file` in JSON format. I.e., the `dataset_file` + would be a newline-delimited JSON file. + + Args: + replay_buffer: A `replay_buffers.TFUniformReplayBuffer` holding + `trajectories.Trajectory` objects. + batch_size: Batch size of environment generated quantities eg. rewards. + dataset_file: File path. Will be overwritten if already exists. + """ + dataset = replay_buffer.as_dataset(sample_batch_size=batch_size) + dataset_size = replay_buffer.num_frames().numpy() + + with open(dataset_file, "w") as f: + for example in dataset.take(count=dataset_size): + traj_dict = build_dict_from_trajectory(example[0]) + f.write(json.dumps(traj_dict) + "\n") + + def load_dataset_into_bigquery( + project_id, + dataset_file, + bigquery_dataset_id, + bigquery_location, + bigquery_table_id): + """Loads training dataset into BigQuery table. + + Loads training dataset of `trajectories.Trajectory` in newline delimited + JSON into a BigQuery dataset and table, using a BigQuery client. + + Args: + project_id: GCP project ID. This is required because otherwise the + BigQuery client will use the ID of the tenant GCP project created as a + result of KFP, which doesn't have proper access to BigQuery. + dataset_file: Path to a JSON file containing the training dataset. + bigquery_dataset_id: A string of the BigQuery dataset ID in the format of + "project.dataset". + bigquery_location: A string of the BigQuery dataset location. + bigquery_table_id: A string of the BigQuery table ID in the format of + "project.dataset.table". + """ + # Construct a BigQuery client object. + client = bigquery.Client(project=project_id) + + # Construct a full Dataset object to send to the API. + dataset = bigquery.Dataset(bigquery_dataset_id) + + # Specify the geographic location where the dataset should reside. + dataset.location = bigquery_location + + # Create the dataset, or get the dataset if it exists. + dataset = client.create_dataset(dataset, exists_ok=True, timeout=30) + + job_config = bigquery.LoadJobConfig( + schema=[ + bigquery.SchemaField("step_type", "INT64", mode="REPEATED"), + bigquery.SchemaField( + "observation", + "RECORD", + mode="REPEATED", + fields=[ + bigquery.SchemaField("observation_batch", "FLOAT64", + "REPEATED") + ]), + bigquery.SchemaField("action", "INT64", mode="REPEATED"), + bigquery.SchemaField("policy_info", "FLOAT64", mode="REPEATED"), + bigquery.SchemaField("next_step_type", "INT64", mode="REPEATED"), + bigquery.SchemaField("reward", "FLOAT64", mode="REPEATED"), + bigquery.SchemaField("discount", "FLOAT64", mode="REPEATED"), + ], + source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON, + ) + + with open(dataset_file, "rb") as source_file: + load_job = client.load_table_from_file( + source_file, bigquery_table_id, job_config=job_config) + + load_job.result() # Wait for the job to complete. + + replay_buffer = generate_simulation_data( + raw_data_path=raw_data_path, + batch_size=batch_size, + rank_k=rank_k, + num_actions=num_actions, + driver_steps=driver_steps) + + write_replay_buffer_to_file( + replay_buffer=replay_buffer, + batch_size=batch_size, + dataset_file=bigquery_tmp_file) + + load_dataset_into_bigquery(project_id, bigquery_tmp_file, bigquery_dataset_id, + bigquery_location, bigquery_table_id) + + outputs = collections.namedtuple( + "Outputs", + ["bigquery_dataset_id", "bigquery_location", "bigquery_table_id"]) + + return outputs(bigquery_dataset_id, bigquery_location, bigquery_table_id) + + def _serialize_str(str_value: str) -> str: + if not isinstance(str_value, str): + raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value)))) + return str_value + + import argparse + _parser = argparse.ArgumentParser(prog='Generate movielens dataset for bigquery', description='Generates BigQuery training data using a MovieLens simulation environment.') + _parser.add_argument("--project-id", dest="project_id", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--raw-data-path", dest="raw_data_path", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--batch-size", dest="batch_size", type=int, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--rank-k", dest="rank_k", type=int, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--num-actions", dest="num_actions", type=int, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--driver-steps", dest="driver_steps", type=int, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--bigquery-tmp-file", dest="bigquery_tmp_file", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--bigquery-dataset-id", dest="bigquery_dataset_id", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--bigquery-location", dest="bigquery_location", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--bigquery-table-id", dest="bigquery_table_id", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=3) + _parsed_args = vars(_parser.parse_args()) + _output_files = _parsed_args.pop("_output_paths", []) + + _outputs = generate_movielens_dataset_for_bigquery(**_parsed_args) + + _output_serializers = [ + _serialize_str, + _serialize_str, + _serialize_str, + + ] + + import os + for idx, output_file in enumerate(_output_files): + try: + os.makedirs(os.path.dirname(output_file)) + except OSError: + pass + with open(output_file, 'w') as f: + f.write(_output_serializers[idx](_outputs[idx])) + args: + - --project-id + - {inputValue: project_id} + - --raw-data-path + - {inputValue: raw_data_path} + - --batch-size + - {inputValue: batch_size} + - --rank-k + - {inputValue: rank_k} + - --num-actions + - {inputValue: num_actions} + - --driver-steps + - {inputValue: driver_steps} + - --bigquery-tmp-file + - {inputValue: bigquery_tmp_file} + - --bigquery-dataset-id + - {inputValue: bigquery_dataset_id} + - --bigquery-location + - {inputValue: bigquery_location} + - --bigquery-table-id + - {inputValue: bigquery_table_id} + - '----output-paths' + - {outputPath: bigquery_dataset_id} + - {outputPath: bigquery_location} + - {outputPath: bigquery_table_id} diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/generator_component.py b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/generator_component.py index 34e1a045a..e23a8020f 100644 --- a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/generator_component.py +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/generator/generator_component.py @@ -259,3 +259,18 @@ def generate_movielens_dataset_for_bigquery( ["bigquery_dataset_id", "bigquery_location", "bigquery_table_id"]) return outputs(bigquery_dataset_id, bigquery_location, bigquery_table_id) + + +if __name__ == "__main__": + from kfp.components import create_component_from_func + + generate_movielens_dataset_for_bigquery_op = create_component_from_func( + func=generate_movielens_dataset_for_bigquery, + base_image="tensorflow/tensorflow:2.5.0", + output_component_file="component.yaml", + packages_to_install=[ + "google-cloud-bigquery==2.20.0", + "tensorflow==2.5.0", + "tf-agents==0.8.0", + ], + ) diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml new file mode 100644 index 000000000..7bfde7703 --- /dev/null +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/component.yaml @@ -0,0 +1,231 @@ +name: Ingest bigquery dataset into tfrecord +description: Ingests data from BigQuery, formats them and outputs TFRecord files. +inputs: +- name: project_id + type: String + description: |- + GCP project ID. This is required because otherwise the BigQuery + client will use the ID of the tenant GCP project created as a result of + KFP, which doesn't have proper access to BigQuery. +- name: bigquery_table_id + type: String + description: |- + A string of the BigQuery table ID in the format of + "project.dataset.table". +- {name: tfrecord_file, type: String, description: Path to file to write the ingestion + result TFRecords.} +- {name: bigquery_max_rows, type: Integer, description: Optional; maximum number of + rows to ingest., optional: true} +outputs: +- {name: tfrecord_file, type: String} +implementation: + container: + image: tensorflow/tensorflow:2.5.0 + command: + - sh + - -c + - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location + 'google-cloud-bigquery==2.20.0' 'tensorflow==2.5.0' || PIP_DISABLE_PIP_VERSION_CHECK=1 + python3 -m pip install --quiet --no-warn-script-location 'google-cloud-bigquery==2.20.0' + 'tensorflow==2.5.0' --user) && "$0" "$@" + - sh + - -ec + - | + program_path=$(mktemp) + printf "%s" "$0" > "$program_path" + python3 -u "$program_path" "$@" + - | + def ingest_bigquery_dataset_into_tfrecord( + project_id, + bigquery_table_id, + tfrecord_file, + bigquery_max_rows = None + ): + """Ingests data from BigQuery, formats them and outputs TFRecord files. + + Serves as the Ingester pipeline component: + 1. Reads data in BigQuery that contains 7 pieces of data: `step_type`, + `observation`, `action`, `policy_info`, `next_step_type`, `reward`, + `discount`. + 2. Packages the data as `tf.train.Example` objects and outputs them as + TFRecord files. + + This function is to be built into a Kubeflow Pipelines (KFP) component. As a + result, this function must be entirely self-contained. This means that the + import statements and helper functions must reside within itself. + + Args: + project_id: GCP project ID. This is required because otherwise the BigQuery + client will use the ID of the tenant GCP project created as a result of + KFP, which doesn't have proper access to BigQuery. + bigquery_table_id: A string of the BigQuery table ID in the format of + "project.dataset.table". + tfrecord_file: Path to file to write the ingestion result TFRecords. + bigquery_max_rows: Optional; maximum number of rows to ingest. + + Returns: + A NamedTuple of the path to the output TFRecord file. + """ + # pylint: disable=g-import-not-at-top + import collections + from typing import Optional + + from google.cloud import bigquery + + import tensorflow as tf + + def read_data_from_bigquery( + project_id, + bigquery_table_id, + bigquery_max_rows): + """Reads data from BigQuery at `bigquery_table_id` and creates an iterator. + + The table contains 7 columns that form `trajectories.Trajectory` objects: + `step_type`, `observation`, `action`, `policy_info`, `next_step_type`, + `reward`, `discount`. + + Args: + project_id: GCP project ID. This is required because otherwise the + BigQuery client will use the ID of the tenant GCP project created as a + result of KFP, which doesn't have proper access to BigQuery. + bigquery_table_id: A string of the BigQuery table ID in the format of + "project.dataset.table". + bigquery_max_rows: Optional; maximum number of rows to fetch. + + Returns: + A row iterator over all data at `bigquery_table_id`. + """ + # Construct a BigQuery client object. + client = bigquery.Client(project=project_id) + + # Get dataset. + query_job = client.query( + f""" + SELECT * FROM {bigquery_table_id} + """ + ) + table = query_job.result(max_results=bigquery_max_rows) + + return table + + def _bytes_feature(tensor): + """Returns a `tf.train.Feature` with bytes from `tensor`. + + Args: + tensor: A `tf.Tensor` object. + + Returns: + A `tf.train.Feature` object containing bytes that represent the content of + `tensor`. + """ + value = tf.io.serialize_tensor(tensor) + if isinstance(value, type(tf.constant(0))): + value = value.numpy() + return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) + + def build_example(data_row): + """Builds a `tf.train.Example` from `data_row` content. + + Args: + data_row: A `bigquery.table.Row` object that contains 7 pieces of data: + `step_type`, `observation`, `action`, `policy_info`, `next_step_type`, + `reward`, `discount`. Each piece of data except `observation` is a 1D + array; `observation` is a 1D array of `{"observation_batch": 1D array}.` + + Returns: + A `tf.train.Example` object holding the same data as `data_row`. + """ + feature = { + "step_type": + _bytes_feature(data_row.get("step_type")), + "observation": + _bytes_feature([ + observation["observation_batch"] + for observation in data_row.get("observation") + ]), + "action": + _bytes_feature(data_row.get("action")), + "policy_info": + _bytes_feature(data_row.get("policy_info")), + "next_step_type": + _bytes_feature(data_row.get("next_step_type")), + "reward": + _bytes_feature(data_row.get("reward")), + "discount": + _bytes_feature(data_row.get("discount")), + } + example_proto = tf.train.Example( + features=tf.train.Features(feature=feature)) + return example_proto + + def write_tfrecords( + tfrecord_file, + table): + """Writes the row data in `table` into TFRecords in `tfrecord_file`. + + Args: + tfrecord_file: Path to file to write the TFRecords. + table: A row iterator over all data to be written. + """ + with tf.io.TFRecordWriter(tfrecord_file) as writer: + for data_row in table: + example = build_example(data_row) + writer.write(example.SerializeToString()) + + table = read_data_from_bigquery( + project_id=project_id, + bigquery_table_id=bigquery_table_id, + bigquery_max_rows=bigquery_max_rows) + + write_tfrecords(tfrecord_file, table) + + outputs = collections.namedtuple( + "Outputs", + ["tfrecord_file"]) + + return outputs(tfrecord_file) + + def _serialize_str(str_value: str) -> str: + if not isinstance(str_value, str): + raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value)))) + return str_value + + import argparse + _parser = argparse.ArgumentParser(prog='Ingest bigquery dataset into tfrecord', description='Ingests data from BigQuery, formats them and outputs TFRecord files.') + _parser.add_argument("--project-id", dest="project_id", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--bigquery-table-id", dest="bigquery_table_id", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--tfrecord-file", dest="tfrecord_file", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--bigquery-max-rows", dest="bigquery_max_rows", type=int, required=False, default=argparse.SUPPRESS) + _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=1) + _parsed_args = vars(_parser.parse_args()) + _output_files = _parsed_args.pop("_output_paths", []) + + _outputs = ingest_bigquery_dataset_into_tfrecord(**_parsed_args) + + _output_serializers = [ + _serialize_str, + + ] + + import os + for idx, output_file in enumerate(_output_files): + try: + os.makedirs(os.path.dirname(output_file)) + except OSError: + pass + with open(output_file, 'w') as f: + f.write(_output_serializers[idx](_outputs[idx])) + args: + - --project-id + - {inputValue: project_id} + - --bigquery-table-id + - {inputValue: bigquery_table_id} + - --tfrecord-file + - {inputValue: tfrecord_file} + - if: + cond: {isPresent: bigquery_max_rows} + then: + - --bigquery-max-rows + - {inputValue: bigquery_max_rows} + - '----output-paths' + - {outputPath: tfrecord_file} diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/ingester_component.py b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/ingester_component.py index 121b538bf..fb43d3591 100644 --- a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/ingester_component.py +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/ingester/ingester_component.py @@ -167,3 +167,17 @@ def ingest_bigquery_dataset_into_tfrecord( ["tfrecord_file"]) return outputs(tfrecord_file) + + +if __name__ == "__main__": + from kfp.components import create_component_from_func + + ingest_bigquery_dataset_into_tfrecord_op = create_component_from_func( + func=ingest_bigquery_dataset_into_tfrecord, + base_image="tensorflow/tensorflow:2.5.0", + output_component_file="component.yaml", + packages_to_install=[ + "google-cloud-bigquery==2.20.0", + "tensorflow==2.5.0", + ], + ) diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml new file mode 100644 index 000000000..e8b376478 --- /dev/null +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/component.yaml @@ -0,0 +1,320 @@ +name: Train reinforcement learning policy +description: Implements off-policy training for a policy on dataset of TFRecord files. +inputs: +- name: training_artifacts_dir + type: String + description: |- + Path to store the Trainer artifacts (trained + policy). +- {name: tfrecord_file, type: String, description: Path to file to write the ingestion + result TFRecords.} +- {name: num_epochs, type: Integer, description: Number of training epochs.} +- name: rank_k + type: Integer + description: |- + Rank for matrix factorization in the MovieLens environment; also + the observation dimension. +- {name: num_actions, type: Integer, description: Number of actions (movie items) + to choose from.} +- {name: tikhonov_weight, type: Float, description: LinUCB Tikhonov regularization + weight of the Trainer.} +- name: agent_alpha + type: Float + description: |- + LinUCB exploration parameter that multiplies the confidence + intervals of the Trainer. +outputs: +- {name: training_artifacts_dir, type: String} +implementation: + container: + image: tensorflow/tensorflow:2.5.0 + command: + - sh + - -c + - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location + 'tensorflow==2.5.0' 'tf-agents==0.8.0' || PIP_DISABLE_PIP_VERSION_CHECK=1 python3 + -m pip install --quiet --no-warn-script-location 'tensorflow==2.5.0' 'tf-agents==0.8.0' + --user) && "$0" "$@" + - sh + - -ec + - | + program_path=$(mktemp) + printf "%s" "$0" > "$program_path" + python3 -u "$program_path" "$@" + - | + def train_reinforcement_learning_policy( + training_artifacts_dir, + tfrecord_file, + num_epochs, + rank_k, + num_actions, + tikhonov_weight, + agent_alpha + ): + """Implements off-policy training for a policy on dataset of TFRecord files. + + The Trainer's task is to submit a remote training job to Vertex AI, with the + training logic of a specified custom training container. The task will be + handled by: `kfp.v2.google.experimental.run_as_aiplatform_custom_job` (which + takes in the component made from this placeholder function) + + This function is to be built into a Kubeflow Pipelines (KFP) component. As a + result, this function must be entirely self-contained. This means that the + import statements and helper functions must reside within itself. + + Args: + training_artifacts_dir: Path to store the Trainer artifacts (trained + policy). + tfrecord_file: Path to file to write the ingestion result TFRecords. + num_epochs: Number of training epochs. + rank_k: Rank for matrix factorization in the MovieLens environment; also + the observation dimension. + num_actions: Number of actions (movie items) to choose from. + tikhonov_weight: LinUCB Tikhonov regularization weight of the Trainer. + agent_alpha: LinUCB exploration parameter that multiplies the confidence + intervals of the Trainer. + + Returns: + A NamedTuple of (`training_artifacts_dir`). + """ + # pylint: disable=g-import-not-at-top + import collections + from typing import Dict, List, NamedTuple # pylint: disable=redefined-outer-name,reimported + + import tensorflow as tf + + from tf_agents import agents + from tf_agents import policies + from tf_agents import trajectories + from tf_agents.bandits.agents import lin_ucb_agent + from tf_agents.policies import policy_saver + from tf_agents.specs import tensor_spec + + import logging + + per_arm = False # Using the non-per-arm version of the MovieLens environment. + + # Mapping from feature name to serialized value + feature_description = { + "step_type": tf.io.FixedLenFeature((), tf.string), + "observation": tf.io.FixedLenFeature((), tf.string), + "action": tf.io.FixedLenFeature((), tf.string), + "policy_info": tf.io.FixedLenFeature((), tf.string), + "next_step_type": tf.io.FixedLenFeature((), tf.string), + "reward": tf.io.FixedLenFeature((), tf.string), + "discount": tf.io.FixedLenFeature((), tf.string), + } + + def _parse_record(raw_record): + """Parses a serialized `tf.train.Example` proto. + + Args: + raw_record: A serialized data record of a `tf.train.Example` proto. + + Returns: + A dict mapping feature names to values as `tf.Tensor` objects of type + string containing serialized protos, following `feature_description`. + """ + return tf.io.parse_single_example(raw_record, feature_description) + + def build_trajectory( + parsed_record, + policy_info): + """Builds a `trajectories.Trajectory` object from `parsed_record`. + + Args: + parsed_record: A dict mapping feature names to values as `tf.Tensor` + objects of type string containing serialized protos. + policy_info: Policy information specification. + + Returns: + A `trajectories.Trajectory` object that contains values as de-serialized + `tf.Tensor` objects from `parsed_record`. + """ + return trajectories.Trajectory( + step_type=tf.expand_dims( + tf.io.parse_tensor(parsed_record["step_type"], out_type=tf.int32), + axis=1), + observation=tf.expand_dims( + tf.io.parse_tensor( + parsed_record["observation"], out_type=tf.float32), + axis=1), + action=tf.expand_dims( + tf.io.parse_tensor(parsed_record["action"], out_type=tf.int32), + axis=1), + policy_info=policy_info, + next_step_type=tf.expand_dims( + tf.io.parse_tensor( + parsed_record["next_step_type"], out_type=tf.int32), + axis=1), + reward=tf.expand_dims( + tf.io.parse_tensor(parsed_record["reward"], out_type=tf.float32), + axis=1), + discount=tf.expand_dims( + tf.io.parse_tensor(parsed_record["discount"], out_type=tf.float32), + axis=1)) + + def train_policy_on_trajectory( + agent, + tfrecord_file, + num_epochs + ): + """Trains the policy in `agent` on the dataset of `tfrecord_file`. + + Parses `tfrecord_file` as `tf.train.Example` objects, packages them into + `trajectories.Trajectory` objects, and trains the agent's policy on these + trajectory objects. + + Args: + agent: A TF-Agents agent that carries the policy to train. + tfrecord_file: Path to the TFRecord file containing the training dataset. + num_epochs: Number of epochs to train the policy. + + Returns: + A NamedTuple of (a trained TF-Agents policy, a dict mapping from + "epoch" to lists of loss values produced at each training step). + """ + raw_dataset = tf.data.TFRecordDataset([tfrecord_file]) + parsed_dataset = raw_dataset.map(_parse_record) + + train_loss = collections.defaultdict(list) + for epoch in range(num_epochs): + for parsed_record in parsed_dataset: + trajectory = build_trajectory(parsed_record, agent.policy.info_spec) + loss, _ = agent.train(trajectory) + train_loss[f"epoch{epoch + 1}"].append(loss.numpy()) + + train_outputs = collections.namedtuple( + "TrainOutputs", + ["policy", "train_loss"]) + return train_outputs(agent.policy, train_loss) + + def execute_training_and_save_policy( + training_artifacts_dir, + tfrecord_file, + num_epochs, + rank_k, + num_actions, + tikhonov_weight, + agent_alpha): + """Executes training for the policy and saves the policy. + + Args: + training_artifacts_dir: Path to store the Trainer artifacts (trained + policy). + tfrecord_file: Path to file to write the ingestion result TFRecords. + num_epochs: Number of training epochs. + rank_k: Rank for matrix factorization in the MovieLens environment; also + the observation dimension. + num_actions: Number of actions (movie items) to choose from. + tikhonov_weight: LinUCB Tikhonov regularization weight of the Trainer. + agent_alpha: LinUCB exploration parameter that multiplies the confidence + intervals of the Trainer. + """ + # Define time step and action specs for one batch. + time_step_spec = trajectories.TimeStep( + step_type=tensor_spec.TensorSpec( + shape=(), dtype=tf.int32, name="step_type"), + reward=tensor_spec.TensorSpec( + shape=(), dtype=tf.float32, name="reward"), + discount=tensor_spec.BoundedTensorSpec( + shape=(), dtype=tf.float32, name="discount", minimum=0., + maximum=1.), + observation=tensor_spec.TensorSpec( + shape=(rank_k,), dtype=tf.float32, + name="observation")) + + action_spec = tensor_spec.BoundedTensorSpec( + shape=(), + dtype=tf.int32, + name="action", + minimum=0, + maximum=num_actions - 1) + + # Define RL agent/algorithm. + agent = lin_ucb_agent.LinearUCBAgent( + time_step_spec=time_step_spec, + action_spec=action_spec, + tikhonov_weight=tikhonov_weight, + alpha=agent_alpha, + dtype=tf.float32, + accepts_per_arm_features=per_arm) + agent.initialize() + logging.info("TimeStep Spec (for each batch):\n%s\n", agent.time_step_spec) + logging.info("Action Spec (for each batch):\n%s\n", agent.action_spec) + + # Perform off-policy training. + policy, _ = train_policy_on_trajectory( + agent=agent, + tfrecord_file=tfrecord_file, + num_epochs=num_epochs) + + # Save trained policy. + saver = policy_saver.PolicySaver(policy) + saver.save(training_artifacts_dir) + + execute_training_and_save_policy( + training_artifacts_dir=training_artifacts_dir, + tfrecord_file=tfrecord_file, + num_epochs=num_epochs, + rank_k=rank_k, + num_actions=num_actions, + tikhonov_weight=tikhonov_weight, + agent_alpha=agent_alpha) + + outputs = collections.namedtuple( + "Outputs", + ["training_artifacts_dir"]) + + return outputs(training_artifacts_dir) + + def _serialize_str(str_value: str) -> str: + if not isinstance(str_value, str): + raise TypeError('Value "{}" has type "{}" instead of str.'.format(str(str_value), str(type(str_value)))) + return str_value + + import argparse + _parser = argparse.ArgumentParser(prog='Train reinforcement learning policy', description='Implements off-policy training for a policy on dataset of TFRecord files.') + _parser.add_argument("--training-artifacts-dir", dest="training_artifacts_dir", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--tfrecord-file", dest="tfrecord_file", type=str, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--num-epochs", dest="num_epochs", type=int, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--rank-k", dest="rank_k", type=int, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--num-actions", dest="num_actions", type=int, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--tikhonov-weight", dest="tikhonov_weight", type=float, required=True, default=argparse.SUPPRESS) + _parser.add_argument("--agent-alpha", dest="agent_alpha", type=float, required=True, default=argparse.SUPPRESS) + _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=1) + _parsed_args = vars(_parser.parse_args()) + _output_files = _parsed_args.pop("_output_paths", []) + + _outputs = train_reinforcement_learning_policy(**_parsed_args) + + _output_serializers = [ + _serialize_str, + + ] + + import os + for idx, output_file in enumerate(_output_files): + try: + os.makedirs(os.path.dirname(output_file)) + except OSError: + pass + with open(output_file, 'w') as f: + f.write(_output_serializers[idx](_outputs[idx])) + args: + - --training-artifacts-dir + - {inputValue: training_artifacts_dir} + - --tfrecord-file + - {inputValue: tfrecord_file} + - --num-epochs + - {inputValue: num_epochs} + - --rank-k + - {inputValue: rank_k} + - --num-actions + - {inputValue: num_actions} + - --tikhonov-weight + - {inputValue: tikhonov_weight} + - --agent-alpha + - {inputValue: agent_alpha} + - '----output-paths' + - {outputPath: training_artifacts_dir} diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/test_trainer_component.py b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/test_trainer_component.py index 0fdc9afb2..d8634b83c 100644 --- a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/test_trainer_component.py +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/test_trainer_component.py @@ -88,7 +88,7 @@ class TestTrainerComponent(unittest.TestCase): def test_training_op_execute_all_steps(self): """Tests that training_op executes all steps.""" - trainer_component.training_op( + trainer_component.train_reinforcement_learning_policy( training_artifacts_dir=TRAINING_ARTIFACTS_DIR, tfrecord_file=TFRECORD_FILE, num_epochs=NUM_EPOCHS, @@ -117,7 +117,7 @@ class TestTrainerComponent(unittest.TestCase): def test_given_zero_epochs_training_op_execute_no_training(self): """Tests that training_op executes zero training with zero num_epochs.""" - trainer_component.training_op( + trainer_component.train_reinforcement_learning_policy( training_artifacts_dir=TRAINING_ARTIFACTS_DIR, tfrecord_file=TFRECORD_FILE, num_epochs=0, @@ -131,7 +131,7 @@ class TestTrainerComponent(unittest.TestCase): def test_given_negative_epochs_training_op_execute_no_training(self): """Tests that training_op executes zero training with negative num_epochs. """ - trainer_component.training_op( + trainer_component.train_reinforcement_learning_policy( training_artifacts_dir=TRAINING_ARTIFACTS_DIR, tfrecord_file=TFRECORD_FILE, num_epochs=-1, @@ -145,7 +145,7 @@ class TestTrainerComponent(unittest.TestCase): def test_given_float_epochs_training_op_raise_exception(self): """Tests that training_op raises an exception for float num_epochs.""" with self.assertRaises(TypeError): - trainer_component.training_op( + trainer_component.train_reinforcement_learning_policy( training_artifacts_dir=TRAINING_ARTIFACTS_DIR, tfrecord_file=TFRECORD_FILE, num_epochs=0.5, diff --git a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/trainer_component.py b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/trainer_component.py index cadb60b11..59c86cd1d 100644 --- a/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/trainer_component.py +++ b/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/trainer/trainer_component.py @@ -16,10 +16,8 @@ # Import for the function return value type. from typing import NamedTuple # pylint: disable=unused-import -from kfp.v2 import components - -def training_op( +def train_reinforcement_learning_policy( training_artifacts_dir: str, tfrecord_file: str, num_epochs: int, @@ -28,7 +26,7 @@ def training_op( tikhonov_weight: float, agent_alpha: float ) -> NamedTuple("Outputs", [ - ("training_artifacts_dir", components.OutputPath), + ("training_artifacts_dir", str), ]): """Implements off-policy training for a policy on dataset of TFRecord files. @@ -249,3 +247,17 @@ def training_op( ["training_artifacts_dir"]) return outputs(training_artifacts_dir) + + +if __name__ == "__main__": + from kfp.components import create_component_from_func + + train_reinforcement_learning_policy_op = create_component_from_func( + func=train_reinforcement_learning_policy, + base_image="tensorflow/tensorflow:2.5.0", + output_component_file="component.yaml", + packages_to_install=[ + "tensorflow==2.5.0", + "tf-agents==0.8.0", + ], + )