qq_lib.batch.pbs

 31class PBSJob(BatchJobInterface):
 32    """
 33    Implementation of BatchJobInterface for PBS.
 34    Stores metadata for a single PBS job.
 35    """
 36
 37    def __init__(self, job_id: str):
 38        """Query the batch system for information about the job with the specified ID."""
 39        self._job_id = job_id
 40        self._info: dict[str, str] = {}
 41
 42        self.update()
 43
 44    def is_empty(self) -> bool:
 45        return not self._info
 46
 47    def get_id(self) -> str:
 48        return self._job_id
 49
 50    def get_account(self) -> str | None:
 51        return None
 52
 53    def update(self) -> None:
 54        # get job info from PBS
 55        command = f"qstat -fxw {self._job_id}"
 56
 57        result = subprocess.run(
 58            ["bash"],
 59            input=command,
 60            text=True,
 61            check=False,
 62            capture_output=True,
 63            errors="replace",
 64        )
 65
 66        if result.returncode != 0:
 67            # if qstat fails, information is empty
 68            logger.debug(
 69                f"qstat failed: no information about job '{self._job_id}' is available: {result.stderr.strip()}"
 70            )
 71            self._info: dict[str, str] = {}
 72        else:
 73            self._info = parse_pbs_dump_to_dictionary(result.stdout)
 74
 75    def get_state(self) -> BatchState:
 76        if not (state := self._info.get("job_state")):
 77            return BatchState.UNKNOWN
 78
 79        # X is used by PBS to indicate finished tasks in unfinished array jobs,
 80        # but qq uses X to indicate failure
 81        if state == "X":
 82            state = "F"
 83
 84        # if the job is finished and the return code is not zero, return FAILED
 85        if state == "F":
 86            exit_code = self.get_exit_code()
 87            # if exit code does not exist, the job never ran and was likely killed
 88            if exit_code is None or exit_code != 0:
 89                return BatchState.FAILED
 90
 91        return BatchState.from_code(state)
 92
 93    def get_comment(self) -> str | None:
 94        return self._info.get("comment")
 95
 96    def get_estimated(self) -> tuple[datetime, str] | None:
 97        if not (raw_time := self._info.get("estimated.start_time")):
 98            logger.debug("No 'estimated.start_time' found.")
 99            return None
100
101        try:
102            time = datetime.strptime(raw_time, CFG.date_formats.pbs)
103            # if the estimated start time is in the past, use the current time
104            if (current_time := datetime.now()) > time:
105                time = current_time
106        except Exception as e:
107            logger.debug(f"Could not parse 'estimated.start_time': {e}.")
108            return None
109
110        if not (raw_vnode := self._info.get("estimated.exec_vnode")):
111            logger.debug("No 'estimated.exec_vnode' found.")
112            return None
113
114        vnodes = []
115        for split in raw_vnode.split("+"):
116            vnodes.append(PBSJob._clean_node_name(split.strip()))
117
118        return (time, " + ".join(vnodes))
119
120    def get_main_node(self) -> str | None:
121        if raw_node := self._info.get("exec_host2"):
122            return PBSJob._clean_node_name(raw_node.split("+")[0].strip())
123
124        return None
125
126    def get_nodes(self) -> list[str] | None:
127        if not (raw_nodes := self._info.get("exec_host2")):
128            return None
129
130        nodes = []
131        for node in raw_nodes.split("+"):
132            nodes.append(PBSJob._clean_node_name(node.strip()))
133
134        return nodes
135
136    def get_short_nodes(self) -> list[str] | None:
137        if not (raw_nodes := self._info.get("exec_host")):
138            return None
139
140        nodes = []
141        for node in raw_nodes.split("+"):
142            nodes.append(PBSJob._clean_node_name(node.strip()))
143
144        return nodes
145
146    def get_name(self) -> str | None:
147        return self._info.get("Job_Name")
148
149    def get_n_cpus(self) -> int | None:
150        return self._get_int_property("Resource_List.ncpus", "the number of CPUs")
151
152    def get_n_gpus(self) -> int | None:
153        return self._get_int_property("Resource_List.ngpus", "the number of GPUs")
154
155    def get_n_nodes(self) -> int | None:
156        return self._get_int_property("Resource_List.nodect", "the number of nodes")
157
158    def get_mem(self) -> Size | None:
159        if not (mem := self._info.get("Resource_List.mem")):
160            logger.debug(
161                f"Could not get information about the amount of memory from the batch system for '{self._job_id}'."
162            )
163            return None
164
165        try:
166            return Size.from_string(mem)
167        except Exception as e:
168            logger.warning(f"Could not parse memory for '{self._job_id}': {e}.")
169            return None
170
171    def get_start_time(self) -> datetime | None:
172        return self._get_datetime_property("stime", "the job start time")
173
174    def get_submission_time(self) -> datetime | None:
175        return self._get_datetime_property("ctime", "the job submission time")
176
177    def get_completion_time(self) -> datetime | None:
178        return self._get_datetime_property("obittime", "the job completion time")
179
180    def get_modification_time(self) -> datetime | None:
181        return (
182            self._get_datetime_property("mtime", "the job modification time")
183            or self.get_submission_time()
184        )
185
186    def get_user(self) -> str | None:
187        if not (user := self._info.get("Job_Owner")):
188            return None
189
190        return user.split("@")[0]
191
192    def get_walltime(self) -> timedelta | None:
193        if not (walltime := self._info.get("Resource_List.walltime")):
194            return None
195
196        try:
197            return hhmmss_to_duration(walltime)
198        except QQError as e:
199            logger.warning(f"Could not parse walltime for '{self._job_id}': {e}.")
200            return None
201
202    def get_queue(self) -> str | None:
203        return self._info.get("queue")
204
205    def get_util_cpu(self) -> int | None:
206        if not (util_cpu := self._info.get("resources_used.cpupercent")):
207            logger.debug(
208                f"Information about CPU utilization is not available for '{self._job_id}'."
209            )
210            return None
211
212        if not (ncpus := self.get_n_cpus()):
213            logger.debug(
214                f"Information about the number of CPUs is not available for '{self._job_id}'."
215            )
216            return None
217
218        try:
219            # PBS report CPU utilization in the same way as `top` - we have to divide by number of CPUs
220            return int(util_cpu) // ncpus
221        except Exception as e:
222            # this catches both invalid util_cpu and invalid get_n_cpus
223            logger.warning(
224                f"Could not parse information about CPU utilization for '{self._job_id}': {e}."
225            )
226            return None
227
228    def get_util_mem(self) -> int | None:
229        if not (util_mem := self._info.get("resources_used.mem")):
230            logger.debug(
231                f"Information about memory utilization is not available for '{self._job_id}'."
232            )
233            return None
234
235        if not (mem := self.get_mem()):
236            logger.debug(
237                f"Information about the amount of memory is not available for '{self._job_id}'."
238            )
239            return None
240
241        try:
242            util_mem_kb = Size.from_string(util_mem).value
243            return int(util_mem_kb / mem.value * 100.0)
244        except Exception as e:
245            logger.warning(
246                f"Could not parse information about memory utilization for '{self._job_id}': {e}."
247            )
248            return None
249
250    def get_exit_code(self) -> int | None:
251        if not (exit := self._info.get("Exit_status")):
252            return None
253
254        try:
255            return int(exit)
256        except Exception as e:
257            logger.warning(f"Could not parse exit code for '{self._job_id}': {e}.")
258            return None
259
260    def get_input_machine(self) -> str | None:
261        return self._info.get("Submit_Host")
262
263    def get_input_dir(self) -> Path | None:
264        if not (env_vars := self._get_env_vars()):
265            logger.debug(
266                f"Could not get list of environment variables for '{self._job_id}'."
267            )
268            return None
269
270        if not (
271            # try qq first
272            # we always try qq first because it provides the most reliable information
273            # PBS sometimes sets the input directory to absolute path (resolving symlinks)
274            # which is not necessarily what we want
275            input_dir := env_vars.get(CFG.env_vars.input_dir)
276            or env_vars.get("PBS_O_WORKDIR")  # if this fails, try PBS
277            or env_vars.get("INF_INPUT_DIR")  # if this fails, try Infinity
278        ):
279            logger.debug(f"Could not obtain input directory for '{self._job_id}'.")
280            return None
281
282        return logical_resolve(Path(input_dir))
283
284    def get_info_file(self) -> Path | None:
285        if not (env_vars := self._get_env_vars()):
286            logger.debug(
287                f"Could not get list of environment variables for '{self._job_id}'."
288            )
289            return None
290
291        if not (info_file := env_vars.get(CFG.env_vars.info_file)):
292            logger.debug(
293                f"Job '{self._job_id}' does not have an assigned qq info file."
294            )
295            return None
296
297        return Path(info_file)
298
299    def to_yaml(self) -> str:
300        # we need to add job id to the start of the dictionary
301        to_dump = {"Job Id": self._job_id} | self._info
302        return yaml.dump(
303            to_dump, default_flow_style=False, sort_keys=False, Dumper=Dumper
304        )
305
306    def get_steps(self) -> Sequence[Self]:
307        # not available for PBS
308        return []
309
310    def get_step_id(self) -> str | None:
311        # no job steps for PBS
312        return None
313
314    def is_array_job(self) -> bool:
315        return (
316            array := self._info.get("array")
317        ) is not None and array.lower() == "true"
318
319    @classmethod
320    def from_dict(cls, job_id: str, info: dict[str, str]) -> Self:
321        """
322        Construct a new instance of PBSJob from a job ID and a dictionary of job information.
323
324        This method bypasses the standard initializer and directly sets the `_job_id` and `_info`
325        attributes of the new instance.
326
327        Args:
328            job_id (str): The unique identifier of the job.
329            info (dict[str, str]): A dictionary containing PBS job metadata as key-value pairs.
330
331        Returns:
332            Self: A new instance of PBSJob.
333
334        Note:
335            This method does not perform any validation or processing of the provided dictionary.
336        """
337        job_info = cls.__new__(cls)
338        job_info._job_id = job_id
339        job_info._info = info
340
341        return job_info
342
343    def get_id_int(self) -> int | None:
344        """
345        Extract the leading numeric portion of the job ID and return it as an integer.
346
347        Returns:
348            int | None: The integer value of the leading digits in the job ID,
349            or `None` if no valid digits are found or conversion fails.
350        """
351        match = re.match(r"\d+", self.get_id())
352        return int(match.group()) if match else None
353
354    def _get_env_vars(self) -> dict[str, str] | None:
355        """
356        Retrieve environment variables associated with the job.
357
358        Returns:
359            dict[str, str] | None: A dictionary of environment variables, or None
360            if no variable list is available.
361        """
362        if not (variable_list := self._info.get("Variable_List")):
363            return None
364
365        return dict(
366            item.split("=", 1) for item in variable_list.split(",") if "=" in item
367        )
368
369    def _get_int_property(self, property: str, property_name: str) -> int | None:
370        """
371        Retrieve an integer property value from the job information.
372
373        If the property is missing or cannot be converted, `None` is returned.
374
375        Args:
376            property (str): The key identifying the property in the job information.
377            property_name (str): A human-readable name of the property for logging.
378
379        Returns:
380            int | None: The integer value of the property, or `None` if unavailable or invalid.
381        """
382        try:
383            return int(self._info[property])
384        except Exception:
385            logger.debug(
386                f"Could not get information about {property_name} from the batch system for '{self._job_id}'."
387            )
388            return None
389
390    def _get_datetime_property(
391        self, property: str, property_name: str
392    ) -> datetime | None:
393        """
394        Retrieve and parse a datetime property from the job information.
395
396        Args:
397            property (str): The key identifying the property in the job information.
398            property_name (str): A human-readable name of the property for logging.
399
400        Returns:
401            datetime | None: A datetime object if parsing succeeds, otherwise None.
402        """
403        if not (raw_datetime := self._info.get(property)):
404            return None
405
406        try:
407            return datetime.strptime(raw_datetime, CFG.date_formats.pbs)
408        except Exception:
409            logger.warning(
410                f"Could not parse information about {property_name} for '{self._job_id}'."
411            )
412            return None
413
414    @staticmethod
415    def _clean_node_name(raw: str) -> str:
416        """
417        Normalize a raw node string to extract the clean hostname.
418
419        Args:
420            raw (str): Raw node string reported by the batch system.
421
422        Returns:
423            str: Cleaned node name.
424        """
425        return raw.split(":", 1)[0].split("/", 1)[0].replace("(", "").replace(")", "")

 90class PBSNode(BatchNodeInterface):
 91    """
 92    Implementation of BatchNodeInterface for PBS.
 93    Stores metadata for a single PBS node.
 94    """
 95
 96    def __init__(self, name: str, server: str | None = None):
 97        self._name = name
 98        self._server = server
 99        self._info: dict[str, str] = {}
100
101        self.update()
102
103    def update(self) -> None:
104        # get node info from PBS
105        command = f"pbsnodes -v {self._name}"
106        if self._server:
107            command += f" -s {self._server}"
108
109        result = subprocess.run(
110            ["bash"],
111            input=command,
112            text=True,
113            check=False,
114            capture_output=True,
115            errors="replace",
116        )
117
118        if result.returncode != 0:
119            raise QQError(f"Node '{self._name}' does not exist.")
120
121        self._info = parse_pbs_dump_to_dictionary(result.stdout)
122
123    def get_name(self) -> str:
124        return self._name
125
126    def get_n_cpus(self) -> int | None:
127        return self._get_int_resource("resources_available.ncpus")
128
129    def get_n_free_cpus(self) -> int | None:
130        return self._get_free_int_resource("ncpus")
131
132    def get_n_gpus(self) -> int | None:
133        return self._get_int_resource("resources_available.ngpus")
134
135    def get_n_free_gpus(self) -> int | None:
136        return self._get_free_int_resource("ngpus")
137
138    def get_cpu_memory(self) -> Size | None:
139        return self._get_size_resource("resources_available.mem")
140
141    def get_free_cpu_memory(self) -> Size | None:
142        return self._get_free_size_resource("mem")
143
144    def get_gpu_memory(self) -> Size | None:
145        return self._get_size_resource("resources_available.gpu_mem")
146
147    def get_free_gpu_memory(self) -> Size | None:
148        return self._get_free_size_resource("gpu_mem")
149
150    def get_local_scratch(self) -> Size | None:
151        return self._get_size_resource("resources_available.scratch_local")
152
153    def get_free_local_scratch(self) -> Size | None:
154        return self._get_free_size_resource("scratch_local")
155
156    def get_ssd_scratch(self) -> Size | None:
157        return self._get_size_resource("resources_available.scratch_ssd")
158
159    def get_free_ssd_scratch(self) -> Size | None:
160        return self._get_free_size_resource("scratch_ssd")
161
162    def get_shared_scratch(self) -> Size | None:
163        return self._get_size_resource("resources_available.scratch_shared")
164
165    def get_free_shared_scratch(self) -> Size | None:
166        return self._get_free_size_resource("scratch_shared")
167
168    def get_properties(self) -> list[str]:
169        return [
170            key.split(".", 1)[-1]
171            for key in self._info
172            if "resources_available" in key and self._info[key] == "True"
173        ]
174
175    def is_available_to_user(self, user: str) -> bool:
176        if not (state := self._info.get("state")):
177            logger.debug(f"Could not get state information for node '{self._name}'.")
178            return False
179
180        if any(
181            disabled_state in state
182            for disabled_state in {"down", "unknown", "unresolvable", "resv-exclusive"}
183        ):
184            return False
185
186        if queue := self._info.get("queue"):
187            return QueuesAvailability.get_or_init(queue, user, self._server)
188
189        return True
190
191    @classmethod
192    def from_dict(cls, name: str, server: str | None, info: dict[str, str]) -> Self:
193        """
194        Construct a new instance of PBSNode from node name and a dictionary of node information.
195
196        Args:
197            name (str): The unique name of the node.
198            server (str | None): Server on which the node is located. If `None`, assumes the current server.
199            info (dict[str, str]): A dictionary containing PBS node metadata as key-value pairs.
200
201        Returns:
202            Self: A new instance of PBSNode.
203
204        Note:
205            This method does not perform any validation or processing of the provided dictionary.
206        """
207        node = cls.__new__(cls)
208        node._name = name
209        node._server = server
210        node._info = info
211
212        return node
213
214    def to_yaml(self) -> str:
215        # we need to add node name to the start of the dictionary
216        to_dump = {"Node": self._name} | self._info
217        return yaml.dump(
218            to_dump, default_flow_style=False, sort_keys=False, Dumper=Dumper
219        )
220
221    def _get_int_resource(self, res: str) -> int | None:
222        """
223        Retrieve an integer-valued resource from the node information.
224
225        Args:
226            res (str): The resource key to retrieve (e.g., "resources_available.ncpus").
227
228        Returns:
229            int: The integer value of the resource, or `None` if unavailable or invalid.
230        """
231        if not (val := self._info.get(res)):
232            return None
233        try:
234            return int(val)
235        except Exception as e:
236            logger.debug(f"Could not parse the value '{val}' of resource '{res}': {e}.")
237            return None
238
239    def _get_free_int_resource(self, res: str) -> int | None:
240        """
241        Compute the number of free units for an integer-valued resource.
242
243        Calculates the difference between the total available (`resources_available.<res>`)
244        and the assigned (`resources_assigned.<res>`) quantities. If the computed
245        difference is negative, returns 0. If the information is not available, returns None.
246
247        Args:
248            res (str): The base resource name (e.g., "ncpus", "ngpus").
249
250        Returns:
251            int | None: The number of unallocated (free) resource units, or None if unavailable.
252        """
253        if not (full := self._get_int_resource(f"resources_available.{res}")):
254            return None
255
256        # if the `resources_assigned` property is missing, we assume it means that there are no resources assigned
257        assigned = self._get_int_resource(f"resources_assigned.{res}") or 0
258
259        if (diff := full - assigned) >= 0:
260            return diff
261
262        return 0
263
264    def _get_size_resource(self, res: str) -> Size | None:
265        """
266        Retrieve a Size resource from the node information.
267
268        Args:
269            res (str): The resource key to retrieve (e.g., "resources_available.mem").
270
271        Returns:
272            Size | None: The parsed Size, or `None` if unavailable or invalid.
273        """
274        if not (val := self._info.get(res)):
275            return None
276
277        try:
278            return Size.from_string(val)
279        except Exception as e:
280            logger.debug(f"Could not parse the value '{val}' of resource '{res}': {e}.")
281            return None
282
283    def _get_free_size_resource(self, res: str) -> Size | None:
284        """
285        Compute the amount of free space for a Size resource.
286
287        Calculates the difference between total available (`resources_available.<res>`)
288        and assigned (`resources_assigned.<res>`) values. If subtraction results in a negative size,
289        returns a zero-size object. If the information are not available, returns None.
290
291        Args:
292            res (str): The base resource name (e.g., "mem", "scratch_local").
293
294        Returns:
295            Size | None: The available (free) size for the resource, or `None` if unavailable.
296        """
297        if not (full := self._get_size_resource(f"resources_available.{res}")):
298            return None
299
300        # if the `resources_assigned` property is missing, we assume it means that there are no resources assigned
301        assigned = self._get_size_resource(f"resources_assigned.{res}") or Size(0, "kb")
302
303        try:
304            return full - assigned
305        except ValueError as e:
306            logger.debug(f"Negative free size of resource '{res}': {e}.")
307            return Size(0, "kb")