Merge pull request #48 from roboflow/fix/license-server-bug

License Server Bug Fix

Author: paulguerrie

.github/workflows/test.yml

@@ -36,8 +36,8 @@ jobs:
         with:
           push: false
           tags: roboflow/roboflow-inference-server-cpu:test
-          cache-from: type=registry,ref=roboflow/roboflow-inference-server-cpu:test
-          cache-to: type=registry,ref=roboflow/roboflow-inference-server-cpu:test,mode=max
+          cache-from: type=registry,ref=roboflow/roboflow-inference-server-cpu:test-cache
+          cache-to: type=registry,ref=roboflow/roboflow-inference-server-cpu:test-cache,mode=max
           platforms: linux/amd64
           file: ./docker/dockerfiles/Dockerfile.onnx.cpu
           outputs: type=docker

inference/core/models/roboflow.py

@@ -28,7 +28,6 @@
     AWS_SECRET_ACCESS_KEY,
     CORE_MODEL_BUCKET,
     INFER_BUCKET,
-    LICENSE_SERVER,
     MODEL_CACHE_DIR,
     ONNXRUNTIME_EXECUTION_PROVIDERS,
     REQUIRED_ONNX_PROVIDERS,
@@ -43,6 +42,7 @@
 from inference.core.utils.image_utils import load_image
 from inference.core.utils.onnx import get_onnxruntime_execution_providers
 from inference.core.utils.preprocess import prepare
+from inference.core.utils.url_utils import ApiUrl
 
 if AWS_ACCESS_KEY_ID and AWS_ACCESS_KEY_ID:
     import boto3
@@ -337,12 +337,9 @@ def get_model_artifacts(self) -> None:
             else:
                 self.log("Downloading model artifacts from Roboflow API")
                 # AWS Keys are not available so we use the API Key to hit the Roboflow API which returns a signed link for downloading model artifacts
-                self.api_url = f"{API_BASE_URL}/ort/{self.endpoint}?api_key={self.api_key}&device={self.device_id}&nocache=true&dynamic=true"
-                if LICENSE_SERVER:
-                    self.api_url = (
-                        f"http://{LICENSE_SERVER}/proxy?url="
-                        + urllib.parse.quote(self.api_url, safe="~()*!'")
-                    )
+                self.api_url = ApiUrl(
+                    f"{API_BASE_URL}/ort/{self.endpoint}?api_key={self.api_key}&device={self.device_id}&nocache=true&dynamic=true"
+                )
                 api_data = get_api_data(self.api_url)
                 if "ort" not in api_data.keys():
                     raise TensorrtRoboflowAPIError(
@@ -359,24 +356,8 @@ def get_model_artifacts(self) -> None:
                 if "colors" in api_data:
                     self.colors = api_data["colors"]
 
-                if LICENSE_SERVER:
-                    license_server_base_url = f"http://{LICENSE_SERVER}/proxy?url="
-                    weights_url = license_server_base_url + urllib.parse.quote(
-                        api_data["model"], safe="~()*!'"
-                    )
-
-                    def get_env_url(api_data):
-                        return license_server_base_url + urllib.parse.quote(
-                            api_data["environment"], safe="~()*!'"
-                        )
-
-                else:
-                    weights_url = api_data["model"]
-
-                    def get_env_url(api_data):
-                        return api_data["environment"]
-
                 t1 = perf_counter()
+                weights_url = ApiUrl(api_data["model"])
                 r = requests.get(weights_url)
                 with self.open_cache(self.weights_file, "wb") as f:
                     f.write(r.content)
@@ -385,7 +366,7 @@ def get_env_url(api_data):
                         "Weights download took longer than 120 seconds, refreshing API request"
                     )
                     api_data = get_api_data(self.api_url)
-                env_url = get_env_url(api_data)
+                env_url = ApiUrl(api_data["environment"])
                 self.environment = requests.get(env_url).json()
                 with open(self.cache_file("environment.json"), "w") as f:
                     json.dump(self.environment, f)
@@ -625,12 +606,9 @@ def download_weights(self) -> None:
                         raise Exception(f"Failed to download model artifacts.")
             else:
                 # AWS Keys are not available so we use the API Key to hit the Roboflow API which returns a signed link for downloading model artifacts
-                self.api_url = f"{API_BASE_URL}/core_model/{self.endpoint}?api_key={self.api_key}&device={self.device_id}&nocache=true"
-                if LICENSE_SERVER:
-                    self.api_url = (
-                        f"http://{LICENSE_SERVER}/proxy?url="
-                        + urllib.parse.quote(self.api_url, safe="~()*!'")
-                    )
+                self.api_url = ApiUrl(
+                    f"{API_BASE_URL}/core_model/{self.endpoint}?api_key={self.api_key}&device={self.device_id}&nocache=true"
+                )
                 api_data = get_api_data(self.api_url)
                 if "weights" not in api_data.keys():
                     raise TensorrtRoboflowAPIError(
@@ -640,13 +618,7 @@ def download_weights(self) -> None:
                 weights_url_keys = api_data["weights"].keys()
 
                 for weights_url_key in weights_url_keys:
-                    if LICENSE_SERVER:
-                        license_server_base_url = f"http://{LICENSE_SERVER}/proxy?url="
-                        weights_url = license_server_base_url + urllib.parse.quote(
-                            api_data["weights"][weights_url_key], safe="~()*!'"
-                        )
-                    else:
-                        weights_url = api_data["weights"][weights_url_key]
+                    weights_url = ApiUrl(api_data["weights"][weights_url_key])
                     t1 = perf_counter()
                     attempts = 0
                     success = False

inference/core/registries/roboflow.py

@@ -7,6 +7,7 @@
 from inference.core.exceptions import DatasetLoadError, WorkspaceLoadError
 from inference.core.models.base import Model
 from inference.core.registries.base import ModelRegistry
+from inference.core.utils.url_utils import ApiUrl
 
 MODEL_TYPE_DEFAULTS = {
     "object-detection": "yolov5v2s",
@@ -72,7 +73,8 @@ def get_model_type(model_id: str, api_key: str) -> str:
             model_type = cache_data["model_type"]
         return project_task_type, model_type
 
-    api_key_info = requests.get("/".join([API_BASE_URL, f"?api_key={api_key}"]))
+    api_url = ApiUrl("/".join([API_BASE_URL, f"?api_key={api_key}"]))
+    api_key_info = requests.get(api_url)
     try:
         api_key_info.raise_for_status()
     except requests.exceptions.HTTPError as e:
@@ -84,12 +86,13 @@ def get_model_type(model_id: str, api_key: str) -> str:
     if workspace_id is None:
         raise WorkspaceLoadError(f"Empty workspace, check your API key")
 
-    dataset_info = requests.get(
+    api_url = ApiUrl(
         "/".join(
             [API_BASE_URL, workspace_id, dataset_id, f"?api_key={api_key}&nocache=true"]
         )
     )
-    version_info = requests.get(
+    dataset_info = requests.get(api_url)
+    api_url = ApiUrl(
         "/".join(
             [
                 API_BASE_URL,
@@ -100,6 +103,7 @@ def get_model_type(model_id: str, api_key: str) -> str:
             ]
         )
     )
+    version_info = requests.get(api_url)
     try:
         dataset_info.raise_for_status()
         version_info.raise_for_status()

inference/core/utils/url_utils.py

@@ -0,0 +1,21 @@
+import urllib
+
+from inference.core.env import LICENSE_SERVER
+
+
+def ProxyUrl(url):
+    """Returns a proxied URL if according to LICENSE_SERVER settings"""
+    return f"http://{LICENSE_SERVER}/proxy?url=" + urllib.parse.quote(
+        url, safe="~()*!'"
+    )
+
+
+def RawUrl(url):
+    """Returns a raw URL"""
+    return url
+
+
+if LICENSE_SERVER:
+    ApiUrl = ProxyUrl
+else:
+    ApiUrl = RawUrl

inference/core/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.8.4"
+__version__ = "0.8.5"
 
 if __name__ == "__main__":
     print(__version__)