qq_lib.batch.slurm

 30class SlurmJob(BatchJobInterface):
 31    """
 32    Implementation of BatchJobInterface for Slurm.
 33    Stores metadata for a single Slurm job.
 34    """
 35
 36    # converts from Slurm state names to qq BatchStates
 37    _STATE_CONVERTER: dict[str, BatchState] = {
 38        "BOOT_FAIL": BatchState.FAILED,
 39        "CANCELLED": BatchState.FAILED,
 40        "COMPLETED": BatchState.FINISHED,
 41        "DEADLINE": BatchState.FAILED,
 42        "FAILED": BatchState.FAILED,
 43        "NODE_FAIL": BatchState.FAILED,
 44        "OUT_OF_MEMORY": BatchState.FAILED,
 45        "PENDING": BatchState.QUEUED,
 46        "PREEMPTED": BatchState.SUSPENDED,
 47        "RUNNING": BatchState.RUNNING,
 48        "SUSPENDED": BatchState.SUSPENDED,
 49        "TIMEOUT": BatchState.FAILED,
 50    }
 51
 52    def __init__(self, job_id: str):
 53        """Query the batch system for information about the job with the specified ID."""
 54        self._job_id = job_id
 55        self._info: dict[str, str] = {}
 56
 57        self.update()
 58
 59    def is_empty(self) -> bool:
 60        return not self._info
 61
 62    def get_id(self) -> str:
 63        return self._job_id
 64
 65    def get_account(self) -> str | None:
 66        return self._info.get("Account")
 67
 68    def update(self) -> None:
 69        # first try `scontrol`
 70        command = f"scontrol show job {self._job_id} -o"
 71        logger.debug(command)
 72
 73        result = subprocess.run(
 74            ["bash"],
 75            input=command,
 76            text=True,
 77            check=False,
 78            capture_output=True,
 79            errors="replace",
 80        )
 81
 82        if result.returncode != 0:
 83            # if scontrol fails, try sacct
 84            logger.debug(
 85                f"scontrol failed for job '{self._job_id}' ({result.stderr.strip()}); trying sacct"
 86            )
 87        else:
 88            self._info: dict[str, str] = parse_slurm_dump_to_dictionary(result.stdout)
 89            return
 90
 91        # if `scontrol` fails, try `sacct`
 92        command = f"sacct --allocations --noheader --parsable2 -j {self._job_id} --format={SACCT_FIELDS} "
 93        logger.debug(command)
 94
 95        result = subprocess.run(
 96            ["bash"],
 97            input=command,
 98            text=True,
 99            check=False,
100            capture_output=True,
101            errors="replace",
102        )
103
104        if result.returncode != 0:
105            # if sacct fails, information is empty
106            logger.debug(
107                f"both scontrol and sacct failed: no information about job '{self._job_id}' is available: {result.stderr.strip()}"
108            )
109            self._info: dict[str, str] = {}
110        else:
111            job: SlurmJob = SlurmJob.from_sacct_string(result.stdout.strip())
112            self._info: dict[str, str] = job._info
113
114    def get_state(self) -> BatchState:
115        if not (raw_state := self._info.get("JobState")):
116            return BatchState.UNKNOWN
117
118        converted_state = SlurmJob._STATE_CONVERTER.get(raw_state) or BatchState.UNKNOWN
119
120        # if the job is queued due to depending on another job, it should be considered "held"
121        if (
122            converted_state == BatchState.QUEUED
123            and (comment := self.get_comment())
124            and "Dependency" in comment
125        ):
126            return BatchState.HELD
127
128        return converted_state
129
130    def get_comment(self) -> str | None:
131        if (reason := self._info.get("Reason")) and reason != "None":
132            return f"Reason: {reason}"
133
134        return None
135
136    def get_estimated(self) -> tuple[datetime, str] | None:
137        # use "StartTime" as an estimate
138        if not (time := self.get_start_time()) or time == "None":
139            return None
140
141        if not (node_list := self._info.get("SchedNodeList")) or "None" in node_list:
142            return None
143
144        return (time, node_list)
145
146    def get_main_node(self) -> str | None:
147        if (main_node := self._info.get("BatchHost")) and "None" not in main_node:
148            return main_node
149
150        # if BatchHost does not exist, use the first node from NodeList
151        if nodes := self.get_nodes():
152            return nodes[0]
153
154        return None
155
156    def get_nodes(self) -> list[str] | None:
157        if (node_list := self._info.get("NodeList")) and "None" not in node_list:
158            return SlurmJob._expand_node_list(node_list)
159
160        return None
161
162    def get_short_nodes(self) -> list[str] | None:
163        # treat all nodes a single node, without expanding
164        # this assumes that get_short_nodes is only used in qq jobs and qq stat
165        if (node_list := self._info.get("NodeList")) and "None" not in node_list:
166            return [node_list]
167
168        return None
169
170    def get_name(self) -> str | None:
171        if not (name := self._info.get("JobName")):
172            logger.debug(f"Could not get job name for '{self._job_id}'.")
173            return None
174
175        return name
176
177    def get_n_cpus(self) -> int | None:
178        min_cpus = (
179            self._get_int_property("MinCPUsNode", "the minimum number of CPUs per node")
180            or 0
181        ) * (self.get_n_nodes() or 0)
182
183        if not (cpus := self._get_int_property("NumCPUs", "the number of CPUs")):
184            return None
185
186        return max(min_cpus, cpus)
187
188    def get_n_gpus(self) -> int | None:
189        tres = self._get_tres()
190        for item in tres.split(","):
191            if item.startswith("gpu") or item.startswith("gres/gpu"):
192                try:
193                    return int(item.split("=")[1])
194                except ValueError as e:
195                    logger.warning(
196                        f"Could not parse the number of GPUs from '{item}': {e}."
197                    )
198                    return None
199
200        return None
201
202    def get_n_nodes(self) -> int | None:
203        return self._get_int_property("NumNodes", "the number of nodes")
204
205    def get_mem(self) -> Size | None:
206        tres = self._get_tres()
207        for item in tres.split(","):
208            if item.startswith("mem="):
209                try:
210                    return Size.from_string(item.split("=", 1)[1])
211                except Exception as e:
212                    logger.warning(f"Could not parse memory for '{self._job_id}': {e}.")
213                    return None
214
215        logger.debug(f"Memory not available for '{self._job_id}'.")
216        return None
217
218    def get_start_time(self) -> datetime | None:
219        return self._get_datetime_property("StartTime", "the job start time")
220
221    def get_submission_time(self) -> datetime | None:
222        return self._get_datetime_property("SubmitTime", "the job submission time")
223
224    def get_completion_time(self) -> datetime | None:
225        # the property EndTime is available for running jobs as well (estimated completion time)
226        # but that should not matter for our purposes
227        return self._get_datetime_property("EndTime", "the job completion time")
228
229    def get_modification_time(self) -> datetime | None:
230        # assuming this is only used for completed jobs
231        return self.get_completion_time() or self.get_submission_time()
232
233    def get_user(self) -> str | None:
234        if not (user := self._info.get("UserId")):
235            logger.debug(f"Could not get user for '{self._job_id}'.")
236            return None
237
238        return user.split("(")[0]
239
240    def get_walltime(self) -> timedelta | None:
241        if not (walltime := self._info.get("TimeLimit")):
242            logger.debug(f"Could not get walltime for '{self._job_id}'.")
243            return None
244
245        try:
246            return dhhmmss_to_duration(walltime)
247        except QQError as e:
248            logger.warning(f"Could not parse walltime for '{self._job_id}': {e}.")
249            return None
250
251    def get_queue(self) -> str | None:
252        if not (queue := self._info.get("Partition")):
253            logger.debug(f"Could not get queue for '{self._job_id}'.")
254            return None
255
256        return queue
257
258    def get_util_cpu(self) -> int | None:
259        # not available in Slurm
260        return None
261
262    def get_util_mem(self) -> int | None:
263        # not available in Slurm
264        return None
265
266    def get_exit_code(self) -> int | None:
267        if not (raw_exit := self._info.get("ExitCode")):
268            return None
269
270        try:
271            # Slurm reports two exit codes; the first one is exit code of the script
272            # the second one is a signal
273            # we return the first non-zero exit code or 0 if both exit codes are 0
274            code, signal = map(int, raw_exit.split(":"))
275            return code or signal
276        except Exception as e:
277            logger.debug(f"Could not parse exit codes '{raw_exit}': {e}.")
278            return None
279
280    def get_input_machine(self) -> str | None:
281        # not available for Slurm
282        return None
283
284    def get_input_dir(self) -> Path | None:
285        if not (raw_dir := self._info.get("WorkDir")):
286            logger.debug(f"Could not obtain input directory for '{self._job_id}'.")
287            return None
288
289        return logical_resolve(Path(raw_dir))
290
291    def get_info_file(self) -> Path | None:
292        if not (input_dir := self.get_input_dir()) or not (name := self.get_name()):
293            return None
294
295        info_file = (input_dir / name).with_suffix(CFG.suffixes.qq_info)
296
297        # we need to check whether the info file actually exists
298        # (or rather if it is available to the user)
299        try:
300            if not info_file.is_file():
301                return None
302        except PermissionError:
303            return None
304
305        return info_file
306
307    def to_yaml(self) -> str:
308        return yaml.dump(
309            self._info, default_flow_style=False, sort_keys=False, Dumper=Dumper
310        )
311
312    def get_steps(self) -> Sequence[Self]:
313        command = f"sacct -j {self._job_id} --parsable2 --format={SACCT_STEP_FIELDS}"
314        logger.debug(command)
315
316        result = subprocess.run(
317            ["bash"],
318            input=command,
319            text=True,
320            check=False,
321            capture_output=True,
322            errors="replace",
323        )
324
325        if result.returncode != 0:
326            logger.debug(f"Could not get steps for a job '{self._job_id}'.")
327            return []
328
329        jobs = []
330        for sacct_string in result.stdout.split("\n"):
331            if sacct_string.strip() == "":
332                continue
333
334            job = SlurmJob._step_from_sacct_string(sacct_string)
335            # only consider job steps with numeric indices
336            if (step_id := job.get_step_id()) and step_id.isnumeric():
337                jobs.append(job)
338
339        return jobs
340
341    def get_step_id(self) -> str | None:
342        try:
343            (_, step) = self._job_id.split(".", maxsplit=1)
344            return step
345        except ValueError:
346            return None
347
348    def is_array_job(self) -> bool:
349        return False
350
351    @classmethod
352    def from_dict(cls, job_id: str, info: dict[str, str]) -> Self:
353        """
354        Construct a new instance of SlurmJob from a job ID and a dictionary of job information.
355
356        This method bypasses the standard initializer and directly sets the `_job_id` and `_info`
357        attributes of the new instance.
358
359        Args:
360            job_id (str): The unique identifier of the job.
361            info (dict[str, str]): A dictionary containing Slurm job metadata as key-value pairs.
362
363        Returns:
364            Self: A new instance of SlurmJob.
365
366        Note:
367            This method does not perform any validation or processing of the provided dictionary.
368        """
369        job_info = cls.__new__(cls)
370        job_info._job_id = job_id
371        job_info._info = info
372
373        return job_info
374
375    @classmethod
376    def from_sacct_string(cls, string: str) -> Self:
377        """
378        Construct a new instance of SlurmJob using a string from sacct.
379
380        Args:
381            string (str): String describing the job properties obtained using sacct.
382
383        Returns:
384            Self: A new instance of SlurmJob.
385        """
386        fields: list[str] = [
387            "JobId",
388            "Account",
389            "JobState",
390            "UserId",
391            "JobName",
392            "Partition",
393            "WorkDir",
394            "AllocCPUs",
395            "ReqCPUs",
396            "AllocTRES",
397            "ReqTRES",
398            "AllocNodes",
399            "ReqNodes",
400            "SubmitTime",
401            "StartTime",
402            "EndTime",
403            "TimeLimit",
404            "NodeList",
405            "Reason",
406            "ExitCode",
407        ]
408
409        split = string.split("|")
410        if len(fields) != len(split):
411            raise QQError(
412                f"Number of items in a sacct string '{string}' ('{len(split)}') does not match the expected number of items ('{len(fields)}'). This is a bug, please report it!"
413            )
414
415        info: dict[str, str] = dict(zip(fields, split))
416
417        # only take the first word from JobState
418        # other words may contain useless additional information
419        info["JobState"] = info["JobState"].split()[0]
420
421        SlurmJob._assign_if_allocated(info, "AllocCPUs", "ReqCPUs", "NumCPUs")
422        SlurmJob._assign_if_allocated(info, "AllocNodes", "ReqNodes", "NumNodes")
423
424        return cls.from_dict(info["JobId"], info)
425
426    @classmethod
427    def _step_from_sacct_string(cls, string: str) -> Self:
428        """
429        Construct a new instance of SlurmJob step using a string from sacct.
430
431        Args:
432            string (str): String describing the job properties obtained using sacct.
433
434        Returns:
435            Self: A new instance of SlurmJob for a job step.
436        """
437        fields: list[str] = [
438            "JobId",
439            "JobState",
440            "StartTime",
441            "EndTime",
442        ]
443
444        split = string.split("|")
445        if len(fields) != len(split):
446            raise QQError(
447                f"Number of items in a sacct string for a slurm step '{string}' ('{len(split)}') does not match the expected number of items ('{len(fields)}'). This is a bug, please report it!"
448            )
449
450        info: dict[str, str] = dict(zip(fields, split))
451
452        # only take the first word from JobState
453        # other words may contain useless additional information
454        info["JobState"] = info["JobState"].split()[0]
455
456        return cls.from_dict(info["JobId"], info)
457
458    def get_ids_for_sorting(self) -> list[int]:
459        """
460        Extract numeric components of the job ID for sorting.
461
462        The method retrieves the leading numeric portion of the job ID, which may
463        contain multiple integer groups separated by underscores. Parsing stops
464        when a non-digit and non-underscore character is encountered.
465
466        Returns:
467            list[int]: A list of integer components extracted from the job ID,
468                or [0] if no valid numeric portion is found.
469        """
470        # get the numerical portion of the job ID (may contain underscores)
471        match = re.match(r"(\d+(?:_\d+)*)", self.get_id())
472        if not match:
473            return [0]
474
475        # split the matched portion into digit groups
476        groups = match.group(1).split("_")
477        return [int(g) for g in groups]
478
479    @staticmethod
480    def _expand_node_list(compact: str) -> list[str]:
481        """
482        Expand a compact Slurm node list expression into individual hostnames.
483
484        This method uses the Slurm `scontrol show hostnames` command to translate
485        a compact node list (e.g., "node[01-03]") into an explicit list of node names.
486        If the expansion fails, the original compact string is returned as a single-element list.
487
488        Args:
489            compact (str): The compact Slurm node list expression to expand.
490
491        Returns:
492            list[str]: A list of fully expanded node hostnames. If expansion fails,
493                returns a list containing the original input string.
494        """
495        command = f"scontrol show hostnames {compact}"
496        logger.debug(command)
497
498        result = subprocess.run(
499            ["bash"],
500            input=command,
501            text=True,
502            check=False,
503            capture_output=True,
504            errors="replace",
505        )
506
507        if result.returncode != 0:
508            logger.warning(
509                f"Could not expand '{compact}' into a list of nodes: {result.stderr.strip()}"
510            )
511            # use unexpanded string
512            return [compact]
513
514        return result.stdout.strip().split("\n")
515
516    def _get_int_property(self, property: str, property_name: str) -> int | None:
517        """
518        Retrieve an integer property value from the job information.
519
520        If the property contains a range (e.g., "MIN-MAX"), only the minimum value
521        is returned. If the property cannot be retrieved or converted to an integer,
522        `None` is returned.
523
524        Args:
525            property (str): The key identifying the property in the job information.
526            property_name (str): A human-readable name of the property for logging.
527
528        Returns:
529            int: The integer value of the property, or `None` if unavailable or invalid.
530        """
531        try:
532            # we split by '-' because pending jobs may have this property shown as MIN-MAX
533            # we show the value of the minimum
534            return int(self._info[property].split("-")[0])
535        except Exception:
536            logger.debug(
537                f"Could not get information about {property_name} from the batch system for '{self._job_id}'."
538            )
539            return None
540
541    def _get_datetime_property(
542        self, property: str, property_name: str
543    ) -> datetime | None:
544        """
545        Retrieve and parse a datetime property from the job information.
546
547        If the property is missing, empty, or marked as unknown, None is returned.
548        A warning is logged if parsing fails.
549
550        Args:
551            property (str): The key identifying the property in the job information.
552            property_name (str): A human-readable name of the property for logging.
553
554        Returns:
555            datetime | None: A datetime object if parsing succeeds, otherwise None.
556        """
557        if not (raw_datetime := self._info.get(property)) or raw_datetime.lower() in [
558            "unknown",
559            "n/a",
560            "none",
561            "",
562        ]:
563            return None
564
565        try:
566            return datetime.strptime(raw_datetime, CFG.date_formats.slurm)
567        except Exception as e:
568            logger.warning(
569                f"Could not parse information about {property_name} for '{self._job_id}': {e}."
570            )
571            return None
572
573    def _get_tres(self) -> str:
574        """
575        Return the AllocTRES property or ReqTRES property, depending on which of them is available.
576        Note that the resources specified in ReqTRES can potentially be different than the resources in AllocTRES.
577        """
578        tres = self._info.get("AllocTRES")
579        if not tres or "null" in tres or "None" in tres or "N/A" in tres:
580            tres = self._info.get("ReqTRES", "")
581
582        return tres
583
584    @staticmethod
585    def _assign_if_allocated(
586        info: dict[str, str], alloc_key: str, req_key: str, target_key: str
587    ) -> None:
588        """
589        Assigns a value to a target key in the `info` dictionary, preferring an allocated value
590        if it exists and is valid; otherwise, falls back to the requested value.
591
592        Args:
593            info (dict[str, str]): The dictionary containing allocation and request data.
594            alloc_key (str): The key for the allocated resource (e.g., "AllocCPUs").
595            req_key (str): The key for the requested resource (e.g., "ReqCPUs").
596            target_key (str): The key under which the resolved value should be stored (e.g., "NumCPUs").
597
598        Notes:
599            - A value is considered invalid if it is `None`, an empty string `""`, or `"0"`.
600            - The function updates `info` in place.
601        """
602        value = info.get(alloc_key)
603        info[target_key] = (
604            value if value not in (None, "None", "", "0") else info.get(req_key, "0")
605        )