From 7aa61048723d44a3da5d92ce085ff31819aa973f Mon Sep 17 00:00:00 2001
From: kk1050 <103404672+kk1050@users.noreply.github.com>
Date: Thu, 7 Jul 2022 17:01:34 +0800
Subject: [PATCH] Add method to check if termination is requested (#811, #1932)

Co-authored-by: kk105 <kkl@m-kabs.hk>
---
 artiq/frontend/artiq_master.py |  1 +
 artiq/master/scheduler.py      | 10 ++++++++++
 artiq/master/worker_impl.py    |  6 ++++++
 3 files changed, 17 insertions(+)

diff --git a/artiq/frontend/artiq_master.py b/artiq/frontend/artiq_master.py
index 99c2fc9f1..da57e6b3d 100755
--- a/artiq/frontend/artiq_master.py
+++ b/artiq/frontend/artiq_master.py
@@ -128,6 +128,7 @@ def main():
         "scheduler_request_termination": scheduler.request_termination,
         "scheduler_get_status": scheduler.get_status,
         "scheduler_check_pause": scheduler.check_pause,
+        "scheduler_check_termination": scheduler.check_termination,
         "ccb_issue": ccb_issue,
     })
     experiment_db.scan_repository_async()
diff --git a/artiq/master/scheduler.py b/artiq/master/scheduler.py
index d6d44acef..eea832a17 100644
--- a/artiq/master/scheduler.py
+++ b/artiq/master/scheduler.py
@@ -490,3 +490,13 @@ class Scheduler:
                     return False
                 return r.priority_key() > run.priority_key()
         raise KeyError("RID not found")
+
+    def check_termination(self, rid):
+        """Returns ``True`` if termination is requested."""
+        for pipeline in self._pipelines.values():
+            if rid in pipeline.pool.runs:
+                run = pipeline.pool.runs[rid]  
+                if run.termination_requested:
+                    return True
+        return False
+        
\ No newline at end of file
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index 97589dc14..33d34ddf8 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -111,6 +111,12 @@ class Scheduler:
             rid = self.rid
         return self._check_pause(rid)
 
+    _check_termination = staticmethod(make_parent_action("scheduler_check_termination"))
+    def check_termination(self, rid=None) -> TBool:
+        if rid is None:
+            rid = self.rid
+        return self._check_termination(rid)
+
     _submit = staticmethod(make_parent_action("scheduler_submit"))
     def submit(self, pipeline_name=None, expid=None, priority=None, due_date=None, flush=False):
         if pipeline_name is None: