Merge pull request #526 from omnivector-solutions/matheushent/hot-fix…

…--PENG-2116

hot fix a bug found in the self update task during QA where the upstream tasks weren't being executed after the package update

jobbergate-agent/CHANGELOG.md

@@ -5,6 +5,7 @@ This file keeps track of all notable changes to jobbergate-agent
 ## Unreleased
 - Added logic to update slurm job status at job submission time [PENG-2193]
 
+- Hot fix regarding the self update task where tasks weren't properly scheduled after the version update
 
 ## 5.0.0a1 -- 2024-04-04
 - Added a task scheduler whose purpose is to self update the agent [PENG-2116]

jobbergate-agent/jobbergate_agent/internals/update.py

@@ -1,11 +1,16 @@
+import re
 import subprocess
 import sys
 
+from apscheduler.job import Job
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
 from loguru import logger
 from pkg_resources import get_distribution
 
 from jobbergate_agent.clients.cluster_api import backend_client as jobbergate_api_client
-from jobbergate_agent.utils.scheduler import schedule_tasks, scheduler
+
+# flake8 doesn't understand the scheduler is used in the self_update_agent function
+from jobbergate_agent.utils.scheduler import schedule_tasks, scheduler  # noqa: F401
 
 
 package_name = "jobbergate_agent"
@@ -25,18 +30,41 @@ def _need_update(current_version: str, upstream_version: str) -> bool:
     """Compare the current version with the upstream version.
 
     In case the current version is the same as the upstream version, return False.
-    As well as, in case the major versions are the same, return False, as we don't want
-    to update across major versions. Otherwise, return True.
-
-    This behaviour allows the agent to update and rollback across all minor and patch versions.
+    If the major versions are different, return False, as updates across major versions are not desired.
+    Otherwise, return True, allowing updates and rollbacks across all minor and patch versions,
+    including handling for pre-release versions ('a' for alpha, 'b' for beta).
     """
-    if current_version == upstream_version:
-        return False
-    current_major, _, _ = map(int, current_version.split("."))
-    upstream_major, _, _ = map(int, upstream_version.split("."))
+    current_major: int | str
+    current_minor: int | str
+    current_patch: int | str
+    upstream_major: int | str
+    upstream_minor: int | str
+    upstream_patch: int | str
+
+    # regular expression to parse version strings: major.minor.patch
+    version_pattern = r"^(\d+)\.(\d+)\.(\d+)$"
+
+    current_match = re.match(version_pattern, current_version)
+    upstream_match = re.match(version_pattern, upstream_version)
+
+    if not current_match or not upstream_match:
+        raise ValueError(
+            f"One of the following versions are improperly formatted: {current_version}, {upstream_version}"
+        )
+
+    current_major, current_minor, current_patch = current_match.groups()
+    upstream_major, upstream_minor, upstream_patch = upstream_match.groups()
+
+    current_major, current_minor, current_patch = map(int, [current_major, current_minor, current_patch])
+    upstream_major, upstream_minor, upstream_patch = map(int, [upstream_major, upstream_minor, upstream_patch])
+
+    # major version check
     if current_major != upstream_major:
         return False
-    return True
+    # minor and patch version check
+    elif current_minor != upstream_minor or current_patch != upstream_patch:
+        return True
+    return False
 
 
 def _update_package(version: str) -> None:
@@ -48,6 +76,8 @@ async def self_update_agent():
 
     In case the agent is updated, the scheduler is shutdown and restarted with the new version.
     """
+    global scheduler
+
     current_version = get_distribution(package_name).version
     upstream_version = await _fetch_upstream_version_info()
     logger.debug(
@@ -60,12 +90,22 @@ async def self_update_agent():
         logger.debug("Shutting down the scheduler...")
         scheduler.shutdown(wait=False)
 
+        logger.debug("Clearing the scheduler jobs...")
+        scheduler_jobs: list[Job] = scheduler.get_jobs()
+        for job in scheduler_jobs:
+            job.remove()
+
         logger.debug(f"Updating {package_name} from version {current_version} to {upstream_version}...")
         _update_package(upstream_version)
         logger.debug("Update completed successfully.")
 
         logger.debug(f"Loading plugins from version {upstream_version}...")
-        schedule_tasks(scheduler)
+        new_scheduler = AsyncIOScheduler()
+        schedule_tasks(new_scheduler)
+        new_scheduler.start()
         logger.debug("Plugins loaded successfully.")
+
+        logger.info("Replacing the scheduler with the new version...")
+        scheduler = new_scheduler
     else:
         logger.debug("No update is required or update crosses a major version divide.")

jobbergate-agent/tests/internals/test_update.py

@@ -59,6 +59,9 @@ def test_need_update(current_version: str, upstream_version: str, expected_resul
         ("1.0", "1.0.1"),  # Improperly formatted current version
         ("1.0.0", "1.0"),  # Improperly formatted upstream version
         ("1", "2"),  # Major version with no minor/patch
+        ("1.0.1a", "1.1.0"),  # Pre-release improperly formatted
+        ("1.0.1", "1.0.10b"),  # Pre-release improperly formatted
+        ("1.0.9a1", "1.0.9a2"),  # Alpha version is not an allowed format
     ],
 )
 def test_need_update__check_improperly_formatted_versions(
@@ -112,7 +115,9 @@ def test_update_package(mocked_sys: mock.MagicMock, mocked_subprocess: mock.Magi
 @mock.patch("jobbergate_agent.internals.update._need_update")
 @mock.patch("jobbergate_agent.internals.update.scheduler")
 @mock.patch("jobbergate_agent.internals.update.schedule_tasks")
+@mock.patch("jobbergate_agent.internals.update.AsyncIOScheduler")
 async def test_self_update_agent(
+    mocked_asyncio_scheduler: mock.MagicMock,
     mocked_schedule_tasks: mock.MagicMock,
     mocked_scheduler: mock.MagicMock,
     mocked_need_update: mock.MagicMock,
@@ -133,6 +138,11 @@ async def test_self_update_agent(
     mocked_need_update.return_value = is_update_available
     mocked_scheduler.shutdown = mock.Mock()
 
+    mocked_new_scheduler = mock.Mock()
+    mocked_asyncio_scheduler.return_value = mocked_new_scheduler
+    mocked_new_scheduler.shutdown = mock.Mock()
+    mocked_new_scheduler.start = mock.Mock()
+
     await self_update_agent()
 
     mocked_get_distribution.assert_called_once_with("jobbergate_agent")
@@ -141,8 +151,17 @@ async def test_self_update_agent(
     if is_update_available:
         mocked_scheduler.shutdown.assert_called_once_with(wait=False)
         mocked_update_package.assert_called_once_with(upstream_version)
-        mocked_schedule_tasks.assert_called_once_with(mocked_scheduler)
+        mocked_schedule_tasks.assert_called_once_with(mocked_new_scheduler)
+        mocked_asyncio_scheduler.assert_called_once_with()
+        mocked_new_scheduler.start.assert_called_once_with()
+
+        # this asserts that the scheduler is updated *in memory* with the new version
+        from jobbergate_agent.internals.update import scheduler
+
+        assert scheduler is mocked_new_scheduler
     else:
         mocked_scheduler.shutdown.assert_not_called()
         mocked_update_package.assert_not_called()
         mocked_schedule_tasks.assert_not_called()
+        mocked_asyncio_scheduler.assert_not_called()
+        mocked_new_scheduler.start.assert_not_called()