Coverage for src/srunx/utils.py: 100%

31 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-24 15:16 +0000

1"""Utility functions for SLURM job management.""" 

2 

3import subprocess 

4 

5from srunx.logging import get_logger 

6from srunx.models import BaseJob, JobStatus 

7 

8logger = get_logger(__name__) 

9 

10 

11def get_job_status(job_id: int) -> BaseJob: 

12 """Get job status and information. 

13 

14 Args: 

15 job_id: SLURM job ID. 

16 

17 Returns: 

18 Job object with current status. 

19 

20 Raises: 

21 subprocess.CalledProcessError: If status query fails. 

22 ValueError: If job information cannot be parsed. 

23 """ 

24 logger.debug(f"Querying status for job {job_id}") 

25 

26 try: 

27 result = subprocess.run( 

28 [ 

29 "sacct", 

30 "-j", 

31 str(job_id), 

32 "--format", 

33 "JobID,JobName,State", 

34 "--noheader", 

35 "--parsable2", 

36 ], 

37 capture_output=True, 

38 text=True, 

39 check=True, 

40 ) 

41 except subprocess.CalledProcessError as e: 

42 logger.error(f"Failed to query job {job_id} status: {e}") 

43 raise 

44 

45 lines = result.stdout.strip().split("\n") 

46 if not lines or not lines[0]: 

47 error_msg = f"No job information found for job {job_id}" 

48 logger.error(error_msg) 

49 raise ValueError(error_msg) 

50 

51 # Parse the first line (main job entry) 

52 job_data = lines[0].split("|") 

53 if len(job_data) < 3: 

54 error_msg = f"Cannot parse job data for job {job_id}" 

55 logger.error(error_msg) 

56 raise ValueError(error_msg) 

57 

58 job_id_str, job_name, status_str = job_data[:3] 

59 logger.debug(f"Job {job_id} status: {status_str}") 

60 

61 # Create job object with available information 

62 job = BaseJob( 

63 name=job_name, 

64 job_id=int(job_id_str), 

65 ) 

66 job.status = JobStatus(status_str) 

67 

68 return job 

69 

70 

71def job_status_msg(job: BaseJob) -> str: 

72 """Generate a formatted status message for a job. 

73 

74 Args: 

75 job: Job object to generate message for. 

76 

77 Returns: 

78 Formatted status message with icons and job information. 

79 """ 

80 icons = { 

81 JobStatus.COMPLETED: "✅", 

82 JobStatus.RUNNING: "🚀", 

83 JobStatus.PENDING: "⌛", 

84 JobStatus.FAILED: "❌", 

85 JobStatus.CANCELLED: "🛑", 

86 JobStatus.TIMEOUT: "⏰", 

87 JobStatus.UNKNOWN: "❓", 

88 } 

89 status_icon = icons.get(job.status, "❓") 

90 job_id_display = job.job_id if job.job_id is not None else "—" 

91 return ( 

92 f"{status_icon} {job.status.name:<12} Job {job.name:<12} (ID: {job_id_display})" 

93 )