Fix some bugs in workflow.

Author: ErjieWu

deepks/io/input/dispatcher.py

@@ -31,6 +31,19 @@ def dispatch_command(config):
         scf_soft = config.get('scf_soft', 'pyscf')
         backend = get_scf_backend(scf_soft)
         backend.collect_stats(**config)
+    elif task_type == 'scf_task':
+        # Low-level per-task SCF runner (called by BatchTask in iterate workflow)
+        from deepks.physics.backends.pyscf.run import main as scf_run_main
+        # Strip orchestration-only keys that pyscf.run.main() doesn't understand
+        scf_config = {k: v for k, v in config.items()
+                      if k not in ('type', 'scf_soft')}
+        scf_run_main(**scf_config)
+    elif task_type == 'train_task':
+        # Low-level training runner (called by BatchTask in iterate workflow)
+        from deepks.ml.train.train import main as train_main
+        train_config = {k: v for k, v in config.items()
+                        if k not in ('type',)}
+        train_main(**train_config)
     elif task_type == 'iterate':
         # Use new iterate workflow
         from deepks.workflows.iterate import run_iterate_workflow

deepks/io/utils.py

@@ -2,6 +2,7 @@
 
 import os
 import shutil
+import warnings
 from collections.abc import Mapping
 from glob import glob
 from itertools import chain
@@ -96,6 +97,15 @@ def save_yaml(data, file_path):
         yaml.dump(data, fp)
 
 
+def dump_yaml_str(data):
+    """Serialize *data* to a YAML string (ruamel.yaml safe mode)."""
+    import io
+    buf = io.StringIO()
+    yaml = YAML(typ='safe', pure=True)
+    yaml.dump(data, buf)
+    return buf.getvalue()
+
+
 def load_array(file):
     ext = os.path.splitext(file)[-1]
     if 'npy' in ext:
@@ -192,3 +202,51 @@ def create_dir(dirname, backup=False):
         os.makedirs(dirname)
     else:
         assert dirname.is_dir(), f'{dirname} is not a dir'
+
+
+# ---------------------------------------------------------------------------
+# Array shape coercion helpers
+# ---------------------------------------------------------------------------
+
+def coerce_box(arr, nframes, fname="box.npy"):
+    """Ensure box array has shape (nframes, 9); accept (nframes, 3, 3)."""
+    if arr.shape == (nframes, 3, 3):
+        warnings.warn(f"{fname}: got shape {arr.shape}, reshaping to ({nframes}, 9).")
+        return arr.reshape(nframes, 9)
+    if arr.shape != (nframes, 9):
+        raise ValueError(f"{fname}: expected shape ({nframes}, 9), got {arr.shape}.")
+    return arr
+
+
+def coerce_energy(arr, nframes, fname="energy.npy"):
+    """Ensure energy array has shape (nframes, 1); accept (nframes,)."""
+    if arr.shape == (nframes,):
+        warnings.warn(f"{fname}: got shape {arr.shape}, reshaping to ({nframes}, 1).")
+        return arr.reshape(nframes, 1)
+    if arr.shape != (nframes, 1):
+        raise ValueError(f"{fname}: expected shape ({nframes}, 1), got {arr.shape}.")
+    return arr
+
+
+def coerce_stress(arr, nframes, fname="stress.npy"):
+    """Ensure stress array has shape (nframes, 6) upper-triangle (xx,xy,xz,yy,yz,zz).
+
+    Accepted input shapes:
+      (nframes, 6)   -- already upper-triangle, returned as-is.
+      (nframes, 3,3) -- full matrix, reshaped then upper-triangle sliced.
+      (nframes, 9)   -- full flat, upper-triangle sliced.
+    """
+    if arr.shape == (nframes, 6):
+        return arr
+    if arr.shape == (nframes, 3, 3):
+        warnings.warn(
+            f"{fname}: got shape {arr.shape}, reshaping to ({nframes}, 9) "
+            f"then taking upper-triangle to ({nframes}, 6)."
+        )
+        arr = arr.reshape(nframes, 9)
+    if arr.shape == (nframes, 9):
+        warnings.warn(f"{fname}: got shape {arr.shape}, taking upper-triangle to ({nframes}, 6).")
+        return arr[:, [0, 1, 2, 4, 5, 8]]
+    raise ValueError(
+        f"{fname}: expected ({nframes},6), ({nframes},9), or ({nframes},3,3), got {arr.shape}."
+    )

deepks/ml/train/train.py

@@ -257,6 +257,6 @@ def main(train_paths, test_paths=None,
     # print("all train time:",end-start)
 
 
-if __name__ == "__main__":
-    from deepks.__main__ import train_cli as cli
-    cli()
+# This module is not intended to be run directly.
+# Training is invoked via PythonTask from the iterate workflow (template.py),
+# which calls main() directly in-process.

deepks/orchestration/scheduler/job/dispatcher.py

@@ -227,7 +227,15 @@ def submit_jobs(self,
                                                  instance_id)
                 job_record.dump()
             else :
-                # finished job, append a None to list
+                # finished job: try to clean up leftover .sub and tag files
+                job_uuid = job_record.get_uuid(cur_hash) if job_record.check_submitted(cur_hash) else None
+                if job_uuid is not None:
+                    try:
+                        context = self.context_fn(work_path, self.session, job_uuid)
+                        context.clean()
+                    except Exception:
+                        pass
+                # append a None to list
                 job_list.append(None)
         assert(len(job_list) == nchunks)
         job_handler = {

deepks/orchestration/scheduler/job/shell.py

@@ -9,32 +9,39 @@ def _default_item(resources, key, value) :
 
 class Shell(Batch) :
 
+    def __init__(self, context, uuid_names=True):
+        super().__init__(context, uuid_names=uuid_names)
+        self.proc = None
+
     def check_status(self) :
         if self.check_finish_tag():
             return JobStatus.finished
         elif self.check_running():
             return JobStatus.running
         else:
             return JobStatus.terminated
-        ## note: check_submitted_tag() only tells us the .sub file exists (was submitted),
-        ## not that the process is still running. Using it here caused an infinite loop
-        ## when a job failed: the .sub file remains, so check_status() kept returning
-        ## 'running' forever instead of 'terminated'.
 
     def check_running(self):
+        # Primary check: use the stored Popen object (reliable, no ps-grep fragility)
+        if self.proc is not None:
+            if self.proc.poll() is None:
+                # process is still alive
+                return True
+            # process has exited; give the finish tag a moment to be flushed
+            time.sleep(2)
+            return False
+        # Fallback for recovered jobs (proc not available): use ps grep
         uuid_names = self.context.job_uuid
-        ## Check if the uuid.sub is running on remote machine
-        cnt = 0
-        ret, stdin, stdout, stderr = self.context.block_call("ps aux | grep %s"%uuid_names)
+        ret, stdin, stdout, stderr = self.context.block_call("ps aux | grep %s" % uuid_names)
         response_list = stdout.read().decode('utf-8').split("\n")
         for response in response_list:
-            if  uuid_names + ".sub" in response:
+            if uuid_names + ".sub" in response:
                 return True
         return False
-    
+
     def exec_sub_script(self, script_str):
         self.context.write_file(self.sub_script_name, script_str)
-        self.proc = self.context.call('cd %s && exec bash %s' % (self.context.remote_root, self.sub_script_name))
+        self.proc = self.context.call('cd %s && bash %s' % (self.context.remote_root, self.sub_script_name))
 
     def default_resources(self, res_) :
         if res_ is None :

deepks/orchestration/workflow/task.py

@@ -190,11 +190,13 @@ class BatchTask(AbstructTask):
     - outlog: the file to redirect stdout to.
     - errlog: the file to redirect stderr to.
     - *_files: the files to forward or backward.
+    - write_files: dict of {filename: content} to write into workdir during preprocess.
     '''
-    def __init__(self, cmds, 
-                 dispatcher=None, resources=None, 
-                 outlog='log', errlog='err', 
+    def __init__(self, cmds,
+                 dispatcher=None, resources=None,
+                 outlog='log', errlog='err',
                  forward_files=None, backward_files=None,
+                 write_files=None,
                  **task_args):
         super().__init__(**task_args)
         self.cmds = check_list(cmds)
@@ -209,10 +211,18 @@ def __init__(self, cmds,
         self.errlog = errlog
         self.forward_files = check_list(forward_files)
         self.backward_files = check_list(backward_files)
-    
+        self.write_files = write_files or {}
+
+    def preprocess(self):
+        super().preprocess()
+        for fname, content in self.write_files.items():
+            fpath = self.workdir / fname
+            fpath.parent.mkdir(parents=True, exist_ok=True)
+            fpath.write_text(content)
+
     def execute(self):
         tdict = self.make_dict(base=self.workdir)
-        self.dispatcher.run_jobs([tdict], group_size=1, work_path='.', 
+        self.dispatcher.run_jobs([tdict], group_size=1, work_path='.',
                                  resources=self.resources, forward_task_deref=True,
                                  outlog=self.outlog, errlog=self.errlog)
 

deepks/physics/backends/abacus/utils.py

@@ -44,26 +44,49 @@ def read_csr(file, dtype=torch.float64):
     all_values = []
     max_iR = 0
     dim = 0
+
+    def _read_tokens(f, n, convert):
+        """Read exactly n tokens from f, skipping comment/blank lines."""
+        tokens = []
+        while len(tokens) < n:
+            line = f.readline()
+            if not line:
+                break
+            s = line.strip()
+            if not s or s.startswith('#'):
+                continue
+            tokens.extend(s.split())
+        return [convert(t) for t in tokens[:n]]
+
     with open(file, 'r') as f:
         dim = int(f.readline().split()[-1])
         num = int(f.readline().split()[-1])
         for _ in range(num):
             nnz = 0
             while nnz == 0:
-                r = f.readline().split()
-                if len(r) == 0:
+                line = f.readline()
+                if not line:
                     break
-                Rx, Ry, Rz, nnz = int(r[0]), int(r[1]), int(r[2]), int(r[3])
+                s = line.strip()
+                if not s or s.startswith('#'):
+                    continue
+                r = s.split()
+                if len(r) < 4:
+                    continue
+                try:
+                    Rx, Ry, Rz, nnz = int(r[0]), int(r[1]), int(r[2]), int(r[3])
+                except ValueError:
+                    continue
                 iRx = R2iR(Rx)
                 iRy = R2iR(Ry)
                 iRz = R2iR(Rz)
                 if max_iR < max(iRx, iRy, iRz):
                     max_iR = max(iRx, iRy, iRz)
             if nnz == 0:
                 break
-            data = [float(x) for x in f.readline().split()]
-            indices = [int(x) for x in f.readline().split()]
-            indptr = [int(x) for x in f.readline().split()]
+            data    = _read_tokens(f, nnz,    float)
+            indices = _read_tokens(f, nnz,    int)
+            indptr  = _read_tokens(f, dim + 1, int)
             matrix = csr_matrix((data, indices, indptr), shape=(dim, dim))
             matrix_coo = matrix.tocoo()
             rows = matrix_coo.row

deepks/physics/backends/pyscf/run.py

@@ -249,6 +249,6 @@ def main(systems, model_file="model.pth", basis='ccpvdz',
             print('group finished')
 
 
-if __name__ == "__main__":
-    from deepks.__main__ import scf_cli as cli
-    cli()
+# This module is not intended to be run directly.
+# SCF calculations are invoked via BatchTask from the iterate workflow (template.py),
+# which calls make_scf_task() using the pyscf backend runner.