qq_lib.jobs

 23class JobsPresenter:
 24    """
 25    Present information about a collection of jobs from the batch system and their statistics.
 26    """
 27
 28    # Mapping of human-readable color names to ANSI escape codes.
 29    _ANSI_COLORS = {
 30        # default
 31        "default": "",
 32        # standard colors
 33        "black": "\033[30m",
 34        "red": "\033[31m",
 35        "green": "\033[32m",
 36        "yellow": "\033[33m",
 37        "blue": "\033[34m",
 38        "magenta": "\033[35m",
 39        "cyan": "\033[36m",
 40        "white": "\033[37m",
 41        # bright colors
 42        "bright_black": "\033[90m",
 43        "bright_red": "\033[91m",
 44        "bright_green": "\033[92m",
 45        "bright_yellow": "\033[93m",
 46        "bright_blue": "\033[94m",
 47        "bright_magenta": "\033[95m",
 48        "bright_cyan": "\033[96m",
 49        "bright_white": "\033[97m",
 50        # other colors
 51        "grey90": "\033[38;5;254m",
 52        "grey70": "\033[38;5;249m",
 53        "grey50": "\033[38;5;244m",
 54        "grey30": "\033[38;5;239m",
 55        "grey10": "\033[38;5;233m",
 56        # bold:
 57        "bold": "\033[1m",
 58        # reset
 59        "reset": "\033[0m",
 60    }
 61
 62    # Table formatting configuration for `tabulate`.
 63    _COMPACT_TABLE = TableFormat(
 64        lineabove=Line("", "", "", ""),
 65        linebelowheader="",
 66        linebetweenrows="",
 67        linebelow=Line("", "", "", ""),
 68        headerrow=("", " ", ""),
 69        datarow=("", " ", ""),
 70        padding=0,
 71        with_header_hide=["lineabove", "linebelow"],
 72    )
 73
 74    def __init__(
 75        self,
 76        batch_system: AnyBatchClass,
 77        jobs: list[BatchJobInterface],
 78        extra: bool,
 79        all: bool,
 80        server: str | None,
 81    ):
 82        """
 83        Initialize the presenter with a list of jobs.
 84
 85        Args:
 86            jobs (list[BatchJobInterface]): List of job information objects
 87                to be presented.
 88            extra (bool): Should show additional info about jobs.
 89            all (bool): Show all jobs, not just queued and running.
 90            server (str | None): Batch server for which the jobs were collected.
 91                `None` = default server.
 92        """
 93        self._batch_system = batch_system
 94        self._jobs = jobs
 95        self._stats = JobsStatistics()
 96        self._extra = extra
 97        self._all = all
 98        self._server = server
 99
100    def create_jobs_info_panel(self, console: Console | None = None) -> Group:
101        """
102        Create a Rich panel displaying job information and statistics.
103
104        Args:
105            console (Console | None): Optional Rich Console instance.
106                If None, a new Console will be created.
107
108        Returns:
109            Group: Rich Group containing the jobs table and stats panel.
110        """
111        console = console or Console()
112
113        jobs_table = self._create_basic_jobs_table()
114        if self._extra:
115            jobs_table = self._insert_extra_info(jobs_table)
116
117        # convert ANSI codes to Rich Text
118        jobs_panel = Text.from_ansi(jobs_table)
119        stats_panel = self._stats.create_stats_panel()
120
121        content = Group(
122            jobs_panel,
123            Text(""),
124            stats_panel,
125        )
126
127        panel = Panel(
128            content,
129            title=Text(
130                "COLLECTED JOBS",
131                style=CFG.jobs_presenter.title_style,
132                justify="center",
133            ),
134            subtitle=Text(
135                f"{self._server}",
136                style=CFG.jobs_presenter.subtitle_style,
137                justify="center",
138            )
139            if self._server
140            else None,
141            border_style=CFG.jobs_presenter.border_style,
142            padding=(1, 1),
143            width=get_panel_width(
144                console, 1, CFG.jobs_presenter.min_width, CFG.jobs_presenter.max_width
145            ),
146            expand=False,
147        )
148
149        return Group(Text(""), panel, Text(""))
150
151    def dump_yaml(self) -> None:
152        """
153        Print the YAML representation of all jobs to stdout.
154        """
155        for job in self._jobs:
156            print(job.to_yaml())
157
158    def _create_basic_jobs_table(self) -> str:
159        """
160        Build a compact tabulated string representation of the job list.
161
162        Returns:
163            str: Tabulated job information with ANSI color codes applied.
164
165        Notes:
166            - Uses `tabulate` with `_COMPACT_TABLE` format because
167              Rich's Table is prohibitively slow for large number of items.
168            - Updates internal job statistics via `self._stats`.
169        """
170        headers = self._get_visible_headers()
171        rows = [self._create_job_row(job, headers) for job in self._jobs]
172
173        return tabulate(
174            rows,
175            headers=self._format_headers(headers),
176            tablefmt=JobsPresenter._COMPACT_TABLE,
177            stralign="center",
178            numalign="center",
179        )
180
181    def _get_visible_headers(self) -> list[str]:
182        """
183        Get list of headers to display based on the batch system configuration.
184
185        Return:
186            list[str]: A list of headers to show.
187        """
188        all_headers = [
189            "S",
190            "Job ID",
191            "User",
192            "Job Name",
193            "Queue",
194            "NCPUs",
195            "NGPUs",
196            "NNodes",
197            "Times",
198            "Node",
199            "%CPU",
200            "%Mem",
201            "Exit" if self._all or CFG.jobs_presenter.columns_to_show else None,
202        ]
203        headers_to_show = (
204            CFG.jobs_presenter.columns_to_show
205            or self._batch_system.jobs_presenter_columns_to_show()
206        )
207        return [h for h in all_headers if h and h in headers_to_show]
208
209    def _create_job_row(self, job: BatchJobInterface, headers: list[str]) -> list[str]:
210        """
211        Create a single row of job data.
212
213        Args:
214            job (BatchJobInterface): Job to show information for.
215            headers (list[str]): List of headers to include in the row
216
217        Returns:
218            list[str]: List of formatted cell values.
219        """
220        state = job.get_state()
221        start_time, end_time = self._get_job_times(job, state)
222
223        # update statistics
224        cpus = job.get_n_cpus() or 0
225        gpus = job.get_n_gpus() or 0
226        nodes = job.get_n_nodes() or 0
227        self._stats.add_job(state, cpus, gpus, nodes)
228
229        if self._server:
230            # show full job ID if we are working with a non-standard server
231            job_id = job.get_id()
232        else:
233            # otherwise, show only the numerical portion of the ID
234            job_id = JobsPresenter._shorten_job_id(job.get_id())
235
236        # build the row
237        row_data: dict[str, str] = {
238            "S": JobsPresenter._color(state.to_code(), state.color),
239            "Job ID": JobsPresenter._main_color(job_id),
240            "User": JobsPresenter._main_color(job.get_user() or ""),
241            "Job Name": JobsPresenter._main_color(
242                JobsPresenter._shorten_job_name(job.get_name() or "")
243            ),
244            "Queue": JobsPresenter._main_color(job.get_queue() or ""),
245            "NCPUs": JobsPresenter._main_color(str(cpus)),
246            "NGPUs": JobsPresenter._main_color(str(gpus)),
247            "NNodes": JobsPresenter._main_color(str(nodes)),
248            "Times": JobsPresenter._format_time(
249                state, start_time, end_time, job.get_walltime()
250            ),
251            "Node": JobsPresenter._format_nodes_or_comment(state, job),
252            "%CPU": JobsPresenter._format_util_cpu(job.get_util_cpu()),
253            "%Mem": JobsPresenter._format_util_mem(job.get_util_mem()),
254            "Exit": JobsPresenter._format_exit_code(job, state) if self._all else "",
255        }
256
257        return [row_data[header] for header in headers if header in row_data]
258
259    @staticmethod
260    def _get_job_times(
261        job: BatchJobInterface, state: BatchState
262    ) -> tuple[datetime | None, datetime | None]:
263        """
264        Get start and end times for a job based on its state.
265
266        Args:
267            job (BatchJobInterface): Job to get times for.
268            state (BatchState): The current job state.
269
270        Returns:
271            tuple[datetime | None, datetime | None]: Tuple of (start_time, end_time).
272        """
273        if state in {BatchState.QUEUED, BatchState.HELD, BatchState.WAITING}:
274            start_time = job.get_submission_time()
275        else:
276            start_time = job.get_start_time() or job.get_submission_time()
277
278        if state in {BatchState.FINISHED, BatchState.FAILED}:
279            end_time = job.get_completion_time() or job.get_modification_time()
280        else:
281            end_time = datetime.now()
282
283        return start_time, end_time
284
285    def _format_headers(self, headers: list[str]) -> list[str]:
286        """
287        Apply formatting to table headers.
288
289        Args:
290            headers (list[str]): List of headers to format.
291
292        Returns:
293            list[str]: List of formatted and colored headers.
294        """
295        return [
296            JobsPresenter._color(
297                header, color=CFG.jobs_presenter.headers_style, bold=True
298            )
299            for header in headers
300        ]
301
302    def _insert_extra_info(self, table: str) -> str:
303        """
304        Augment a formatted job table with additional information about each job.
305
306        Lines where job attributes are missing are skipped.
307
308        Args:
309            table (str): The formatted table string containing one line per job.
310
311        Returns:
312            str: A new table string including the extra job information lines.
313        """
314        split_table = table.splitlines()
315        table_with_extra_info = split_table[0] + "\n"
316
317        for line, job in zip(split_table[1:], self._jobs):
318            table_with_extra_info += line + "\n"
319
320            if input_machine := job.get_input_machine():
321                table_with_extra_info += JobsPresenter._color(
322                    f" >   Input machine:   {input_machine}\n",
323                    CFG.jobs_presenter.extra_info_style,
324                )
325
326            if input_dir := job.get_input_dir():
327                table_with_extra_info += JobsPresenter._color(
328                    f" >   Input directory: {str(input_dir)}\n",
329                    CFG.jobs_presenter.extra_info_style,
330                )
331
332            if comment := job.get_comment():
333                table_with_extra_info += JobsPresenter._color(
334                    f" >   Comment:         {comment}\n",
335                    CFG.jobs_presenter.extra_info_style,
336                )
337            table_with_extra_info += "\n"
338
339        return table_with_extra_info
340
341    @staticmethod
342    def _format_time(
343        state: BatchState,
344        start_time: datetime | None,
345        end_time: datetime | None,
346        walltime: timedelta | None,
347    ) -> str:
348        """
349        Format the job running time, queued time or completion time with color coding.
350
351        Args:
352            state (BatchState): Current job state.
353            start_time (datetime | None): Job submission or start time.
354            end_time (datetime | None): Job completion or current time.
355            walltime (timedelta | None): Scheduled walltime for the job.
356
357        Returns:
358            str: ANSI-colored string representing elapsed or finished time.
359        """
360        # return an empty string if any of the required times is missing
361        if start_time is None or end_time is None or walltime is None:
362            return ""
363
364        match state:
365            case BatchState.UNKNOWN | BatchState.SUSPENDED:
366                return ""
367            case BatchState.FAILED | BatchState.FINISHED:
368                return JobsPresenter._color(
369                    end_time.strftime(CFG.date_formats.standard), color=state.color
370                )
371            case (
372                BatchState.HELD
373                | BatchState.QUEUED
374                | BatchState.WAITING
375                | BatchState.MOVING
376            ):
377                return JobsPresenter._color(
378                    format_duration_wdhhmmss(end_time - start_time),
379                    color=state.color,
380                )
381            case BatchState.RUNNING | BatchState.EXITING:
382                run_time = end_time - start_time
383                return JobsPresenter._color(
384                    format_duration_wdhhmmss(run_time),
385                    color=CFG.jobs_presenter.strong_warning_style
386                    if run_time > walltime
387                    else state.color,
388                ) + JobsPresenter._main_color(
389                    f" / {format_duration_wdhhmmss(walltime)}"
390                )
391
392        return Text("")
393
394    @staticmethod
395    def _format_util_cpu(util: int | None) -> str:
396        """
397        Format CPU utilization with color coding.
398
399        Args:
400            util (int | None): CPU usage percentage.
401
402        Returns:
403            str: ANSI-colored string representation of CPU utilization,
404                 or empty string if `util` is None.
405        """
406        if util is None:
407            return ""
408
409        if util > 100:
410            color = CFG.jobs_presenter.strong_warning_style
411        elif util >= 80:
412            color = CFG.jobs_presenter.main_style
413        elif util >= 60:
414            color = CFG.jobs_presenter.mild_warning_style
415        else:
416            color = CFG.jobs_presenter.strong_warning_style
417
418        return JobsPresenter._color(str(util), color=color)
419
420    @staticmethod
421    def _format_util_mem(util: int | None) -> str:
422        """
423        Format memory utilization with color coding.
424
425        Args:
426            util (int | None): Memory usage percentage.
427
428        Returns:
429            str: ANSI-colored string representation of memory utilization,
430                 or empty string if `util` is None.
431        """
432        if util is None:
433            return ""
434
435        if util < 90:
436            color = CFG.jobs_presenter.main_style
437        elif util < 100:
438            color = CFG.jobs_presenter.mild_warning_style
439        else:
440            color = CFG.jobs_presenter.strong_warning_style
441
442        return JobsPresenter._color(str(util), color=color)
443
444    @staticmethod
445    def _format_exit_code(job: BatchJobInterface, state: BatchState) -> str:
446        """
447        Get formatted exit code if the job is completed and color it appropriately.
448
449        The color of the exit code is set based on the state of the job,
450        not on the value of the exit code.
451
452        If the job is not completed, returns an empty string.
453
454        Args:
455            job (BatchJobInterface): Job to get the exit code for.
456            state (BatchState): The current job state.
457
458        Returns:
459            str: ANSI-colored exit code. Empty string if the job is not completed
460            or the exit code is undefined.
461        """
462        if (exit_code := job.get_exit_code()) is None:
463            return ""
464
465        match state:
466            case BatchState.FINISHED:
467                return JobsPresenter._main_color(str(exit_code))
468            case BatchState.FAILED:
469                return JobsPresenter._color(
470                    str(exit_code), color=CFG.jobs_presenter.strong_warning_style
471                )
472            case _:
473                return ""
474
475    @staticmethod
476    def _format_nodes_or_comment(state: BatchState, job: BatchJobInterface) -> str:
477        """
478        Format node information or an estimated runtime comment.
479
480        Args:
481            state (BatchState): Current job state.
482            job (BatchJobInterface): Job information object.
483
484        Returns:
485            str: ANSI-colored string for working node(s) or estimated start,
486                 or an empty string if neither information is available.
487        """
488        if nodes := job.get_short_nodes():
489            return JobsPresenter._main_color(
490                JobsPresenter._shorten_nodes(" + ".join(nodes)),
491            )
492
493        if state in {BatchState.FINISHED, BatchState.FAILED}:
494            return ""
495
496        if estimated := job.get_estimated():
497            truncated_nodes = JobsPresenter._shorten_nodes(estimated[1])
498            return JobsPresenter._color(
499                f"{truncated_nodes} in {format_duration_wdhhmmss(estimated[0] - datetime.now()).rsplit(':', 1)[0]}",
500                color=state.color,
501            )
502
503        return ""
504
505    @staticmethod
506    def _shorten_job_id(job_id: str) -> str:
507        """
508        Shorten the job ID to its primary component (before the first dot).
509
510        Args:
511            job_id (str): Full job identifier.
512
513        Returns:
514            str: Shortened job ID.
515        """
516        return job_id.split(".", 1)[0]
517
518    @staticmethod
519    def _shorten_job_name(job_name: str) -> str:
520        """
521        Truncate a job name if it exceeds the maximum allowed display length.
522
523        Args:
524            job_name (str): The original job name string.
525
526        Returns:
527            str: The possibly shortened job name. If the original name length is
528                less than or equal to the configured limit, it is returned unchanged.
529        """
530        if len(job_name) > CFG.jobs_presenter.max_job_name_length:
531            return f"{job_name[: CFG.jobs_presenter.max_job_name_length]}…"
532
533        return job_name
534
535    @staticmethod
536    def _shorten_nodes(nodes: str) -> str:
537        """
538        Truncate a list of nodes if it exceeds the maximum allowed display length.
539
540        Args:
541            nodes (str): The original nodes string.
542
543        Returns:
544            str: The possibly shortened list of nodes. If the original string length
545                is less than or equal to the configured limit, it is returned unchanged.
546        """
547        if len(nodes) > CFG.jobs_presenter.max_nodes_length:
548            return f"{nodes[: CFG.jobs_presenter.max_nodes_length]}…"
549
550        return nodes
551
552    @staticmethod
553    def _color(string: str, color: str | None = None, bold: bool = False) -> str:
554        """
555        Apply ANSI color codes and optional bold styling to a string.
556
557        Args:
558            string (str): The string to colorize.
559            color (str | None): Optional color.
560            bold (bool): Whether to apply bold formatting.
561
562        Returns:
563            str: ANSI-colored and optionally bolded string.
564        """
565        return f"{JobsPresenter._ANSI_COLORS['bold'] if bold else ''}{JobsPresenter._ANSI_COLORS[color] if color else ''}{string}{JobsPresenter._ANSI_COLORS['reset'] if color or bold else ''}"
566
567    @staticmethod
568    def _main_color(string: str, bold: bool = False) -> str:
569        """
570        Apply the main presenter color with optional bold styling.
571
572        Args:
573            string (str): String to format.
574            bold (bool): Whether to apply bold formatting.
575
576        Returns:
577            str: ANSI-colored string in the main presenter color.
578        """
579        return JobsPresenter._color(string, CFG.jobs_presenter.main_style, bold)
580
581    @staticmethod
582    def _secondary_color(string: str, bold: bool = False) -> str:
583        """
584        Apply the secondary presenter color with optional bold styling.
585
586        Args:
587            string (str): String to format.
588            bold (bool): Whether to apply bold formatting.
589
590        Returns:
591            Text: ANSI-colored Rich Text object in secondary color.
592        """
593        return JobsPresenter._color(string, CFG.jobs_presenter.secondary_style, bold)