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

1import subprocess 

2 

3from srunx.logging import get_logger 

4from srunx.models import BaseJob, JobStatus 

5 

6logger = get_logger(__name__) 

7 

8 

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

10 """Get job status and information. 

11 

12 Args: 

13 job_id: SLURM job ID. 

14 

15 Returns: 

16 Job object with current status. 

17 

18 Raises: 

19 subprocess.CalledProcessError: If status query fails. 

20 ValueError: If job information cannot be parsed. 

21 """ 

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

23 

24 try: 

25 result = subprocess.run( 

26 [ 

27 "sacct", 

28 "-j", 

29 str(job_id), 

30 "--format", 

31 "JobID,JobName,State", 

32 "--noheader", 

33 "--parsable2", 

34 ], 

35 capture_output=True, 

36 text=True, 

37 check=True, 

38 ) 

39 except subprocess.CalledProcessError as e: 

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

41 raise 

42 

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

44 if not lines or not lines[0]: 

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

46 logger.error(error_msg) 

47 raise ValueError(error_msg) 

48 

49 # Parse the first line (main job entry) 

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

51 if len(job_data) < 3: 

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

53 logger.error(error_msg) 

54 raise ValueError(error_msg) 

55 

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

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

58 

59 # Create job object with available information 

60 job = BaseJob( 

61 name=job_name, 

62 job_id=int(job_id_str), 

63 status=JobStatus(status_str), 

64 ) 

65 

66 return job