UI config api changes

dataproc_jupyter_plugin/controllers/vertex/vertex.py

@@ -0,0 +1,42 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import aiohttp
+import tornado
+from jupyter_server.base.handlers import APIHandler
+
+from dataproc_jupyter_plugin import credentials
+from dataproc_jupyter_plugin.services.vertex import vertex
+
+
+class UIConfigController(APIHandler):
+    @tornado.web.authenticated
+    async def get(self):
+        """Returns available ui config"""
+        try:
+            async with aiohttp.ClientSession() as client_session:
+                client = vertex.Client(
+                    await credentials.get_cached(), self.log, client_session
+                )
+                configs = await client.list_uiconfig()
+                response = []
+                for config in configs:
+                    env = config.dict()
+                    response.append(env)
+                self.finish(json.dumps(response))
+        except Exception as e:
+            self.log.exception(f"Error fetching ui config: {str(e)}")
+            self.finish({"error": str(e)})

dataproc_jupyter_plugin/handlers.py

@@ -38,6 +38,7 @@
     dataproc,
     executor,
 )
+from dataproc_jupyter_plugin.controllers.vertex import vertex
 
 _region_not_set_error = """GCP region not set in gcloud.
 
@@ -209,6 +210,7 @@ def full_path(name):
         "bigQueryPreview": bigquery.PreviewController,
         "bigQueryProjectsList": bigquery.ProjectsController,
         "bigQuerySearch": bigquery.SearchController,
+        "api/vertex/uiConfig": vertex.UIConfigController
     }
     handlers = [(full_path(name), handler) for name, handler in handlersMap.items()]
     web_app.add_handlers(host_pattern, handlers)

dataproc_jupyter_plugin/services/vertex/vertex.py

@@ -0,0 +1,77 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# TODO: modify code
+from typing import List
+
+from dataproc_jupyter_plugin import urls
+from dataproc_jupyter_plugin.commons.constants import (
+    VERTEX_SERVICE_NAME,
+    CONTENT_TYPE,
+)
+
+
+class Client:
+    def __init__(self, credentials, log, client_session):
+        self.log = log
+        if not (
+            ("access_token" in credentials)
+            and ("project_id" in credentials)
+            and ("region_id" in credentials)
+        ):
+            self.log.exception("Missing required credentials")
+            raise ValueError("Missing required credentials")
+        self._access_token = credentials["access_token"]
+        self.project_id = credentials["project_id"]
+        self.region_id = credentials["region_id"]
+        self.client_session = client_session
+
+    def create_headers(self):
+        return {
+            "Content-Type": CONTENT_TYPE,
+            "Authorization": f"Bearer {self._access_token}",
+        }
+
+    async def list_uiconfig(self) -> List[]:
+        try:
+            uiconfig = []
+            # vertex_url = await urls.gcp_service_url(VERTEX_SERVICE_NAME)
+            # api_endpoint = f"{self.region_id}-{vertex_url}/ui/projects/{self.project_id}/locations/{self.region_id}/uiConfig"
+            api_endpoint = "https://us-central1-aiplatform.googleapis.com/ui/projects/bngy-prod-colab/locations/us-central1/uiConfig"
+
+            headers = self.create_headers()
+            async with self.client_session.get(
+                api_endpoint, headers=headers
+            ) as response:
+                if response.status == 200:
+                    resp = await response.json()
+                    if not resp:
+                        return uiconfig
+                    else:
+                        for machineconfig in resp.get("notebookRuntimeConfig").get("machineConfigs"):
+                            ramBytes_in_gb = int(machineconfig.get('ramBytes')) / 1000000000
+                            formatted_config = {
+                                "machineType": f"{machineconfig.get('machineType')} ({machineconfig.get('cpuCount')} CPUs, {ramBytes_in_gb} GB RAM)",
+                                "acceleratorConfigs": machineconfig.get("acceleratorConfigs")
+                            }
+                            uiconfig.append(formatted_config)
+                        return uiconfig
+                else:
+                    self.log.exception("Error listing ui config")
+                    raise Exception(
+                        f"Error getting vertex ui config: {response.reason} {await response.text()}"
+                    )
+        except Exception as e:
+            self.log.exception(f"Error fetching ui config: {str(e)}")
+            return {"Error fetching ui config": str(e)}