Switch from requests to aiohttp

I would like us to eventually move to Tornado (from fastapi)
to be consistent with what we use in the JupyterHub ecosystem.
One part of that is to make all network stuff async, and this
moves us from `requests` to `aiohttp`, making everything async
as needed as we go.

This isn't complete, a WIP.

Author: yuvipanda

pyproject.toml

@@ -31,3 +31,6 @@ escapism = { git = "https://github.com/jupyterhub/escapism", tag = "1.0.1" }
 [pytest]
 log_cli = true
 log_cli_level = "INFO"
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"

src/jupyterhub_cost_monitoring/app.py

@@ -1,5 +1,6 @@
 from datetime import timedelta
 
+import aiohttp
 import requests
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import Response
@@ -27,6 +28,7 @@
 app = FastAPI()
 app.add_middleware(MetricsMiddleware)
 logger = get_logger(__name__)
+client_session = aiohttp.ClientSession()
 
 
 @app.get("/")
@@ -256,7 +258,7 @@ def total_costs_per_group(
 
 
 @app.get("/costs-per-user")
-def costs_per_user(
+async def costs_per_user(
     from_date: str | None = Query(
         None, alias="from", description="Start date in YYYY-MM-DDTHH:MMZ format"
     ),
@@ -315,22 +317,16 @@ def costs_per_user(
     # Get per-user costs by combining AWS costs with Prometheus usage data
     results = []
     for ug in usergroup:
-        try:
-            per_user_costs = query_total_costs_per_user(
-                date_range, hub, component, user, ug, limit
-            )
-        except requests.exceptions.HTTPError as e:
-            response = e.response
-            raise HTTPException(status_code=response.status_code, detail=response.text)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"{e}")
+        per_user_costs = await query_total_costs_per_user(
+            client_session, date_range, hub, component, user, ug, limit
+        )
         results.extend(per_user_costs)
 
     return results
 
 
 @app.get("/total-usage")
-def total_usage(
+async def total_usage(
     from_date: str | None = Query(
         None, alias="from", description="Start date in YYYY-MM-DDTHH:MMZ format"
     ),
@@ -359,7 +355,7 @@ def total_usage(
         user = None
 
     try:
-        return query_usage(date_range, hub, component, user)
+        return await query_usage(client_session, date_range, hub, component, user)
     except requests.exceptions.HTTPError as e:
         response = e.response
         raise HTTPException(status_code=response.status_code, detail=response.text)

src/jupyterhub_cost_monitoring/query_cost_aws.py

@@ -6,8 +6,8 @@
 import functools
 from pprint import pformat
 
+import aiohttp
 import boto3
-import requests
 
 from .cache import ttl_lru_cache
 from .const_cost_aws import (
@@ -514,7 +514,8 @@ def query_total_costs_per_component(
 
 
 @ttl_lru_cache(seconds_to_live=3600)
-def query_total_costs_per_user(
+async def query_total_costs_per_user(
+    client: aiohttp.ClientSession,
     date_range: DateRange,
     hub: str = None,
     component: str = None,
@@ -556,15 +557,13 @@ def query_total_costs_per_user(
     # Get user usage percentages from Prometheus using the same DateRange object
     # This ensures we query the same logical date range for both AWS and Prometheus,
     # accounting for their different date range semantics (exclusive vs inclusive)
-    try:
-        usage_shares = query_usage(
-            date_range,
-            hub_name=hub,
-            component_name=component,
-            user_name=user,
-        )
-    except requests.exceptions.ConnectionError:
-        raise
+    usage_shares = await query_usage(
+        client,
+        date_range,
+        hub_name=hub,
+        component_name=component,
+        user_name=user,
+    )
     results = []
     for entry in usage_shares:
         d = entry["date"]
@@ -577,7 +576,7 @@ def query_total_costs_per_user(
             )  # Adjust usage share to cost
             results.append(entry)
     results = [x for x in results if x["hub"] != "binder"]  # Exclude binder hubs
-    user_groups = query_user_groups(date_range, hub, user)
+    user_groups = await query_user_groups(client, date_range, hub, user)
     seen = set()
     list_groups = []
     # Ensure uniquely keyed entries when double-counting group costs
@@ -630,7 +629,8 @@ def query_total_costs_per_user(
 
 
 @ttl_lru_cache(seconds_to_live=3600)
-def query_total_costs_per_group(
+async def query_total_costs_per_group(
+    client: aiohttp.ClientSession,
     date_range: DateRange,
 ):
     """
@@ -642,11 +642,7 @@ def query_total_costs_per_group(
     Returns:
         List of dicts with keys: date, usergroup and cost.
     """
-    try:
-        results = query_total_costs_per_user(date_range=date_range)
-    except Exception as e:
-        logger.exception(f"HTTP request failed: {e}")
-        raise
+    results = await query_total_costs_per_user(client, date_range=date_range)
     response = {}
     for r in results:
         key = (r["date"], r["usergroup"])

src/jupyterhub_cost_monitoring/query_usage.py

@@ -6,8 +6,8 @@
 from collections import defaultdict
 from datetime import datetime, timedelta, timezone
 
+import aiohttp
 import escapism
-import requests
 from yarl import URL
 
 from .cache import ttl_lru_cache
@@ -23,7 +23,9 @@
 prometheus_password = os.environ.get("PROMETHEUS_PASSWORD", "")
 
 
-def query_prometheus(query: str, date_range: DateRange, step: str) -> requests.Response:
+async def query_prometheus(
+    client: aiohttp.ClientSession, query: str, date_range: DateRange, step: str
+) -> dict:
     """
     Query the Prometheus server with the given query over a date range.
 
@@ -42,9 +44,7 @@ def query_prometheus(query: str, date_range: DateRange, step: str) -> requests.R
         scheme="http", host=prometheus_host, port=prometheus_port
     )
     if prometheus_username != "" and prometheus_password != "":
-        prometheus_auth = requests.auth.HTTPBasicAuth(
-            prometheus_username, prometheus_password
-        )
+        prometheus_auth = aiohttp.BasicAuth(prometheus_username, prometheus_password)
     else:
         prometheus_auth = None
     parameters = {
@@ -53,15 +53,16 @@ def query_prometheus(query: str, date_range: DateRange, step: str) -> requests.R
         "end": to_date,
         "step": step,
     }
-    query_api = URL(prometheus_api.with_path("/api/v1/query_range"))
-    with requests.get(query_api, params=parameters, auth=prometheus_auth) as response:
+    query_api = prometheus_api.with_path("/api/v1/query_range").with_query(parameters)
+    async with client.get(query_api, auth=prometheus_auth) as response:
         logger.info(f"Querying Prometheus: {response.url}")
         response.raise_for_status()
-        result = response.json()
+        result = await response.json()
         return result
 
 
-def query_usage(
+async def query_usage(
+    client: aiohttp.ClientSession,
     date_range: DateRange,
     hub_name: str | None,
     component_name: str | None,
@@ -84,23 +85,18 @@ def query_usage(
     if component_name is None:
         # Query all components defined in USAGE_MAP
         for component, params in USAGE_MAP.items():
-            try:
-                response = query_prometheus(
-                    params["query"], date_range, step=params["step"]
-                )
-            except requests.exceptions.RequestException:
-                raise
+            response = await query_prometheus(
+                client, params["query"], date_range, step=params["step"]
+            )
             result.extend(_process_response(response, component))
     else:
         # Query specific component only
-        try:
-            response = query_prometheus(
-                USAGE_MAP[component_name]["query"],
-                date_range,
-                step=USAGE_MAP[component_name]["step"],
-            )
-        except requests.exceptions.RequestException:
-            raise
+        response = await query_prometheus(
+            client,
+            USAGE_MAP[component_name]["query"],
+            date_range,
+            step=USAGE_MAP[component_name]["step"],
+        )
         result.extend(_process_response(response, component_name))
     # Calculate daily cost factors from absolute usage totals)
     result = _calculate_daily_cost_factors(result, hub_name=hub_name)
@@ -111,9 +107,9 @@ def query_usage(
 
 
 def _process_response(
-    response: requests.Response,
+    response: dict,
     component_name: str,
-) -> dict:
+) -> list[dict]:
     """
     Process the response from the Prometheus server to extract absolute usage data.
 
@@ -256,7 +252,8 @@ def _calculate_daily_cost_factors(
 
 
 @ttl_lru_cache(seconds_to_live=3600)
-def query_user_groups(
+async def query_user_groups(
+    client: aiohttp.ClientSession,
     hub_name: str | None = None,
     user_name: str | None = None,
     group_name: str | None = None,
@@ -266,17 +263,13 @@ def query_user_groups(
     """
     now_date = get_now_date() - timedelta(days=1)
     date_range = DateRange(start_date=now_date, end_date=now_date)
-    try:
-        response = query_prometheus(USER_GROUP_INFO, date_range, step="1d")
-    except requests.exceptions.RequestException as e:
-        logger.exception(f"HTTP request failed: {e}")
-        raise
+    response = await query_prometheus(client, USER_GROUP_INFO, date_range, step="1d")
     result = _process_user_groups(response, hub_name, user_name, group_name)
     return result
 
 
 def _process_user_groups(
-    response: requests.Response,
+    response: dict,
     hub_name: str | None = None,
     user_name: str | None = None,
     group_name: str | None = None,
@@ -306,16 +299,13 @@ def _process_user_groups(
 
 
 @ttl_lru_cache(seconds_to_live=3600)
-def query_users_with_multiple_groups(
+async def query_users_with_multiple_groups(
+    client: aiohttp.ClientSession,
     date_range: DateRange,
     hub_name: str | None = None,
     user_name: str | None = None,
 ) -> list[dict]:
-    try:
-        response = query_user_groups(hub_name=hub_name, user_name=user_name)
-    except requests.exceptions.RequestException as e:
-        logger.exception(f"HTTP request failed: {e}")
-        raise
+    response = await query_user_groups(client, hub_name=hub_name, user_name=user_name)
     grouped = defaultdict(
         lambda: {"username": None, "hub": None, "usergroups": [], "has_multiple": False}
     )
@@ -340,16 +330,13 @@ def query_users_with_multiple_groups(
 
 
 @ttl_lru_cache(seconds_to_live=3600)
-def query_users_with_no_groups(
+async def query_users_with_no_groups(
+    client: aiohttp.ClientSession,
     date_range: DateRange,
     hub_name: str | None = None,
     user_name: str | None = None,
 ) -> list[dict]:
-    try:
-        response = query_user_groups(hub_name=hub_name, user_name=user_name)
-    except requests.exceptions.RequestException as e:
-        logger.exception(f"HTTP request failed: {e}")
-        raise
+    response = await query_user_groups(client, hub_name=hub_name, user_name=user_name)
     grouped = defaultdict(lambda: {"username": None, "hub": None})
     for entry in response:
         key = (entry["username"], entry["hub"])

tests/conftest.py

@@ -3,6 +3,7 @@
 from datetime import datetime, timezone
 from unittest.mock import patch
 
+import aiohttp
 import boto3
 import pytest
 from botocore.stub import Stubber
@@ -14,6 +15,12 @@
 # Usage and cost data fixtures for test_cost.py
 
 
+@pytest.fixture(scope="session")
+async def client_session():
+    async with aiohttp.ClientSession() as client:
+        yield client
+
+
 @pytest.fixture(scope="function")
 def input_data_usage():
     with open("tests/data/test_data_usage.json") as f: