DIRACCommon 9.0.0a66__py3-none-any.whl → 9.0.0a67__py3-none-any.whl

DIRACCommon/WorkloadManagementSystem/Client/JobStatus.py

@@ -0,0 +1,95 @@
+"""
+This module contains constants and lists for the possible job states.
+"""
+
+from DIRACCommon.Core.Utilities.StateMachine import State, StateMachine
+
+#:
+SUBMITTING = "Submitting"
+#:
+RECEIVED = "Received"
+#:
+CHECKING = "Checking"
+#:
+STAGING = "Staging"
+#:
+SCOUTING = "Scouting"
+#:
+WAITING = "Waiting"
+#:
+MATCHED = "Matched"
+#: The Rescheduled status is effectively never stored in the DB.
+#: It could be considered a "virtual" status, and might even be dropped.
+RESCHEDULED = "Rescheduled"
+#:
+RUNNING = "Running"
+#:
+STALLED = "Stalled"
+#:
+COMPLETING = "Completing"
+#:
+DONE = "Done"
+#:
+COMPLETED = "Completed"
+#:
+FAILED = "Failed"
+#:
+DELETED = "Deleted"
+#:
+KILLED = "Killed"
+
+#: Possible job states
+JOB_STATES = [
+    SUBMITTING,
+    RECEIVED,
+    CHECKING,
+    SCOUTING,
+    STAGING,
+    WAITING,
+    MATCHED,
+    RESCHEDULED,
+    RUNNING,
+    STALLED,
+    COMPLETING,
+    DONE,
+    COMPLETED,
+    FAILED,
+    DELETED,
+    KILLED,
+]
+
+# Job States when the payload work has finished
+JOB_FINAL_STATES = [DONE, COMPLETED, FAILED, KILLED]
+
+# WMS internal job States indicating the job object won't be updated
+JOB_REALLY_FINAL_STATES = [DELETED]
+
+
+class JobsStateMachine(StateMachine):
+    """Jobs state machine"""
+
+    def __init__(self, state):
+        """c'tor
+        Defines the state machine transactions
+        """
+        super().__init__(state)
+
+        # States transitions
+        self.states = {
+            DELETED: State(15),  # final state
+            KILLED: State(14, [DELETED], defState=KILLED),
+            FAILED: State(13, [RESCHEDULED, DELETED], defState=FAILED),
+            DONE: State(12, [DELETED], defState=DONE),
+            COMPLETED: State(11, [DONE, FAILED], defState=COMPLETED),
+            COMPLETING: State(10, [DONE, FAILED, COMPLETED, STALLED, KILLED], defState=COMPLETING),
+            STALLED: State(9, [RUNNING, FAILED, KILLED], defState=STALLED),
+            RUNNING: State(8, [STALLED, DONE, FAILED, RESCHEDULED, COMPLETING, KILLED, RECEIVED], defState=RUNNING),
+            RESCHEDULED: State(7, [WAITING, RECEIVED, DELETED, FAILED, KILLED], defState=RESCHEDULED),
+            MATCHED: State(6, [RUNNING, FAILED, RESCHEDULED, KILLED], defState=MATCHED),
+            WAITING: State(5, [MATCHED, RESCHEDULED, DELETED, KILLED], defState=WAITING),
+            STAGING: State(4, [CHECKING, WAITING, FAILED, KILLED], defState=STAGING),
+            SCOUTING: State(3, [CHECKING, FAILED, STALLED, KILLED], defState=SCOUTING),
+            CHECKING: State(2, [SCOUTING, STAGING, WAITING, RESCHEDULED, FAILED, DELETED, KILLED], defState=CHECKING),
+            RECEIVED: State(1, [SCOUTING, CHECKING, STAGING, WAITING, FAILED, DELETED, KILLED], defState=RECEIVED),
+            SUBMITTING: State(0, [RECEIVED, CHECKING, DELETED, KILLED], defState=SUBMITTING),  # initial state
+        }

DIRACCommon/WorkloadManagementSystem/Client/__init__.py

@@ -0,0 +1 @@
+"""DIRACCommon WorkloadManagementSystem client utilities"""

DIRACCommon/WorkloadManagementSystem/DB/JobDBUtils.py

@@ -0,0 +1,170 @@
+"""Stateless JobDB utilities extracted from DIRAC for DIRACCommon"""
+
+from __future__ import annotations
+
+import base64
+import zlib
+from typing import TypedDict
+
+from DIRACCommon.Core.Utilities.ReturnValues import S_OK, DOKReturnType, S_ERROR
+from DIRACCommon.Core.Utilities.DErrno import EWMSSUBM
+from DIRACCommon.WorkloadManagementSystem.Client import JobStatus
+from DIRACCommon.WorkloadManagementSystem.Client.JobState.JobManifest import JobManifest, JobManifestConfig
+
+
+def compressJDL(jdl):
+    """Return compressed JDL string."""
+    return base64.b64encode(zlib.compress(jdl.encode(), -1)).decode()
+
+
+def extractJDL(compressedJDL):
+    """Return decompressed JDL string."""
+    # the starting bracket is guaranteeed by JobManager.submitJob
+    # we need the check to be backward compatible
+    if isinstance(compressedJDL, bytes):
+        if compressedJDL.startswith(b"["):
+            return compressedJDL.decode()
+    else:
+        if compressedJDL.startswith("["):
+            return compressedJDL
+    return zlib.decompress(base64.b64decode(compressedJDL)).decode()
+
+
+def fixJDL(jdl: str) -> str:
+    """Fix possible lack of brackets in JDL"""
+    # 1.- insert original JDL on DB and get new JobID
+    # Fix the possible lack of the brackets in the JDL
+    if jdl.strip()[0].find("[") != 0:
+        jdl = "[" + jdl + "]"
+    return jdl
+
+
+class CheckAndPrepareJobConfig(TypedDict):
+    """Dictionary type for defining the information checkAndPrepareJob needs from the CS"""
+
+    inputDataPolicyForVO: str
+    softwareDistModuleForVO: str
+    defaultCPUTimeForOwnerGroup: int
+    getDIRACPlatform: callable[list[str], DOKReturnType[list[str]]]
+
+
+def checkAndPrepareJob(
+    jobID, classAdJob, classAdReq, owner, ownerGroup, jobAttrs, vo, *, config: CheckAndPrepareJobConfig
+):
+    error = ""
+
+    jdlOwner = classAdJob.getAttributeString("Owner")
+    jdlOwnerGroup = classAdJob.getAttributeString("OwnerGroup")
+    jdlVO = classAdJob.getAttributeString("VirtualOrganization")
+
+    # The below is commented out since this is always overwritten by the submitter IDs
+    # but the check allows to findout inconsistent client environments
+    if jdlOwner and jdlOwner != owner:
+        error = "Wrong Owner in JDL"
+    elif jdlOwnerGroup and jdlOwnerGroup != ownerGroup:
+        error = "Wrong Owner Group in JDL"
+    elif jdlVO and jdlVO != vo:
+        error = "Wrong Virtual Organization in JDL"
+
+    classAdJob.insertAttributeString("Owner", owner)
+    classAdJob.insertAttributeString("OwnerGroup", ownerGroup)
+
+    if vo:
+        classAdJob.insertAttributeString("VirtualOrganization", vo)
+
+    classAdReq.insertAttributeString("Owner", owner)
+    classAdReq.insertAttributeString("OwnerGroup", ownerGroup)
+    if vo:
+        classAdReq.insertAttributeString("VirtualOrganization", vo)
+
+    if config["inputDataPolicyForVO"] and not classAdJob.lookupAttribute("InputDataModule"):
+        classAdJob.insertAttributeString("InputDataModule", config["inputDataPolicyForVO"])
+
+    if config["softwareDistModuleForVO"] and not classAdJob.lookupAttribute("SoftwareDistModule"):
+        classAdJob.insertAttributeString("SoftwareDistModule", config["softwareDistModuleForVO"])
+
+    # priority
+    priority = classAdJob.getAttributeInt("Priority")
+    if priority is None:
+        priority = 0
+    classAdReq.insertAttributeInt("UserPriority", priority)
+
+    # CPU time
+    cpuTime = classAdJob.getAttributeInt("CPUTime")
+    if cpuTime is None:
+        cpuTime = config["defaultCPUTimeForOwnerGroup"]
+    classAdReq.insertAttributeInt("CPUTime", cpuTime)
+
+    # platform(s)
+    platformList = classAdJob.getListFromExpression("Platform")
+    if platformList:
+        result = config["getDIRACPlatform"](platformList)
+        if not result["OK"]:
+            return result
+        if result["Value"]:
+            classAdReq.insertAttributeVectorString("Platforms", result["Value"])
+        else:
+            error = "OS compatibility info not found"
+    if error:
+        retVal = S_ERROR(EWMSSUBM, error)
+        retVal["JobId"] = jobID
+        retVal["Status"] = JobStatus.FAILED
+        retVal["MinorStatus"] = error
+
+        jobAttrs["Status"] = JobStatus.FAILED
+
+        jobAttrs["MinorStatus"] = error
+        return retVal
+    return S_OK()
+
+
+def checkAndAddOwner(
+    jdl: str, owner: str, ownerGroup: str, *, job_manifest_config: JobManifestConfig
+) -> DOKReturnType[JobManifest]:
+    jobManifest = JobManifest()
+    res = jobManifest.load(jdl)
+    if not res["OK"]:
+        return res
+
+    jobManifest.setOptionsFromDict({"Owner": owner, "OwnerGroup": ownerGroup})
+    res = jobManifest.check(config=job_manifest_config)
+    if not res["OK"]:
+        return res
+
+    return S_OK(jobManifest)
+
+
+def createJDLWithInitialStatus(
+    classAdJob, classAdReq, jdl2DBParameters, jobAttrs, initialStatus, initialMinorStatus, *, modern=False
+):
+    """
+    :param modern: if True, store boolean instead of string for VerifiedFlag (used by diracx only)
+    """
+    priority = classAdJob.getAttributeInt("Priority")
+    if priority is None:
+        priority = 0
+    jobAttrs["UserPriority"] = priority
+
+    for jdlName in jdl2DBParameters:
+        # Defaults are set by the DB.
+        jdlValue = classAdJob.getAttributeString(jdlName)
+        if jdlValue:
+            jobAttrs[jdlName] = jdlValue
+
+    jdlValue = classAdJob.getAttributeString("Site")
+    if jdlValue:
+        if jdlValue.find(",") != -1:
+            jobAttrs["Site"] = "Multiple"
+        else:
+            jobAttrs["Site"] = jdlValue
+
+    jobAttrs["VerifiedFlag"] = True if modern else "True"
+
+    jobAttrs["Status"] = initialStatus
+
+    jobAttrs["MinorStatus"] = initialMinorStatus
+
+    reqJDL = classAdReq.asJDL()
+    classAdJob.insertAttributeInt("JobRequirements", reqJDL)
+
+    return classAdJob.asJDL()

DIRACCommon/WorkloadManagementSystem/DB/__init__.py

@@ -0,0 +1 @@
+"""DIRACCommon WorkloadManagementSystem DB utilities"""

DIRACCommon/WorkloadManagementSystem/Utilities/JobModel.py

@@ -0,0 +1,236 @@
+""" This module contains the JobModel class, which is used to validate the job description """
+
+# pylint: disable=no-self-argument, no-self-use, invalid-name, missing-function-docstring
+
+from collections.abc import Iterable
+from typing import Annotated, Any, Callable, ClassVar, Self, TypeAlias, TypedDict
+
+from pydantic import BaseModel, BeforeValidator, ConfigDict, field_validator, model_validator
+
+from DIRACCommon.Core.Utilities.ReturnValues import DErrorReturnType
+
+
+# HACK: Convert appropriate iterables into sets
+def default_set_validator(value):
+    if value is None:
+        return set()
+    elif not isinstance(value, Iterable):
+        return value
+    elif isinstance(value, (str, bytes, bytearray)):
+        return value
+    else:
+        return set(value)
+
+
+CoercibleSetStr: TypeAlias = Annotated[set[str], BeforeValidator(default_set_validator)]
+
+
+class BaseJobDescriptionModelConfg(TypedDict):
+    """Dictionary type for defining the information JobDescriptionModel needs from the CS"""
+
+    # Default values
+    cpuTime: int
+    priority: int
+    # Bounds
+    minCPUTime: int
+    maxCPUTime: int
+    allowedJobTypes: list[str]
+    maxInputDataFiles: int
+    minNumberOfProcessors: int
+    maxNumberOfProcessors: int
+    minPriority: int
+    maxPriority: int
+    possibleLogLevels: list[str]
+    sites: DErrorReturnType[list[str]]
+
+
+class BaseJobDescriptionModel(BaseModel):
+    """Base model for the job description (not parametric)"""
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    # This must be overridden in subclasses
+    _config_builder: ClassVar[Callable[[], BaseJobDescriptionModelConfg] | None] = None
+
+    @model_validator(mode="before")
+    def injectDefaultValues(cls, values: dict[str, Any]) -> dict[str, Any]:
+        if cls._config_builder is None:
+            raise NotImplementedError("You must define a _config_builder class attribute")
+        config = cls._config_builder()
+        values.setdefault("cpuTime", config["cpuTime"])
+        values.setdefault("priority", config["priority"])
+        return values
+
+    arguments: str = ""
+    bannedSites: CoercibleSetStr = set()
+    # TODO: This should use a field factory
+    cpuTime: int
+    executable: str
+    executionEnvironment: dict = None
+    gridCE: str = ""
+    inputSandbox: CoercibleSetStr = set()
+    inputData: CoercibleSetStr = set()
+    inputDataPolicy: str = ""
+    jobConfigArgs: str = ""
+    jobGroup: str = ""
+    jobType: str = "User"
+    jobName: str = "Name"
+    # TODO: This should be an StrEnum
+    logLevel: str = "INFO"
+    # TODO: This can't be None with this type hint
+    maxNumberOfProcessors: int = None
+    minNumberOfProcessors: int = 1
+    outputData: CoercibleSetStr = set()
+    outputPath: str = ""
+    outputSandbox: CoercibleSetStr = set()
+    outputSE: str = ""
+    platform: str = ""
+    # TODO: This should use a field factory
+    priority: int
+    sites: CoercibleSetStr = set()
+    stderr: str = "std.err"
+    stdout: str = "std.out"
+    tags: CoercibleSetStr = set()
+    extraFields: dict[str, Any] = {}
+
+    @field_validator("cpuTime")
+    def checkCPUTimeBounds(cls, v):
+        minCPUTime = cls._config_builder()["minCPUTime"]
+        maxCPUTime = cls._config_builder()["maxCPUTime"]
+        if not minCPUTime <= v <= maxCPUTime:
+            raise ValueError(f"cpuTime out of bounds (must be between {minCPUTime} and {maxCPUTime})")
+        return v
+
+    @field_validator("executable")
+    def checkExecutableIsNotAnEmptyString(cls, v: str):
+        if not v:
+            raise ValueError("executable must not be an empty string")
+        return v
+
+    @field_validator("jobType")
+    def checkJobTypeIsAllowed(cls, v: str):
+        allowedTypes = cls._config_builder()["allowedJobTypes"]
+        if v not in allowedTypes:
+            raise ValueError(f"jobType '{v}' is not allowed for this kind of user (must be in {allowedTypes})")
+        return v
+
+    @field_validator("inputData")
+    def checkInputDataDoesntContainDoubleSlashes(cls, v):
+        if v:
+            for lfn in v:
+                if lfn.find("//") > -1:
+                    raise ValueError("Input data contains //")
+        return v
+
+    @field_validator("inputData")
+    def addLFNPrefixIfStringStartsWithASlash(cls, v: set[str]):
+        if v:
+            v = {lfn.strip() for lfn in v if lfn.strip()}
+            v = {f"LFN:{lfn}" if lfn.startswith("/") else lfn for lfn in v}
+
+            for lfn in v:
+                if not lfn.startswith("LFN:/"):
+                    raise ValueError("Input data files must start with LFN:/")
+        return v
+
+    @model_validator(mode="after")
+    def checkNumberOfInputDataFiles(self) -> Self:
+        if self.inputData:
+            maxInputDataFiles = self._config_builder()["maxInputDataFiles"]
+            if self.jobType == "User" and len(self.inputData) >= maxInputDataFiles:
+                raise ValueError(f"inputData contains too many files (must contain at most {maxInputDataFiles})")
+        return self
+
+    @field_validator("inputSandbox")
+    def checkLFNSandboxesAreWellFormated(cls, v: set[str]):
+        for inputSandbox in v:
+            if inputSandbox.startswith("LFN:") and not inputSandbox.startswith("LFN:/"):
+                raise ValueError("LFN files must start by LFN:/")
+        return v
+
+    @field_validator("logLevel")
+    def checkLogLevelIsValid(cls, v: str):
+        v = v.upper()
+        possibleLogLevels = cls._config_builder()["possibleLogLevels"]
+        if v not in possibleLogLevels:
+            raise ValueError(f"Log level {v} not in {possibleLogLevels}")
+        return v
+
+    @field_validator("minNumberOfProcessors")
+    def checkMinNumberOfProcessorsBounds(cls, v):
+        minNumberOfProcessors = cls._config_builder()["minNumberOfProcessors"]
+        maxNumberOfProcessors = cls._config_builder()["maxNumberOfProcessors"]
+        if not minNumberOfProcessors <= v <= maxNumberOfProcessors:
+            raise ValueError(
+                f"minNumberOfProcessors out of bounds (must be between {minNumberOfProcessors} and {maxNumberOfProcessors})"
+            )
+        return v
+
+    @field_validator("maxNumberOfProcessors")
+    def checkMaxNumberOfProcessorsBounds(cls, v):
+        minNumberOfProcessors = cls._config_builder()["minNumberOfProcessors"]
+        maxNumberOfProcessors = cls._config_builder()["maxNumberOfProcessors"]
+        if not minNumberOfProcessors <= v <= maxNumberOfProcessors:
+            raise ValueError(
+                f"maxNumberOfProcessors out of bounds (must be between {minNumberOfProcessors} and {maxNumberOfProcessors})"
+            )
+        return v
+
+    @model_validator(mode="after")
+    def checkThatMaxNumberOfProcessorsIsGreaterThanMinNumberOfProcessors(self) -> Self:
+        if self.maxNumberOfProcessors:
+            if self.maxNumberOfProcessors < self.minNumberOfProcessors:
+                raise ValueError("maxNumberOfProcessors must be greater than minNumberOfProcessors")
+        return self
+
+    @model_validator(mode="after")
+    def addTagsDependingOnNumberOfProcessors(self) -> Self:
+        if self.minNumberOfProcessors == self.maxNumberOfProcessors:
+            self.tags.add(f"{self.minNumberOfProcessors}Processors")
+        if self.minNumberOfProcessors > 1:
+            self.tags.add("MultiProcessor")
+        return self
+
+    @field_validator("sites")
+    def checkSites(cls, v: set[str]):
+        if v:
+            res = cls._config_builder()["sites"]
+            if not res["OK"]:
+                raise ValueError(res["Message"])
+            invalidSites = v - set(res["Value"]).union({"ANY"})
+            if invalidSites:
+                raise ValueError(f"Invalid sites: {' '.join(invalidSites)}")
+        return v
+
+    @model_validator(mode="after")
+    def checkThatSitesAndBannedSitesAreNotMutuallyExclusive(self) -> Self:
+        if self.sites and self.bannedSites:
+            while self.bannedSites:
+                self.sites.discard(self.bannedSites.pop())
+            if not self.sites:
+                raise ValueError("sites and bannedSites are mutually exclusive")
+        return self
+
+    @field_validator("priority")
+    def checkPriorityBounds(cls, v):
+        minPriority = cls._config_builder()["minPriority"]
+        maxPriority = cls._config_builder()["maxPriority"]
+        if not minPriority <= v <= maxPriority:
+            raise ValueError(f"priority out of bounds (must be between {minPriority} and {maxPriority})")
+        return v
+
+
+class JobDescriptionModel(BaseJobDescriptionModel):
+    """Model for the job description (non parametric job with user credentials, i.e server side)"""
+
+    owner: str
+    ownerGroup: str
+    vo: str
+
+    @model_validator(mode="after")
+    def checkLFNMatchesREGEX(self) -> Self:
+        if self.inputData:
+            for lfn in self.inputData:
+                if not lfn.startswith(f"LFN:/{self.vo}/"):
+                    raise ValueError(f"Input data not correctly specified (must start with LFN:/{self.vo}/)")
+        return self

DIRACCommon/WorkloadManagementSystem/Utilities/JobStatusUtility.py

@@ -0,0 +1,93 @@
+"""Stateless job status utility functions"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+from DIRACCommon.Core.Utilities.ReturnValues import S_OK
+from DIRACCommon.Core.Utilities.TimeUtilities import toEpoch, fromString
+from DIRACCommon.WorkloadManagementSystem.Client.JobStatus import RUNNING, JOB_FINAL_STATES, JobsStateMachine
+
+
+def getStartAndEndTime(startTime, endTime, updateTimes, timeStamps, statusDict):
+    """Get start and end times from job status updates
+
+    :param startTime: current start time
+    :param endTime: current end time
+    :param updateTimes: list of update times
+    :param timeStamps: list of (timestamp, status) tuples
+    :param statusDict: dictionary mapping update times to status dictionaries
+    :return: tuple of (newStartTime, newEndTime)
+    """
+    newStat = ""
+    firstUpdate = toEpoch(fromString(updateTimes[0]))
+    for ts, st in timeStamps:
+        if firstUpdate >= ts:
+            newStat = st
+    # Pick up start and end times from all updates
+    for updTime in updateTimes:
+        sDict = statusDict[updTime]
+        newStat = sDict.get("Status", newStat)
+
+        if not startTime and newStat == RUNNING:
+            # Pick up the start date when the job starts running if not existing
+            startTime = updTime
+        elif not endTime and newStat in JOB_FINAL_STATES:
+            # Pick up the end time when the job is in a final status
+            endTime = updTime
+
+    return startTime, endTime
+
+
+def getNewStatus(
+    jobID: int,
+    updateTimes: list[datetime],
+    lastTime: datetime,
+    statusDict: dict[datetime, Any],
+    currentStatus,
+    force: bool,
+    log,
+):
+    """Get new job status from status updates
+
+    :param jobID: job ID
+    :param updateTimes: list of update times
+    :param lastTime: last update time
+    :param statusDict: dictionary mapping update times to status dictionaries
+    :param currentStatus: current job status
+    :param force: whether to force status update without state machine validation
+    :param log: logger object
+    :return: S_OK((status, minor, application)) or S_ERROR
+    """
+    status = ""
+    minor = ""
+    application = ""
+    # Get the last status values looping on the most recent upupdateTimes in chronological order
+    for updTime in [dt for dt in updateTimes if dt >= lastTime]:
+        sDict = statusDict[updTime]
+        log.debug(f"\tTime {updTime} - Statuses {str(sDict)}")
+        status = sDict.get("Status", currentStatus)
+        # evaluate the state machine if the status is changing
+        if not force and status != currentStatus:
+            res = JobsStateMachine(currentStatus).getNextState(status)
+            if not res["OK"]:
+                return res
+            newStat = res["Value"]
+            # If the JobsStateMachine does not accept the candidate, don't update
+            if newStat != status:
+                # keeping the same status
+                log.error(
+                    f"Job Status Error: {jobID} can't move from {currentStatus} to {status}: using {newStat}",
+                )
+                status = newStat
+                sDict["Status"] = newStat
+                # Change the source to indicate this is not what was requested
+                source = sDict.get("Source", "")
+                sDict["Source"] = source + "(SM)"
+            # at this stage status == newStat. Set currentStatus to this new status
+            currentStatus = newStat
+
+        minor = sDict.get("MinorStatus", minor)
+        application = sDict.get("ApplicationStatus", application)
+    return S_OK((status, minor, application))