74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
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 |