Add tenant id, relax token format requirements (#75)

grafana · Jul 29, 2024 · e9c811d · e9c811d
1 parent 12e001d
commit e9c811d
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 36 deletions.
diff --git a/docker-compose-jupyter.yml b/docker-compose-jupyter.yml
@@ -8,4 +8,4 @@ services:
     ports:
       - "8888:8888"
     environment:
-      - GF_AI_TRAINING_CREDS=83bcaff6228b39bbe431af5e19fb4368e2a03dd3@ai-training-api:8000
+      - GF_AI_TRAINING_CREDS=83bcaff6228b39bbe431af5e19fb4368e2a03dd3@1337@http://ai-training-api:8000
diff --git a/docker-compose-pytorch.yml b/docker-compose-pytorch.yml
@@ -8,4 +8,4 @@ services:
     ports:
       - "8888:8888"
     environment:
-      - GF_AI_TRAINING_CREDS=83bcaff6228b39bbe431af5e19fb4368e2a03dd3@ai-training-api:8000
+      - GF_AI_TRAINING_CREDS=83bcaff6228b39bbe431af5e19fb4368e2a03dd3@1337@http://ai-training-api:8000
diff --git a/o11y/smoke.ipynb b/o11y/smoke.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 1,
    "id": "07827d16-8a84-4280-ab52-cece5b6dbaf1",
    "metadata": {},
    "outputs": [
@@ -40,17 +40,13 @@
       "    Uninstalling charset-normalizer-3.3.2:\n",
       "      Successfully uninstalled charset-normalizer-3.3.2\n",
       "  Attempting uninstall: certifi\n",
-      "    Found existing installation: certifi 2024.7.4\n",
-      "    Uninstalling certifi-2024.7.4:\n",
-      "      Successfully uninstalled certifi-2024.7.4\n",
+      "    Found existing installation: certifi 2024.6.2\n",
+      "    Uninstalling certifi-2024.6.2:\n",
+      "      Successfully uninstalled certifi-2024.6.2\n",
       "  Attempting uninstall: requests\n",
       "    Found existing installation: requests 2.32.3\n",
       "    Uninstalling requests-2.32.3:\n",
       "      Successfully uninstalled requests-2.32.3\n",
-      "  Attempting uninstall: o11y\n",
-      "    Found existing installation: o11y 0.1.0\n",
-      "    Uninstalling o11y-0.1.0:\n",
-      "      Successfully uninstalled o11y-0.1.0\n",
       "Successfully installed certifi-2024.7.4 charset-normalizer-3.3.2 idna-3.7 o11y-0.1.0 requests-2.32.3 urllib3-2.2.2\n"
      ]
     }
@@ -77,7 +73,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 2,
    "id": "b3a03c39-891a-4524-a258-36f1a7104dcf",
    "metadata": {},
    "outputs": [],
@@ -87,15 +83,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 3,
    "id": "fc49064e-9974-4d7f-9ef1-c08d9fb06087",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "83bcaff6228b39bbe431af5e19fb4368e2a03dd3@ai-training-api:8000\n"
+      "83bcaff6228b39bbe431af5e19fb4368e2a03dd3@1337@http://ai-training-api:8000\n"
      ]
     }
    ],
@@ -109,7 +105,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 4,
    "id": "e6a85341-29d2-4e19-b7e9-429c5c7f6d95",
    "metadata": {
     "scrolled": true
@@ -156,16 +152,28 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 5,
    "id": "2a41c1dd-e3d2-4f0e-95cb-d57a3832d6e0",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "token_in_request\n",
+      "83bcaff6228b39bbe431af5e19fb4368e2a03dd3\n",
+      "url_in_request\n",
+      "http://ai-training-api:8000/api/v1/process/new\n",
       "Finished run with metrics ('train', 'test')\n",
+      "token_in_request\n",
+      "83bcaff6228b39bbe431af5e19fb4368e2a03dd3\n",
+      "url_in_request\n",
+      "http://ai-training-api:8000/api/v1/process/new\n",
       "Finished run with metrics ('train', 'val')\n",
+      "token_in_request\n",
+      "83bcaff6228b39bbe431af5e19fb4368e2a03dd3\n",
+      "url_in_request\n",
+      "http://ai-training-api:8000/api/v1/process/new\n",
       "Finished run with metrics ('train', 'test', 'val')\n"
      ]
     }

diff --git a/o11y/src/o11y/_internal/client.py b/o11y/src/o11y/_internal/client.py
@@ -15,52 +15,43 @@ class Client:
     def __init__(self):
         self.process_uuid = None
         self.user_metadata = None
+        self.url = None
+        self.token = None
+        self.tenant_id = None
         # We are going to assume that the user has set the credentials in the environment
         # There are other flows but it's the easiest one
         login_string = os.environ.get('GF_AI_TRAINING_CREDS')
         self.set_credentials(login_string)
 
     def set_credentials(self, login_string):
-
         if not login_string or type(login_string) != str:
             logger.error("No login string provided, please set GF_AI_TRAINING_CREDS environment variable")
             return False
-        # Count @ characters in the login string, should be 1
-        if login_string.count("@") != 1:
+        # Count @ characters in the login string, should be 2
+        if login_string.count("@") != 2:
             logger.error("Invalid login string format")
             return False
 
-        token, url = login_string.split("@")
-        url = "http://" + url
-        # Check that the token is exactly 40 characters of hex
-        if len(token) != 40 or not all(c in "0123456789abcdef" for c in token):
-            logger.error("Invalid token format")
-            return "Invalid token format"
-        # Disabled until I can figure out how to make this work with service names
-        # This is both a dev env problem and a problem when running in e.g. k8s
-        # # Validate that the url is a url
-        # if not validate_url(url):
-        #     logger.error("Invalid url format")
-        #     return "Invalid url format"
+        token, tenant_id, url = login_string.split("@")
+        if not url.startswith("http://") and not url.startswith("https://"):
+            url = "http://" + url
 
         self.url = url
         self.token = token
+        self.tenant_id = tenant_id
         return True
 
-    # Returns a boolean indicating if the process was successfully registered
-    # POST /api/v1/process/new
-    # Takes a JSON object with a key “user_metadata” containing a dictionary of metadata
-    # Returns a JSON object with a key “process_uuid” containing the UUID of the process
     def register_process(self, data):
         # If the process is currently registered, clear everything from it
         if self.process_uuid:
             self.process_uuid = None
             self.user_metadata = None
 
         headers = {
-            'Authorization': f'Bearer {self.token}',
+            'Authorization': f'Bearer {self.tenant_id}:{self.token}',
             'Content-Type': 'application/json'
         }
+
         response = requests.post(f'{self.url}/api/v1/process/new', headers=headers, data=json.dumps(data))
         if response.status_code != 200:
             logging.error(f'Failed to register with error: {response.text}')
@@ -80,7 +71,7 @@ def update_metadata(self, process_uuid, user_metadata):
             logging.error("No process registered, unable to update metadata")
             return False
         headers = {
-            'Authorization': f'Bearer {self.token}',
+            'Authorization': f'Bearer {self.tenant_id}:{self.token}',
             'Content-Type': 'application/json'
         }
         data = {