Coverage for src/srunx/utils.py: 100%
31 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 21:22 +0900
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 21:22 +0900
1"""Utility functions for SLURM job management."""
3import subprocess
5from srunx.logging import get_logger
6from srunx.models import BaseJob, JobStatus
8logger = get_logger(__name__)
11def get_job_status(job_id: int) -> BaseJob:
12 """Get job status and information.
14 Args:
15 job_id: SLURM job ID.
17 Returns:
18 Job object with current status.
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}")
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
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)
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)
58 job_id_str, job_name, status_str = job_data[:3]
59 logger.debug(f"Job {job_id} status: {status_str}")
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)
68 return job
71def job_status_msg(job: BaseJob) -> str:
72 """Generate a formatted status message for a job.
74 Args:
75 job: Job object to generate message for.
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 )