Coverage for src/srunx/callbacks.py: 86%

29 statements  

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

1"""Callback system for job state notifications.""" 

2 

3from slack_sdk import WebhookClient 

4 

5from srunx.models import JobType, Workflow 

6from srunx.utils import job_status_msg 

7 

8 

9class Callback: 

10 """Base callback class for job state notifications.""" 

11 

12 def on_job_submitted(self, job: JobType) -> None: 

13 """Called when a job is submitted to SLURM. 

14 

15 Args: 

16 job: Job that was submitted. 

17 """ 

18 pass 

19 

20 def on_job_completed(self, job: JobType) -> None: 

21 """Called when a job completes successfully. 

22 

23 Args: 

24 job: Job that completed. 

25 """ 

26 pass 

27 

28 def on_job_failed(self, job: JobType) -> None: 

29 """Called when a job fails. 

30 

31 Args: 

32 job: Job that failed. 

33 """ 

34 pass 

35 

36 def on_job_running(self, job: JobType) -> None: 

37 """Called when a job starts running. 

38 

39 Args: 

40 job: Job that started running. 

41 """ 

42 pass 

43 

44 def on_job_cancelled(self, job: JobType) -> None: 

45 """Called when a job is cancelled. 

46 

47 Args: 

48 job: Job that was cancelled. 

49 """ 

50 pass 

51 

52 def on_workflow_started(self, workflow: Workflow) -> None: 

53 """Called when a workflow starts. 

54 

55 Args: 

56 workflow: Workflow that started. 

57 """ 

58 pass 

59 

60 def on_workflow_completed(self, workflow: Workflow) -> None: 

61 """Called when a workflow completes. 

62 

63 Args: 

64 workflow: Workflow that completed. 

65 """ 

66 pass 

67 

68 

69class SlackCallback(Callback): 

70 """Callback that sends notifications to Slack via webhook.""" 

71 

72 def __init__(self, webhook_url: str): 

73 """Initialize Slack callback. 

74 

75 Args: 

76 webhook_url: Slack webhook URL for sending notifications. 

77 """ 

78 self.client = WebhookClient(webhook_url) 

79 

80 def on_job_submitted(self, job: JobType) -> None: 

81 """Send a message to Slack. 

82 

83 Args: 

84 job: Job that completed. 

85 message: Message to send. 

86 """ 

87 self.client.send( 

88 text="Job submitted", 

89 blocks=[ 

90 { 

91 "type": "section", 

92 "text": { 

93 "type": "mrkdwn", 

94 "text": f"`🌋 {'SUBMITTED':<12} Job {job.name:<12} (ID: {job.job_id})`", 

95 }, 

96 } 

97 ], 

98 ) 

99 

100 def on_job_completed(self, job: JobType) -> None: 

101 """Send completion notification to Slack. 

102 

103 Args: 

104 job: Job that completed. 

105 """ 

106 self.client.send( 

107 text="Job completed", 

108 blocks=[ 

109 { 

110 "type": "section", 

111 "text": {"type": "mrkdwn", "text": f"`{job_status_msg(job)}`"}, 

112 } 

113 ], 

114 ) 

115 

116 def on_job_failed(self, job: JobType) -> None: 

117 """Send failure notification to Slack. 

118 

119 Args: 

120 job: Job that failed. 

121 """ 

122 self.client.send( 

123 text="Job failed", 

124 blocks=[ 

125 { 

126 "type": "section", 

127 "text": {"type": "mrkdwn", "text": f"`{job_status_msg(job)}`"}, 

128 } 

129 ], 

130 ) 

131 

132 def on_workflow_completed(self, workflow: Workflow) -> None: 

133 """Send completion notification to Slack. 

134 

135 Args: 

136 workflow: Workflow that completed. 

137 """ 

138 self.client.send( 

139 text="Workflow completed", 

140 blocks=[ 

141 { 

142 "type": "section", 

143 "text": { 

144 "type": "mrkdwn", 

145 "text": f"🎉 Workflow {workflow.name} completed🎉", 

146 }, 

147 } 

148 ], 

149 )