Merge pull request #1183 from roboflow/feat/remove-uuid-validator-from-request-correlation-middleware-for-gcp-and-dd

Feat/remove UUID validator from request correlation middleware for gcp and dd

Author: grzegorz-roboflow

inference/core/env.py

@@ -231,6 +231,9 @@
 # Whether is's GCP serverless service
 GCP_SERVERLESS = str2bool(os.getenv("GCP_SERVERLESS", "False"))
 
+# Header where correlaction ID for logging is stored
+CORRELACTION_ID_HEADER = os.getenv("CORRELACTION_ID_HEADER", "X-Request-ID")
+
 # Flag to enable legacy route, default is True
 LEGACY_ROUTE_ENABLED = str2bool(os.getenv("LEGACY_ROUTE_ENABLED", True))
 

inference/core/interfaces/http/http_api.py

@@ -5,6 +5,7 @@
 from functools import partial, wraps
 from time import sleep
 from typing import Any, Dict, List, Optional, Union
+from uuid import uuid4
 
 import asgi_correlation_id
 import uvicorn
@@ -108,6 +109,8 @@
     CORE_MODEL_TROCR_ENABLED,
     CORE_MODEL_YOLO_WORLD_ENABLED,
     CORE_MODELS_ENABLED,
+    CORRELACTION_ID_HEADER,
+    DEDICATED_DEPLOYMENT_ID,
     DEDICATED_DEPLOYMENT_WORKSPACE_URL,
     DISABLE_WORKFLOW_ENDPOINTS,
     DOCKER_SOCKET_PATH,
@@ -630,7 +633,17 @@ async def on_shutdown():
                 strip_dirs=False,
                 sort_by="cumulative",
             )
-        app.add_middleware(asgi_correlation_id.CorrelationIdMiddleware)
+        if DEDICATED_DEPLOYMENT_ID or GCP_SERVERLESS:
+            app.add_middleware(
+                asgi_correlation_id.CorrelationIdMiddleware,
+                header_name=CORRELACTION_ID_HEADER,
+                update_request_header=True,
+                generator=lambda: uuid4().hex,
+                validator=lambda a: True,
+                transformer=lambda a: a,
+            )
+        else:
+            app.add_middleware(asgi_correlation_id.CorrelationIdMiddleware)
 
         if METRICS_ENABLED: