diff --git a/scripts/cleanup_jobs.py b/scripts/cleanup_jobs.py index 1f946ead..2b8f3d9f 100644 --- a/scripts/cleanup_jobs.py +++ b/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 diff --git a/scripts_test/requirements.txt b/scripts_test/requirements.txt index 47f8db36..bb94ba1e 100644 --- a/scripts_test/requirements.txt +++ b/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 diff --git a/scripts_test/test_cleanup_jobs.py b/scripts_test/test_cleanup_jobs.py index f99fbc0f..e7b5fad6 100644 --- a/scripts_test/test_cleanup_jobs.py +++ b/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