From 060706da709dbe9bf69afc8d84c0b30be2db4c5f Mon Sep 17 00:00:00 2001
From: Jusong Yu <jusong.yeu@gmail.com>
Date: Tue, 17 Dec 2024 09:06:57 +0100
Subject: [PATCH] More explicit error handling for operate_calcjobs when
 connection reset

---
 src/aiida_sssp_workflow/workflows/__init__.py     | 15 +++++++++------
 .../workflows/evaluate/_phonon_frequencies.py     |  3 ++-
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/aiida_sssp_workflow/workflows/__init__.py b/src/aiida_sssp_workflow/workflows/__init__.py
index 6528042a..388240bf 100644
--- a/src/aiida_sssp_workflow/workflows/__init__.py
+++ b/src/aiida_sssp_workflow/workflows/__init__.py
@@ -31,12 +31,15 @@ def on_terminated(self):
             self.report(f"{type(self)}: remote folders will not be cleaned")
             return
 
-        # The clen_wordir is modified to not clean the remote folder of the caching node.
-        cleaned_calcs = operate_calcjobs(
-            self.node, operator=clean_workdir, all_same_nodes=False
-        )
-
-        if cleaned_calcs:
+        try:
+            cleaned_calcs = operate_calcjobs(
+                self.node, operator=clean_workdir, all_same_nodes=False
+            )
+        except RuntimeError as exc:
+            self.logger.warning(
+                f"clean remote workdir folder {self.inputs.clean_workir} failed: {exc}"
+            )
+        else:
             self.report(
                 f"cleaned remote folders of calculations: {' '.join(map(str, cleaned_calcs))}"
             )
diff --git a/src/aiida_sssp_workflow/workflows/evaluate/_phonon_frequencies.py b/src/aiida_sssp_workflow/workflows/evaluate/_phonon_frequencies.py
index b761a7d4..019b9f8e 100644
--- a/src/aiida_sssp_workflow/workflows/evaluate/_phonon_frequencies.py
+++ b/src/aiida_sssp_workflow/workflows/evaluate/_phonon_frequencies.py
@@ -3,6 +3,7 @@
 WorkChain calculate phonon frequencies at Gamma
 """
 
+from builtins import RuntimeError
 from aiida import orm
 from aiida.common import NotExistentAttributeError
 from aiida.engine import ToContext, while_
@@ -198,7 +199,7 @@ def finalize(self):
                 cleaned_calcs = operate_calcjobs(
                     self.node, operator=clean_workdir, all_same_nodes=False
                 )
-            except ConnectionError as exc:
+            except RuntimeError as exc:
                 self.logger.warning(
                     f"clean remote workdir folder {self.inputs.clean_workir} failed: {exc}"
                 )