Coverage for sparkle/CLI/help/jobs.py: 54%

1"""File to help with RunRunner jobs.""" 

2from pathlib import Path 

3 

4from tabulate import tabulate 

5 

6from runrunner.base import Status 

7from runrunner.slurm import SlurmRun 

8 

9from sparkle.platform.cli_types import TEXT 

10 

11 

12def get_runs_from_file(path: Path, print_error: bool = False) -> list[SlurmRun]: 

13 """Retrieve all run objects from file storage. 

14 

15 Args: 

16 path: Path object where to look recursively for the files. 

17 

18 Returns: 

19 List of all found SlumRun objects. 

20 """ 

21 runs = [] 

22 for file in path.rglob("*.json"): 

23 # TODO: RunRunner should be adapted to have more general methods for runs 

24 # So this method can work for both local and slurm 

25 try: 

26 run_obj = SlurmRun.from_file(file) 

27 runs.append(run_obj) 

28 except Exception as ex: 

29 # Not a (correct) RunRunner JSON file 

30 if print_error: 

31 print(f"[WARNING] Could not load file: {file}. Exception: {ex}") 

32 return runs 

33 

34 

35def create_jobs_table(jobs: list[SlurmRun], markup: bool = True) -> str: 

36 """Create a table of jobs. 

37 

38 Args: 

39 runs: List of SlurmRun objects. 

40 markup: By default some mark up will be applied to the table. 

41 If false, a more plain version will be created. 

42 

43 Returns: 

44 A table of jobs as a string. 

45 """ 

46 job_table = [["RunId", "Name", "Quality of Service", "Partition", "Status", 

47 "Dependencies", "Finished Jobs", "Run Time"]] 

48 for job in jobs: 

49 # Count number of jobs that have finished 

50 finished_jobs_count = sum(1 for status in job.all_status 

51 if status == Status.COMPLETED) 

52 if markup: # Format job.status 

53 status_text = \ 

54 TEXT.format_text([TEXT.BOLD], job.status) \ 

55 if job.status == Status.RUNNING else \ 

56 (TEXT.format_text([TEXT.ITALIC], job.status) 

57 if job.status == Status.COMPLETED else job.status) 

58 else: 

59 status_text = job.status 

60 job_table.append( 

61 [job.run_id, 

62 job.name, 

63 job.qos, 

64 job.partition, 

65 status_text, 

66 "None" if len(job.dependencies) == 0 else ", ".join(job.dependencies), 

67 f"{finished_jobs_count}/{len(job.all_status)}", 

68 job.runtime]) 

69 if markup: 

70 return tabulate(job_table, headers="firstrow", tablefmt="grid") 

71 return job_table