nep框架搭建

nep_auto/utils/runner.py

@@ -0,0 +1,74 @@
+import subprocess
+import os
+import time
+import logging
+
+
+class CommandRunner:
+    def __init__(self, machine_config):
+        """
+        :param machine_config: config/machine.yaml 中 'systems' -> 'current_system' 对应的内容
+        """
+        self.config = machine_config
+        self.logger = logging.getLogger("NEP_Auto")
+        self.mode = self.config.get("type", "local")  # local 或 slurm
+
+    def run(self, tool_name, cwd=".", wait=True, extra_args=""):
+        """
+        核心运行方法
+        :param tool_name: machine.yaml 中 tools 下的键名 (如 'gpumd', 'vasp')
+        :param cwd: 执行命令的工作目录
+        :param wait: 是否等待命令结束 (True: 阻塞, False: 后台运行)
+        :param extra_args: 附加在命令后的参数
+        """
+        # 1. 获取工具配置
+        tool_conf = self.config.get("tools", {}).get(tool_name)
+        if not tool_conf:
+            self.logger.error(f"❌ 找不到工具配置: {tool_name}")
+            raise ValueError(f"Tool {tool_name} not defined in machine.yaml")
+
+        cmd = tool_conf.get("command")
+        env_setup = tool_conf.get("env_setup", "")
+
+        # 2. 组装命令 (Local 模式)
+        if self.mode == "local":
+            full_cmd = f"{cmd} {extra_args}"
+
+            # 如果有环境加载脚本，用 && 连接
+            if env_setup:
+                full_cmd = f"{env_setup} && {full_cmd}"
+
+            self.logger.info(f"⚙️ [Local] Executing: {full_cmd}")
+            self.logger.info(f"   📂 Workdir: {cwd}")
+
+            try:
+                # 使用 bash 执行以支持 source 命令
+                process = subprocess.Popen(
+                    full_cmd,
+                    shell=True,
+                    cwd=cwd,
+                    executable="/bin/bash",
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True
+                )
+
+                if wait:
+                    stdout, stderr = process.communicate()
+                    if process.returncode != 0:
+                        self.logger.error(f"❌ Execution failed (Code {process.returncode})")
+                        self.logger.error(f"Stderr: {stderr}")
+                        raise RuntimeError(f"Command failed: {full_cmd}")
+                    return True
+                else:
+                    return process  # 返回进程对象供监控
+
+            except Exception as e:
+                self.logger.error(f"❌ Runner Error: {str(e)}")
+                raise
+
+        # 3. Slurm 模式 (预留接口，暂未实现具体逻辑)
+        elif self.mode == "slurm":
+            self.logger.warning("⚠️ Slurm mode not fully implemented yet.")
+            # 这里未来会生成 sbatch 脚本并提交
+            return False