Migrate to using Traitlets

yuvipanda · yuvipanda · commit 7bc03ae06d1b · 2026-07-22T16:14:11.000-07:00
This lets us eventually generalize the codebase to both
other cloud providers, as well as customize the tags that are
looked at
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,6 +13,7 @@ dependencies = [
     "prometheus-client~=0.24",
     "requests>=2.32.4",
     "yarl>=1.20.1",
+    "traitlets"
 ]
 
 [dependency-groups]
diff --git a/src/jupyterhub_cost_monitoring/prometheus.py b/src/jupyterhub_cost_monitoring/prometheus.py
@@ -0,0 +1,377 @@
+"""
+Query the Prometheus server to get usage of JupyterHub resources.
+"""
+
+import os
+from collections import defaultdict
+from datetime import datetime, timedelta, timezone
+
+import escapism
+import requests
+from yarl import URL
+
+from .cache import ttl_lru_cache
+from .const_usage import USAGE_MAP, USER_GROUP_INFO
+from .date_utils import DateRange, get_now_date
+from .logs import get_logger
+
+logger = get_logger(__name__)
+
+prometheus_host = os.environ.get("SUPPORT_PROMETHEUS_SERVER_SERVICE_HOST", "localhost")
+prometheus_port = int(os.environ.get("SUPPORT_PROMETHEUS_SERVER_SERVICE_PORT", 9090))
+prometheus_username = os.environ.get("PROMETHEUS_USERNAME", "")
+prometheus_password = os.environ.get("PROMETHEUS_PASSWORD", "")
+
+class Prometheus(LoggingConfigurable):
+    def query(self, query: str, date_range: DateRange, step: str) -> requests.Response:
+        """
+        Query the Prometheus server with the given query over a date range.
+
+        Args:
+            query: The Prometheus query string
+            date_range: DateRange object containing the time period for the query
+            step: The query resolution step duration
+
+        Returns:
+            JSON response from Prometheus API
+        """
+        # Use Prometheus-formatted dates (inclusive date range with ISO timestamps)
+        from_date, to_date = date_range.prometheus_range
+
+        prometheus_api = URL.build(
+            scheme="http", host=prometheus_host, port=prometheus_port
+        )
+        if prometheus_username != "" and prometheus_password != "":
+            prometheus_auth = requests.auth.HTTPBasicAuth(
+                prometheus_username, prometheus_password
+            )
+        else:
+            prometheus_auth = None
+        parameters = {
+            "query": query,
+            "start": from_date,
+            "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:
+            logger.info(f"Querying Prometheus: {response.url}")
+            response.raise_for_status()
+            result = response.json()
+            return result
+
+
+    def query_usage(
+        self,
+        date_range: DateRange,
+        hub_name: str | None,
+        component_name: str | None,
+        user_name: str | None,
+    ) -> list[dict]:
+        """
+        Query usage cost factors per user from the Prometheus server.
+
+        Returns daily usage cost factors (0-1) for each user, where cost factors represent
+        each user's share of total resource usage and sum to 1 across all users
+        within each date/hub/component combination.
+
+        Args:
+            date_range: DateRange object containing the time period for the query
+            hub_name: Optional name of the hub to filter results.
+            component_name: Optional name of the component to filter results.
+            user_name: Optional name of the user to filter results.
+        """
+        result = []
+        if component_name is None:
+            # Query all components defined in USAGE_MAP
+            for component, params in USAGE_MAP.items():
+                try:
+                    response = self.query(
+                        params["query"], date_range, step=params["step"]
+                    )
+                except requests.exceptions.RequestException:
+                    raise
+                result.extend(self._process_response(response, component))
+        else:
+            # Query specific component only
+            try:
+                response = self.query(
+                    USAGE_MAP[component_name]["query"],
+                    date_range,
+                    step=USAGE_MAP[component_name]["step"],
+                )
+            except requests.exceptions.RequestException:
+                raise
+            result.extend(self._process_response(response, component_name))
+        # Calculate daily cost factors from absolute usage totals)
+        result = self._calculate_daily_cost_factors(result, hub_name=hub_name)
+        # sort the result by date
+        result.sort(key=lambda x: (x["date"], x["component"], x["hub"], x["user"]))
+        result = self._filter_json(result, hub=hub_name, user=user_name)
+        return result
+
+
+    def _process_response(
+        self,
+        response: requests.Response,
+        component_name: str,
+    ) -> dict:
+        """
+        Process the response from the Prometheus server to extract absolute usage data.
+
+        Converts the time series data into a list of usage records, then pivots by date
+        and sums the absolute usage values across time steps within each date.
+
+        If the component_name is home storage, then rename the escaped username used for the directory to the unescaped version.
+        """
+        result = []
+        for data in response["data"]["result"]:
+            hub = data["metric"]["namespace"]
+            user = data["metric"]["username"]
+            date = [
+                datetime.fromtimestamp(value[0], tz=timezone.utc).strftime("%Y-%m-%d")
+                for value in data["values"]
+            ]
+            usage = [float(value[1]) for value in data["values"]]
+            result.append(
+                {
+                    "hub": hub,
+                    "component": component_name,
+                    "user": user,
+                    "date": date,
+                    "value": usage,
+                }
+            )
+        pivoted_result = self._pivot_response_dict(result)
+        processed_result = self._sum_absolute_usage_by_date(pivoted_result)
+
+        if component_name == "home storage":
+            for entry in processed_result:
+                if "shared" not in entry["user"]:
+                    try:
+                        entry["user"] = escapism.unescape(entry["user"], escape_char="-")
+                    except ValueError:
+                        logger.warning(
+                            f"Could not unescape username {entry['user']} for home storage component."
+                        )
+                        continue
+        return processed_result
+
+
+    def _filter_json(self, result: list[dict], **filters):
+        return [
+            item
+            for item in result
+            if all(filters[k] is None or item.get(k) == filters[k] for k in filters)
+        ]
+
+
+    def _pivot_response_dict(self, result: list[dict]) -> list[dict]:
+        """
+        Pivot the response dictionary to have top-level keys as dates.
+        """
+        pivot = []
+        for entry in result:
+            for date, value in zip(entry["date"], entry["value"]):
+                pivot.append(
+                    {
+                        "date": date,
+                        "user": entry["user"],
+                        "hub": entry["hub"],
+                        "component": entry["component"],
+                        "value": value,
+                    }
+                )
+        return pivot
+
+
+    def _sum_absolute_usage_by_date(self, result: list[dict]) -> list[dict]:
+        """
+        Sum the absolute usage values by date.
+
+        The Prometheus queries can return multiple absolute usage values per day.
+        We sum across all entries within each date to get the total daily usage for each user.
+        """
+        sums = defaultdict(float)
+
+        for entry in result:
+            key = (
+                entry["date"],
+                entry["user"],
+                entry["hub"],
+                entry["component"],
+            )
+            sums[key] += entry["value"]
+
+        return [
+            {
+                "date": date,
+                "user": user,
+                "hub": hub,
+                "component": component,
+                "value": total,
+            }
+            for (date, user, hub, component), total in sums.items()
+        ]
+
+
+    def _calculate_daily_cost_factors(
+        self, result: list[dict], hub_name: str | None = None
+    ) -> list[dict]:
+        """
+        Calculate daily usage cost factors from absolute usage values.
+
+        Converts absolute usage values to cost factors by dividing each user's usage
+        by the total usage for all users within the appropriate grouping.
+
+        If hub_name is None: cost factors are calculated across all hubs for each date/component
+        If hub_name is specified: cost factors are calculated per hub for each date/component
+
+        This ensures that cost factors sum to 1 for the appropriate grouping.
+        """
+        # Calculate total usage for the appropriate grouping
+        totals = defaultdict(float)
+        for entry in result:
+            if hub_name is None:
+                # When no specific hub requested, calculate totals across all hubs
+                key = (entry["date"], entry["component"])
+            else:
+                # When specific hub requested, calculate totals per hub
+                key = (entry["date"], entry["hub"], entry["component"])
+            totals[key] += entry["value"]
+
+        # Convert absolute values to cost factors
+        for entry in result:
+            if hub_name is None:
+                # When no specific hub requested, use cross-hub totals
+                key = (entry["date"], entry["component"])
+            else:
+                # When specific hub requested, use per-hub totals
+                key = (entry["date"], entry["hub"], entry["component"])
+
+            total = totals[key]
+            if total > 0:
+                entry["value"] = entry["value"] / total
+            else:
+                entry["value"] = 0.0
+        return result
+
+
+    @ttl_lru_cache(seconds_to_live=3600)
+    def query_user_groups(
+        self,
+        hub_name: str | None = None,
+        user_name: str | None = None,
+        group_name: str | None = None,
+    ) -> list[dict]:
+        """
+        Get user group information from the Prometheus server for the most recent day.
+        """
+        now_date = get_now_date() - timedelta(days=1)
+        date_range = DateRange(start_date=now_date, end_date=now_date)
+        try:
+            response = self.query(USER_GROUP_INFO, date_range, step="1d")
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"HTTP request failed: {e}")
+            raise
+        result = self._process_user_groups(response, hub_name, user_name, group_name)
+        return result
+
+
+    def _process_user_groups(
+        self,
+        response: requests.Response,
+        hub_name: str | None = None,
+        user_name: str | None = None,
+        group_name: str | None = None,
+    ) -> list[dict]:
+        """
+        Process the response from the Prometheus server to extract user group information. Note that only the most recent date of user group membership is used.
+        """
+        result = []
+        unique_keys = set()
+        for data in response["data"]["result"]:
+            hub = data["metric"]["namespace"]
+            user = data["metric"]["username"]
+            user_escaped = data["metric"]["username_escaped"]
+            group = data["metric"]["usergroup"]
+            key = (hub, user, user_escaped, group)
+            if key not in unique_keys:
+                unique_keys.add(key)
+                result.append(
+                    {
+                        "hub": hub,
+                        "username": user,
+                        "username_escaped": user_escaped,
+                        "usergroup": group,
+                    }
+                )
+        return result
+
+
+    @ttl_lru_cache(seconds_to_live=3600)
+    def query_users_with_multiple_groups(
+        self,
+        date_range: DateRange,
+        hub_name: str | None = None,
+        user_name: str | None = None,
+    ) -> list[dict]:
+        try:
+            response = self.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
+        grouped = defaultdict(
+            lambda: {"username": None, "hub": None, "usergroups": [], "has_multiple": False}
+        )
+        for entry in response:
+            k = (entry["username"], entry["hub"])
+            g = grouped[k]
+            g["username"] = entry["username"]
+            g["hub"] = entry["hub"]
+            if entry["usergroup"] == "multiple":
+                g["has_multiple"] = True
+                continue
+            g["usergroups"].append(entry["usergroup"])
+        result = []
+        for v in grouped.values():
+            if v["has_multiple"]:
+                for group in v["usergroups"]:
+                    result.append(
+                        {"username": v["username"], "hub": v["hub"], "usergroup": group}
+                    )
+
+        return result
+
+
+    @ttl_lru_cache(seconds_to_live=3600)
+    def query_users_with_no_groups(
+        self,
+        date_range: DateRange,
+        hub_name: str | None = None,
+        user_name: str | None = None,
+    ) -> list[dict]:
+        try:
+            response = self.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
+        grouped = defaultdict(lambda: {"username": None, "hub": None})
+        for entry in response:
+            key = (entry["username"], entry["hub"])
+            if grouped[key]["username"] is None:
+                grouped[key]["username"] = entry["username"]
+                grouped[key]["hub"] = entry["hub"]
+                if entry["usergroup"] == "none":
+                    logger.debug(
+                        f"User {entry['username']} in hub {entry['hub']} has no groups."
+                    )
+                    grouped[key]["has_none"] = True
+                else:
+                    grouped[key]["has_none"] = False
+        result = [
+            {"username": v["username"], "hub": v["hub"]}
+            for v in grouped.values()
+            if v["has_none"]
+        ]
+        return result
diff --git a/src/jupyterhub_cost_monitoring/query_usage.py b/src/jupyterhub_cost_monitoring/query_usage.py

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@ dependencies = [`
`13`	`13`	`"prometheus-client~=0.24",`
`14`	`14`	`"requests>=2.32.4",`
`15`	`15`	`"yarl>=1.20.1",`
	`16`	`+ "traitlets"`
`16`	`17`	`]`
`17`	`18`
`18`	`19`	`[dependency-groups]`