Skip to content

Commit 52077bb

Browse files
committed
Depend on cachetools instead of implementing our own cache
More robust (but pure TTL) cache set up, with same 128 maxsize. https://cachetools.readthedocs.io/en/stable/ has more options we can switch to if needed.
1 parent 1fb3ef2 commit 52077bb

5 files changed

Lines changed: 19 additions & 48 deletions

File tree

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ dependencies = [
1313
"prometheus-client~=0.24",
1414
"requests>=2.32.4",
1515
"yarl>=1.20.1",
16+
"cachetools"
1617
]
1718

1819
[dependency-groups]

src/jupyterhub_cost_monitoring/cache.py

Lines changed: 0 additions & 31 deletions
This file was deleted.

src/jupyterhub_cost_monitoring/query_cost_aws.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,11 @@
66
import functools
77
from pprint import pformat
88

9+
from cachetools.func import ttl_cache
910
import boto3
1011
import requests
1112

12-
from .cache import ttl_lru_cache
13+
from .cache import ttl_cache
1314
from .const_cost_aws import (
1415
FILTER_ATTRIBUTABLE_COSTS,
1516
FILTER_CORE_COSTS,
@@ -64,7 +65,7 @@ def query_aws_cost_explorer(metrics, granularity, from_date, to_date, filter, gr
6465
return response
6566

6667

67-
@ttl_lru_cache(seconds_to_live=3600)
68+
@ttl_cache(ttl=3600)
6869
def query_hub_names(date_range: DateRange):
6970
"""
7071
Query hub names from AWS Cost Explorer within the given date range.
@@ -87,7 +88,7 @@ def query_hub_names(date_range: DateRange):
8788
return hub_names
8889

8990

90-
@ttl_lru_cache(seconds_to_live=3600)
91+
@ttl_cache(ttl=3600)
9192
def query_total_costs(date_range: DateRange):
9293
"""
9394
Query total costs from AWS Cost Explorer for the given date range.
@@ -119,7 +120,7 @@ def query_total_costs(date_range: DateRange):
119120
return processed_response
120121

121122

122-
@ttl_lru_cache(seconds_to_live=3600)
123+
@ttl_cache(ttl=3600)
123124
def _query_total_costs(date_range: DateRange, add_attributable_costs_filter):
124125
"""
125126
Internal function to query total costs from AWS Cost Explorer.
@@ -169,7 +170,7 @@ def _query_total_costs(date_range: DateRange, add_attributable_costs_filter):
169170
return processed_response
170171

171172

172-
@ttl_lru_cache(seconds_to_live=3600)
173+
@ttl_cache(ttl=3600)
173174
def query_total_costs_per_hub(date_range: DateRange):
174175
"""
175176
Query total costs per hub from AWS Cost Explorer for the given date range.
@@ -378,7 +379,7 @@ def _process_core_costs(entries_by_date, core_cost_response):
378379
logger.debug(f"Added new core entry for {date}: {core_cost:.2f}")
379380

380381

381-
@ttl_lru_cache(seconds_to_live=3600)
382+
@ttl_cache(ttl=3600)
382383
def query_total_costs_per_component(
383384
date_range: DateRange,
384385
hub_name: str = None,
@@ -513,7 +514,7 @@ def query_total_costs_per_component(
513514
return final_response
514515

515516

516-
@ttl_lru_cache(seconds_to_live=3600)
517+
@ttl_cache(ttl=3600)
517518
def query_total_costs_per_user(
518519
date_range: DateRange,
519520
hub: str = None,
@@ -629,7 +630,7 @@ def query_total_costs_per_user(
629630
return results
630631

631632

632-
@ttl_lru_cache(seconds_to_live=3600)
633+
@ttl_cache(ttl=3600)
633634
def query_total_costs_per_group(
634635
date_range: DateRange,
635636
):

src/jupyterhub_cost_monitoring/query_usage.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
import requests
1111
from yarl import URL
1212

13-
from .cache import ttl_lru_cache
13+
from .cache import ttl_cache
1414
from .const_usage import USAGE_MAP, USER_GROUP_INFO
1515
from .date_utils import DateRange, get_now_date
1616
from .logs import get_logger
@@ -255,7 +255,7 @@ def _calculate_daily_cost_factors(
255255
return result
256256

257257

258-
@ttl_lru_cache(seconds_to_live=3600)
258+
@ttl_cache(ttl=3600)
259259
def query_user_groups(
260260
hub_name: str | None = None,
261261
user_name: str | None = None,
@@ -305,7 +305,7 @@ def _process_user_groups(
305305
return result
306306

307307

308-
@ttl_lru_cache(seconds_to_live=3600)
308+
@ttl_cache(ttl=3600)
309309
def query_users_with_multiple_groups(
310310
date_range: DateRange,
311311
hub_name: str | None = None,
@@ -339,7 +339,7 @@ def query_users_with_multiple_groups(
339339
return result
340340

341341

342-
@ttl_lru_cache(seconds_to_live=3600)
342+
@ttl_cache(ttl=3600)
343343
def query_users_with_no_groups(
344344
date_range: DateRange,
345345
hub_name: str | None = None,

tests/test_date_utils.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
import pytest
1212

13-
from src.jupyterhub_cost_monitoring.cache import ttl_lru_cache
13+
from src.jupyterhub_cost_monitoring.cache import ttl_cache
1414
from src.jupyterhub_cost_monitoring.date_utils import (
1515
DateRange,
1616
ensure_utc_datetime,
@@ -229,14 +229,14 @@ def test_timezone_handling_in_parsing(self):
229229

230230

231231
class TestCacheIntegration:
232-
"""Test DateRange compatibility with ttl_lru_cache."""
232+
"""Test DateRange compatibility with ttl_cache."""
233233

234234
def test_daterange_as_cache_key(self):
235235
"""Test that DateRange objects work as cache keys."""
236236

237237
call_count = 0
238238

239-
@ttl_lru_cache(seconds_to_live=300)
239+
@ttl_cache(ttl=300)
240240
def cached_function(date_range: DateRange) -> str:
241241
nonlocal call_count
242242
call_count += 1
@@ -261,7 +261,7 @@ def test_cache_miss_with_different_ranges(self):
261261
"""Test that different DateRange objects result in cache misses."""
262262
call_count = 0
263263

264-
@ttl_lru_cache(seconds_to_live=300)
264+
@ttl_cache(ttl=300)
265265
def cached_function(date_range: DateRange) -> str:
266266
nonlocal call_count
267267
call_count += 1
@@ -286,7 +286,7 @@ def test_cache_hit_with_same_dates_different_times(self):
286286
"""Test that DateRange objects with same dates but different times result in cache hits."""
287287
call_count = 0
288288

289-
@ttl_lru_cache(seconds_to_live=300)
289+
@ttl_cache(ttl=300)
290290
def cached_function(date_range: DateRange) -> str:
291291
nonlocal call_count
292292
call_count += 1

0 commit comments

Comments
 (0)