Use a single httpx client for each worker thread (#862)

Author: HansKallekleiv

backend_py/primary/poetry.lock

@@ -0,0 +1,49 @@
+from typing import Optional
+import logging
+
+import httpx
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class HTTPXAsyncClientWrapper:
+    """Global async client wrapper for HTTPX."""
+
+    _instance: Optional["HTTPXAsyncClientWrapper"] = None
+    _async_client: Optional[httpx.AsyncClient] = None
+
+    def __new__(cls) -> "HTTPXAsyncClientWrapper":
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    @property
+    def client(self) -> httpx.AsyncClient:
+        """Get the async client instance."""
+        if self._async_client is None:
+            raise RuntimeError("HTTPXAsyncClientWrapper not initialized. Call start() first.")
+        return self._async_client
+
+    def start(self) -> None:
+        """Instantiate the client. Call from the FastAPI startup hook."""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient()
+            LOGGER.info(f"httpx AsyncClient instantiated. Id {id(self._async_client)}")
+
+    async def stop(self) -> None:
+        """Gracefully shutdown. Call from FastAPI shutdown hook."""
+        if self._async_client is not None:
+            LOGGER.info(
+                f"httpx async_client.is_closed: {self._async_client.is_closed}. " f"Id: {id(self._async_client)}"
+            )
+            await self._async_client.aclose()
+            LOGGER.info(
+                f"httpx async_client.is_closed: {self._async_client.is_closed}. " f"Id: {id(self._async_client)}"
+            )
+            self._async_client = None
+            LOGGER.info("httpx AsyncClient closed")
+
+
+# Create a singleton instance of the async client
+httpx_async_client = HTTPXAsyncClientWrapper()

backend_py/primary/primary/httpx_client.py

@@ -38,18 +38,21 @@
 from primary.utils.logging_setup import ensure_console_log_handler_is_configured, setup_normal_log_levels
 
 from . import config
-
+from .httpx_client import httpx_async_client
 
 ensure_console_log_handler_is_configured()
 setup_normal_log_levels()
 
 # temporarily set some loggers to DEBUG
 # logging.getLogger().setLevel(logging.DEBUG)
 logging.getLogger("primary.services.sumo_access").setLevel(logging.DEBUG)
-logging.getLogger("primary.services.user_session_manager").setLevel(logging.DEBUG)
+logging.getLogger("primary.services.smda_access").setLevel(logging.DEBUG)
+logging.getLogger("primary.services.ssdl_access").setLevel(logging.DEBUG)
 logging.getLogger("primary.services.user_grid3d_service").setLevel(logging.DEBUG)
 logging.getLogger("primary.routers.grid3d").setLevel(logging.DEBUG)
 logging.getLogger("primary.routers.dev").setLevel(logging.DEBUG)
+# logging.getLogger("uvicorn.error").setLevel(logging.DEBUG)
+# logging.getLogger("uvicorn.access").setLevel(logging.DEBUG)
 
 LOGGER = logging.getLogger(__name__)
 
@@ -71,6 +74,17 @@ def custom_generate_unique_id(route: APIRoute) -> str:
     LOGGER.warning("Skipping telemetry configuration, APPLICATIONINSIGHTS_CONNECTION_STRING env variable not set.")
 
 
+# Start the httpx client on startup and stop it on shutdown of the app
+@app.on_event("startup")
+async def startup_event() -> None:
+    httpx_async_client.start()
+
+
+@app.on_event("shutdown")
+async def shutdown_event() -> None:
+    await httpx_async_client.stop()
+
+
 # The tags we add here will determine the name of the frontend api service for our endpoints as well as
 # providing some grouping when viewing the openapi documentation.
 app.include_router(explore_router, tags=["explore"])

backend_py/primary/primary/main.py

@@ -2,9 +2,10 @@
 from typing import Mapping
 from urllib.parse import urljoin
 
-# Using the same http client as sumo
 import httpx
 
+from primary.httpx_client import httpx_async_client
+
 
 class GraphApiAccess:
     def __init__(self, access_token: str):
@@ -15,12 +16,12 @@ def _make_headers(self) -> Mapping[str, str]:
         return {"Authorization": f"Bearer {self._access_token}"}
 
     async def _request(self, url: str) -> httpx.Response:
-        async with httpx.AsyncClient() as client:
-            response = await client.get(
-                url,
-                headers=self._make_headers(),
-            )
-            return response
+
+        response = await httpx_async_client.client.get(
+            url,
+            headers=self._make_headers(),
+        )
+        return response
 
     async def get_user_profile_photo(self, user_id: str) -> str | None:
         request_url = urljoin(self.base_url, "me/photo/$value" if user_id == "me" else f"users/{user_id}/photo/$value")

backend_py/primary/primary/services/graph_access/graph_access.py

@@ -1,11 +1,10 @@
 import logging
 from typing import List
 
-import httpx
-
 from webviz_pkg.core_utils.perf_timer import PerfTimer
 
 from primary import config
+from primary.httpx_client import httpx_async_client
 from primary.services.service_exceptions import ServiceRequestError, Service
 
 LOGGER = logging.getLogger(__name__)
@@ -28,35 +27,34 @@ async def smda_get_request(access_token: str, endpoint: str, params: dict) -> Li
     timer = PerfTimer()
     single_lap_timer = PerfTimer()
 
-    async with httpx.AsyncClient() as client:
-        results: List[dict] = []
-        page: int = 1
-        while True:
-            response = await client.get(urlstring, params=params, headers=headers, timeout=60)
-            LOGGER.info(f"TIME SMDA fetch '{endpoint}', page {page}, took {single_lap_timer.lap_s():.2f} seconds")
-            page += 1
-            if response.status_code == 200:
-                result = response.json()["data"]["results"]
-                if not result:
-                    raise ServiceRequestError(f"No data found for endpoint: '{endpoint}'", Service.SMDA)
-
-                results.extend(result)
-
-                next_request = response.json()["data"]["next"]
-                if next_request is None:
-                    break
-                params["_next"] = next_request
-            elif response.status_code == 404:
-                LOGGER.error(f"{str(response.status_code) } {endpoint} either does not exists or can not be found")
-                raise ServiceRequestError(
-                    f"[{str(response.status_code)}] '{endpoint}' either does not exists or can not be found",
-                    Service.SMDA,
-                )
-            else:
-                raise ServiceRequestError(
-                    f"[{str(response.status_code)}] Cannot fetch data from endpoint: '{endpoint}'", Service.SMDA
-                )
-
-        LOGGER.info(f"TIME SMDA fetch '{endpoint}' took {timer.lap_s():.2f} seconds")
+    results: List[dict] = []
+    page: int = 1
+    while True:
+        response = await httpx_async_client.client.get(urlstring, params=params, headers=headers, timeout=60)
+        LOGGER.debug(f"TIME SMDA fetch '{endpoint}', page {page}, took {single_lap_timer.lap_s():.2f} seconds")
+        page += 1
+        if response.status_code == 200:
+            result = response.json()["data"]["results"]
+            if not result:
+                raise ServiceRequestError(f"No data found for endpoint: '{endpoint}'", Service.SMDA)
+
+            results.extend(result)
+
+            next_request = response.json()["data"]["next"]
+            if next_request is None:
+                break
+            params["_next"] = next_request
+        elif response.status_code == 404:
+            LOGGER.error(f"{str(response.status_code) } {endpoint} either does not exists or can not be found")
+            raise ServiceRequestError(
+                f"[{str(response.status_code)}] '{endpoint}' either does not exists or can not be found",
+                Service.SMDA,
+            )
+        else:
+            raise ServiceRequestError(
+                f"[{str(response.status_code)}] Cannot fetch data from endpoint: '{endpoint}'", Service.SMDA
+            )
+
+    LOGGER.debug(f"TIME SMDA fetch '{endpoint}' took {timer.lap_s():.2f} seconds")
 
     return results

backend_py/primary/primary/services/smda_access/_smda_get_request.py

@@ -1,11 +1,10 @@
 import logging
 from typing import List, Optional
 
-import httpx
-
 from webviz_pkg.core_utils.perf_timer import PerfTimer
 
 from primary import config
+from primary.httpx_client import httpx_async_client
 from primary.services.service_exceptions import (
     Service,
     InvalidDataError,
@@ -30,23 +29,22 @@ async def ssdl_get_request(access_token: str, endpoint: str, params: Optional[di
     }
     timer = PerfTimer()
 
-    async with httpx.AsyncClient() as client:
-        response = await client.get(urlstring, params=params, headers=headers, timeout=60)
-        results = []
-        if response.status_code == 200:
-            results = response.json()
-
-        elif response.status_code == 401:
-            raise AuthorizationError("Unauthorized access to SSDL", Service.SSDL)
-        elif response.status_code == 403:
-            raise AuthorizationError("Forbidden access to SSDL", Service.SSDL)
-        elif response.status_code == 404:
-            raise InvalidDataError(f"Endpoint {endpoint} either does not exists or can not be found", Service.SSDL)
-
-        # Capture other errors
-        else:
-            raise InvalidParameterError(f"Can not fetch data from endpoint {endpoint}", Service.SSDL)
-
-        print(f"TIME SSDL fetch {endpoint} took {timer.lap_s():.2f} seconds")
-        LOGGER.debug(f"TIME SSDL fetch {endpoint} took {timer.lap_s():.2f} seconds")
+    response = await httpx_async_client.client.get(urlstring, params=params, headers=headers, timeout=60)
+    results = []
+    if response.status_code == 200:
+        results = response.json()
+
+    elif response.status_code == 401:
+        raise AuthorizationError("Unauthorized access to SSDL", Service.SSDL)
+    elif response.status_code == 403:
+        raise AuthorizationError("Forbidden access to SSDL", Service.SSDL)
+    elif response.status_code == 404:
+        raise InvalidDataError(f"Endpoint {endpoint} either does not exists or can not be found", Service.SSDL)
+
+    # Capture other errors
+    else:
+        raise InvalidParameterError(f"Can not fetch data from endpoint {endpoint}", Service.SSDL)
+
+    print(f"TIME SSDL fetch {endpoint} took {timer.lap_s():.2f} seconds")
+    LOGGER.debug(f"TIME SSDL fetch {endpoint} took {timer.lap_s():.2f} seconds")
     return results

backend_py/primary/primary/services/ssdl_access/_ssdl_get_request.py

@@ -1,20 +1,51 @@
 import logging
+from typing import Any
 
 from fmu.sumo.explorer.explorer import SumoClient, Pit
 from fmu.sumo.explorer.objects import CaseCollection, Case
 from webviz_pkg.core_utils.perf_timer import PerfTimer
+from webviz_pkg.core_utils.perf_metrics import PerfMetrics
 
 from primary import config
+from primary.httpx_client import httpx_async_client
 from primary.services.service_exceptions import Service, NoDataError, MultipleDataMatchesError
 
 LOGGER = logging.getLogger(__name__)
 
 
+class SynchronousMethodCallError(Exception):
+    """Custom error for when synchronous methods are called instead of async."""
+
+
+class FakeHTTPXClient:
+    """A fake HTTPX client to ensure we use async methods instead of sync ones.
+    This is needed as we do not want to allow any synchronous HTTP calls in the primary service.
+    Ideally this should be handled by the SumoClient. https://github.com/equinor/fmu-sumo/issues/369"""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:  # pylint: disable=unused-argument
+        self._error_msg = "🚫 Do not use a synchronous http class!. Use the async http class instead. "
+
+    def __getattr__(self, name: str) -> None:
+        """Catch any synchronous method calls and raise a helpful error."""
+        async_methods = {"get", "post", "put", "patch", "delete", "head", "options"}
+        if name in async_methods:
+            raise SynchronousMethodCallError(self._error_msg.format(method=name))
+        raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'")
+
+
 def create_sumo_client(access_token: str) -> SumoClient:
+    timer = PerfMetrics()
     if access_token == "DUMMY_TOKEN_FOR_TESTING":  # nosec bandit B105
         sumo_client = SumoClient(env=config.SUMO_ENV, interactive=False)
     else:
-        sumo_client = SumoClient(env=config.SUMO_ENV, token=access_token, interactive=False)
+        sumo_client = SumoClient(
+            env=config.SUMO_ENV,
+            token=access_token,
+            http_client=FakeHTTPXClient(),
+            async_http_client=httpx_async_client.client,
+        )
+    timer.record_lap("create_sumo_client()")
+    LOGGER.debug(f"{timer.to_string()}ms")
     return sumo_client
 
 

backend_py/primary/primary/services/sumo_access/_helpers.py

@@ -40,9 +40,7 @@ async def get_case_name_async(self) -> str:
     async def get_iterations_async(self) -> list[IterationInfo]:
         case: Case = await self._get_or_create_case_obj()
 
-        # Stick with the sync version for now, since there is a bug in the async version of SumoExplorer
-        # See: https://github.com/equinor/fmu-sumo/issues/326
-        iterations = case.iterations
+        iterations = await case.iterations_async
 
         iter_info_arr: list[IterationInfo] = []
         for iteration in iterations:

backend_py/primary/primary/services/sumo_access/case_inspector.py

@@ -43,7 +43,7 @@ async def get_group_tree_table(self, realization: Optional[int]) -> Optional[pd.
             if await table_collection.length_async() > 1:
                 raise ValueError("Multiple tables found.")
 
-            group_tree_df = table_collection[0].to_pandas()
+            group_tree_df = await table_collection[0].to_pandas_async()
 
             _validate_group_tree_df(group_tree_df)
 

backend_py/primary/primary/services/sumo_access/group_tree_access.py

@@ -53,9 +53,10 @@ async def get_realization_table_async(
             iteration=self._iteration_name,
             realization=realization,
         )
-        if not table_collection:
+        table_length = await table_collection.length_async()
+        if table_length == 0:
             raise ValueError(f"No table found for {table_schema=}")
-        if len(table_collection) > 1:
+        if table_length > 1:
             raise ValueError(f"Multiple tables found for {table_schema=}")
 
         sumo_table = await table_collection.getitem_async(0)