Merge pull request #125 from felix5572/master

Author: njzjz

dpdispatcher/submission.py

@@ -156,14 +156,19 @@ def run_submission(self, *, exit_on_submit=False, clean=True):
         if not self.belonging_jobs:
             self.generate_jobs()
         self.try_recover_from_json()
+        self.update_submission_state()
         if self.check_all_finished():
             dlog.info('info:check_all_finished: True')
         else:
             dlog.info('info:check_all_finished: False')
             self.upload_jobs()
             self.handle_unexpected_submission_state()
             self.submission_to_json()
-        time.sleep(1)
+            time.sleep(1)
+            self.update_submission_state()
+            self.check_all_finished()
+            self.handle_unexpected_submission_state()
+
         while not self.check_all_finished():
             if exit_on_submit is True:
                 dlog.info(f"submission succeeded: {self.submission_hash}")
@@ -179,6 +184,7 @@ def run_submission(self, *, exit_on_submit=False, clean=True):
                 dlog.debug(self.serialize())
                 raise e
             else:
+                self.update_submission_state()
                 self.handle_unexpected_submission_state()
             finally:
                 pass
@@ -189,7 +195,7 @@ def run_submission(self, *, exit_on_submit=False, clean=True):
             self.clean_jobs()
         return self.serialize()
 
-    def get_submission_state(self):
+    def update_submission_state(self):
         """check whether all the jobs in the submission.
 
         Notes
@@ -201,7 +207,7 @@ def get_submission_state(self):
                 # finished job will be finished for ever, skip
                 continue
             job.get_job_state()
-            dlog.debug(f"debug:get_submission_state: job: {job.job_hash}, {job.job_id}, {repr(job.job_state)}")
+            dlog.debug(f"debug:update_submission_state: job: {job.job_hash}, {job.job_id}, {job.job_state}")
         # self.submission_to_json()
 
     def handle_unexpected_submission_state(self):
@@ -217,9 +223,10 @@ def handle_unexpected_submission_state(self):
             self.submission_to_json()
             raise RuntimeError(
                 f"Meet errors will handle unexpected submission state.\n"
-                f"Debug information: remote_root=={self.remote_root}.\n"
+                f"Debug information: remote_root=={self.machine.context.remote_root}.\n"
                 f"Debug information: submission_hash=={self.submission_hash}.\n"
                 f"Please check the dirs and scripts in remote_root"
+                f"The job information mentioned above may help"
             ) from e
 
     # not used here, submitting job is in handle_unexpected_submission_state.
@@ -231,14 +238,16 @@ def handle_unexpected_submission_state(self):
     #         job.submit_job()
     #     self.get_submission_state()
 
+    # def update_submi
+
     def check_all_finished(self):
         """check whether all the jobs in the submission.
 
         Notes
         -----
         This method will not handle unexpected job state in the submission.
         """
-        self.get_submission_state()
+        # self.update_submission_state()
         if any( (job.job_state in  [JobStatus.terminated, JobStatus.unknown] ) for job in self.belonging_jobs):
             self.submission_to_json()
         if any( (job.job_state in  [JobStatus.running,
@@ -294,7 +303,7 @@ def clean_jobs(self):
         self.machine.context.clean()
 
     def submission_to_json(self):
-        # self.get_submission_state()
+        # self.update_submission_state()
         write_str = json.dumps(self.serialize(), indent=4, default=str)
         submission_file_name = "{submission_hash}.json".format(submission_hash=self.submission_hash)
         self.machine.context.write_file(submission_file_name, write_str=write_str)
@@ -532,11 +541,10 @@ def handle_unexpected_job_state(self):
                 raise RuntimeError(f"job:{self.job_hash} {self.job_id} failed {self.fail_count} times.job_detail:{self}")
             self.submit_job()
             dlog.info("job:{job_hash} re-submit after terminated; new job_id is {job_id}".format(job_hash=self.job_hash, job_id=self.job_id))
+            time.sleep(0.2)
             self.get_job_state()
-            dlog.info("job:{job_hash} job_id:{job_id} after re-submitting; the state now is {job_state}".format(
-                job_hash=self.job_hash,
-                job_id=self.job_id,
-                job_state=JobStatus(self.job_state)))
+            dlog.info(f"job:{self.job_hash} job_id:{self.job_id} after re-submitting; the state now is {repr(self.job_state)}")
+            self.handle_unexpected_job_state()
 
         if job_state == JobStatus.unsubmitted:
             dlog.info(f"job: {self.job_hash} unsubmitted; submit it")
@@ -758,3 +766,5 @@ def arginfo():
         resources_format = Argument("resources", dict, resources_args)
         return resources_format
 
+
+# %%

tests/graph.pb

@@ -0,0 +1,14 @@
+{
+    "machine":{
+        "batch_type": "Shell",
+        "context_type": "LazyLocalContext",
+        "local_root": "./test_shell_trival_dir"
+    },
+    "resources":{
+        "number_node": 1,
+        "cpu_per_node": 4,
+        "gpu_per_node": 0,
+        "queue_name": "CPU",
+        "group_size": 2
+    }
+}

tests/jsons/machine_lazylocal_shell.json

@@ -51,9 +51,9 @@ def test_submission_to_json(self):
         pass
 
     @patch('dpdispatcher.Submission.submission_to_json')
-    @patch('dpdispatcher.Submission.get_submission_state')
-    def test_check_all_finished(self, patch_get_submission_state, patch_submission_to_json):
-        patch_get_submission_state = MagicMock(return_value=None)
+    @patch('dpdispatcher.Submission.update_submission_state')
+    def test_check_all_finished(self, patch_update_submission_state, patch_submission_to_json):
+        patch_update_submission_state = MagicMock(return_value=None)
         patch_submission_to_json = MagicMock(return_value=None)
 
         self.submission.belonging_jobs[0].job_state = JobStatus.running

tests/test_class_submission.py

@@ -48,6 +48,41 @@ def test_shell_trival(self):
             f2 = os.path.join('test_shell_trival_dir/', 'parent_dir/', dir, 'out.txt')
             self.assertEqual(get_file_md5(f1), get_file_md5(f2))
 
+    def test_shell_fail(self):
+        with open('jsons/machine_local_shell.json', 'r') as f:
+            machine_dict = json.load(f)
+
+        machine = Machine(**machine_dict['machine'])
+        resources = Resources(**machine_dict['resources'])
+
+        task = Task(command='cat mock_fail_task.txt && exit 1',
+            task_work_path='./',
+            forward_files=[], 
+            backward_files=['out.txt'],
+            outlog='out.txt')
+
+        task_list = [task,]
+
+        submission = Submission(work_base='fail_dir/',
+            machine=machine,
+            resources=resources,
+            forward_common_files=[],
+            backward_common_files=[],
+            task_list=task_list
+        )
+        with self.assertRaises(RuntimeError):
+            submission.run_submission()
+
+    def test_shell_recover(self):
+        with open('jsons/machine_lazylocal_shell.json', 'r') as f:
+            machine_dict = json.load(f)
+
+        machine = Machine(**machine_dict['machine'])
+        resources = Resources(**machine_dict['resources'])
+
+        pass
+
+
     @classmethod
     def tearDownClass(cls):
         shutil.rmtree('tmp_shell_trival_dir/')

tests/test_shell_trival.py

@@ -0,0 +1 @@
+# mock file for unittest; test when dpdispatcher meets fail task

tests/test_shell_trival_dir/fail_dir/mock_fail_task.txt

@@ -1,15 +1,47 @@
 import os,sys,json,glob,shutil,uuid,getpass
 import unittest
-from pathlib import Path
+import pathlib
+from paramiko.ssh_exception import NoValidConnectionsError
+from paramiko.ssh_exception import SSHException
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 __package__ = 'tests'
 from .context import SSHContext, SSHSession
-from .context import setUpModule
+from .context import Machine
+from .sample_class import SampleClass
 
 class TestSSHContext(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        with open('jsons/machine_ali_ehpc.json', 'r') as f:
+            mdata = json.load(f)
+        try:
+            cls.machine = Machine.load_from_dict(mdata['machine'])
+        except SSHException:
+            raise unittest.SkipTest("SSHException ssh cannot connect")
+        cls.submission = SampleClass.get_sample_submission()
+        cls.submission.bind_machine(cls.machine)
+        cls.submission_hash = cls.submission.submission_hash
+        file_list = ['bct-1/log.lammps', 'bct-2/log.lammps', 'bct-3/log.lammps', 'bct-4/log.lammps']
+        for file in file_list:
+            cls.machine.context.write_file(file, '# mock log')
+    
     def setUp(self):
-        self.tmp_local_root = 'test_context_dir/'
-        self.tmp_remote_root = 'tmp_ssh_context_dir/'
+        self.context = self.__class__.machine.context
+
+    def test_ssh_session(self):
+        self.assertIsInstance(
+            self.__class__.machine.context.ssh_session, SSHSession
+        )
+
+    def test_upload(self):
+        self.context.upload(self.__class__.submission)
+        check_file_list = ['graph.pb', 'bct-1/conf.lmp', 'bct-4/input.lammps']
+        for file in check_file_list:
+            self.assertTrue(self.context.check_file_exists(os.path.join(self.context.remote_root, file)))
+
+    def test_download(self):
+        self.context.download(self.__class__.submission)
+