on_kill() operator added.

docs/python_airflow_operator.md

@@ -81,6 +81,22 @@ Runs an Armada job and calls the job_service_client for polling.
 
 
 
+#### on_kill()
+Stops the JobService from listening to the JobSet and cancels the jobs.
+
+
+* **Returns**
+
+    None
+
+
+
+* **Return type**
+
+    None
+
+
+
 #### render_template_fields(context, jinja_env=None)
 Template all attributes listed in *self.template_fields*.
 

third_party/airflow/armada/operators/armada.py

@@ -151,6 +151,26 @@ def _get_lookout_url(self, job_id: str) -> str:
             return ""
         return self.lookout_url_template.replace("<job_id>", job_id)
 
+    def on_kill(self) -> None:
+        """
+        Stops the JobService from listening to the JobSet and cancels the jobs.
+
+        :return: None
+        """
+        try:
+            if self.job_set_id and self.queue:
+                # Cancel the jobs using the Armada client
+                self.armada_client.cancel_job(
+                    job_set_id=self.job_set_id, queue=self.queue
+                )
+                armada_logger.info(
+                    "Queue %s and JobSetId %s has been cancelled.",
+                    self.queue,
+                    self.job_set_id,
+                )
+        except Exception as e:
+            armada_logger.warning("Error during job cancellation: %s", str(e))
+
     def render_template_fields(
         self,
         context: Context,

third_party/airflow/tests/unit/test_armada_operator.py

@@ -1,8 +1,17 @@
+from armada.operators.armada import ArmadaOperator
+
+from armada_client import ArmadaClient
+from armada.operators.armada import ArmadaOperator
+from unittest.mock import patch, Mock
+
 import copy
 from unittest.mock import patch, Mock
 
 import grpc
 import pytest
+import unittest
+
+
 
 from armada.jobservice import jobservice_pb2
 from armada.operators.armada import ArmadaOperator
@@ -45,6 +54,35 @@ def test_get_lookout_url(lookout_url_template, job_id, expected_url):
     assert operator._get_lookout_url(job_id) == expected_url
 
 
+class TestJobService(unittest.TestCase):
+    @patch.object(JobService, "cancel_jobs")
+    def test_on_kill(self, mock_cancel_jobs):
+
+        operator = ArmadaOperator(
+            name="test_operator",
+            armada_channel_args={},
+            job_service_channel_args={},
+            armada_queue="test_queue",
+            job_request_items=[],
+        )
+
+        operator.job_set_id = "test_job_set_id"
+        operator.queue = "test_queue"
+
+        armada_client_instance = mock_armada_client.return_value
+        armada_client_instance.cancel_job.return_value = None
+
+        operator.on_kill()
+
+        armada_client_instance.cancel_job.assert_called_once_with(
+            job_set_id="test_job_set_id", queue="test_queue"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
+
+
 def test_deepcopy_operator():
     armada_channel_args = {"target": "127.0.0.1:50051"}
     job_service_channel_args = {"target": "127.0.0.1:60003"}