Feature rescheduleJob method

src/diracx/db/jobs/db.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
+import logging
 from datetime import datetime, timezone
 from typing import Any
 
-from sqlalchemy import func, insert, select, update
+from sqlalchemy import delete, func, insert, select, update
 
 from diracx.core.exceptions import InvalidQueryError
+from diracx.core.models import ScalarSearchOperator, ScalarSearchSpec
 from diracx.core.utils import JobStatus
 
 from ..utils import BaseDB, apply_search_filters
 from .schema import Base as JobDBBase
-from .schema import InputData, JobJDLs, Jobs
+from .schema import InputData, JobJDLs, JobParameters, Jobs, OptimizerParameters
 
 
 class JobDB(BaseDB):
@@ -22,6 +24,11 @@
     # to find a way to make it dynamic
     jdl2DBParameters = ["JobName", "JobType", "JobGroup"]
 
+    # TODO: set maxRescheduling value from CS
+    # maxRescheduling = self.getCSOption("MaxRescheduling", 3)
+    # For now:
+    maxRescheduling = 3
+
     async def summary(self, group_by, search) -> list[dict[str, str | int]]:
         columns = [Jobs.__table__.columns[x] for x in group_by]
 
@@ -141,6 +148,20 @@
         )
         await self.conn.execute(stmt)
 
+    async def getJobJDL(self, job_id: int, original: bool = False) -> str:
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import extractJDL
+
+        if original:
+            stmt = select(JobJDLs.OriginalJDL).where(JobJDLs.JobID == job_id)
+        else:
+            stmt = select(JobJDLs.JDL).where(JobJDLs.JobID == job_id)
+
+        jdl = (await self.conn.execute(stmt)).scalar_one()
+        if jdl:
+            jdl = extractJDL(jdl)
+
+        return jdl
+
     async def insert(
         self,
         jdl,
@@ -193,7 +214,7 @@
 
             return {
                 "JobID": job_id,
-                "Status": JobStatus.FAILED,
+                "Status": JobStatus.Failed,
                 "MinorStatus": "Error in JDL syntax",
             }
 
@@ -242,3 +263,149 @@
             "MinorStatus": initial_minor_status,
             "TimeStamp": datetime.now(tz=timezone.utc),
         }
+
+    async def delete_job_parameters(self, job_id: int):
+        stmt = delete(JobParameters).where(JobParameters.JobID == job_id)
+        await self.conn.execute(stmt)
+
+    async def delete_job_optimizer_parameters(self, job_id: int):
+        stmt = delete(OptimizerParameters).where(OptimizerParameters.JobID == job_id)
+        await self.conn.execute(stmt)
+
+    async def rescheduleJob(self, job_id) -> dict[str, Any]:
+        """Reschedule given job"""
+        from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
+        from DIRAC.Core.Utilities.ReturnValues import SErrorException
+        from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus
+
+        result = await self.search(
+            parameters=[
+                "Status",
+                "MinorStatus",
+                "VerifiedFlag",
+                "RescheduleCounter",
+                "Owner",
+                "OwnerGroup",
+            ],
+            search=[
+                ScalarSearchSpec(
+                    parameter="JobID", operator=ScalarSearchOperator.EQUAL, value=job_id
+                )
+            ],
+            sorts=[],
+        )
+        if not result:
+            raise ValueError(f"Job {job_id} not found.")
+
+        jobAttrs = result[0]
+
+        if "VerifiedFlag" not in jobAttrs:
+            raise ValueError(f"Job {job_id} not found in the system")
+
+        if not jobAttrs["VerifiedFlag"]:
+            raise ValueError(
+                f"Job {job_id} not Verified: Status {jobAttrs['Status']}, Minor Status: {jobAttrs['MinorStatus']}"
+            )
+
+        reschedule_counter = int(jobAttrs["RescheduleCounter"]) + 1
+
+        # TODO: update maxRescheduling:
+        # self.maxRescheduling = self.getCSOption("MaxRescheduling", self.maxRescheduling)
+
+        if reschedule_counter > self.maxRescheduling:
+            logging.warn(f"Job {job_id}: Maximum number of reschedulings is reached.")
+            self.setJobAttributes(
+                job_id,
+                {
+                    "Status": JobStatus.Failed,
+                    "MinorStatus": "Maximum of reschedulings reached",
+                },
+            )
+            raise ValueError(
+                f"Maximum number of reschedulings is reached: {self.maxRescheduling}"
+            )
+
+        new_job_attributes = {"RescheduleCounter": reschedule_counter}
+
+        # TODO: get the job parameters from JobMonitoringClient
+        # result = JobMonitoringClient().getJobParameters(jobID)
+        # if result["OK"]:
+        #     parDict = result["Value"]
+        #     for key, value in parDict.get(jobID, {}).items():
+        #         result = self.setAtticJobParameter(jobID, key, value, rescheduleCounter - 1)
+        #         if not result["OK"]:
+        #             break
+
+        # Delete job in JobParameters and OptimizerParameters:
+        await self.delete_job_parameters(job_id)
+        await self.delete_job_optimizer_parameters(job_id)
+
+        job_jdl = await self.getJobJDL(job_id, original=True)
+        if not job_jdl.strip().startswith("["):
+            job_jdl = f"[{job_jdl}]"
+
+        classAdJob = ClassAd(job_jdl)
+        classAdReq = ClassAd("[]")
+        retVal = {}
+        retVal["JobID"] = job_id
+
+        classAdJob.insertAttributeInt("JobID", job_id)
+
+        try:
+            result = self._checkAndPrepareJob(
+                job_id,
+                classAdJob,
+                classAdReq,
+                jobAttrs["Owner"],
+                jobAttrs["OwnerGroup"],
+                new_job_attributes,
+                classAdJob.getAttributeString("VirtualOrganization"),
+            )
+        except SErrorException as e:
+            raise ValueError(e) from e
+
+        priority = classAdJob.getAttributeInt("Priority")
+        if priority is None:
+            priority = 0
+        jobAttrs["UserPriority"] = priority
+
+        siteList = classAdJob.getListFromExpression("Site")
+        if not siteList:
+            site = "ANY"
+        elif len(siteList) > 1:
+            site = "Multiple"
+        else:
+            site = siteList[0]
+
+        jobAttrs["Site"] = site
+
+        jobAttrs["Status"] = JobStatus.RECEIVED
+
+        jobAttrs["MinorStatus"] = JobMinorStatus.RESCHEDULED
+
+        jobAttrs["ApplicationStatus"] = "Unknown"
+
+        jobAttrs["ApplicationNumStatus"] = 0
+
+        jobAttrs["LastUpdateTime"] = str(datetime.utcnow())
+
+        jobAttrs["RescheduleTime"] = str(datetime.utcnow())
+
+        reqJDL = classAdReq.asJDL()
+        classAdJob.insertAttributeInt("JobRequirements", reqJDL)
+
+        jobJDL = classAdJob.asJDL()
+
+        # Replace the JobID placeholder if any
+        jobJDL = jobJDL.replace("%j", str(job_id))
+
+        result = self.setJobJDL(job_id, jobJDL)
+
+        result = self.setJobAttributes(job_id, jobAttrs)
+
+        retVal["InputData"] = classAdJob.lookupAttribute("InputData")
+        retVal["RescheduleCounter"] = reschedule_counter
+        retVal["Status"] = JobStatus.RECEIVED
+        retVal["MinorStatus"] = JobMinorStatus.RESCHEDULED
+
+        return retVal

src/diracx/routers/job_manager/__init__.py

@@ -273,6 +273,59 @@
     return [{"job_id": job.job_id, "status": job.status} for job in job_update]
 
 
+@router.post("/reschedule")
+async def reschedule_bulk_jobs(
+    job_ids: Annotated[list[int], Query()],
+    job_db: JobDB,
+    user_info: Annotated[UserInfo, Depends(verify_dirac_token)],
+):
+    rescheduled_jobs = []
+    # TODO:
+    # validJobList, invalidJobList, nonauthJobList, ownerJobList = self.jobPolicy.evaluateJobRights(
+    #        jobList, RIGHT_RESCHEDULE
+    #    )
+    # For the moment all jobs are valid:
+    valid_job_list = job_ids
+    for job_id in valid_job_list:
+        # TODO: delete job in TaskQueueDB
+        # self.taskQueueDB.deleteJob(jobID)
+        result = job_db.rescheduleJob(job_id)
+        # TODO: add reschedule logging in JobLoggingDB
+        # self.jobLoggingDB.addLoggingRecord(
+        #         result["JobID"],
+        #         status=result["Status"],
+        #         minorStatus=result["MinorStatus"],
+        #         applicationStatus="Unknown",
+        #         source="JobManager",
+        #     )
+        if result:
+            rescheduled_jobs.append(job_id)
+    # To uncomment when jobPolicy is setup:
+    # if invalid_job_list or non_auth_job_list:
+    #     logging.error("Some jobs failed to reschedule")
+    #     if invalid_job_list:
+    #         logging.info(f"Invalid jobs: {invalid_job_list}")
+    #     if non_auth_job_list:
+    #         logging.info(f"Non authorized jobs: {nonauthJobList}")
+
+    # TODO: send jobs to OtimizationMind
+    #  self.__sendJobsToOptimizationMind(validJobList)
+    return rescheduled_jobs
+
+
+@router.post("/{job_id}/reschedule")
+async def reschedule_single_job(
+    job_id: int,
+    job_db: JobDB,
+    user_info: Annotated[UserInfo, Depends(verify_dirac_token)],
+):
+    try:
+        result = await job_db.rescheduleJob(job_id)
+    except ValueError as e:
+        raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail=str(e)) from e
+    return result
+
+
 EXAMPLE_SEARCHES = {
     "Show all": {
         "summary": "Show all",

tests/routers/test_job_manager.py

@@ -1,6 +1,7 @@
 from http import HTTPStatus
 
 import pytest
+from fastapi.testclient import TestClient
 
 TEST_JDL = """
     Arguments = "jobDescription.xml -o LogLevel=INFO";
@@ -208,3 +209,19 @@ def test_user_cannot_submit_multiple_jdl_if_at_least_one_of_them_is_parametric(
 def test_user_without_the_normal_user_property_cannot_submit_job(admin_user_client):
     res = admin_user_client.post("/jobs/", json=[TEST_JDL])
     assert res.status_code == HTTPStatus.FORBIDDEN, res.json()
+
+
+def test_insert_and_reschedule(normal_user_client: TestClient):
+    job_definitions = [TEST_JDL]
+    r = normal_user_client.post("/jobs/", json=job_definitions)
+    assert r.status_code == 200, r.json()
+    assert len(r.json()) == len(job_definitions)
+
+    submitted_job_ids = sorted([job_dict["JobID"] for job_dict in r.json()])
+
+    # Test /jobs/reschedule
+    r = normal_user_client.post(
+        "/jobs/reschedule",
+        params={"job_ids": submitted_job_ids},
+    )
+    assert r.status_code == 200, r.json()