dma-cli 0.1.0__py3-none-any.whl

dma_cli/__init__.py

@@ -0,0 +1,3 @@
+"""DMA CLI —— 任务提交命令行工具"""
+
+__version__ = "0.1.0"

dma_cli/__main__.py

@@ -0,0 +1,6 @@
+"""支持 `python -m dma_cli`（在 tools/dma-cli 目录下运行）。"""
+
+from .cli import main
+
+if __name__ == "__main__":
+    main()

dma_cli/calc.py

@@ -0,0 +1,42 @@
+"""内置默认任务脚本。
+
+"""
+import argparse
+import time
+import sys
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num1", type=float, default=1)
+    parser.add_argument("--operator", choices=["+", "-"], default="+")
+    parser.add_argument("--num2", type=float, default=1)
+    args = parser.parse_args()
+
+    # 打开 run.log 文件并将 print 输出重定向追加至日志
+    with open("run.log", "a", encoding="utf-8") as f:
+        f.write(f"--- 任务启动时间: {time.strftime('%Y-%m-%d %H:%M:%S')} ---\n")
+        f.flush()
+
+        total_seconds = 10 * 60
+        interval = 5
+        elapsed = 0
+        while elapsed < total_seconds:
+            msg = "我大概运行10分钟，你可以先检查其他功能\n"
+            sys.stdout.write(msg)
+            sys.stdout.flush()
+            f.write(msg)
+            f.flush()
+            time.sleep(interval)
+            elapsed += interval
+
+        result = args.num1 + args.num2 if args.operator == "+" else args.num1 - args.num2
+        res_msg = f"计算结果: {args.num1} {args.operator} {args.num2} = {result}\n"
+        sys.stdout.write(res_msg)
+        sys.stdout.flush()
+        f.write(res_msg)
+        f.flush()
+
+
+if __name__ == "__main__":
+    main()

dma_cli/cli.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+"""DMA CLI —— 任务提交命令行工具入口。
+
+用法:
+    dma --help
+    dma login
+    dma whoami
+    dma project list
+    dma project use 3948
+    dma sku list --kind cpu
+    dma image list
+    dma job submit --machine_type c2_m2_cpu --image_address <url> --command "python calc.py"
+    dma job submit -i job.example.json
+    dma job submit --batch_file batch.json
+"""
+import argparse
+import os
+import sys
+
+from .client import DmaClient
+from .resources import (
+    handle_project_list, handle_project_use, handle_project_current,
+    handle_sku_list,
+    handle_image_list,
+    handle_jobs, handle_job_rm,
+)
+from .submit import handle_submit
+
+# 不强制鉴权的命令
+_NO_AUTH = {"login", "whoami"}
+
+
+def cmd_login(client, args):
+    creds = client.login(args.email)
+    print(f"✅ 登录成功: email={creds.get('email')} user_id={creds.get('user_id')} uexid={creds.get('uexid')}")
+
+
+def _fmt_project(p):
+    if not p:
+        return "未设置"
+    return f"{p['projectId']} | {p['name']}"
+
+
+def cmd_whoami(client, args):
+    u = client.current_user()
+    print(f"base_url:  {client.base_url}")
+    print(f"email:     {u['email']}")
+    print(f"user_id:   {u['user_id']}")
+    print(f"uexid:     {u['uexid']}")
+    print(f"token:     {'已配置' if client.is_authed() else '未配置'} (来源: {u['token_source']})")
+    print(f"当前项目:   {_fmt_project(client.current_project())}")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(prog="dma", description="DMA 任务提交命令行工具")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    # login
+    sp = sub.add_parser("login", help="验证码登录，缓存 token/userId/uexid")
+    sp.add_argument("--email", default=None, help="登录邮箱（默认用 config.yaml 的 auth.email）")
+    sp.set_defaults(func=cmd_login)
+
+    # whoami
+    sp = sub.add_parser("whoami", help="显示当前用户与当前项目")
+    sp.set_defaults(func=cmd_whoami)
+
+    # project
+    sp = sub.add_parser("project", help="项目管理")
+    ps = sp.add_subparsers(dest="subcmd", required=True)
+    ps.add_parser("list", help="列项目").set_defaults(func=handle_project_list)
+    sp_use = ps.add_parser("use", help="设为当前项目")
+    sp_use.add_argument("project_id", type=int)
+    sp_use.set_defaults(func=handle_project_use)
+    ps.add_parser("current", help="显示当前项目").set_defaults(func=handle_project_current)
+
+    # sku (仅查询，不持久化)
+    sp = sub.add_parser("sku", help="机型查询（参考）")
+    ps = sp.add_subparsers(dest="subcmd", required=True)
+    sp_list = ps.add_parser("list", help="列机型")
+    sp_list.add_argument("--kind", choices=["cpu", "gpu"], default="cpu", help="机型类型（默认 cpu）")
+    sp_list.set_defaults(func=handle_sku_list)
+
+    # image (软件目录，仅查询，不持久化)
+    sp = sub.add_parser("image", help="镜像(软件目录)查询")
+    ps = sp.add_subparsers(dest="subcmd", required=True)
+    ps.add_parser("list", help="列软件及其版本").set_defaults(func=handle_image_list)
+
+    # job (任务操作)
+    sp = sub.add_parser("job", help="任务操作")
+    ps = sp.add_subparsers(dest="subcmd", required=True)
+
+    # job submit
+    sp_submit = ps.add_parser("submit", help="提交任务（单个或批量）")
+    sp_submit.add_argument("-i", "--config_file", default=None, help="配置 JSON 文件路径")
+    sp_submit.add_argument("--input_directory", default="./", help="输入文件目录（目录下所有文件上传）")
+    sp_submit.add_argument("-g", "--job_group_id", type=int, default=0, help="Job Group ID（0=自动创建）")
+    sp_submit.add_argument("--project_id", type=int, default=0, help="项目 ID（0=用当前选择）")
+    sp_submit.add_argument("-n", "--job_name", default=None, help="任务名称")
+    sp_submit.add_argument("--name_prefix", default=None, help="批量任务名前缀")
+    sp_submit.add_argument("-m", "--image_address", default=None, help="镜像地址（自动查询匹配）")
+    sp_submit.add_argument("-t", "--machine_type", default=None, help="机器类型名（自动查询匹配）")
+    sp_submit.add_argument("-c", "--command", default=None, help="执行命令（未指定时取软件默认）")
+    sp_submit.add_argument("-r", "--result_path", default=None, help="结果路径")
+    sp_submit.add_argument("-l", "--log_file", default=None, help="日志文件名")
+    sp_submit.add_argument("--max_run_time", type=int, default=0, help="最大运行时间（分钟，0=用默认）")
+    sp_submit.add_argument("--max_reschedule_times", type=int, default=0, help="最大重调度次数")
+    sp_submit.add_argument("--nnode", type=int, default=1, help="节点数")
+    sp_submit.add_argument("--out_files", nargs="+", default=None, help="输出文件列表")
+    sp_submit.add_argument("--dataset_path", nargs="+", default=None, help="数据集路径列表")
+    sp_submit.add_argument("--count", type=int, default=1, help="批量提交任务数")
+    sp_submit.add_argument("--concurrency", type=int, default=5, help="批量并发数")
+    sp_submit.add_argument("--batch_file", default=None, help="批量任务 JSON 数组文件路径")
+    sp_submit.add_argument("--batch_input_dir", default=None, help="批量目录扫描路径")
+    sp_submit.add_argument("--no_check", action="store_true", help="跳过 check_machine_type 预检")
+    sp_submit.add_argument("--dry_run", action="store_true", help="只打印 snapshot body 不提交")
+    sp_submit.add_argument("--poll", action="store_true", help="提交后轮询状态")
+    sp_submit.set_defaults(func=handle_submit)
+
+    # jobs (隐藏命令)
+    sp_jobs = ps.add_parser("jobs", help=argparse.SUPPRESS)
+    sp_jobs.add_argument("--status", type=int, default=-1, help="-1=全部 0=待运行 1=运行中 2=成功 3=失败 4=取消 5=排队 6=超时")
+    sp_jobs.add_argument("--project_id", type=int, default=None, help="按项目过滤")
+    sp_jobs.set_defaults(func=handle_jobs)
+
+    # job rm (隐藏命令)
+    sp_rm = ps.add_parser("rm", help=argparse.SUPPRESS)
+    sp_rm.add_argument("job_id", help="任务 ID")
+    sp_rm.set_defaults(func=handle_job_rm)
+
+    return p
+
+
+def main():
+    parser = build_parser()
+    args = parser.parse_args()
+    client = DmaClient()
+    if args.cmd not in _NO_AUTH and not client.is_authed():
+        print("❌ 未登录/未配 token，请先 `dma login` 或在 config.yaml / 环境变量 DMA_TOKEN 填写")
+        sys.exit(1)
+    try:
+        args.func(client, args)
+    except RuntimeError as e:
+        print(f"❌ {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

dma_cli/client.py

@@ -0,0 +1,245 @@
+"""DmaClient: 配置/状态加载、鉴权头、登录、GET/POST、重试。
+
+接口形态移植自压测脚本（prepare_users.py 登录、cleanup_jobs.py 鉴权头与 job/list），
+改为 requests 同步版。仅引用接口形态，不修改任何压测文件。
+"""
+import base64
+import json
+import os
+import time
+from pathlib import Path
+from typing import Optional
+
+import requests
+import urllib3
+import yaml
+
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+from .state import load_session, load_creds, save_creds, creds_path
+
+# ---- 配置文件定位 ----
+_PKG_DIR = Path(__file__).parent
+# 配置文件查找顺序：当前目录 -> 包目录
+CONFIG_FILE = Path.cwd() / "config.yaml" if (Path.cwd() / "config.yaml").exists() else _PKG_DIR / "config.yaml"
+EXAMPLE_CONFIG = _PKG_DIR / "config.example.yaml"
+
+DEFAULT_CONFIG = {
+    "base_url": "https://szc.test.dp.tech",
+    "auth": {"email": "", "captcha": "952754", "token": "", "uexid": "", "user_id": ""},
+    "defaults": {
+        "biz_platform": 2,
+        "name_prefix": "cli-job",
+        "cmd": "python calc.py",
+        "max_run_time": 15,
+        "max_run_time_unit": "d",
+        "concurrency": 5,
+    },
+}
+
+# HAR 实测鉴权头（登录来源标注）
+X_BOHRIUM_PARAMS = "eyJsb2dpblNvdXJjZSI6ImVtYWlsX2NvZGUiLCJ1dG1Tb3VyY2UiOiIifQ=="
+LOGIN_PATH = "/platform-gateway/api/account/login/email_code"
+
+# 任务状态映射（移植自 cleanup_jobs.py）
+STATUS_MAP = {
+    -1: "异常/未知",
+    0: "待运行",
+    1: "运行中",
+    2: "运行成功",
+    3: "运行失败",
+    4: "已取消",
+    5: "排队中",
+    6: "超时",
+}
+
+
+def _load_config() -> dict:
+    """加载 config.yaml（缺失则用 config.example.yaml），与内置默认浅合并。"""
+    cfg = {k: (v.copy() if isinstance(v, dict) else v) for k, v in DEFAULT_CONFIG.items()}
+    src = CONFIG_FILE if CONFIG_FILE.exists() else EXAMPLE_CONFIG
+    try:
+        with open(src, "r", encoding="utf-8") as f:
+            user_cfg = yaml.safe_load(f) or {}
+    except FileNotFoundError:
+        user_cfg = {}
+    for k in cfg:
+        if k not in user_cfg:
+            continue
+        if isinstance(cfg[k], dict) and isinstance(user_cfg[k], dict):
+            cfg[k].update(user_cfg[k])
+        else:
+            cfg[k] = user_cfg[k]
+    return cfg
+
+
+class DmaClient:
+    def __init__(self, http_session=None):
+        # requests.Session 非线程安全；批量并发时每线程传入独立 Session
+        self._http = http_session or requests.Session()
+        self.config = _load_config()
+        # 环境变量覆盖（DMA_BASE_URL / DMA_TOKEN）
+        env_base = os.environ.get("DMA_BASE_URL")
+        if env_base:
+            self.config["base_url"] = env_base
+        env_token = os.environ.get("DMA_TOKEN")
+        if env_token:
+            self.config["auth"]["token"] = env_token
+        self.session = load_session()
+        self.creds = load_creds()
+
+    # ---- 基本属性 ----
+    @property
+    def base_url(self) -> str:
+        return str(self.config["base_url"]).rstrip("/")
+
+    @property
+    def defaults(self) -> dict:
+        return self.config.get("defaults", {})
+
+    def reload_session(self) -> dict:
+        self.session = load_session()
+        return self.session
+
+    # ---- 凭据 ----
+    def _token(self) -> str:
+        return self.config["auth"].get("token") or self.creds.get("token") or ""
+
+    def _uexid(self) -> str:
+        u = self.config["auth"].get("uexid") or self.creds.get("uexid")
+        if u:
+            return u
+        # 留空则从 JWT 解析 fe_<userId>
+        token = self._token()
+        if token:
+            _, uexid = self._parse_jwt(token)
+            return uexid
+        return ""
+
+    def _user_id(self) -> str:
+        u = self.config["auth"].get("user_id") or self.creds.get("user_id")
+        if u:
+            return u
+        token = self._token()
+        if token:
+            uid, _ = self._parse_jwt(token)
+            return uid
+        return ""
+
+    def is_authed(self) -> bool:
+        return bool(self._token() and self._uexid())
+
+    def current_user(self) -> dict:
+        return {
+            "email": self.creds.get("email") or self.config["auth"].get("email") or "",
+            "user_id": self._user_id(),
+            "uexid": self._uexid(),
+            "token_source": "config" if self.config["auth"].get("token") else ("cache" if self.creds.get("token") else "none"),
+        }
+
+    def auth_headers(self) -> dict:
+        """通用鉴权头（HAR 实测）。JSON body 时由调用方补 Content-Type。"""
+        return {
+            "Authorization": f"Bearer {self._token()}",
+            "Uexid": self._uexid(),
+            "Accept": "application/json, text/plain, */*",
+            "Content-Language": "zh-cn",
+            "X-Bohrium-Params": X_BOHRIUM_PARAMS,
+        }
+
+    # ---- 登录 ----
+    def login(self, email: str = None) -> dict:
+        """验证码登录，解析 JWT 缓存 token/user_id/uexid。"""
+        email = email or self.config["auth"].get("email")
+        if not email:
+            raise RuntimeError("未配置 auth.email，无法验证码登录（请在 config.yaml 填写或改用 token）")
+        captcha = self.config["auth"].get("captcha") or "952754"
+        url = self.base_url + LOGIN_PATH
+        payload = {
+            "channel": "gkx-talent-universe",
+            "businessLine": "bohrium",
+            "email": email,
+            "emailCode": captcha,
+            "code": captcha,
+        }
+        resp = self._http.post(url, json=payload, timeout=15, verify=False)
+        data = resp.json()
+        if data.get("code") != 0:
+            raise RuntimeError(
+                f"登录失败: code={data.get('code')} msg={data.get('message') or data.get('msg')}"
+            )
+        token = (data.get("data") or {}).get("token") or (data.get("data") or {}).get("accessToken")
+        if not token:
+            raise RuntimeError("登录返回无 token")
+        user_id, uexid = self._parse_jwt(token)
+        self.creds = {"token": token, "user_id": user_id, "uexid": uexid, "email": email}
+        save_creds(self.creds)
+        return self.creds
+
+    @staticmethod
+    def _parse_jwt(token: str):
+        """从 JWT payload 解析 userId，推导 uexid = fe_<userId>（不校验签名）。"""
+        try:
+            payload_b64 = token.split(".")[1]
+            payload_b64 += "=" * (-len(payload_b64) % 4)
+            payload = json.loads(base64.b64decode(payload_b64))
+            uid = str((payload.get("identity") or {}).get("userId", ""))
+            uexid = f"fe_{uid}" if uid else "fe_0"
+            return uid, uexid
+        except Exception:
+            return "", "fe_0"
+
+    # ---- 通用请求（带重试）----
+    def _request(self, method, path, *, params=None, json_body=None, data=None,
+                 headers=None, timeout=30, retries=3):
+        url = self.base_url + path if path.startswith("/") else path
+        h = self.auth_headers()
+        if headers:
+            h.update(headers)
+        if json_body is not None and "Content-Type" not in h:
+            h["Content-Type"] = "application/json;charset=UTF-8"
+        last_err = None
+        for attempt in range(retries):
+            try:
+                r = self._http.request(
+                    method, url, headers=h, params=params, json=json_body,
+                    data=data, timeout=timeout, verify=False,
+                )
+                try:
+                    return r.json()
+                except ValueError:
+                    return {"code": -1, "msg": f"非 JSON 响应: HTTP {r.status_code}", "_raw": r.text[:200]}
+            except (requests.ConnectionError, requests.Timeout) as e:
+                last_err = e
+                if attempt < retries - 1:
+                    time.sleep(1.0 * (attempt + 1))
+        raise RuntimeError(f"请求失败（{retries} 次重试后）: {last_err}")
+
+    def get(self, path, params=None, timeout=30):
+        return self._request("GET", path, params=params, timeout=timeout)
+
+    def post(self, path, data=None, json_body=None, headers=None, timeout=30):
+        """通用 POST 方法，支持 data 或 json_body。"""
+        return self._request("POST", path, data=data, json_body=json_body, headers=headers, timeout=timeout)
+
+    def post_json(self, path, body=None, headers=None, timeout=30):
+        # body 为 None 时发空 JSON 对象 {}
+        return self._request("POST", path, json_body=(body if body is not None else {}), headers=headers, timeout=timeout)
+
+    def post_raw(self, path, data, headers=None, timeout=30):
+        """发送原始 body（如 upload 步骤的文件原文），调用方需在 headers 指定 Content-Type。"""
+        return self._request("POST", path, data=data, headers=headers, timeout=timeout)
+
+    def post_empty(self, path, headers=None, timeout=30):
+        """发送空 body（如 commit / del）。"""
+        return self._request("POST", path, headers=headers, timeout=timeout)
+
+    # ---- 当前选择快捷访问 ----
+    def current_project(self) -> Optional[dict]:
+        return self.session.get("project")
+
+    def current_sku(self) -> Optional[dict]:
+        return self.session.get("sku")
+
+    def current_software(self) -> Optional[dict]:
+        return self.session.get("software")