xflow.framework 0.1.0__py3-none-any.whl

xflow/framework/node.py

@@ -0,0 +1,392 @@
+# Copyright (c) 2025-2026, zhaowcheng <zhaowcheng@163.com>
+
+"""
+节点。
+"""
+import os
+import shutil
+import inspect
+import tempfile
+
+from pathlib import PurePosixPath, Path
+
+from typing import Generator, Dict, Optional, Union, Literal
+from contextlib import contextmanager
+
+from xflow.framework.ssh import SSHConnection
+from xflow.framework.container import ContainerConnection
+from xflow.framework.ssh import CommandResult
+from xflow.framework.pipeline import Pipeline
+
+
+class Node(object):
+    """
+    节点基类。
+    """
+    def __init__(
+        self,
+        name: str,
+        user: str,
+        bwd: str,
+        conn: Union[SSHConnection, ContainerConnection],
+        envs: Optional[Dict[str, str]] = None
+    ):
+        """
+        :param name: 名称。
+        :param user: 用户名。
+        :param bwd: 基础工作目录。
+        :param envs: 环境变量。
+        """
+        self.__name = name
+        self.__user = user
+        self.__envs = envs
+        self.__bwd: PurePosixPath = PurePosixPath(bwd)
+        self.__cwd: PurePosixPath = PurePosixPath('')
+        self.__conn = conn
+        self.__nixenv = {}
+
+    @property
+    def name(self) -> str:
+        """
+        节点名。
+        """
+        return self.__name
+    
+    @property
+    def bwd(self) -> PurePosixPath:
+        """
+        基础工作目录。
+        """
+        return self.__bwd
+    
+    @property
+    def cwd(self) -> PurePosixPath:
+        """
+        当前工作目录（与 pipeline 相关）。
+        """
+        # 如果已指定绝对路径，则直接返回，否则自动拼接。
+        if self.__cwd.is_absolute():
+            return self.__cwd
+        p = self.pipeline
+        return self.bwd.joinpath(p.name, f'{p.buildid}', self.__cwd)
+    
+    @property
+    def scriptdir(self) -> PurePosixPath:
+        """
+        脚本存放目录。
+        """
+        return self.cwd.joinpath('scripts')
+    
+    @property
+    def pipeline(self) -> 'Pipeline':
+        """
+        当前 pipeline。
+        """
+        for frame_info in inspect.stack():
+            frame = frame_info.frame
+            for obj in frame.f_locals.values():
+                if hasattr(obj, '__class__') and issubclass(obj.__class__, Pipeline):
+                    return obj
+                
+    @property
+    def is_native(self) -> bool:
+        """
+        是否原生节点。
+        """
+        return self.__class__.__name__ == 'NativeNode'
+    
+    @property
+    def is_container(self) -> bool:
+        """
+        是否容器节点。
+        """
+        return self.__class__.__name__ == 'ContainerNode'
+
+    @property
+    def existed(self) -> bool:
+        """
+        是否一个已存在的容器（即不是从镜像临时创建的）。
+        """
+        return self.__conn.existed
+
+    def remove(self, force: bool = False) -> None:
+        """
+        删除本容器。
+
+        :param force: 是否强制删除正在运行的容器。
+        """
+        self.__conn.remove(force=force)
+                
+    def mkcwd(self) -> None:
+        """
+        创建当前工作目录。
+        """
+        self.__conn.exec(f'mkdir -p {self.cwd}')
+
+    def rmcwd(self) -> None:
+        """
+        删除当前工作目录。
+        """
+        self.__conn.exec(f'rm -rf {self.cwd}')
+    
+    @contextmanager
+    def dir(self, path: str | PurePosixPath) -> Generator[None, None, None]:
+        """
+        切换工作目录。
+
+        >>> with dir('/my/workdir'):     # doctest: +SKIP
+        ...     d = exec('pwd')          # doctest: +SKIP
+        ...                              # doctest: +SKIP
+        >>> d                            # doctest: +SKIP
+        '/my/workdir'                    # doctest: +SKIP
+        """
+        try:
+            self.__cwd = PurePosixPath(path)
+            yield
+        finally:
+            self.__cwd = PurePosixPath('')
+
+    @contextmanager
+    def nixenv(
+        self, 
+        flake: str | PurePosixPath,
+        system: Optional[Literal['x86_64-linux', 
+                                 'aarch64-linux', 
+                                 'loongarch64-linux', 
+                                 'mips64el-linux']] = None,
+        name: str = 'default',
+        options: Optional[str] = None
+    ) -> Generator[None, None, None]:
+        """
+        使用 nix develop 进入一个 nix shell 环境。
+
+        :param flake: flake（目录）路径。
+        :param system: 系统平台，为空则为当前系统平台。
+        :param name: 环境名称。
+        :param options: nix develop 命令选项。
+        """
+        try:
+            self.__nixenv.update(
+                {
+                    'flake': flake,
+                    'system': system,
+                    'name': name,
+                    'options': options
+                }
+            )
+            yield
+        finally:
+            self.__nixenv.clear()
+            
+    def exec(
+        self,
+        cmd: str,
+        envs: Optional[Dict[str, str]] = None
+    ) -> CommandResult:
+        """
+        执行命令。
+
+        :param cmd: 被执行的命令。
+        :param envs: 环境变量。
+        :return: 命令输出。
+
+        :raises: 
+            `CommandError` -- 命令返回码不为 0。
+
+        >>> exec('ls /home')  # successful                              # doctest: +SKIP
+        >>> exec('ls /errpath')  # CommandError                         # doctest: +SKIP
+        """
+        if self.__nixenv:
+            flake = self.__nixenv['flake']
+            system = self.__nixenv['system']
+            name = self.__nixenv['name']
+            options = self.__nixenv['options']
+            if system:
+                attr = f'devShells.{system}.{name}'
+            else:
+                attr = name
+            options = (options or '') + f' --log-format raw -c {cmd}'
+            cmd = f'nix develop {flake}#{attr} {options.strip()}'
+        with self.__conn.dir(self.cwd):
+            return self.__conn.exec(cmd, envs=envs)
+        
+    def exec_script(
+        self,
+        script: str | Path,
+        argstr: str = '',
+        envs: Optional[Dict[str, str]] = None
+    ) -> CommandResult:
+        """
+        上传并执行脚本。
+
+        :param script: 本地脚本路径（相对项目目录）。
+        :param argstr: 脚本参数。
+        :param envs: 环境变量。
+        :return: 脚本输出。
+
+        :raises: 
+            `CommandError` -- 返回码不为 0。
+
+        >>> exec_script('scripts/my_script.sh')                         # doctest: +SKIP
+        """
+        lscript = Path(script).absolute()
+        rscript = self.scriptdir.joinpath(lscript.name)
+        rdir = rscript.parent
+        if not self.exists(rscript):
+            if not self.exists(rdir):
+                self.__conn.exec(f'mkdir -p {rdir}')
+            self.putfile(lscript, rdir)
+            self.__conn.exec(f'chmod +x {rscript}')
+        return self.exec(f'{rscript} {argstr}', envs=envs)
+        
+    def getfile(
+        self, 
+        rfile: Union[str, PurePosixPath], 
+        ldir: Union[str, Path]
+    ) -> None:
+        """
+        从远端下载文件 `rfile` 到本地目录 `ldir`。
+        
+        :param rfile: 远端文件。
+        :param ldir: 本地目录。
+
+        >>> getfile('/tmp/myfile', '/home')  # /home/myfile
+        >>> getfile('/tmp/myfile', 'D:\\')  # D:\\myfile
+        """
+        self.__conn.getfile(rfile, ldir)
+    
+    def putfile(
+        self, 
+        lfile: Union[str, Path], 
+        rdir: Union[str, PurePosixPath]
+    ) -> None:
+        """
+        上传本地文件 `lfile` 到远端目录 `rdir`。
+
+        :param lfile: 本地文件。
+        :param rdir: 远端目录。
+
+        >>> putfile('/home/myfile', '/tmp')  # /tmp/myfile
+        >>> putfile('D:\\myfile', '/tmp')  # /tmp/myfile
+        """
+        return self.__conn.putfile(lfile, rdir)
+    
+    def exists(self, path: Union[str, PurePosixPath]) -> bool:
+        """
+        检查远端路径是否存在。
+        """
+        return self.__conn.exists(path)
+    
+    def write(self, text: str, rfile: Union[str, PurePosixPath]) -> None:
+        """
+        写入一个文本文件到远端路径 `path`。
+
+        :param text: 要写入的文本。
+        :param rfile: 远端文件路径。
+        """
+        rfile = PurePosixPath(rfile)
+        filename = rfile.name
+        rdir = rfile.parent
+        tmpdir = Path(tempfile.mkdtemp())
+        lfile = tmpdir.joinpath(filename)
+        with open(lfile, 'w') as f:
+            f.write(text)
+        self.putfile(lfile, rdir)
+        shutil.rmtree(tmpdir)
+    
+    def git(
+        self,
+        repourl: str,
+        revision: Optional[str] = None,
+        directory: Optional[str | PurePosixPath] = None,
+        options: Optional[str] = None
+    ) -> None:
+        """
+        git clone 一个仓库。
+
+        :param repourl: 仓库地址。
+        :param revision: 分支、tag 或 commit。
+        :param directory: 克隆的目标目录。
+        :param options: git clone 命令选项。
+        """
+        humanish = repourl.split('/')[-1].replace('.git', '')
+        directory = directory or humanish
+        self.exec(f'git clone {options or ""} {repourl} {directory}')
+        if revision:
+            with self.dir(directory):
+                self.exec(f'git checkout {revision}')
+
+
+class NativeNode(Node):
+    """
+    原生节点（非容器）。
+    """
+    def __init__(
+        self,
+        name: str,
+        ip: str,
+        sshport: int,
+        user: str,
+        password: str,
+        bwd: str,
+        envs: Optional[Dict[str, str]] = None
+    ):
+        """
+        :param name: 名称。
+        :param ip: IP 地址。
+        :param sshport: SSH 端口。
+        :param user: 用户名。
+        :param password: 用户密码。
+        :param bwd: 基础工作目录。
+        :param envs: 环境变量。
+        """
+        conn = SSHConnection(ip, user, password, port=sshport, envs=envs)
+        super().__init__(name, user, bwd, conn, envs=envs)
+
+
+class ContainerNode(Node):
+    """
+    容器节点。
+    """
+    def __init__(
+        self,
+        name: str,
+        ip: str,
+        port: int,
+        bwd: str,
+        user: Optional[str] = None,
+        container: Optional[str] = None,
+        image: Optional[str] = None,
+        runargs: Optional[Dict] = None,
+        cacert: Optional[str] = None,
+        clientcert: Optional[str] = None,
+        clientkey: Optional[str] = None,
+        envs: Optional[Dict[str, str]] = None
+    ):
+        """
+        :param name: 名称。
+        :param ip: IP 地址。
+        :param port: Docker 服务端口。
+        :param user: 用户名。
+        :param bwd: 基础工作目录。
+        :param container: 容器名，和 `image` 为互斥参数，必须且只能指定其中一个。
+        :param image: 镜像名，和 `container` 为互斥参数，必须且只能指定其中一个，
+                      如果该参数指定，则每次连接用该镜像创建一个新容器。
+        :param runargs: 使用镜像启动容器时的参数。
+        :param cacert: CA 证书。
+        :param clientcert: 客户端证书，当指定 `clientkey` 参数时必须同时指定该参数。
+        :param clientkey: 客户端私钥，当指定 `clientcert` 参数时必须同时指定该参数。
+        :param envs: 环境变量。
+        """
+        conn = ContainerConnection(ip, 
+                                   port=port, 
+                                   user=user, 
+                                   name=container, 
+                                   image=image, 
+                                   runargs=runargs,
+                                   cacert=cacert, 
+                                   clientcert=clientcert, 
+                                   clientkey=clientkey, 
+                                   envs=envs)
+        super().__init__(name, user, bwd, conn, envs=envs)
+    

xflow/framework/pipeline.py

@@ -0,0 +1,189 @@
+# Copyright (c) 2025-2026, zhaowcheng <zhaowcheng@163.com>
+
+"""
+流水线模块。
+"""
+
+import re
+import traceback
+
+import click
+
+from typing import List, Literal, Any, Optional, Iterable, TYPE_CHECKING, Union, get_args, get_origin
+from pathlib import Path
+
+from filelock import FileLock
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+
+if TYPE_CHECKING:
+    from xflow.framework.env import Env
+
+TResult = Literal['FAILED', 'SUCCESSFUL']
+
+
+class Pipeline(object):
+    """
+    流水线。
+    """ 
+    class Option(FieldInfo):
+        """
+        流水线参数。
+        """
+        def __init__(
+            self, 
+            desc: Optional[str] = None,
+            default: Any = None,
+            choices: Optional[Iterable] = None
+        ):
+            """
+            :param desc: 描述。
+            :param default: 默认值。
+            :param choices: 枚举值。
+            """
+            kwargs = {}
+            if desc is not None:
+                kwargs['description'] = desc
+            if default is not None:
+                kwargs['default'] = default
+            if choices is not None:
+                kwargs.update(
+                    {
+                        'json_schema_extra': {
+                            'typed-settings': {
+                                'click': {
+                                    'type': click.Choice(choices)
+                                }
+                            }
+                        }
+                    }
+                )
+            super().__init__(**kwargs)
+           
+
+    class Options(BaseModel):
+        """
+        流水线参数表。
+        """
+        def __init__(self, **data):
+            data = self.__set_default_values(**data)
+            super().__init__(**data)
+
+        def __set_default_values(self, **data) -> dict:
+            """
+            为 Optional 类型但未设置 default 的参数补充默认值 None。
+            """
+            for field_name, field_info in self.__class__.model_fields.items():
+                field_type = field_info.annotation
+                origin = get_origin(field_type)
+                args = get_args(field_type)
+                if origin is Union and type(None) in args and field_name not in data:
+                    data[field_name] = None
+            return data
+
+    def __init__(
+        self, 
+        projdir: str | Path,
+        env: 'Env',
+        nodename: str, 
+        options: Options
+    ):
+        """
+        :param projdir: 项目目录。
+        :param env: 环境信息。
+        :param nodename: 执行节点名称。
+        :param options: 流水线参数。
+        """
+        self.projdir = Path(projdir)
+        self.bwd = self.projdir.joinpath('workdir')  # base workdir
+        self.env = env
+        self.node = self.env.get_node(nodename)
+        self.options = options
+        self.result: TResult = None
+        self.buildid: int = None
+
+    @property
+    def name(self) -> str:
+        """
+        名称。
+        """
+        return self.__class__.__name__
+    
+    @property
+    def cwd(self) -> Path:
+        """
+        当前工作目录。
+        """
+        return self.bwd.joinpath(self.name, f'{self.buildid}')
+    
+    def setup(self) -> None:
+        """
+        前置步骤，如果失败将不会执行任何 stage，直接执行 teardown。
+        """
+        # 获取 buildid
+        print(f'Getting BuildID: ', end='')
+        parent = self.bwd.joinpath(self.name)
+        parent.mkdir(parents=True, exist_ok=True)
+        idfile = parent.joinpath('buildid.txt')
+        idfile.touch()
+        with FileLock(f'{idfile}.lock'):
+            with open(idfile, 'r+') as f:
+                newid = int(f.read() or 0) + 1
+            with open(idfile, 'w') as f:
+                f.write(f'{newid}')
+        self.buildid = newid
+        print(self.buildid)
+
+        # 打印参数
+        print(f'options: {self.options}')
+
+        # 创建远端工作目录
+        self.node.mkcwd()
+
+        # 创建本地工作目录
+        self.cwd.mkdir(parents=True, exist_ok=True) 
+    
+    def stage1(self) -> None:
+        """
+        阶段 1，一个 pipeline 至少包含一个 stage1，增加则按照
+        stage2, stage3, ... 这样依次增加，执行时也将按此顺序。
+        """
+        raise NotImplementedError
+
+    @property
+    def stages(self) -> List[str]:
+        """
+        所有阶段名称。
+        """
+        return sorted([n for n in dir(self.__class__) if re.match(r'stage\d+', n)])
+
+    def run(self) -> TResult:
+        """
+        执行 pipeline。
+        """
+        title = lambda t: print(t.center(80, '='))
+        try:
+            title('setup')
+            self.setup()
+            for stage in self.stages:
+                title(stage)
+                getattr(self, stage)()
+            self.result = 'SUCCESSFUL'
+        except:
+            self.result = 'FAILED'
+            traceback.print_exc()
+        finally:
+            title('teardown')
+            self.teardown()
+        return self.result
+            
+    def teardown(self) -> None:
+        """
+        后置步骤，不管成功与否，都会执行该步骤。
+        """
+        if self.result == 'SUCCESSFUL':
+            if self.node.is_container and not self.node.existed:
+                self.node.remove(force=True)
+            else:
+                self.node.rmcwd()
+