follow symlink in LocalContext downloading (#201)

* follow symlink in LocalContext downloading to prevent broken symbolic links are downloaded * add tests * use copyfile for symlinks * remove realpath * fix typo * add tests for common_file
deepmodeling · Jun 22, 2022 · ddecb7c · ddecb7c
1 parent 0edfa7c
commit ddecb7c
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/dpdispatcher/local_context.py b/dpdispatcher/local_context.py
@@ -230,7 +230,11 @@ def download(self,
                         pass
                     elif (os.path.exists(rfile)) and (not os.path.exists(lfile)) :
                         # trivial case, download happily
-                        shutil.move(rfile, lfile)
+                        # for links, copy instead of moving (default behavior of copyfile is following symlinks)
+                        if not os.path.islink(rfile):
+                            shutil.move(rfile, lfile)
+                        else:
+                            shutil.copyfile(rfile, lfile)
                     elif (os.path.exists(rfile)) and (os.path.exists(lfile)) :
                         # both exists, replace!
                         dlog.info('find existing %s, replacing by %s' % (lfile, rfile))
@@ -276,7 +280,10 @@ def download(self,
                     pass
                 elif (os.path.exists(rfile)) and (not os.path.exists(lfile)) :
                     # trivial case, download happily
-                    shutil.move(rfile, lfile)
+                    if not os.path.islink(rfile):
+                        shutil.move(rfile, lfile)
+                    else:
+                        shutil.copyfile(rfile, lfile)
                 elif (os.path.exists(rfile)) and (os.path.exists(lfile)) :
                     dlog.info(f"both exist rfile:{rfile}; lfile:{lfile}")
                     # both exists, replace!

diff --git a/tests/test_local_context.py b/tests/test_local_context.py
@@ -273,3 +273,33 @@ def test_download_replace_old_files(self):
         md5_new = get_file_md5(target_file)
         self.assertNotEqual(md5_old, md5_new)
 
+    def test_download_symlink(self):
+        task1 = MagicMock(
+            task_work_path='bct-1/',
+            backward_files=['input.lammps.symlink']
+        )
+        submission = MagicMock(work_base='0_md/',
+            belonging_tasks=[task1],
+            backward_common_files=['graph.pb.symlink'],
+            submission_hash='0_md/')
+        os.symlink(os.path.abspath(os.path.join(self.tmp_remote_root, "0_md", "bct-1", "input.lammps")), os.path.join(self.tmp_remote_root, "0_md", "bct-1", "input.lammps.symlink"))
+        os.symlink(os.path.abspath(os.path.join(self.tmp_remote_root, "0_md", "graph.pb")), os.path.join(self.tmp_remote_root, "0_md", "graph.pb.symlink"))
+
+        self.local_context.bind_submission(submission)
+        self.local_context.download(
+            submission)
+        self.local_context.clean()
+        task_file = os.path.join(
+            self.tmp_local_root,
+            '0_md',
+            'bct-1',
+            'input.lammps.symlink',
+        )
+        common_file = os.path.join(
+            self.tmp_local_root,
+            '0_md',
+            'graph.pb.symlink',
+        )
+        self.assertTrue(os.path.isfile(task_file))
+        self.assertTrue(os.path.isfile(common_file))
+