feat(agent): add task to report status to the API (#539)

jobbergate-agent/CHANGELOG.md

@@ -4,6 +4,7 @@ This file keeps track of all notable changes to jobbergate-agent
 
 ## Unreleased
 
+- Added a new task to report at the expected interval that the agent is up and running [ASP-4602]
 - Patched cwd issues on remote job submission to avoid permission denied errors on the folder
 
 ## 5.0.0 -- 2024-04-18

jobbergate-agent/jobbergate_agent/jobbergate/report_health.py

@@ -0,0 +1,13 @@
+from loguru import logger
+
+from jobbergate_agent.clients.cluster_api import backend_client as jobbergate_api_client
+from jobbergate_agent.utils.exception import JobbergateApiError
+from jobbergate_agent.utils.logging import log_error
+
+
+async def report_health_status(interval: int) -> None:
+    """Ping the API to report the agent's status."""
+    logger.debug("Reporting status to the API")
+    with JobbergateApiError.handle_errors("Failed to report agent status", do_except=log_error):
+        response = await jobbergate_api_client.put("jobbergate/clusters/status", params={"interval": interval})
+        response.raise_for_status()

jobbergate-agent/jobbergate_agent/tasks.py

@@ -8,6 +8,7 @@
 
 from jobbergate_agent.clients.cluster_api import backend_client
 from jobbergate_agent.internals.update import self_update_agent
+from jobbergate_agent.jobbergate.report_health import report_health_status
 from jobbergate_agent.jobbergate.submit import submit_pending_jobs
 from jobbergate_agent.jobbergate.update import update_active_jobs
 from jobbergate_agent.settings import SETTINGS
@@ -38,6 +39,16 @@ def pending_submissions_task(scheduler: BaseScheduler) -> Job:
     return scheduler.add_job(submit_pending_jobs, "interval", seconds=SETTINGS.TASK_JOBS_INTERVAL_SECONDS)
 
 
+def status_report_task(scheduler: BaseScheduler) -> Job:
+    """
+    Schedule a task to report the status.
+    """
+    seconds_between_calls = SETTINGS.TASK_JOBS_INTERVAL_SECONDS
+    return scheduler.add_job(
+        report_health_status, "interval", seconds=seconds_between_calls, kwargs={"interval": seconds_between_calls}
+    )
+
+
 @logger_wraps()
 async def trigger_garbage_collections(interval_between_calls: int = 60) -> None:
     """Trigger maintenance tasks on the Jobbergate API."""

jobbergate-agent/pyproject.toml

@@ -48,6 +48,7 @@ jg-run = "jobbergate_agent.main:main"
 [tool.poetry.plugins.'jobbergate_agent.tasks']
 active-jobs = 'jobbergate_agent.tasks:active_submissions_task'
 pending-jobs = 'jobbergate_agent.tasks:pending_submissions_task'
+report-status = 'jobbergate_agent.tasks:status_report_task'
 garbage-collection = 'jobbergate_agent.tasks:garbage_collection_task'
 self-update = 'jobbergate_agent.tasks:self_update_task'
 

jobbergate-agent/tests/jobbergate/test_report_health.py

@@ -0,0 +1,36 @@
+import httpx
+import pytest
+import respx
+
+from jobbergate_agent.jobbergate.report_health import report_health_status
+from jobbergate_agent.settings import SETTINGS
+from jobbergate_agent.utils.exception import JobbergateApiError
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_access_token")
+async def test_report_health_status__success():
+    interval = 60
+    async with respx.mock:
+        status_route = respx.put(f"{SETTINGS.BASE_API_URL}/jobbergate/clusters/status")
+        status_route.mock(return_value=httpx.Response(status_code=202))
+
+        await report_health_status(interval)
+
+    assert status_route.call_count == 1
+    assert dict(status_route.calls.last.request.url.params) == {"interval": str(interval)}
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_access_token")
+async def test_report_health_status__failure():
+    interval = 60
+    async with respx.mock:
+        status_route = respx.put(f"{SETTINGS.BASE_API_URL}/jobbergate/clusters/status")
+        status_route.mock(return_value=httpx.Response(status_code=500))
+
+        with pytest.raises(JobbergateApiError):
+            await report_health_status(interval)
+
+    assert status_route.call_count == 1
+    assert dict(status_route.calls.last.request.url.params) == {"interval": str(interval)}

jobbergate-agent/tests/jobbergate/test_submit.py

@@ -52,7 +52,6 @@ def user_mapper():
 
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("mock_access_token")
-@pytest.mark.asyncio
 async def test_retrieve_submission_file__success():
     """
     Test that the ``retrieve_submission_file()`` function can retrieve a submission file

jobbergate-agent/tests/utils.py/test_plugin.py

@@ -9,6 +9,7 @@
     garbage_collection_task,
     pending_submissions_task,
     self_update_task,
+    status_report_task,
 )
 from jobbergate_agent.utils.plugin import load_plugins
 from jobbergate_agent.utils.user_mapper import SingleUserMapper
@@ -19,6 +20,7 @@ def test_discover_tasks__success():
     expected_result = {
         "active-jobs": active_submissions_task,
         "pending-jobs": pending_submissions_task,
+        "report-status": status_report_task,
         "garbage-collection": garbage_collection_task,
         "self-update": self_update_task,
     }

jobbergate-agent/tests/utils.py/test_scheduler.py

@@ -54,6 +54,7 @@ def test_scheduler_end_to_end(tweak_settings):
         "active-jobs",
         "garbage-collection",
         "pending-jobs",
+        "report-status",
         "self-update",
     }