fix milatools_usage_report.py (#141)

* fix milatools_usage_report.py
mila-iqia · Nov 21, 2024 · 85bd044 · 85bd044
1 parent 8a4ca09
commit 85bd044
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/examples/milatools_usage_report.py b/examples/milatools_usage_report.py
@@ -36,8 +36,8 @@
 from pandas.core.indexes.datetimes import DatetimeIndex
 from typing_extensions import TypeGuard
 
+from sarc.client.job import _jobs_collection
 from sarc.config import MTL
-from sarc.jobs.job import jobs_collection
 
 logger = get_logger(__name__)
 
@@ -126,7 +126,7 @@ def _get_cache_dir():
 
 def _get_all_clusters(start_date: datetime, end_date: datetime):
     cache_dir = _get_cache_dir()
-    job_db: pymongo.collection.Collection = jobs_collection().get_collection()
+    job_db: pymongo.collection.Collection = _jobs_collection().get_collection()
 
     if (
         all_clusters_file := cache_dir / f"all_clusters_{start_date}_{end_date}.pkl"
@@ -428,7 +428,7 @@ def _get_unique_users(
         assert _is_iterable_of(all_users, str) and isinstance(all_users, set)
         return milatools_users, all_users
 
-    job_structured_db = jobs_collection()
+    job_structured_db = _jobs_collection()
     job_db: pymongo.collection.Collection = job_structured_db.get_collection()
 
     _period_filter = _get_filter(start_date, end_date, cluster, name=None)