hpc-task 0.0.1__tar.gz

hpc_task-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 pj.ren
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

hpc_task-0.0.1/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: hpc_task
+Version: 0.0.1
+Summary: HPC Task python library.
+Home-page: https://gitee.com/pjren/hpc_task
+Author: Renpj
+Author-email: 0403114076@163.com
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: paramiko
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary
+
+# HPC Task
+ 
+Python package for easy HPC task management based on paramiko.
+
+## Installation
+
+pip install -U hpc_task
+
+**Requirements**
+* paramiko
+
+## Usage
+
+See tests
+
+## TODO

hpc_task-0.0.1/README.md

@@ -0,0 +1,16 @@
+# HPC Task
+ 
+Python package for easy HPC task management based on paramiko.
+
+## Installation
+
+pip install -U hpc_task
+
+**Requirements**
+* paramiko
+
+## Usage
+
+See tests
+
+## TODO

hpc_task-0.0.1/hpc_task/hpc.py

@@ -0,0 +1,147 @@
+import os
+
+import paramiko
+from paramiko import SSHClient, SSHException
+
+
+class HPCTask:
+    def __init__(self, workdir='.'):
+        self._ssh_jump = None
+        self.jobid = None
+        self.ssh_client = None
+        self.workdir = workdir
+
+    def connect(self, target_host, jump_host=None):
+        """
+        开启队列
+        :return: jobid
+        """
+        # 建立 ssh 连接
+        # 跳板机连接信息
+        if self.ssh_client is None:
+            try:
+                # 使用隧道连接到目标服务器
+                target_client = SSHClient()
+                target_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+                # 首先连接到跳板机
+                if jump_host is not None:
+                    # 创建SSH客户端
+                    jump_client = SSHClient()
+                    jump_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+                    jump_client.connect(**jump_host)
+                    # 在跳板机上创建到目标服务器的隧道
+                    transport = jump_client.get_transport()
+                    dest_addr = (target_host['hostname'], target_host['port'])
+                    local_addr = ('127.0.0.1', 0)  # 本地任意端口
+                    channel = transport.open_channel('direct-tcpip', dest_addr, local_addr)
+                    target_client.connect('127.0.0.1',
+                                          port=channel.getpeername()[1],
+                                          username=target_host['username'],
+                                          password=target_host['password'],
+                                          sock=channel)
+                    self._ssh_jump = jump_client
+                else:
+                    target_client.connect(**target_host)
+
+                self.ssh_client = target_client
+                print("Connect to SSH success.")
+            except SSHException as e:
+                raise e
+
+    def prerun(self):
+        # 提交任务，占据节点
+        # TODO: 命令改成配置，适配常见 hpc
+
+        if self.ssh_client is None:
+            raise RuntimeError('ssh client is not connected')
+        commands = [f'mkdir -p {self.workdir}',
+                    f'cd {self.workdir}',
+                    f'cp ~/bin/hpc_job.chess hpc_job.chess',
+                    f'bsub < hpc_job.chess']
+        stdin, stdout, stderr = self.ssh_client.exec_command(';'.join(commands))
+        jobid = stdout.read().decode().strip()  # "Job <688518> is submitted to queue <proj>."
+        jobid = jobid.split()[1].lstrip('<').rstrip('>')
+        self.jobid = jobid
+        return stdin, stdout, stderr
+
+    def postrun(self):
+        """
+        说明：关闭任务节点占用
+        bkill JOBID是从任务头部开始杀, KILL 信号会传递到子进程
+        pkill gosh-remote 是直接杀
+        二者可能相同, 也可能不相同. 取决于 bsub 时如何定义的.
+        通常 bsub 是用一个 script 调 gosh-remote, 这时二者就不同了.
+        那个主调script 可能会做信号处理, 会顺着 gosh-remote的调用进程下去, 逐一 KILL.
+        gosh-remote我不记得是否有信号处理的逻辑, 得做实验确认一下.
+
+        :return:
+
+        TODO: 命令改成配置，适配常见 hpc
+        """
+
+        if self.ssh_client is None:
+            raise RuntimeError('ssh client is not connected')
+        commands = ['sleep 1', 'pkill gosh-remote', f'bkill {self.jobid}']
+        stdin, stdout, stderr = self.ssh_client.exec_command(';'.join(commands))
+        return stdin, stdout, stderr
+
+    @property
+    def status(self):
+        """
+        查询作业状态：排队，运行，结束
+        # TODO: 统一不同的排队系统状态码
+        :return:
+        """
+        if self.ssh_client is None:
+            raise RuntimeError('ssh client is not connected')
+        commands = [f'bjobs -noheader {self.jobid}',]
+        stdin, stdout, stderr = self.ssh_client.exec_command(';'.join(commands))
+        stat = stdout.read().decode().strip().split()  # 688559  renpeng RUN   proj       Khpcserver0 72*Knode44  scheduler  Sep  8 10:55
+        if len(stat) > 5:
+            return stat[2]
+        else:
+            return "UNKNOWN"
+
+    def submit(self):
+        """
+
+        :return: None
+        """
+        return None
+
+    def upload(self):
+        """
+        TODO: 使用 rsync
+        :return: file sync status
+        """
+        if self.ssh_client is None:
+            raise RuntimeError('ssh client is not connected')
+
+        # 创建服务器目录，如果不存在
+        self.ssh_client.exec_command(f'if [ ! -d "{self.workdir}" ]; then mkdir -p {self.workdir};fi')
+        sftp_client = self.ssh_client.open_sftp()
+        filelist = os.listdir(self.workdir)  # TODO: 递归所有文件夹
+        for filename in filelist:
+            sftp_client.put(os.path.join(self.workdir,filename), os.path.join(self.workdir,filename), confirm=False)
+        sftp_client.close()
+        return None
+
+    def download(self):
+        if self.ssh_client is None:
+            raise RuntimeError('ssh client is not connected')
+
+        sftp_client = self.ssh_client.open_sftp()
+        for filename in sftp_client.listdir(self.workdir):  # TODO: 递归所有文件夹
+            sftp_client.get(os.path.join(self.workdir,filename), os.path.join(self.workdir,filename))
+        sftp_client.close()
+        return None
+
+    def close(self):
+        """
+        关闭队列
+        """
+        if self.ssh_client is not None:
+            self.ssh_client.close()
+        if self._ssh_jump is not None:
+            self._ssh_jump.close()
+        return None

hpc_task-0.0.1/hpc_task.egg-info/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: hpc_task
+Version: 0.0.1
+Summary: HPC Task python library.
+Home-page: https://gitee.com/pjren/hpc_task
+Author: Renpj
+Author-email: 0403114076@163.com
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: paramiko
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary
+
+# HPC Task
+ 
+Python package for easy HPC task management based on paramiko.
+
+## Installation
+
+pip install -U hpc_task
+
+**Requirements**
+* paramiko
+
+## Usage
+
+See tests
+
+## TODO

hpc_task-0.0.1/hpc_task.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+LICENSE
+README.md
+setup.py
+hpc_task/__init__.py
+hpc_task/hpc.py
+hpc_task.egg-info/PKG-INFO
+hpc_task.egg-info/SOURCES.txt
+hpc_task.egg-info/dependency_links.txt
+hpc_task.egg-info/requires.txt
+hpc_task.egg-info/top_level.txt
+tests/test_task.py

hpc_task-0.0.1/hpc_task.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

hpc_task-0.0.1/hpc_task.egg-info/requires.txt

@@ -0,0 +1 @@
+paramiko

hpc_task-0.0.1/hpc_task.egg-info/top_level.txt

@@ -0,0 +1 @@
+hpc_task

hpc_task-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

hpc_task-0.0.1/setup.py

@@ -0,0 +1,26 @@
+from setuptools import setup
+
+with open("README.md", "r", encoding='utf-8') as f:
+    long_description = f.read()
+
+install_requires = [
+    'paramiko',
+]
+
+setup(
+    name='hpc_task',
+    version='0.0.1',
+    packages=['hpc_task'],
+    url='https://gitee.com/pjren/hpc_task',
+    license='MIT',
+    author='Renpj',
+    author_email='0403114076@163.com',
+    description='HPC Task python library.',
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    install_requires=install_requires,
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "Operating System :: OS Independent",
+    ],
+)

hpc_task-0.0.1/tests/test_task.py

@@ -0,0 +1,94 @@
+import os
+
+import pytest
+import ase
+from ase.calculators.vasp import Vasp
+from hpc_task.hpc import HPCTask
+
+jump_host = {
+    'hostname': os.getenv("JUMP_HOST_IP"),
+    'port': int(os.getenv("JUMP_HOST_PORT")),
+    'username': os.getenv("JUMP_HOST_USER"),
+    'password': os.getenv("JUMP_HOST_PASS"),  # 建议使用密钥认证而非密码
+}
+
+target_host = {
+    'hostname': os.getenv("TARGET_HOST_IP"),
+    'port': int(os.getenv("TARGET_HOST_PORT")),
+    'username': os.getenv("TARGET_HOST_USER"),
+    'password': os.getenv("TARGET_HOST_PASS"),  # 建议使用密钥认证而非密码
+}
+
+class TestHPCTask:
+    def setup_method(self):
+        pass
+
+    def teardown_method(self):
+        pass
+
+
+    def test_connect(self):
+        hpc = HPCTask()
+        hpc.connect(target_host, jump_host)
+        stdin, stdout, stderr = hpc.ssh_client.exec_command('hostname')
+        print(stdout.read().decode().strip())
+        hpc.close()
+
+    def test_pre_post_run(self):
+        workdir = 'test_hpc_run'
+        hpc = HPCTask(workdir=workdir)
+        hpc.connect(target_host, jump_host)
+        stdin, stdout, stderr = hpc.prerun()
+        print(hpc.jobid)
+        print(f"Job {hpc.jobid} status: {hpc.status}")
+        stdin, stdout, stderr = hpc.postrun()
+        print(stdout.read().decode().strip())
+
+    def test_upload(self):
+        workdir = 'test_hpc_run'
+        hpc = HPCTask(workdir=workdir)
+        hpc.connect(target_host, jump_host)
+        # 在本地创建 workdir
+        if not os.path.exists(workdir):
+            os.makedirs(workdir)
+        # 新建一个空文件
+        with open(os.path.join(workdir, 'test.xyz'), 'w') as f:
+            f.write('test')
+        hpc.upload()
+
+    def test_download(self):
+        workdir = 'test_hpc_run'
+        hpc = HPCTask(workdir=workdir)
+        hpc.connect(target_host, jump_host)
+        if not os.path.exists(workdir):
+            os.makedirs(workdir)
+        hpc.download()
+
+    def test_vasp_calc(self):
+        atoms = ase.Atoms('N2', positions=[(0.,0.,0.),(1.4,0.,0.)],cell=[10,10,10], pbc=True)
+        workdir = 'test_vasp_calc'
+        curdir = os.path.abspath('.')
+        python_command = [   # TODO: 写成一个python 的 template
+            "from hpc_task.hpc import HPCTask",
+            "import os",
+            f"os.chdir('{curdir}')",
+            f"hpc = HPCTask(workdir='{workdir}')",
+            f"hpc.connect({target_host}, {jump_host})",
+            "hpc.upload()",
+            f"stdin,stdout,stderr = hpc.ssh_client.exec_command('cd {workdir};mpirun -np 2 vasp544_std')",
+            "print(stdout.read().decode().strip())",
+            "hpc.download()",
+            "hpc.close()",
+        ]
+        calc_command = f'python -c "{';'.join(python_command)}"'
+        calc = Vasp(xc='pbe',
+                    command=calc_command,
+                    directory=workdir,
+                    gamma=True,
+                    encut=400,
+                    lwave=False,
+                    lcharg=False,
+                    )
+        atoms.calc = calc
+        e = atoms.get_potential_energy()
+        print(f"Energy of N2 is {e} eV.")