Merge pull request #212 from notoraptor/cw-560-update-cleanup-jobs

[CW-560] Faire en sorte que le script de ménage des jobs accepte une durée (en jours)

scripts/cleanup_jobs.py

@@ -6,7 +6,7 @@
 
 import argparse
 import sys
-from datetime import datetime
+from datetime import datetime, timedelta
 from slurm_state.mongo_client import get_mongo_client
 from slurm_state.config import get_config
 
@@ -42,6 +42,15 @@ def main(arguments: list):
             "If specified, script will delete all jobs older than given date."
         ),
     )
+    group.add_argument(
+        "-t",
+        "--days",
+        type=int,
+        help=(
+            "Time (in days) of older job to keep. "
+            "If specified, script will delete all jobs updated before latest <--days> days."
+        ),
+    )
     parser.add_argument(
         "--debug",
         action="store_true",
@@ -59,6 +68,10 @@ def main(arguments: list):
         keep_n_most_recent_jobs(args.jobs)
     elif args.jobs_per_user is not None:
         keep_n_most_recent_jobs_per_user(args.jobs_per_user)
+    elif args.days is not None:
+        print(f"Keeping jobs updated in latest {args.days} days.")
+        older_date = datetime.now() - timedelta(days=args.days)
+        keep_jobs_from_date(older_date)
     else:
         if args.date.count("-") == 2:
             date_format = "%Y-%m-%d"
@@ -150,7 +163,7 @@ def keep_n_most_recent_jobs_per_user(n: int):
 
 
 def keep_jobs_from_date(date: datetime):
-    print(f"Keeping jobs starting from: {date}")
+    print(f"Keeping jobs updated since: {date}")
     mc = _get_db()
     db_jobs = mc["jobs"]
     nb_total_jobs = db_jobs.count_documents({})
@@ -161,7 +174,7 @@ def keep_jobs_from_date(date: datetime):
         db_jobs.find({"cw.last_slurm_update": {"$lt": date.timestamp()}})
     )
     if not jobs_to_delete:
-        print(f"No job found before {date}, nothing to do.")
+        print(f"No job updated before {date}, nothing to do.")
         return
 
     # Delete jobs

scripts_test/requirements.txt

@@ -7,5 +7,6 @@ packaging==23.2
 pluggy==1.3.0
 pymongo==4.6.3
 pytest==7.4.2
+pytest-freezegun==0.4.2
 toml==0.10.2
 tomli==2.0.1

scripts_test/test_cleanup_jobs.py

@@ -21,7 +21,9 @@ def __init__(self):
         client = MongoClient(get_config("mongo.connection_string"))
         mc = client[db_name]
 
-        base_datetime = datetime.now()
+        base_datetime = datetime(
+            year=2020, month=1, day=1, hour=0, minute=0, second=0, microsecond=0
+        )
 
         # Create fake users, intended to be checked when cleaning jobs per user.
         fake_users = [
@@ -274,3 +276,40 @@ def _fmt_date(d: datetime):
         # all jobs should be deleted.
         remaining_jobs = ctx.check_user_props()
         assert len(remaining_jobs) == 0
+
+
+@pytest.mark.freeze_time("2024-01-01 00:00:00")
+def test_keep_jobs_after_a_time():
+    with CleanupTestContext() as ctx:
+        too_old_date = ctx.base_datetime - timedelta(days=1)
+        inbound_date_1 = ctx.base_datetime + timedelta(days=15)
+        inbound_date_2 = ctx.base_datetime + timedelta(days=60)
+        new_date = ctx.base_datetime + timedelta(days=ctx.NB_JOBS)
+
+        current_date = datetime.now()
+        too_old_time = (current_date - too_old_date).days
+        inbound_time_1 = (current_date - inbound_date_1).days
+        inbound_time_2 = (current_date - inbound_date_2).days
+        new_time = (current_date - new_date).days
+
+        jobs = ctx.check_user_props()
+
+        cleanup_jobs(["-t", str(too_old_time)])
+        # Date is too old, so no job should be deleted
+        assert jobs == ctx.check_user_props()
+
+        cleanup_jobs(["--days", str(inbound_time_1)])
+        remaining_jobs = ctx.check_user_props()
+        assert len(remaining_jobs) == 100 - 15
+        assert jobs[15:] == remaining_jobs
+
+        cleanup_jobs(["-t", str(inbound_time_2)])
+        remaining_jobs = ctx.check_user_props()
+        assert len(remaining_jobs) == 100 - 60
+        assert jobs[60:] == remaining_jobs
+
+        cleanup_jobs(["-t", str(new_time)])
+        # With a date more recent than latest job,
+        # all jobs should be deleted.
+        remaining_jobs = ctx.check_user_props()
+        assert len(remaining_jobs) == 0