Coverage for src/srunx/callbacks.py: 86%
29 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 15:10 +0000
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 15:10 +0000
1"""Callback system for job state notifications."""
3from slack_sdk import WebhookClient
5from srunx.models import JobType, Workflow
6from srunx.utils import job_status_msg
9class Callback:
10 """Base callback class for job state notifications."""
12 def on_job_submitted(self, job: JobType) -> None:
13 """Called when a job is submitted to SLURM.
15 Args:
16 job: Job that was submitted.
17 """
18 pass
20 def on_job_completed(self, job: JobType) -> None:
21 """Called when a job completes successfully.
23 Args:
24 job: Job that completed.
25 """
26 pass
28 def on_job_failed(self, job: JobType) -> None:
29 """Called when a job fails.
31 Args:
32 job: Job that failed.
33 """
34 pass
36 def on_job_running(self, job: JobType) -> None:
37 """Called when a job starts running.
39 Args:
40 job: Job that started running.
41 """
42 pass
44 def on_job_cancelled(self, job: JobType) -> None:
45 """Called when a job is cancelled.
47 Args:
48 job: Job that was cancelled.
49 """
50 pass
52 def on_workflow_started(self, workflow: Workflow) -> None:
53 """Called when a workflow starts.
55 Args:
56 workflow: Workflow that started.
57 """
58 pass
60 def on_workflow_completed(self, workflow: Workflow) -> None:
61 """Called when a workflow completes.
63 Args:
64 workflow: Workflow that completed.
65 """
66 pass
69class SlackCallback(Callback):
70 """Callback that sends notifications to Slack via webhook."""
72 def __init__(self, webhook_url: str):
73 """Initialize Slack callback.
75 Args:
76 webhook_url: Slack webhook URL for sending notifications.
77 """
78 self.client = WebhookClient(webhook_url)
80 def on_job_submitted(self, job: JobType) -> None:
81 """Send a message to Slack.
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 )
100 def on_job_completed(self, job: JobType) -> None:
101 """Send completion notification to Slack.
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 )
116 def on_job_failed(self, job: JobType) -> None:
117 """Send failure notification to Slack.
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 )
132 def on_workflow_completed(self, workflow: Workflow) -> None:
133 """Send completion notification to Slack.
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 )