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

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

2

3 import subprocess

4

5 from srunx.logging import get_logger

6 from srunx.models import BaseJob, JobStatus

7

8 logger = get_logger(__name__)

9

10

11 def 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

71 def 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 )