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

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

2

3 from slack_sdk import WebhookClient

4

5 from srunx.models import JobType, Workflow

6 from srunx.utils import job_status_msg

7

8

9 class 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

69 class 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 )