From f7aa93af172b3996e6de02f5e4140053dd4a96cc Mon Sep 17 00:00:00 2001
From: SchrodingersCattt <gmy721212@163.com>
Date: Mon, 20 Jul 2026 11:19:09 +0000
Subject: [PATCH 1/9] feat: add 'on_success' clean strategy to preserve workdir
 on failure
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Extend run_submission(clean=...) to accept string strategies in addition
to the existing bool:
- True / 'always': always clean remote workdir (backward compatible default)
- False / 'never': never clean
- 'on_success': only clean when ALL jobs finished successfully;
  preserve remote workdir on failure for post-mortem debugging

This is especially useful when debugging LAMMPS/DP-train failures on
remote clusters — previously clean=True would delete stderr/log files
before they could be inspected.

Add _should_clean() helper and comprehensive unit tests.

Backward compatible: True/False behavior is unchanged.
---
 dpdispatcher/submission.py   | 56 +++++++++++++++++++++++++++-
 tests/test_clean_strategy.py | 72 ++++++++++++++++++++++++++++++++++++
 2 files changed, 127 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_clean_strategy.py

diff --git a/dpdispatcher/submission.py b/dpdispatcher/submission.py
index da55a03c..5f19b304 100644
--- a/dpdispatcher/submission.py
+++ b/dpdispatcher/submission.py
@@ -210,6 +210,21 @@ def run_submission(
         Forth, wait until the tasks in the submission finished and download the result file to local directory.
         If dry_run is True, submission will be uploaded but not be executed and exit.
         If exit_on_submit is True, submission will exit.
+
+        Parameters
+        ----------
+        dry_run : bool
+            If True, only upload without execution.
+        exit_on_submit : bool
+            If True, exit after submission without waiting.
+        clean : bool or str
+            Controls whether to clean remote working directory after completion.
+            - True or "always": always clean (default, backward compatible)
+            - False or "never": never clean
+            - "on_success": only clean when all jobs finished successfully;
+              preserve remote workdir on failure for debugging.
+        check_interval : int
+            Seconds between status polling iterations.
         """
         assert self.resources is not None
         if not self.belonging_jobs:
@@ -261,10 +276,49 @@ def run_submission(
         self.handle_unexpected_submission_state()
         self.try_download_result()
         self.submission_to_json()
-        if clean:
+
+        # Determine whether to clean remote workdir
+        should_clean = self._should_clean(clean)
+        if should_clean:
             self.clean_jobs()
+        elif clean == "on_success":
+            dlog.info(
+                "clean='on_success': some jobs did not finish successfully, "
+                "preserving remote workdir for debugging at: "
+                f"{self.machine.context.remote_root}"
+            )
         return self.serialize()
 
+    def _should_clean(self, clean) -> bool:
+        """Determine whether remote workdir should be cleaned.
+
+        Parameters
+        ----------
+        clean : bool or str
+            - True or "always": always clean
+            - False or "never": never clean
+            - "on_success": clean only when all jobs finished successfully
+
+        Returns
+        -------
+        bool
+            Whether to perform clean.
+        """
+        if clean is True or clean == "always":
+            return True
+        if clean is False or clean == "never":
+            return False
+        if clean == "on_success":
+            return all(
+                job.job_state == JobStatus.finished for job in self.belonging_jobs
+            )
+        # Unknown clean value — treat as True for backward compatibility
+        dlog.warning(
+            f"Unknown clean strategy '{clean}', treating as True. "
+            f"Valid options: True, False, 'always', 'never', 'on_success'."
+        )
+        return True
+
     def try_download_result(self):
         start_time = time.time()
         retry_interval = 60  # retry every 1 minute
diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
new file mode 100644
index 00000000..5ed39be0
--- /dev/null
+++ b/tests/test_clean_strategy.py
@@ -0,0 +1,72 @@
+"""Test PR3: clean strategy (always / never / on_success) in Submission.run_submission()."""
+
+import os
+import sys
+import unittest
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from dpdispatcher.submission import Submission
+from dpdispatcher.utils.job_status import JobStatus
+
+
+class TestShouldClean(unittest.TestCase):
+    """Unit tests for Submission._should_clean() logic."""
+
+    def _make_submission_with_jobs(self, job_states):
+        """Create a Submission with mocked jobs in given states."""
+        submission = Submission.__new__(Submission)
+        submission.belonging_jobs = []
+        for state in job_states:
+            job = MagicMock()
+            job.job_state = state
+            submission.belonging_jobs.append(job)
+        return submission
+
+    def test_clean_true_always_cleans(self):
+        """clean=True (legacy) should always return True."""
+        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.terminated])
+        self.assertTrue(sub._should_clean(True))
+
+    def test_clean_false_never_cleans(self):
+        """clean=False (legacy) should always return False."""
+        sub = self._make_submission_with_jobs([JobStatus.finished])
+        self.assertFalse(sub._should_clean(False))
+
+    def test_clean_always_string(self):
+        """clean='always' behaves same as True."""
+        sub = self._make_submission_with_jobs([JobStatus.terminated])
+        self.assertTrue(sub._should_clean("always"))
+
+    def test_clean_never_string(self):
+        """clean='never' behaves same as False."""
+        sub = self._make_submission_with_jobs([JobStatus.finished])
+        self.assertFalse(sub._should_clean("never"))
+
+    def test_on_success_all_finished(self):
+        """clean='on_success' with all jobs finished → should clean."""
+        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.finished])
+        self.assertTrue(sub._should_clean("on_success"))
+
+    def test_on_success_some_terminated(self):
+        """clean='on_success' with some terminated jobs → should NOT clean."""
+        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.terminated])
+        self.assertFalse(sub._should_clean("on_success"))
+
+    def test_on_success_all_terminated(self):
+        """clean='on_success' with all terminated → should NOT clean."""
+        sub = self._make_submission_with_jobs([JobStatus.terminated, JobStatus.terminated])
+        self.assertFalse(sub._should_clean("on_success"))
+
+    def test_unknown_strategy_warns_and_cleans(self):
+        """Unknown clean value should warn and default to True."""
+        sub = self._make_submission_with_jobs([JobStatus.finished])
+        with self.assertLogs("dpdispatcher", level="WARNING") as cm:
+            result = sub._should_clean("invalid_value")
+        self.assertTrue(result)
+        self.assertTrue(any("Unknown clean strategy" in msg for msg in cm.output))
+
+
+if __name__ == "__main__":
+    unittest.main()

From 0842e7a19b3f53c537551e2de026eb87aa9e6763 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 20 Jul 2026 11:42:16 +0000
Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_clean_strategy.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
index 5ed39be0..d4d1ad13 100644
--- a/tests/test_clean_strategy.py
+++ b/tests/test_clean_strategy.py
@@ -3,7 +3,7 @@
 import os
 import sys
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
@@ -26,7 +26,9 @@ def _make_submission_with_jobs(self, job_states):
 
     def test_clean_true_always_cleans(self):
         """clean=True (legacy) should always return True."""
-        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.terminated])
+        sub = self._make_submission_with_jobs(
+            [JobStatus.finished, JobStatus.terminated]
+        )
         self.assertTrue(sub._should_clean(True))
 
     def test_clean_false_never_cleans(self):
@@ -51,12 +53,16 @@ def test_on_success_all_finished(self):
 
     def test_on_success_some_terminated(self):
         """clean='on_success' with some terminated jobs → should NOT clean."""
-        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.terminated])
+        sub = self._make_submission_with_jobs(
+            [JobStatus.finished, JobStatus.terminated]
+        )
         self.assertFalse(sub._should_clean("on_success"))
 
     def test_on_success_all_terminated(self):
         """clean='on_success' with all terminated → should NOT clean."""
-        sub = self._make_submission_with_jobs([JobStatus.terminated, JobStatus.terminated])
+        sub = self._make_submission_with_jobs(
+            [JobStatus.terminated, JobStatus.terminated]
+        )
         self.assertFalse(sub._should_clean("on_success"))
 
     def test_unknown_strategy_warns_and_cleans(self):

From 1dee3ea9a9f74f2d32366e93c82892ad42099da2 Mon Sep 17 00:00:00 2001
From: SchrodingersCattt <gmy721212@163.com>
Date: Mon, 20 Jul 2026 22:39:41 +0000
Subject: [PATCH 3/9] fix: on_success must not clean when ratio_unfinished
 kills jobs

Address PR review: remove_unfinished_tasks() mutates killed jobs' state
to JobStatus.finished, which deceived _should_clean('on_success') into
cleaning the remote workdir even though tasks were deliberately killed.

Fix: use a while/else pattern to track whether all jobs genuinely
completed (all_jobs_genuinely_finished flag). The flag is True only when
the while loop exits naturally (all jobs really finished), and False
when the ratio_unfinished early-exit path triggers break.

_should_clean() now takes all_genuinely_finished as a parameter instead
of inspecting (potentially mutated) job states.

Add integration test TestCleanWithRatioUnfinished that mocks the full
run_submission path with ratio_unfinished triggering early exit, and
verifies clean_jobs is NOT called with clean='on_success'.
---
 dpdispatcher/submission.py   |  24 +++++--
 tests/test_clean_strategy.py | 135 ++++++++++++++++++++++++++++-------
 2 files changed, 129 insertions(+), 30 deletions(-)

diff --git a/dpdispatcher/submission.py b/dpdispatcher/submission.py
index 5f19b304..781a9292 100644
--- a/dpdispatcher/submission.py
+++ b/dpdispatcher/submission.py
@@ -248,6 +248,11 @@ def run_submission(
             self.handle_unexpected_submission_state()
 
         ratio_unfinished = self.resources.strategy["ratio_unfinished"]
+        # Track whether all jobs genuinely succeeded (before any state mutation).
+        # remove_unfinished_tasks() rewrites killed jobs' state to "finished",
+        # which would fool _should_clean("on_success"). We capture the real
+        # outcome here: True only if the loop exits normally (all finished).
+        all_jobs_genuinely_finished = False
         while not self.check_all_finished():
             if exit_on_submit is True:
                 dlog.info(f"submission succeeded: {self.submission_hash}")
@@ -273,12 +278,16 @@ def run_submission(
                 self.handle_unexpected_submission_state()
             finally:
                 pass
+        else:
+            # Loop exited normally (check_all_finished() was True from the start
+            # or became True without hitting the ratio_unfinished early-exit).
+            all_jobs_genuinely_finished = True
         self.handle_unexpected_submission_state()
         self.try_download_result()
         self.submission_to_json()
 
         # Determine whether to clean remote workdir
-        should_clean = self._should_clean(clean)
+        should_clean = self._should_clean(clean, all_jobs_genuinely_finished)
         if should_clean:
             self.clean_jobs()
         elif clean == "on_success":
@@ -289,7 +298,7 @@ def run_submission(
             )
         return self.serialize()
 
-    def _should_clean(self, clean) -> bool:
+    def _should_clean(self, clean, all_genuinely_finished: bool = True) -> bool:
         """Determine whether remote workdir should be cleaned.
 
         Parameters
@@ -297,7 +306,12 @@ def _should_clean(self, clean) -> bool:
         clean : bool or str
             - True or "always": always clean
             - False or "never": never clean
-            - "on_success": clean only when all jobs finished successfully
+            - "on_success": clean only when all jobs genuinely finished
+              (not killed by ratio_unfinished early-exit)
+        all_genuinely_finished : bool
+            Whether all jobs completed successfully without intervention.
+            When ratio_unfinished triggers remove_unfinished_tasks(), this
+            is False even though job states have been mutated to "finished".
 
         Returns
         -------
@@ -309,9 +323,7 @@ def _should_clean(self, clean) -> bool:
         if clean is False or clean == "never":
             return False
         if clean == "on_success":
-            return all(
-                job.job_state == JobStatus.finished for job in self.belonging_jobs
-            )
+            return all_genuinely_finished
         # Unknown clean value — treat as True for backward compatibility
         dlog.warning(
             f"Unknown clean strategy '{clean}', treating as True. "
diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
index d4d1ad13..6425d7e7 100644
--- a/tests/test_clean_strategy.py
+++ b/tests/test_clean_strategy.py
@@ -3,7 +3,7 @@
 import os
 import sys
 import unittest
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
@@ -25,45 +25,47 @@ def _make_submission_with_jobs(self, job_states):
         return submission
 
     def test_clean_true_always_cleans(self):
-        """clean=True (legacy) should always return True."""
-        sub = self._make_submission_with_jobs(
-            [JobStatus.finished, JobStatus.terminated]
-        )
-        self.assertTrue(sub._should_clean(True))
+        """clean=True (legacy) should always return True regardless of job states."""
+        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.terminated])
+        self.assertTrue(sub._should_clean(True, all_genuinely_finished=False))
 
     def test_clean_false_never_cleans(self):
         """clean=False (legacy) should always return False."""
         sub = self._make_submission_with_jobs([JobStatus.finished])
-        self.assertFalse(sub._should_clean(False))
+        self.assertFalse(sub._should_clean(False, all_genuinely_finished=True))
 
     def test_clean_always_string(self):
         """clean='always' behaves same as True."""
         sub = self._make_submission_with_jobs([JobStatus.terminated])
-        self.assertTrue(sub._should_clean("always"))
+        self.assertTrue(sub._should_clean("always", all_genuinely_finished=False))
 
     def test_clean_never_string(self):
         """clean='never' behaves same as False."""
         sub = self._make_submission_with_jobs([JobStatus.finished])
-        self.assertFalse(sub._should_clean("never"))
+        self.assertFalse(sub._should_clean("never", all_genuinely_finished=True))
 
-    def test_on_success_all_finished(self):
-        """clean='on_success' with all jobs finished → should clean."""
+    def test_on_success_genuinely_finished(self):
+        """clean='on_success' with all_genuinely_finished=True → should clean."""
         sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.finished])
-        self.assertTrue(sub._should_clean("on_success"))
+        self.assertTrue(sub._should_clean("on_success", all_genuinely_finished=True))
 
-    def test_on_success_some_terminated(self):
-        """clean='on_success' with some terminated jobs → should NOT clean."""
-        sub = self._make_submission_with_jobs(
-            [JobStatus.finished, JobStatus.terminated]
-        )
-        self.assertFalse(sub._should_clean("on_success"))
+    def test_on_success_not_genuinely_finished(self):
+        """clean='on_success' with all_genuinely_finished=False → should NOT clean.
 
-    def test_on_success_all_terminated(self):
-        """clean='on_success' with all terminated → should NOT clean."""
-        sub = self._make_submission_with_jobs(
-            [JobStatus.terminated, JobStatus.terminated]
-        )
-        self.assertFalse(sub._should_clean("on_success"))
+        This covers the ratio_unfinished path where remove_unfinished_tasks()
+        kills jobs and mutates their state to 'finished', but the submission
+        did not genuinely succeed.
+        """
+        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.finished])
+        # Even though all job_states are "finished" (mutated by remove_unfinished_tasks),
+        # we explicitly know not all jobs genuinely completed.
+        self.assertFalse(sub._should_clean("on_success", all_genuinely_finished=False))
+
+    def test_on_success_default_genuinely_finished(self):
+        """Default all_genuinely_finished=True for backward compat (normal exit path)."""
+        sub = self._make_submission_with_jobs([JobStatus.finished])
+        # When called without the second arg, defaults to True
+        self.assertTrue(sub._should_clean("on_success"))
 
     def test_unknown_strategy_warns_and_cleans(self):
         """Unknown clean value should warn and default to True."""
@@ -74,5 +76,90 @@ def test_unknown_strategy_warns_and_cleans(self):
         self.assertTrue(any("Unknown clean strategy" in msg for msg in cm.output))
 
 
+class TestCleanWithRatioUnfinished(unittest.TestCase):
+    """Integration test: on_success should NOT clean when ratio_unfinished triggers early exit.
+
+    This tests the full run_submission path where remove_unfinished_tasks()
+    is triggered, mutating job states to 'finished'. The clean='on_success'
+    strategy must still recognize that some tasks were killed (not genuinely
+    successful) and preserve the remote workdir.
+    """
+
+    def test_ratio_unfinished_prevents_clean_on_success(self):
+        """ratio_unfinished early-exit → clean='on_success' must NOT clean."""
+        sub = Submission.__new__(Submission)
+        sub.belonging_jobs = []
+        sub.belonging_tasks = []
+        sub.submission_hash = "test_hash"
+
+        # Mock machine and context
+        sub.machine = MagicMock()
+        sub.machine.context.remote_root = "/tmp/fake_remote"
+        sub.machine.context.local_root = "/tmp/fake_local"
+
+        # Mock resources with ratio_unfinished > 0
+        sub.resources = MagicMock()
+        sub.resources.strategy = {"ratio_unfinished": 0.5}
+        sub.resources.wait_time = 0
+
+        # Create 4 jobs: 2 finished, 2 running (will be killed by ratio_unfinished)
+        for i in range(4):
+            job = MagicMock()
+            job.job_hash = f"job_{i}"
+            job.job_id = str(i)
+            if i < 2:
+                job.job_state = JobStatus.finished
+            else:
+                job.job_state = JobStatus.running
+            sub.belonging_jobs.append(job)
+
+        # Create tasks matching job states
+        for i in range(4):
+            task = MagicMock()
+            if i < 2:
+                task.task_state = JobStatus.finished
+            else:
+                task.task_state = JobStatus.running
+            sub.belonging_tasks.append(task)
+
+        # Mock methods called by run_submission
+        sub.generate_jobs = MagicMock()
+        sub.try_recover_from_json = MagicMock()
+        sub.upload_jobs = MagicMock()
+        sub.handle_unexpected_submission_state = MagicMock()
+        sub.submission_to_json = MagicMock()
+        sub.try_download_result = MagicMock()
+        sub.clean_jobs = MagicMock()
+        sub.serialize = MagicMock(return_value={})
+
+        # Make update_submission_state a no-op (states are manually set)
+        sub.update_submission_state = MagicMock()
+
+        # check_all_finished is called multiple times in run_submission:
+        #   1. Line 234: if self.check_all_finished() — initial check
+        #   2. Line 246: self.check_all_finished() — after upload (return discarded)
+        #   3. Line 256: while not self.check_all_finished() — loop condition
+        #   Inside loop: ratio_unfinished triggers break
+        call_count = [0]
+
+        def mock_check_all_finished():
+            call_count[0] += 1
+            # Calls 1-3: not all finished (jobs 2,3 still running)
+            if call_count[0] <= 3:
+                return False
+            return True
+
+        sub.check_all_finished = mock_check_all_finished
+
+        # check_ratio_unfinished returns True (triggers early exit)
+        sub.check_ratio_unfinished = MagicMock(return_value=True)
+
+        # Run with clean="on_success"
+        sub.run_submission(clean="on_success", check_interval=0)
+
+        # clean_jobs should NOT have been called because not all genuinely finished
+        sub.clean_jobs.assert_not_called()
+
+
 if __name__ == "__main__":
     unittest.main()

From 864e64e73972db369a9742a20cdefc4947fec3e5 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 20 Jul 2026 22:39:54 +0000
Subject: [PATCH 4/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_clean_strategy.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
index 6425d7e7..e1363c5c 100644
--- a/tests/test_clean_strategy.py
+++ b/tests/test_clean_strategy.py
@@ -3,7 +3,7 @@
 import os
 import sys
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
@@ -26,7 +26,9 @@ def _make_submission_with_jobs(self, job_states):
 
     def test_clean_true_always_cleans(self):
         """clean=True (legacy) should always return True regardless of job states."""
-        sub = self._make_submission_with_jobs([JobStatus.finished, JobStatus.terminated])
+        sub = self._make_submission_with_jobs(
+            [JobStatus.finished, JobStatus.terminated]
+        )
         self.assertTrue(sub._should_clean(True, all_genuinely_finished=False))
 
     def test_clean_false_never_cleans(self):

From 03692c1423aaeb90bad3b1dad278abb121371264 Mon Sep 17 00:00:00 2001
From: SchrodingersCattt <gmy721212@163.com>
Date: Mon, 20 Jul 2026 22:47:34 +0000
Subject: [PATCH 5/9] fix: add Union type hint for clean param, raise
 ValueError on invalid strategy

Address CodeRabbit review comments:

1. Add Union[bool, str] type hint to _should_clean(clean=...) parameter,
   import Union from typing (Python 3.7+ compatible).

2. Replace warning+return True on unknown clean value with raise ValueError.
   A typo like 'on_sucess' would previously silently clean the workdir
   (the most destructive option), defeating the purpose of on_success.
   Now it fails loudly at the call site.

Update test to assertRaises(ValueError) instead of assertLogs(WARNING).
---
 dpdispatcher/submission.py   | 19 ++++++++++++-------
 tests/test_clean_strategy.py | 12 ++++++------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/dpdispatcher/submission.py b/dpdispatcher/submission.py
index 781a9292..0285b341 100644
--- a/dpdispatcher/submission.py
+++ b/dpdispatcher/submission.py
@@ -9,7 +9,7 @@
 import time
 import uuid
 from hashlib import sha1
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import yaml
 from dargs.dargs import Argument, Variant
@@ -298,12 +298,14 @@ def run_submission(
             )
         return self.serialize()
 
-    def _should_clean(self, clean, all_genuinely_finished: bool = True) -> bool:
+    def _should_clean(
+        self, clean: Union[bool, str], all_genuinely_finished: bool = True
+    ) -> bool:
         """Determine whether remote workdir should be cleaned.
 
         Parameters
         ----------
-        clean : bool or str
+        clean : Union[bool, str]
             - True or "always": always clean
             - False or "never": never clean
             - "on_success": clean only when all jobs genuinely finished
@@ -317,6 +319,11 @@ def _should_clean(self, clean, all_genuinely_finished: bool = True) -> bool:
         -------
         bool
             Whether to perform clean.
+
+        Raises
+        ------
+        ValueError
+            If clean is not a recognized strategy.
         """
         if clean is True or clean == "always":
             return True
@@ -324,12 +331,10 @@ def _should_clean(self, clean, all_genuinely_finished: bool = True) -> bool:
             return False
         if clean == "on_success":
             return all_genuinely_finished
-        # Unknown clean value — treat as True for backward compatibility
-        dlog.warning(
-            f"Unknown clean strategy '{clean}', treating as True. "
+        raise ValueError(
+            f"Unknown clean strategy '{clean}'. "
             f"Valid options: True, False, 'always', 'never', 'on_success'."
         )
-        return True
 
     def try_download_result(self):
         start_time = time.time()
diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
index 6425d7e7..973c8877 100644
--- a/tests/test_clean_strategy.py
+++ b/tests/test_clean_strategy.py
@@ -67,13 +67,13 @@ def test_on_success_default_genuinely_finished(self):
         # When called without the second arg, defaults to True
         self.assertTrue(sub._should_clean("on_success"))
 
-    def test_unknown_strategy_warns_and_cleans(self):
-        """Unknown clean value should warn and default to True."""
+    def test_unknown_strategy_raises(self):
+        """Unknown clean value should raise ValueError (fail loudly, not silently clean)."""
         sub = self._make_submission_with_jobs([JobStatus.finished])
-        with self.assertLogs("dpdispatcher", level="WARNING") as cm:
-            result = sub._should_clean("invalid_value")
-        self.assertTrue(result)
-        self.assertTrue(any("Unknown clean strategy" in msg for msg in cm.output))
+        with self.assertRaises(ValueError) as ctx:
+            sub._should_clean("invalid_value")
+        self.assertIn("Unknown clean strategy", str(ctx.exception))
+        self.assertIn("invalid_value", str(ctx.exception))
 
 
 class TestCleanWithRatioUnfinished(unittest.TestCase):

From d09f1a306967a769d01fc0265c9046f16c48c642 Mon Sep 17 00:00:00 2001
From: SchrodingersCattt <gmy721212@163.com>
Date: Tue, 4 Aug 2026 09:50:27 +0000
Subject: [PATCH 6/9] fix: fail-fast validation of clean strategy before
 submission

Reject unknown clean strategies (e.g. typos like 'on_sucess') immediately at the
top of run_submission(), before try_recover_from_json() / upload_jobs() /
scheduler submission. Previously the ValueError was only raised at the end of
run_submission(), after hours of HPC work had already been performed.

Addresses reviewer feedback from njzjz-bot (P2).
---
 dpdispatcher/submission.py   |  2 ++
 tests/test_clean_strategy.py | 23 +++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/dpdispatcher/submission.py b/dpdispatcher/submission.py
index 0285b341..8a5ee438 100644
--- a/dpdispatcher/submission.py
+++ b/dpdispatcher/submission.py
@@ -227,6 +227,8 @@ def run_submission(
             Seconds between status polling iterations.
         """
         assert self.resources is not None
+        # Fail-fast: reject invalid clean strategies before recovery/upload/submission.
+        self._should_clean(clean, all_genuinely_finished=False)
         if not self.belonging_jobs:
             self.generate_jobs()
         self.try_recover_from_json()
diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
index f6c67e35..9806dbf3 100644
--- a/tests/test_clean_strategy.py
+++ b/tests/test_clean_strategy.py
@@ -78,6 +78,29 @@ def test_unknown_strategy_raises(self):
         self.assertIn("invalid_value", str(ctx.exception))
 
 
+class TestInvalidStrategyFailsFast(unittest.TestCase):
+    """Invalid clean strategy must raise ValueError BEFORE upload_jobs() is called."""
+
+    def test_invalid_strategy_raises_before_upload(self):
+        """Passing an invalid clean strategy raises ValueError without calling upload_jobs."""
+        from unittest.mock import patch
+
+        sub = Submission.__new__(Submission)
+        sub.belonging_jobs = []
+        sub.belonging_tasks = []
+        sub.submission_hash = "test_hash"
+        sub.machine = MagicMock()
+        sub.resources = MagicMock()
+        sub.resources.strategy = {"ratio_unfinished": 0.0}
+        sub.resources.wait_time = 0
+
+        with patch.object(Submission, "upload_jobs") as mock_upload:
+            with self.assertRaises(ValueError) as ctx:
+                sub.run_submission(clean="on_sucess", check_interval=0)
+            self.assertIn("Unknown clean strategy", str(ctx.exception))
+            mock_upload.assert_not_called()
+
+
 class TestCleanWithRatioUnfinished(unittest.TestCase):
     """Integration test: on_success should NOT clean when ratio_unfinished triggers early exit.
 

From 8229974ef29d7401097484d9da99c79c73bfdc58 Mon Sep 17 00:00:00 2001
From: SchrodingersCattt <gmy721212@163.com>
Date: Tue, 4 Aug 2026 10:02:35 +0000
Subject: [PATCH 7/9] fix: add type hint to clean param, fix async truthiness
 check

- Annotate run_submission(clean: Union[bool, str] = True) for Python 3.7+ compat
- Fix async_run_submission: use _should_clean() instead of truthiness check
  ('never' is truthy but should not trigger the warning)

Addresses coderabbit review comments.
---
 dpdispatcher/submission.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/dpdispatcher/submission.py b/dpdispatcher/submission.py
index 8a5ee438..5249473d 100644
--- a/dpdispatcher/submission.py
+++ b/dpdispatcher/submission.py
@@ -201,7 +201,12 @@ def bind_machine(self, machine):
         return self
 
     def run_submission(
-        self, *, dry_run=False, exit_on_submit=False, clean=True, check_interval=30
+        self,
+        *,
+        dry_run=False,
+        exit_on_submit=False,
+        clean: Union[bool, str] = True,
+        check_interval=30,
     ):
         """Main method to execute the submission.
         First, check whether old Submission exists on the remote machine, and try to recover from it.
@@ -392,9 +397,10 @@ async def async_run_submission(self, **kwargs):
         May raise Error if pass `clean=True` explicitly when submit to pbs or slurm.
         """
         kwargs = {**{"clean": False}, **kwargs}
-        if kwargs["clean"]:
+        if self._should_clean(kwargs["clean"]):
             dlog.warning(
-                "Using async submission with `clean=True`, job may fail in queue system"
+                "Using async submission with a clean strategy that can delete "
+                "the remote workdir. Jobs may fail in queue systems."
             )
         loop = asyncio.get_event_loop()
         wrapped_submission = functools.partial(self.run_submission, **kwargs)

From 32d75f08855b9372761da5eddba6b4407e74eb46 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Sun, 30 Aug 2026 01:37:33 +0800
Subject: [PATCH 8/9] fix: preserve results after download failure

Make result-download success part of on_success cleanup, and complete the affected public type annotations.

Coding-Agent: Codex

Codex-Version: codex-cli 0.151.0

Model: gpt-5.6-sol

Reasoning-Effort: xhigh
---
 dpdispatcher/submission.py   | 56 ++++++++++++++++++++++--------------
 tests/test_clean_strategy.py | 53 ++++++++++++++++++++++++++++++++--
 2 files changed, 85 insertions(+), 24 deletions(-)

diff --git a/dpdispatcher/submission.py b/dpdispatcher/submission.py
index ebf994ce..b90fa499 100644
--- a/dpdispatcher/submission.py
+++ b/dpdispatcher/submission.py
@@ -203,11 +203,11 @@ def bind_machine(self, machine):
     def run_submission(
         self,
         *,
-        dry_run=False,
-        exit_on_submit=False,
+        dry_run: bool = False,
+        exit_on_submit: bool = False,
         clean: Union[bool, str] = True,
-        check_interval=30,
-    ):
+        check_interval: int = 30,
+    ) -> Dict[str, Any]:
         """Main method to execute the submission.
         First, check whether old Submission exists on the remote machine, and try to recover from it.
         Second, upload the local files to the remote machine where the tasks to be executed.
@@ -224,6 +224,7 @@ def run_submission(
             If True, exit after submission without waiting.
         clean : bool or str
             Controls whether to clean remote working directory after completion.
+
             - True or "always": always clean (default, backward compatible)
             - False or "never": never clean
             - "on_success": only clean when all jobs finished successfully;
@@ -288,7 +289,10 @@ def run_submission(
                 # The loop condition became false without ratio-based early exit.
                 all_jobs_genuinely_finished = True
             self.handle_unexpected_submission_state()
-            self.try_download_result()
+            results_downloaded = self.try_download_result()
+            all_jobs_genuinely_finished = (
+                all_jobs_genuinely_finished and results_downloaded
+            )
         finally:
             # Cover recovery, initial submission, polling, and final download
             # failures so exhausted retries always preserve diagnostics.
@@ -388,7 +392,16 @@ def try_download_error_info(self) -> None:
                         f"Could not download error file for job {job.job_hash}: {e}"
                     )
 
-    def try_download_result(self):
+    def try_download_result(self) -> bool:
+        """Download result files, retrying transient failures for up to 24 hours.
+
+        Returns
+        -------
+        bool
+            Whether all result files were downloaded successfully. A false
+            result prevents ``clean='on_success'`` from deleting the only
+            remaining remote copy after retry exhaustion.
+        """
         start_time = time.time()
         retry_interval = 60  # retry every 1 minute
         success = False
@@ -412,6 +425,7 @@ def try_download_result(self):
                 else:  # > 24 h
                     dlog.info("Maximum retries time reached. Exiting.")
                     break
+        return success
 
     async def async_run_submission(self, **kwargs):
         """Async interface of run_submission.
@@ -690,13 +704,13 @@ class Task:
 
     def __init__(
         self,
-        command,
-        task_work_path,
+        command: str,
+        task_work_path: str,
         forward_files: Optional[Sequence[str]] = None,
         backward_files: Optional[Sequence[str]] = None,
-        outlog="log",
-        errlog="err",
-    ):
+        outlog: Optional[str] = "log",
+        errlog: Optional[str] = "err",
+    ) -> None:
         self.command = command
         self.task_work_path = task_work_path
         # Detach task state from caller-owned lists and constructor defaults.
@@ -1191,25 +1205,25 @@ class Resources:
 
     def __init__(
         self,
-        number_node,
-        cpu_per_node,
-        gpu_per_node,
-        queue_name,
-        group_size,
+        number_node: int,
+        cpu_per_node: int,
+        gpu_per_node: int,
+        queue_name: str,
+        group_size: int,
         *,
         custom_flags: Optional[Sequence[str]] = None,
         strategy: Optional[Dict[str, Any]] = None,
-        para_deg=1,
+        para_deg: int = 1,
         module_unload_list: Optional[Sequence[str]] = None,
-        module_purge=False,
+        module_purge: bool = False,
         module_list: Optional[Sequence[str]] = None,
         source_list: Optional[Sequence[str]] = None,
         envs: Optional[Dict[str, Any]] = None,
         prepend_script: Optional[Sequence[str]] = None,
         append_script: Optional[Sequence[str]] = None,
-        wait_time=0,
-        **kwargs,
-    ):
+        wait_time: int = 0,
+        **kwargs: Any,
+    ) -> None:
         self.number_node = number_node
         self.cpu_per_node = cpu_per_node
         self.gpu_per_node = gpu_per_node
diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
index 9806dbf3..677a8e67 100644
--- a/tests/test_clean_strategy.py
+++ b/tests/test_clean_strategy.py
@@ -3,7 +3,7 @@
 import os
 import sys
 import unittest
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
@@ -83,8 +83,6 @@ class TestInvalidStrategyFailsFast(unittest.TestCase):
 
     def test_invalid_strategy_raises_before_upload(self):
         """Passing an invalid clean strategy raises ValueError without calling upload_jobs."""
-        from unittest.mock import patch
-
         sub = Submission.__new__(Submission)
         sub.belonging_jobs = []
         sub.belonging_tasks = []
@@ -101,6 +99,55 @@ def test_invalid_strategy_raises_before_upload(self):
             mock_upload.assert_not_called()
 
 
+class TestDownloadResult(unittest.TestCase):
+    """Result-download status must distinguish success from retry exhaustion."""
+
+    def test_successful_download_returns_true(self):
+        sub = Submission.__new__(Submission)
+        sub.download_jobs = MagicMock()
+
+        self.assertTrue(sub.try_download_result())
+        sub.download_jobs.assert_called_once_with()
+
+    def test_retry_exhaustion_returns_false(self):
+        sub = Submission.__new__(Submission)
+        sub.download_jobs = MagicMock(side_effect=OSError("temporary failure"))
+
+        with patch("dpdispatcher.submission.time.time", side_effect=[0, 86400]), patch(
+            "dpdispatcher.submission.time.sleep"
+        ) as mock_sleep:
+            self.assertFalse(sub.try_download_result())
+
+        sub.download_jobs.assert_called_once_with()
+        mock_sleep.assert_not_called()
+
+    def test_retry_exhaustion_prevents_on_success_cleanup(self):
+        sub = Submission.__new__(Submission)
+        sub.belonging_jobs = [MagicMock(job_state=JobStatus.finished)]
+        sub.belonging_tasks = []
+        sub.submission_hash = "test_hash"
+        sub.machine = MagicMock()
+        sub.machine.context.remote_root = "/tmp/fake_remote"
+        sub.resources = MagicMock()
+        sub.resources.strategy = {"ratio_unfinished": 0.0}
+        sub.resources.wait_time = 0
+
+        sub.try_recover_from_json = MagicMock()
+        sub.update_submission_state = MagicMock()
+        sub.check_all_finished = MagicMock(return_value=True)
+        sub.handle_unexpected_submission_state = MagicMock()
+        sub.try_download_result = MagicMock(return_value=False)
+        sub.try_download_error_info = MagicMock()
+        sub.submission_to_json = MagicMock()
+        sub.clean_jobs = MagicMock()
+        sub.serialize = MagicMock(return_value={})
+
+        sub.run_submission(clean="on_success", check_interval=0)
+
+        sub.try_download_result.assert_called_once_with()
+        sub.clean_jobs.assert_not_called()
+
+
 class TestCleanWithRatioUnfinished(unittest.TestCase):
     """Integration test: on_success should NOT clean when ratio_unfinished triggers early exit.
 

From 44bb3b1e7cadc5d6fa9d4b07a0baac09a176f480 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Sun, 30 Aug 2026 01:55:48 +0800
Subject: [PATCH 9/9] fix(test): isolate retry clock from logging

Mock the submission logger so older Python logging implementations cannot consume the retry test's time side effects.

Coding-Agent: Codex

Codex-Version: codex-cli 0.151.0

Model: gpt-5.6-sol

Reasoning-Effort: xhigh
---
 tests/test_clean_strategy.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/test_clean_strategy.py b/tests/test_clean_strategy.py
index 677a8e67..4f690522 100644
--- a/tests/test_clean_strategy.py
+++ b/tests/test_clean_strategy.py
@@ -113,9 +113,12 @@ def test_retry_exhaustion_returns_false(self):
         sub = Submission.__new__(Submission)
         sub.download_jobs = MagicMock(side_effect=OSError("temporary failure"))
 
-        with patch("dpdispatcher.submission.time.time", side_effect=[0, 86400]), patch(
-            "dpdispatcher.submission.time.sleep"
-        ) as mock_sleep:
+        # Mock the module logger as well as the clock: older Python logging
+        # implementations call ``time.time()`` while formatting this expected
+        # retry error, which would otherwise consume the test clock values.
+        with patch("dpdispatcher.submission.dlog"), patch(
+            "dpdispatcher.submission.time.time", side_effect=[0, 86400]
+        ), patch("dpdispatcher.submission.time.sleep") as mock_sleep:
             self.assertFalse(sub.try_download_result())
 
         sub.download_jobs.assert_called_once_with()