Coverage for src/srunx/workflows/tasks.py: 100%

1"""Prefect tasks for SLURM workflow management."""

2

3 from prefect import task

4 from srunx.client import Slurm

5 from srunx.logging import get_logger

6 from srunx.models import BaseJob, Job, ShellJob

7

8 logger = get_logger(__name__)

9

10

11@task

12 def submit_and_monitor_job(

13 job: Job | ShellJob, poll_interval: int = 30

14) -> Job | ShellJob:

15 """Submit a SLURM job and monitor until completion.

16

17 This Prefect task handles the complete lifecycle of a SLURM job:

18 submission, monitoring, and completion verification.

19

20 Args:

21 job: Job configuration.

22 poll_interval: Status polling interval in seconds.

23

24 Returns:

25 Completed Job instance.

26

27 Raises:

28 RuntimeError: If the SLURM job fails.

29 subprocess.CalledProcessError: If job operations fail.

30 """

31 logger.info(f"Starting SLURM job submission and monitoring for '{job.name}'")

32 client = Slurm()

33

34 # Submit the job

35 submitted_job = client.run(job)

36 logger.info(f"Job '{submitted_job.name}' submitted with ID {submitted_job.job_id}")

37

38 # Wait for completion

39 completed_job = client.monitor(submitted_job, poll_interval=poll_interval)

40 logger.info(f"Job '{completed_job.name}' (ID: {completed_job.job_id}) completed")

41

42 assert isinstance(completed_job, Job | ShellJob)

43

44 return completed_job

45

46

47@task

48 def submit_job_async(job: Job | ShellJob) -> Job | ShellJob:

49 """Submit a SLURM job without waiting for completion.

50

51 Args:

52 job: Job configuration.

53

54 Returns:

55 Submitted Job instance with job_id.

56 """

57 logger.info(f"Submitting async SLURM job '{job.name}'")

58 client = Slurm()

59 submitted_job = client.run(job)

60 logger.info(

61 f"Async job '{submitted_job.name}' submitted with ID {submitted_job.job_id}"

62 )

63 assert isinstance(submitted_job, Job | ShellJob)

64 return submitted_job

65

66

67@task

68 def wait_for_job(job_id: int, poll_interval: int = 30) -> BaseJob:

69 """Wait for a job to complete.

70

71 Args:

72 job_id: SLURM job ID.

73 poll_interval: Polling interval in seconds.

74

75 Returns:

76 Completed job object.

77 """

78 logger.info(f"Waiting for job {job_id} to complete")

79 client = Slurm()

80 completed_job = client.monitor(job_id, poll_interval)

81 logger.info(f"Job {job_id} completed")

82 return completed_job