Fix endpoint issues in pyglove colab. Should resolve #1044

PiperOrigin-RevId: 603124975

Author: xingyousong

docs/advanced_topics/pyglove/vizier_as_backend.ipynb

@@ -49,7 +49,7 @@
         "\n",
         "import pyglove as pg\n",
         "from vizier import pyglove as pg_vizier\n",
-        "from vizier.service import vizier_server"
+        "from vizier.service import servers"
       ]
     },
     {
@@ -136,7 +136,7 @@
         "id": "zIbLEGi6prpm"
       },
       "source": [
-        "Alternatively, if using a remote server, the endpoint can be specified as well:"
+        "**Alternatively**, if using a remote server, the endpoint can be specified as well:"
       ]
     },
     {
@@ -147,7 +147,7 @@
       },
       "outputs": [],
       "source": [
-        "server = vizier_server.DefaultVizierServer(host=hostname)  # Normally hosted on a remote machine.\n",
+        "server = servers.DefaultVizierServer()  # Normally hosted on a remote machine.\n",
         "pg_vizier.init('my_study', vizier_endpoint=server.endpoint)"
       ]
     },
@@ -169,15 +169,15 @@
       },
       "outputs": [],
       "source": [
-        "num_workers = 10\n",
+        "NUM_WORKERS = 10\n",
         "\n",
         "\n",
         "def work_fn(worker_id):\n",
         "  print(f\"Worker ID: {worker_id}\")\n",
         "  for value, feedback in pg.sample(\n",
         "      search_space,\n",
         "      algorithm=algorithm,\n",
-        "      num_examples=num_trials // num_workers,\n",
+        "      num_examples=num_trials // NUM_WORKERS,\n",
         "      name=\"worker_run\",\n",
         "  ):\n",
         "    reward = evaluator(value)\n",
@@ -217,7 +217,7 @@
       "outputs": [],
       "source": [
         "with multiprocessing.pool.ThreadPool(num_workers) as pool:\n",
-        "  pool.map(work_fn, range(num_workers))"
+        "  pool.map(work_fn, range(NUM_WORKERS))"
       ]
     },
     {
@@ -238,7 +238,7 @@
       "outputs": [],
       "source": [
         "processes = []\n",
-        "for worker_id in range(num_workers):\n",
+        "for worker_id in range(NUM_WORKERS):\n",
         "  p = multiprocessing.Process(target=work_fn, args=(worker_id,))\n",
         "  p.start()\n",
         "  processes.append(p)\n",
@@ -265,7 +265,7 @@
       "outputs": [],
       "source": [
         "# Server Machine\n",
-        "server = vizier_server.DefaultVizierServer(host=hostname)"
+        "server = servers.DefaultVizierServer()"
       ]
     },
     {