DIRAC 9.0.0a42__py3-none-any.whl → 9.0.7__py3-none-any.whl

DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py

@@ -28,6 +28,7 @@ from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getCESiteMapping
 from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Utilities import DErrno
 from DIRAC.Core.Utilities.MySQL import _quotedList
+from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
 from DIRAC.FrameworkSystem.Client.Logger import contextLogger
 from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus
 from DIRAC.WorkloadManagementSystem.Client import PilotStatus
@@ -401,9 +402,18 @@ AND SubmissionTime < DATE_SUB(UTC_TIMESTAMP(),INTERVAL %d DAY)"
             if result["Value"]:
                 return int(result["Value"][0][0])
             return 0
-        refString = ",".join(["'" + ref + "'" for ref in pilotRef])
-        req = f"SELECT PilotID from PilotAgents WHERE PilotJobReference in ( {refString} )"
-        result = self._query(req)
+
+        sqlCmd = "CREATE TEMPORARY TABLE to_select_PilotAgents (PilotID VARCHAR(255) NOT NULL, PRIMARY KEY (PilotID)) ENGINE=MEMORY;"
+        returnValueOrRaise(self._update(sqlCmd))
+        try:
+            sqlCmd = "INSERT INTO to_select_PilotAgents (PilotID) VALUES ( %s )"
+            returnValueOrRaise(self._updatemany(sqlCmd, [(p,) for p in pilotRef]))
+            sqlCmd = "SELECT PilotID FROM PilotAgents JOIN to_select_PilotAgents USING (PilotID)"
+            result = self._query(sqlCmd)
+        finally:
+            sqlCmd = "DROP TEMPORARY TABLE to_select_PilotAgents"
+            returnValueOrRaise(self._update(sqlCmd))
+
         if not result["OK"]:
             return []
         if result["Value"]:
@@ -437,9 +447,10 @@ AND SubmissionTime < DATE_SUB(UTC_TIMESTAMP(),INTERVAL %d DAY)"
         """Get IDs of Jobs that were executed by a pilot"""
         cmd = "SELECT pilotID,JobID FROM JobToPilotMapping "
         if isinstance(pilotID, list):
-            cmd = cmd + " WHERE pilotID IN (%s)" % ",".join(["%s" % x for x in pilotID])
+            pilotIDs_string = ",".join(str(int(x)) for x in pilotID)
+            cmd = f"{cmd} WHERE pilotID IN ({pilotIDs_string})"
         else:
-            cmd = cmd + f" WHERE pilotID = {pilotID}"
+            cmd = f"{cmd} WHERE pilotID = {pilotID}"
 
         result = self._query(cmd)
         if not result["OK"]:

DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.sql

@@ -46,7 +46,7 @@ CREATE TABLE `PilotAgents` (
   KEY `PilotJobReference` (`PilotJobReference`),
   KEY `Status` (`Status`),
   KEY `Statuskey` (`GridSite`,`DestinationSite`,`Status`)
-) ENGINE=InnoDB DEFAULT CHARSET=latin1;
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
 
 
 DROP TABLE IF EXISTS `JobToPilotMapping`;
@@ -56,7 +56,7 @@ CREATE TABLE `JobToPilotMapping` (
   `StartTime` DATETIME NOT NULL,
   KEY `JobID` (`JobID`),
   KEY `PilotID` (`PilotID`)
-) ENGINE=InnoDB DEFAULT CHARSET=latin1;
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
 
 DROP TABLE IF EXISTS `PilotOutput`;
 CREATE TABLE `PilotOutput` (
@@ -64,4 +64,4 @@ CREATE TABLE `PilotOutput` (
   `StdOutput` MEDIUMTEXT,
   `StdError` MEDIUMTEXT,
   PRIMARY KEY (`PilotID`)
-) ENGINE=InnoDB DEFAULT CHARSET=latin1;
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;

DIRAC/WorkloadManagementSystem/DB/SandboxMetadataDB.py

@@ -1,10 +1,11 @@
 """ SandboxMetadataDB class is a front-end to the metadata for sandboxes
 """
-from DIRAC import S_ERROR, S_OK, gLogger
+from DIRAC import S_ERROR, S_OK
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Security import Properties
 from DIRAC.Core.Utilities import List
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue, returnValueOrRaise
 
 
 class SandboxMetadataDB(DB):
@@ -63,7 +64,7 @@ class SandboxMetadataDB(DB):
                 "Type": "VARCHAR(64) NOT NULL",
             },
             "Indexes": {"Entity": ["EntityId"], "SBIndex": ["SBId"]},
-            "UniqueIndexes": {"Mapping": ["SBId", "EntityId", "Type"]},
+            "PrimaryKey": ["SBId", "EntityId", "Type"],
         }
 
         for tableName in self.__tablesDesc:
@@ -219,60 +220,27 @@ class SandboxMetadataDB(DB):
             return result
         return S_OK(assigned)
 
-    def __filterEntitiesByRequester(self, entitiesList, requesterName, requesterGroup):
+    @convertToReturnValue
+    def unassignEntities(self, entities: list):
         """
-        Given a list of entities and a requester, return the ones that the requester is allowed to modify
-        """
-        sqlCond = ["s.OwnerId=o.OwnerId", "s.SBId=e.SBId"]
-        requesterProps = Registry.getPropertiesForEntity(requesterGroup, name=requesterName)
-        if Properties.JOB_ADMINISTRATOR in requesterProps:
-            # Do nothing, just ensure it doesn't fit in the other cases
-            pass
-        elif Properties.JOB_SHARING in requesterProps:
-            sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
-        elif Properties.NORMAL_USER in requesterProps:
-            sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
-            sqlCond.append(f"o.Owner='{requesterName}'")
-        else:
-            return S_ERROR("Not authorized to access sandbox")
-        for i in range(len(entitiesList)):
-            entitiesList[i] = self._escapeString(entitiesList[i])["Value"]
-        if len(entitiesList) == 1:
-            sqlCond.append(f"e.EntityId = {entitiesList[0]}")
-        else:
-            sqlCond.append(f"e.EntityId in ( {', '.join(entitiesList)} )")
-        sqlCmd = "SELECT DISTINCT e.EntityId FROM `sb_EntityMapping` e, `sb_SandBoxes` s, `sb_Owners` o WHERE"
-        sqlCmd = f"{sqlCmd} {' AND '.join(sqlCond)}"
-        result = self._query(sqlCmd)
-        if not result["OK"]:
-            return result
-        return S_OK([row[0] for row in result["Value"]])
-
-    def unassignEntities(self, entities, requesterName, requesterGroup):
-        """
-        Unassign jobs to sandboxes
+        Unassign entities to sandboxes. Entities are a list of strings, e.g. ['job:1234', 'job:5678'].
 
         :param list entities: list of entities to unassign
         """
-        updated = 0
         if not entities:
-            return S_OK()
-        result = self.__filterEntitiesByRequester(entities, requesterName, requesterGroup)
-        if not result["OK"]:
-            gLogger.error("Cannot filter entities", result["Message"])
-            return result
-        ids = result["Value"]
-        if not ids:
-            return S_OK(0)
-        sqlCmd = "DELETE FROM `sb_EntityMapping` WHERE EntityId in ( %s )" % ", ".join(
-            ["'%s'" % str(eid) for eid in ids]
-        )
-        result = self._update(sqlCmd)
-        if not result["OK"]:
-            gLogger.error("Cannot unassign entities", result["Message"])
-        else:
-            updated += 1
-        return S_OK(updated)
+            return None
+
+        sqlCmd = "CREATE TEMPORARY TABLE to_delete_EntityId (EntityId VARCHAR(128) NOT NULL, PRIMARY KEY (EntityId)) ENGINE=MEMORY;"
+        returnValueOrRaise(self._update(sqlCmd))
+        try:
+            sqlCmd = "INSERT INTO to_delete_EntityId (EntityId) VALUES ( %s )"
+            returnValueOrRaise(self._updatemany(sqlCmd, [(e,) for e in entities]))
+            sqlCmd = "DELETE m from `sb_EntityMapping` m JOIN to_delete_EntityId t USING (EntityId)"
+            returnValueOrRaise(self._update(sqlCmd))
+        finally:
+            sqlCmd = "DROP TEMPORARY TABLE to_delete_EntityId"
+            returnValueOrRaise(self._update(sqlCmd))
+        return 1
 
     def getSandboxesAssignedToEntity(self, entityId, requesterName, requesterGroup, requestedVO):
         """
@@ -315,22 +283,37 @@ class SandboxMetadataDB(DB):
             "TIMESTAMPDIFF( DAY, LastAccessTime, UTC_TIMESTAMP() ) >= %d" % self.__assignedSBGraceDays,
             f"! Assigned AND TIMESTAMPDIFF( DAY, LastAccessTime, UTC_TIMESTAMP() ) >= {self.__unassignedSBGraceDays}",
         ]
-        sqlCmd = f"SELECT SBId, SEName, SEPFN FROM `sb_SandBoxes` WHERE ( {' ) OR ( '.join(sqlCond)} )"
+        # Exclude sandboxes that are in S3 as those are handled by DiracX
+        sqlCmd = f"SELECT SBId, SEName, SEPFN FROM `sb_SandBoxes` WHERE SEPFN not like '/S3/%' AND (( {' ) OR ( '.join(sqlCond)} ))"
         return self._query(sqlCmd)
 
+    @convertToReturnValue
     def deleteSandboxes(self, SBIdList):
         """
-        Delete sandboxes
+        Delete sandboxes using a temporary table for efficiency and consistency.
         """
-        sqlSBList = ", ".join([str(sbid) for sbid in SBIdList])
-        for table in ("sb_SandBoxes", "sb_EntityMapping"):
-            sqlCmd = f"DELETE FROM `{table}` WHERE SBId IN ( {sqlSBList} )"
-            result = self._update(sqlCmd)
-            if not result["OK"]:
-                return result
+        if not SBIdList:
+            return S_OK()
+        # Create temporary table
+        sqlCmd = "CREATE TEMPORARY TABLE to_delete_SBId (SBId INTEGER(10) UNSIGNED NOT NULL, PRIMARY KEY (SBId)) ENGINE=MEMORY;"
+        returnValueOrRaise(self._update(sqlCmd))
+        try:
+            # Insert SBIds into temporary table
+            sqlCmd = "INSERT INTO to_delete_SBId (SBId) VALUES (%s)"
+            returnValueOrRaise(self._updatemany(sqlCmd, [(sbid,) for sbid in SBIdList]))
+            # Delete from sb_EntityMapping first (to respect FK constraints if any)
+            sqlCmd = "DELETE FROM `sb_EntityMapping` WHERE SBId IN (SELECT SBId FROM to_delete_SBId)"
+            returnValueOrRaise(self._update(sqlCmd))
+            # Delete from sb_SandBoxes
+            sqlCmd = "DELETE FROM `sb_SandBoxes` WHERE SBId IN (SELECT SBId FROM to_delete_SBId)"
+            returnValueOrRaise(self._update(sqlCmd))
+        finally:
+            # Drop temporary table
+            sqlCmd = "DROP TEMPORARY TABLE to_delete_SBId"
+            returnValueOrRaise(self._update(sqlCmd))
         return S_OK()
 
-    def getSandboxId(self, SEName, SEPFN, requesterName, requesterGroup, field="SBId", requesterDN=None):
+    def getSandboxId(self, SEName, SEPFN, requesterName, requesterGroup, field="SBId"):
         """
         Get the sandboxId if it exists
 
@@ -350,7 +333,7 @@ class SandboxMetadataDB(DB):
             "s.OwnerId=o.OwnerId",
         ]
         sqlCmd = f"SELECT s.{field} FROM `sb_SandBoxes` s, `sb_Owners` o WHERE"
-        requesterProps = Registry.getPropertiesForEntity(requesterGroup, name=requesterName, dn=requesterDN)
+        requesterProps = Registry.getPropertiesForEntity(requesterGroup, name=requesterName)
         if Properties.JOB_ADMINISTRATOR in requesterProps or Properties.JOB_MONITOR in requesterProps:
             # Do nothing, just ensure it doesn't fit in the other cases
             pass
@@ -370,24 +353,3 @@ class SandboxMetadataDB(DB):
         if not data:
             return S_ERROR("No sandbox matches the requirements")
         return S_OK(data[0][0])
-
-    def getSandboxOwner(self, SEName, SEPFN, requesterDN, requesterGroup):
-        """get the sandbox owner, if such sandbox exists
-
-        :param SEName: name of the StorageElement
-        :param SEPFN: PFN of the Sandbox
-        :param requesterDN: host DN used as credentials
-        :param requesterGroup: group used to use as credentials (should be 'hosts')
-
-        :returns: S_OK with tuple (owner, ownerGroup, VO)
-        """
-        if not (res := self.getSandboxId(SEName, SEPFN, None, requesterGroup, "OwnerId", requesterDN=requesterDN))[
-            "OK"
-        ]:
-            return res
-
-        if not (
-            res := self._query(f"SELECT `Owner`, `OwnerGroup`, `VO` FROM `sb_Owners` WHERE `OwnerId` = {res['Value']}")
-        )["OK"]:
-            return res
-        return S_OK(res["Value"][0])

DIRAC/WorkloadManagementSystem/DB/StatusUtils.py

@@ -0,0 +1,125 @@
+from DIRAC import S_ERROR, S_OK, gLogger
+from DIRAC.StorageManagementSystem.DB.StorageManagementDB import StorageManagementDB
+from DIRAC.WorkloadManagementSystem.Client import JobStatus
+from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
+from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
+from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE, RIGHT_KILL
+from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import _filterJobStateTransition
+
+
+def _deleteJob(jobID, force=False):
+    """Set the job status to "Deleted"
+    and remove the pilot that ran and its logging info if the pilot is finished.
+
+    :param int jobID: job ID
+    :return: S_OK()/S_ERROR()
+    """
+    if not (result := JobDB().setJobStatus(jobID, JobStatus.DELETED, "Checking accounting", force=force))["OK"]:
+        gLogger.warn("Failed to set job Deleted status", result["Message"])
+        return result
+
+    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
+        gLogger.warn("Failed to delete job from the TaskQueue")
+
+    # if it was the last job for the pilot
+    result = PilotAgentsDB().getPilotsForJobID(jobID)
+    if not result["OK"]:
+        gLogger.error("Failed to get Pilots for JobID", result["Message"])
+        return result
+    for pilot in result["Value"]:
+        res = PilotAgentsDB().getJobsForPilot(pilot)
+        if not res["OK"]:
+            gLogger.error("Failed to get jobs for pilot", res["Message"])
+            return res
+        if not res["Value"]:  # if list of jobs for pilot is empty, delete pilot
+            result = PilotAgentsDB().getPilotInfo(pilotID=pilot)
+            if not result["OK"]:
+                gLogger.error("Failed to get pilot info", result["Message"])
+                return result
+            ret = PilotAgentsDB().deletePilot(result["Value"]["PilotJobReference"])
+            if not ret["OK"]:
+                gLogger.error("Failed to delete pilot from PilotAgentsDB", ret["Message"])
+                return ret
+
+    return S_OK()
+
+
+def _killJob(jobID, sendKillCommand=True, force=False):
+    """Kill one job
+
+    :param int jobID: job ID
+    :param bool sendKillCommand: send kill command
+
+    :return: S_OK()/S_ERROR()
+    """
+    if sendKillCommand:
+        if not (result := JobDB().setJobCommand(jobID, "Kill"))["OK"]:
+            gLogger.warn("Failed to set job Kill command", result["Message"])
+            return result
+
+    gLogger.info("Job marked for termination", jobID)
+    if not (result := JobDB().setJobStatus(jobID, JobStatus.KILLED, "Marked for termination", force=force))["OK"]:
+        gLogger.warn("Failed to set job Killed status", result["Message"])
+    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
+        gLogger.warn("Failed to delete job from the TaskQueue", result["Message"])
+
+    return S_OK()
+
+
+def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
+    """Kill (== set the status to "KILLED") or delete (== set the status to "DELETED") jobs as necessary
+
+    :param str right: RIGHT_KILL or RIGHT_DELETE
+
+    :return: S_OK()/S_ERROR()
+    """
+    badIDs = []
+
+    killJobList = []
+    deleteJobList = []
+    if validJobList:
+        result = JobDB().getJobsAttributes(validJobList, ["Status"])
+        if not result["OK"]:
+            return result
+        jobStates = result["Value"]
+
+        # Get the jobs allowed to transition to the Killed state
+        killJobList.extend(_filterJobStateTransition(jobStates, JobStatus.KILLED))
+
+        if right == RIGHT_DELETE:
+            # Get the jobs allowed to transition to the Deleted state
+            deleteJobList.extend(_filterJobStateTransition(jobStates, JobStatus.DELETED))
+
+        for jobID in killJobList:
+            result = _killJob(jobID, force=force)
+            if not result["OK"]:
+                badIDs.append(jobID)
+
+        for jobID in deleteJobList:
+            result = _deleteJob(jobID, force=force)
+            if not result["OK"]:
+                badIDs.append(jobID)
+
+        # Look for jobs that are in the Staging state to send kill signal to the stager
+        stagingJobList = [jobID for jobID, sDict in jobStates.items() if sDict["Status"] == JobStatus.STAGING]
+
+        if stagingJobList:
+            stagerDB = StorageManagementDB()
+            gLogger.info("Going to send killing signal to stager as well!")
+            result = stagerDB.killTasksBySourceTaskID(stagingJobList)
+            if not result["OK"]:
+                gLogger.warn("Failed to kill some Stager tasks", result["Message"])
+
+    if nonauthJobList or badIDs:
+        result = S_ERROR("Some jobs failed deletion")
+        if nonauthJobList:
+            gLogger.warn("Non-authorized JobIDs won't be deleted", str(nonauthJobList))
+            result["NonauthorizedJobIDs"] = nonauthJobList
+        if badIDs:
+            gLogger.warn("JobIDs failed to be deleted", str(badIDs))
+            result["FailedJobIDs"] = badIDs
+        return result
+
+    jobsList = killJobList if right == RIGHT_KILL else deleteJobList
+    return S_OK(jobsList)

DIRAC/WorkloadManagementSystem/DB/tests/Test_JobDB.py

@@ -28,7 +28,7 @@ def test_getInputData(jobDB: JobDB):
     """Test the getInputData method from JobDB"""
     # Arrange
     jobDB._escapeString = MagicMock(return_value=S_OK())
-    jobDB._query = MagicMock(return_value=S_OK((("/vo/user/lfn1",), ("LFN:/vo/user/lfn2",))))
+    jobDB._query = MagicMock(return_value=S_OK([(1234, "/vo/user/lfn1"), (1234, "LFN:/vo/user/lfn2")]))
 
     # Act
     res = jobDB.getInputData(1234)

DIRAC/WorkloadManagementSystem/DB/tests/Test_StatusUtils.py

@@ -0,0 +1,28 @@
+""" unit test (pytest) of JobAdministration module
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# sut
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
+
+
+@pytest.mark.parametrize(
+    "jobIDs_list, right",
+    [
+        ([], "Kill"),
+        ([], "Delete"),
+        (1, "Kill"),
+        ([1, 2], "Kill"),
+    ],
+)
+def test___kill_delete_jobs(mocker, jobIDs_list, right):
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.JobDB", MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.TaskQueueDB", MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.PilotAgentsDB", MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.StorageManagementDB", MagicMock())
+
+    res = kill_delete_jobs(right, jobIDs_list)
+    assert res["OK"]

DIRAC/WorkloadManagementSystem/Executor/JobSanity.py

@@ -1,6 +1,6 @@
 """ The Job Sanity executor assigns sandboxes to the job """
 from DIRAC import S_OK
-from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient
+from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
 from DIRAC.WorkloadManagementSystem.Executor.Base.OptimizerExecutor import OptimizerExecutor
 
 
@@ -15,7 +15,6 @@ class JobSanity(OptimizerExecutor):
     @classmethod
     def initializeOptimizer(cls):
         """Initialize specific parameters for JobSanityAgent."""
-        cls.sandboxClient = SandboxStoreClient(useCertificates=True, smdb=True)
         return S_OK()
 
     def optimizeJob(self, jid, jobState):
@@ -48,6 +47,7 @@ class JobSanity(OptimizerExecutor):
         ownerName = manifest.getOption("Owner")
         ownerGroup = manifest.getOption("OwnerGroup")
         isbList = manifest.getOption("InputSandbox", [])
+        vo = manifest.getOption("VirtualOrganization")
         sbsToAssign = []
         for isb in isbList:
             if isb.startswith("SB:"):
@@ -55,8 +55,9 @@ class JobSanity(OptimizerExecutor):
         numSBsToAssign = len(sbsToAssign)
         if not numSBsToAssign:
             return S_OK(0)
-        self.jobLog.info("Assigning sandboxes", f"({numSBsToAssign} on behalf of {ownerName}@{ownerGroup})")
-        result = self.sandboxClient.assignSandboxesToJob(jobState.jid, sbsToAssign, ownerName, ownerGroup)
+        self.jobLog.info("Assigning sandboxes", f"({numSBsToAssign} on behalf of {ownerName}@{ownerGroup}@{vo})")
+        eId = f"Job:{jobState.jid}"
+        result = SandboxMetadataDB().assignSandboxesToEntities({eId: sbsToAssign}, ownerName, ownerGroup)
         if not result["OK"]:
             self.jobLog.error("Could not assign sandboxes in the SandboxStore")
             return result

DIRAC/WorkloadManagementSystem/Executor/JobScheduling.py

@@ -549,6 +549,10 @@ class JobScheduling(OptimizerExecutor):
         if not result["OK"]:
             return result
 
+        result = jobState.commitChanges()
+        if not result["OK"]:
+            self.jobLog.error("Could not save changes for job", result["Message"])
+
         return S_OK(stageLFNs)
 
     def __updateSharedSESites(self, jobManifest, stageSite, stagedLFNs, opData):

DIRAC/WorkloadManagementSystem/FutureClient/JobStateUpdateClient.py

@@ -1,11 +1,17 @@
-import importlib
+from __future__ import annotations
+
 import functools
+import time
 from datetime import datetime, timezone
-
+from typing import TYPE_CHECKING
 
 from DIRAC.Core.Security.DiracX import DiracXClient, FutureClient, addRPCStub
-from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue, returnValueOrRaise
 from DIRAC.Core.Utilities.TimeUtilities import fromString
+from DIRAC.WorkloadManagementSystem.Client import JobStatus
+
+if TYPE_CHECKING:
+    from diracx.client.models import JobCommand
 
 
 def stripValueIfOK(func):
@@ -21,30 +27,49 @@ def stripValueIfOK(func):
     def wrapper(*args, **kwargs):
         result = func(*args, **kwargs)
         if result.get("OK"):
-            assert result.pop("Value") is None, "Value should be None if OK"
+            if result.get("Value") is None:
+                result.pop("Value")
         return result
 
     return wrapper
 
 
 class JobStateUpdateClient(FutureClient):
-    @stripValueIfOK
     @convertToReturnValue
     def sendHeartBeat(self, jobID: str | int, dynamicData: dict, staticData: dict):
-        print("HACK: This is a no-op until we decide what to do")
+        """Send a heartbeat from a Job.
+
+        The behaviour of this function is not strictly the same as in legacy
+        DIRAC. Most notably, in legacy DIRAC the heartbeat always overrides the
+        job status to Running whereas in DiracX the job state machine is still
+        respected. Additionally, DiracX updates the job logging information
+        when status transitions occur as a result of the heartbeat.
+        """
+
+        with DiracXClient() as api:
+            body = {jobID: dynamicData | staticData}
+            if len(body[jobID]) != len(dynamicData) + len(staticData):
+                raise NotImplementedError(f"Duplicate keys: {dynamicData=} {staticData=}")
+            commands: list[JobCommand] = api.jobs.add_heartbeat(body)
+        # Legacy DIRAC returns a dictionary of {command: arguments}
+        result = {}
+        for command in commands:
+            if command.job_id != jobID:
+                raise NotImplementedError(f"Job ID mismatch: {jobID=} {command.job_id=}")
+            result[command.command] = command.arguments
+        return result
 
     @stripValueIfOK
     @convertToReturnValue
     def setJobApplicationStatus(self, jobID: str | int, appStatus: str, source: str = "Unknown"):
         statusDict = {
-            "application_status": appStatus,
+            "ApplicationStatus": appStatus,
         }
         if source:
             statusDict["Source"] = source
         with DiracXClient() as api:
-            api.jobs.set_single_job_status(
-                jobID,
-                {datetime.now(tz=timezone.utc): statusDict},
+            api.jobs.set_job_statuses(
+                {jobID: {datetime.now(tz=timezone.utc): statusDict}},
             )
 
     @stripValueIfOK
@@ -52,34 +77,29 @@ class JobStateUpdateClient(FutureClient):
     def setJobAttribute(self, jobID: str | int, attribute: str, value: str):
         with DiracXClient() as api:
             if attribute == "Status":
-                api.jobs.set_single_job_status(
-                    jobID,
-                    {datetime.now(tz=timezone.utc): {"status": value}},
+                return api.jobs.set_job_statuses(
+                    {jobID: {datetime.now(tz=timezone.utc): {"Status": value}}},
                 )
             else:
-                api.jobs.set_single_job_properties(jobID, {attribute: value})
-
-    @stripValueIfOK
-    @convertToReturnValue
-    def setJobFlag(self, jobID: str | int, flag: str):
-        with DiracXClient() as api:
-            api.jobs.set_single_job_properties(jobID, {flag: True})
+                return api.jobs.patch_metadata({jobID: {attribute: value}})
 
     @stripValueIfOK
     @convertToReturnValue
     def setJobParameter(self, jobID: str | int, name: str, value: str):
-        print("HACK: This is a no-op until we decide what to do")
+        with DiracXClient() as api:
+            api.jobs.patch_metadata({jobID: {name: value}})
 
     @stripValueIfOK
     @convertToReturnValue
     def setJobParameters(self, jobID: str | int, parameters: list):
-        print("HACK: This is a no-op until we decide what to do")
+        with DiracXClient() as api:
+            api.jobs.patch_metadata({jobID: {k: v for k, v in parameters}})
 
     @stripValueIfOK
     @convertToReturnValue
     def setJobSite(self, jobID: str | int, site: str):
         with DiracXClient() as api:
-            api.jobs.set_single_job_properties(jobID, {"Site": site})
+            api.jobs.patch_metadata({jobID: {"Site": site}})
 
     @stripValueIfOK
     @convertToReturnValue
@@ -102,9 +122,8 @@ class JobStateUpdateClient(FutureClient):
         if datetime_ is None:
             datetime_ = datetime.utcnow()
         with DiracXClient() as api:
-            api.jobs.set_single_job_status(
-                jobID,
-                {fromString(datetime_).replace(tzinfo=timezone.utc): statusDict},
+            api.jobs.set_job_statuses(
+                {jobID: {fromString(datetime_).replace(tzinfo=timezone.utc): statusDict}},
                 force=force,
             )
 
@@ -114,7 +133,7 @@ class JobStateUpdateClient(FutureClient):
     def setJobStatusBulk(self, jobID: str | int, statusDict: dict, force=False):
         statusDict = {fromString(k).replace(tzinfo=timezone.utc): v for k, v in statusDict.items()}
         with DiracXClient() as api:
-            api.jobs.set_job_status_bulk(
+            api.jobs.set_job_statuses(
                 {jobID: statusDict},
                 force=force,
             )
@@ -122,13 +141,36 @@ class JobStateUpdateClient(FutureClient):
     @stripValueIfOK
     @convertToReturnValue
     def setJobsParameter(self, jobsParameterDict: dict):
-        print("HACK: This is a no-op until we decide what to do")
+        with DiracXClient() as api:
+            updates = {job_id: {k: v} for job_id, (k, v) in jobsParameterDict.items()}
+            api.jobs.patch_metadata(updates)
 
     @stripValueIfOK
     @convertToReturnValue
-    def unsetJobFlag(self, jobID: str | int, flag: str):
-        with DiracXClient() as api:
-            api.jobs.set_single_job_properties(jobID, {flag: False})
-
     def updateJobFromStager(self, jobID: str | int, status: str):
-        raise NotImplementedError("TODO")
+        if status == "Done":
+            jobStatus = JobStatus.CHECKING
+            minorStatus = "JobScheduling"
+        else:
+            jobStatus = None
+            minorStatus = "Staging input files failed"
+
+        trials = 10
+        query = [{"parameter": "JobID", "operator": "eq", "value": jobID}]
+        with DiracXClient() as api:
+            for i in range(trials):
+                result = api.jobs.search(parameters=["Status"], search=query)
+                if not result:
+                    return None
+                if result[0]["Status"] == JobStatus.STAGING:
+                    break
+                time.sleep(1)
+            else:
+                return f"Job is not in Staging after {trials} seconds"
+
+            retVal = self.setJobStatus(jobID, status=jobStatus, minorStatus=minorStatus, source="StagerSystem")
+            # As there might not be a value (see stripValueIfOK), only call
+            # returnValueOrRaise if the return value is not OK
+            if not retVal["OK"]:  # pylint: disable=unsubscriptable-object
+                returnValueOrRaise(retVal)
+            return None if i == 0 else f"Found job in Staging after {i} seconds"