Merge branch 'main' into vizier_sdk

notebooks/community/ml_ops/stage2/get_started_vertex_vizier.ipynb

@@ -84,7 +84,8 @@
         "The steps performed include:\n",
         "\n",
         "- Hyperparameter tuning with Random algorithm.\n",
-        "- Hyperparameter tuning with Vizier (Bayesian) algorithm."
+        "- Hyperparameter tuning with Vizier (Bayesian) algorithm.\n",
+        "- Suggesting trials and updating results for Vizier study"
       ]
     },
     {
@@ -187,7 +188,8 @@
         "if IS_WORKBENCH_NOTEBOOK:\n",
         "    USER_FLAG = \"--user\"\n",
         "\n",
-        "! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG -q"
+        "! pip3 install --upgrade $USER_FLAG -q google-cloud-aiplatform \\\n",
+        "                                       google-vizier==0.0.4"
       ]
     },
     {
@@ -329,25 +331,32 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "timestamp"
+        "id": "06571eb4063b"
       },
       "source": [
-        "#### Timestamp\n",
+        "#### UUID\n",
         "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial."
+        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "timestamp"
+        "id": "4e166d927e36"
       },
       "outputs": [],
       "source": [
-        "from datetime import datetime\n",
+        "import random\n",
+        "import string\n",
         "\n",
-        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
+        "\n",
+        "# Generate a uuid of a specifed length(default=8)\n",
+        "def generate_uuid(length: int = 8) -> str:\n",
+        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
+        "\n",
+        "\n",
+        "UUID = generate_uuid()"
       ]
     },
     {
@@ -358,7 +367,7 @@
       "source": [
         "### Authenticate your Google Cloud account\n",
         "\n",
-        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. Skip this step.\n",
+        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. \n",
         "\n",
         "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
         "\n",
@@ -446,7 +455,7 @@
       "outputs": [],
       "source": [
         "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + UUID"
       ]
     },
     {
@@ -509,7 +518,8 @@
       },
       "outputs": [],
       "source": [
-        "import google.cloud.aiplatform as aip"
+        "import google.cloud.aiplatform as aip\n",
+        "from google.cloud.aiplatform.vizier import Study, pyvizier"
       ]
     },
     {
@@ -534,35 +544,6 @@
         "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "aip_constants"
-      },
-      "source": [
-        "#### Vertex AI constants\n",
-        "\n",
-        "Setup up the following constants for Vertex AI:\n",
-        "\n",
-        "- `API_ENDPOINT`: The Vertex AI API service endpoint for `Dataset`, `Model`, `Job`, `Pipeline` and `Endpoint` services.\n",
-        "- `PARENT`: The Vertex AI location root path for `Dataset`, `Model`, `Job`, `Pipeline` and `Endpoint` resources."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "aip_constants"
-      },
-      "outputs": [],
-      "source": [
-        "# API service endpoint\n",
-        "API_ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
-        "\n",
-        "# Vertex location root path for your dataset, model and endpoint resources\n",
-        "PARENT = \"projects/\" + PROJECT_ID + \"/locations/\" + REGION"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -626,7 +607,7 @@
         "if os.getenv(\"IS_TESTING_TF\"):\n",
         "    TF = os.getenv(\"IS_TESTING_TF\")\n",
         "else:\n",
-        "    TF = \"2.1\".replace(\".\", \"-\")\n",
+        "    TF = \"2.5\".replace(\".\", \"-\")\n",
         "\n",
         "if TF[0] == \"2\":\n",
         "    if TRAIN_GPU:\n",
@@ -1031,7 +1012,7 @@
       },
       "outputs": [],
       "source": [
-        "JOB_NAME = \"custom_job_\" + TIMESTAMP\n",
+        "JOB_NAME = \"custom_job_\" + UUID\n",
         "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, JOB_NAME)\n",
         "\n",
         "if not TRAIN_NGPU or TRAIN_NGPU < 2:\n",
@@ -1094,9 +1075,7 @@
       },
       "outputs": [],
       "source": [
-        "job = aip.CustomJob(\n",
-        "    display_name=\"boston_\" + TIMESTAMP, worker_pool_specs=worker_pool_spec\n",
-        ")"
+        "job = aip.CustomJob(display_name=\"boston_\" + UUID, worker_pool_specs=worker_pool_spec)"
       ]
     },
     {
@@ -1128,7 +1107,7 @@
         "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
         "\n",
         "hpt_job = aip.HyperparameterTuningJob(\n",
-        "    display_name=\"boston_\" + TIMESTAMP,\n",
+        "    display_name=\"boston_\" + UUID,\n",
         "    custom_job=job,\n",
         "    metric_spec={\n",
         "        \"val_loss\": \"minimize\",\n",
@@ -1309,7 +1288,7 @@
       "outputs": [],
       "source": [
         "job = aip.CustomJob(\n",
-        "    display_name=\"boston_\" + TIMESTAMP,\n",
+        "    display_name=\"boston_\" + UUID,\n",
         "    worker_pool_specs=worker_pool_spec,\n",
         "    base_output_dir=MODEL_DIR,\n",
         ")"
@@ -1344,7 +1323,7 @@
         "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
         "\n",
         "hpt_job = aip.HyperparameterTuningJob(\n",
-        "    display_name=\"boston_\" + TIMESTAMP,\n",
+        "    display_name=\"boston_\" + UUID,\n",
         "    custom_job=job,\n",
         "    metric_spec={\n",
         "        \"val_loss\": \"minimize\",\n",
@@ -1513,22 +1492,25 @@
         "id": "vizier_client"
       },
       "source": [
-        "### Create Vizier client\n",
+        "### Specify the algorithm used to suggest trial parameters\n",
         "\n",
-        "Create a client side connection to the Vertex AI Vizier service."
+        "First, you create a `StudyConfig`, and specify the algorithm to suggest the next trial.\n",
+        "\n",
+        "    GRID_SEARCH: grid search\n",
+        "    RANDOM_SEARCH: random search\n",
+        "    ALGORIGTHM_UNSPECIFIED: Vizier bayesian algorithm"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "vizier_client"
+        "id": "d7dd26490358"
       },
       "outputs": [],
       "source": [
-        "vizier_client = aip.gapic.VizierServiceClient(\n",
-        "    client_options=dict(api_endpoint=API_ENDPOINT)\n",
-        ")"
+        "problem = pyvizier.StudyConfig()\n",
+        "problem.algorithm = pyvizier.Algorithm.RANDOM_SEARCH"
       ]
     },
     {
@@ -1543,7 +1525,15 @@
         "\n",
         "In the following example, the goal is to maximize y = x^2 with x in the range of \\[-10. 10\\]. This example has only one parameter and uses an easily calculated function to help demonstrate how to use Vizier.\n",
         "\n",
-        "First, you will create the study using the `create_study()` method."
+        "First, you specify the metrics to minimize or maximize in the study as a list to the property `metric_information`. Then you specify the parameters to the study using the `add_XXX_params()` method for the corresponding data type:\n",
+        "\n",
+        "    - add_bool_param\n",
+        "    - add_categorical_param\n",
+        "    - add_discrete_param\n",
+        "    - add_float_param\n",
+        "    - add_int_param\n",
+        "\n",
+        "You create the study using the `create_or_load()` method."
       ]
     },
     {
@@ -1554,28 +1544,19 @@
       },
       "outputs": [],
       "source": [
-        "STUDY_DISPLAY_NAME = \"xpow2\" + TIMESTAMP\n",
+        "STUDY_DISPLAY_NAME = \"xpow2\" + UUID\n",
         "\n",
-        "param_x = {\n",
-        "    \"parameter_id\": \"x\",\n",
-        "    \"double_value_spec\": {\"min_value\": -10.0, \"max_value\": 10.0},\n",
-        "}\n",
+        "problem.metric_information.append(\n",
+        "    pyvizier.MetricInformation(name=\"y\", goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE)\n",
+        ")\n",
         "\n",
-        "metric_y = {\"metric_id\": \"y\", \"goal\": \"MAXIMIZE\"}\n",
+        "params = problem.search_space.select_root()\n",
+        "params.add_float_param(\"x\", -10.0, 10.0, scale_type=pyvizier.ScaleType.LINEAR)\n",
         "\n",
-        "study = {\n",
-        "    \"display_name\": STUDY_DISPLAY_NAME,\n",
-        "    \"study_spec\": {\n",
-        "        \"algorithm\": \"RANDOM_SEARCH\",\n",
-        "        \"parameters\": [param_x],\n",
-        "        \"metrics\": [metric_y],\n",
-        "    },\n",
-        "}\n",
+        "study = Study.create_or_load(display_name=STUDY_DISPLAY_NAME, problem=problem)\n",
         "\n",
-        "study = vizier_client.create_study(parent=PARENT, study=study)\n",
         "STUDY_NAME = study.name\n",
-        "\n",
-        "print(STUDY_NAME)"
+        "print(\"STUDY_NAME: {}\".format(STUDY_NAME))"
       ]
     },
     {
@@ -1586,9 +1567,7 @@
       "source": [
         "### Get Vizier study\n",
         "\n",
-        "You can get a study using the method `get_study()`, with the following key/value pairs:\n",
-        "\n",
-        "- `name`: The name of the study."
+        "You can get a study using the method `list()`."
       ]
     },
     {
@@ -1599,9 +1578,8 @@
       },
       "outputs": [],
       "source": [
-        "study = vizier_client.get_study({\"name\": STUDY_NAME})\n",
-        "\n",
-        "print(study)"
+        "studies = Study.list()\n",
+        "print(studies[0].gca_resource)"
       ]
     },
     {
@@ -1612,11 +1590,9 @@
       "source": [
         "### Get suggested trial\n",
         "\n",
-        "Next, query the Vizier service for a suggested trial(s) using the method `suggest_trials`, with the following key/value pairs:\n",
+        "Next, query the Vizier service for a suggested trial(s) using the method `suggest()`, with the following key/value pairs:\n",
         "\n",
-        "- `parent`: The name of the study.\n",
-        "- `suggestion_count`: The number of trials to suggest.\n",
-        "- `client_id`: blah\n",
+        "- `count`: The number of trials to suggest.\n",
         "\n",
         "This call is a long running operation. The method `result()` from the response object will wait until the call has completed."
       ]
@@ -1625,18 +1601,13 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "vizier_suggest_trial"
+        "id": "11ff2c4562cb"
       },
       "outputs": [],
       "source": [
         "SUGGEST_COUNT = 1\n",
-        "CLIENT_ID = \"1001\"\n",
         "\n",
-        "response = vizier_client.suggest_trials(\n",
-        "    {\"parent\": STUDY_NAME, \"suggestion_count\": SUGGEST_COUNT, \"client_id\": CLIENT_ID}\n",
-        ")\n",
-        "\n",
-        "trials = response.result().trials\n",
+        "trials = study.suggest(count=SUGGEST_COUNT)\n",
         "\n",
         "print(trials)\n",
         "\n",
@@ -1679,12 +1650,10 @@
       "source": [
         "RESULT = 0.01\n",
         "\n",
-        "vizier_client.add_trial_measurement(\n",
-        "    {\n",
-        "        \"trial_name\": TRIAL_ID,\n",
-        "        \"measurement\": {\"metrics\": [{\"metric_id\": \"y\", \"value\": RESULT}]},\n",
-        "    }\n",
-        ")"
+        "measurement = pyvizier.Measurement()\n",
+        "measurement.metrics[\"y\"] = RESULT\n",
+        "\n",
+        "trials[0].add_measurement(measurement)"
       ]
     },
     {
@@ -1695,7 +1664,7 @@
       "source": [
         "### Delete the Vizier study\n",
         "\n",
-        "The method 'delete_study()' will delete the study."
+        "The method 'delete()' will delete the study."
       ]
     },
     {
@@ -1706,7 +1675,7 @@
       },
       "outputs": [],
       "source": [
-        "vizier_client.delete_study({\"name\": STUDY_NAME})"
+        "study.delete()"
       ]
     },
     {