zyworkflow 0.0.1__py3-none-any.whl

zyworkflow/example/train_client.py

@@ -0,0 +1,301 @@
+#!/usr/bin/env python3
+import json
+import time
+import requests
+import argparse
+from typing import Optional, Literal
+from zyworkflow.utils.logger_config import setup_train_client_logger
+
+logger = setup_train_client_logger()
+
+
+class TrainClient:
+    def __init__(self, base_url: str = "http://localhost:8001"):
+        self.base_url = base_url.rstrip('/')
+        self.train_endpoint = f"{self.base_url}/train"
+
+    def start_training(
+        self,
+        algo_type: Literal["bnn"],
+        action_type: Literal["pick", "place"],
+        task_id: str,
+        ability_id: str,
+        dataset_id: str,
+        model_id: str,
+        batch_size: int = 48,
+        seq_len: int = 4,
+        action_chunk: int = 8,
+        lr: float = 1e-4,
+        num_epochs: int = 500,
+        start_epoch: int = 0,
+        lambda_joints: float = 10.0,
+        lambda_grip: float = 5.0,
+        lambda_success: float = 2.0,
+        log_path: Optional[str] = None,
+        ckpt_dir: Optional[str] = None,
+        success_mode: str = "within_horizon",
+        report_url: Optional[str] = None,
+    ) -> dict:
+        payload = {
+            "algo_type": algo_type,
+            "action_type": action_type,
+            "task_id": task_id,
+            "ability_id": ability_id,
+            "dataset_id": dataset_id,
+            "model_id": model_id,
+            "batch_size": batch_size,
+            "seq_len": seq_len,
+            "action_chunk": action_chunk,
+            "lr": lr,
+            "num_epochs": num_epochs,
+            "start_epoch": start_epoch,
+            "lambda_joints": lambda_joints,
+            "lambda_grip": lambda_grip,
+            "lambda_success": lambda_success,
+            "success_mode": success_mode,
+        }
+
+        if log_path is not None:
+            payload["log_path"] = log_path
+        if ckpt_dir is not None:
+            payload["ckpt_dir"] = ckpt_dir
+        if report_url is not None:
+            payload["report_url"] = report_url
+
+        try:
+            logger.info(f"正在向 {self.train_endpoint} 发送训练请求...")
+            logger.debug(f"请求参数: {json.dumps(payload, indent=2, ensure_ascii=False)}")
+
+            response = requests.post(self.train_endpoint, json=payload, timeout=30)
+            response.raise_for_status()
+
+            result = response.json()
+            logger.info(f"训练任务已启动: {result.get('task_name')}")
+            print("训练任务已启动!")
+            print(f"任务名称: {result.get('task_name')}")
+            print(f"状态: {result.get('status')}")
+            print(f"消息: {result.get('message')}")
+
+            return result
+
+        except requests.exceptions.RequestException as e:
+            logger.error(f"请求失败: {e}")
+            print(f"请求失败: {e}")
+            if hasattr(e, 'response') and e.response is not None:
+                try:
+                    error_detail = e.response.json()
+                    logger.error(f"错误详情: {json.dumps(error_detail, indent=2, ensure_ascii=False)}")
+                    print(f"错误详情: {json.dumps(error_detail, indent=2, ensure_ascii=False)}")
+                except Exception:
+                    logger.error(f"响应内容: {e.response.text}")
+                    print(f"响应内容: {e.response.text}")
+            raise
+
+    def get_training_status(self, task_id: str, ability_id: str, model_id: str) -> dict:
+        try:
+            url = f"{self.base_url}/train/status/{task_id}/{ability_id}/{model_id}"
+            logger.debug(f"查询训练状态: {task_id}-{ability_id}-{model_id}")
+            response = requests.get(url, timeout=10)
+            response.raise_for_status()
+            status = response.json()
+            logger.debug(f"训练状态: {status}")
+            return status
+
+        except requests.exceptions.RequestException as e:
+            logger.error(f"查询状态失败: {e}")
+            print(f"查询状态失败: {e}")
+            if hasattr(e, 'response') and e.response is not None:
+                try:
+                    error_detail = e.response.json()
+                    logger.error(f"错误详情: {json.dumps(error_detail, indent=2, ensure_ascii=False)}")
+                    print(f"错误详情: {json.dumps(error_detail, indent=2, ensure_ascii=False)}")
+                except Exception:
+                    logger.error(f"响应内容: {e.response.text}")
+                    print(f"响应内容: {e.response.text}")
+            raise
+
+    def stop_training(self, task_id: str, ability_id: str, model_id: str) -> dict:
+        try:
+            url = f"{self.base_url}/train/stop/{task_id}/{ability_id}/{model_id}"
+            logger.info(f"发送停止训练请求: {task_id}-{ability_id}-{model_id}")
+            response = requests.post(url, timeout=10)
+            response.raise_for_status()
+            result = response.json()
+            logger.info(f"停止训练结果: {result}")
+            return result
+
+        except requests.exceptions.RequestException as e:
+            logger.error(f"停止训练失败: {e}")
+            print(f"停止训练失败: {e}")
+            if hasattr(e, 'response') and e.response is not None:
+                try:
+                    error_detail = e.response.json()
+                    logger.error(f"错误详情: {json.dumps(error_detail, indent=2, ensure_ascii=False)}")
+                    print(f"错误详情: {json.dumps(error_detail, indent=2, ensure_ascii=False)}")
+                except Exception:
+                    logger.error(f"响应内容: {e.response.text}")
+                    print(f"响应内容: {e.response.text}")
+            raise
+
+    def monitor_training(self, task_id: str, ability_id: str, model_id: str, interval: int = 10):
+        task_name = f"{task_id}-{ability_id}-{model_id}"
+        logger.info(f"开始监控训练任务: {task_name}, 查询间隔: {interval} 秒")
+        print(f"\n开始监控训练任务: {task_name}")
+        print(f"查询间隔: {interval} 秒")
+        print("-" * 60)
+
+        while True:
+            try:
+                status = self.get_training_status(task_id, ability_id, model_id)
+                status_msg = f"状态: {status.get('status')} | 消息: {status.get('message', 'N/A')}"
+                logger.info(f"[{task_name}] {status_msg}")
+                print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {status_msg}")
+
+                if status.get('status') in ['completed', 'failed', 'stopped']:
+                    logger.info(f"训练任务结束: {task_name}, 状态: {status.get('status')}")
+                    print("\n" + "=" * 60)
+                    print(f"训练任务结束: {status.get('status')}")
+                    print(f"最终消息: {status.get('message', 'N/A')}")
+                    break
+
+                time.sleep(interval)
+
+            except KeyboardInterrupt:
+                logger.warning("监控已中断")
+                print("\n\n监控已中断")
+                break
+            except Exception as e:
+                logger.error(f"监控出错: {e}")
+                print(f"监控出错: {e}")
+                time.sleep(interval)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="训练客户端 - 向 FastAPI 服务发送训练请求")
+
+    parser.add_argument("--url", type=str, default="http://localhost:8001",
+                        help="API 服务地址 (默认: http://localhost:8001)")
+    parser.add_argument("--algo-type", type=str, default="bnn", choices=["bnn"],
+                        help="算法类型 (默认: bnn)")
+    parser.add_argument("--action-type", type=str, default="pick", choices=["pick", "place"],
+                        help="动作类型 (默认: pick)")
+    parser.add_argument("--task-id", type=str, default=None,
+                        help="任务ID（启动训练/查询状态/停止训练时必填）")
+    parser.add_argument("--ability-id", type=str, default=None,
+                        help="原子动作ID（启动训练/查询状态/停止训练时必填）")
+    parser.add_argument("--dataset-id", type=str, default=None,
+                        help="数据集ID（启动训练时必填）")
+    parser.add_argument("--model-id", type=str, default=None,
+                        help="模型ID（启动训练/查询状态/停止训练时必填）")
+
+    parser.add_argument("--batch_size", type=int, default=48,
+                        help="批次大小 (默认: 48)")
+    parser.add_argument("--seq_len", type=int, default=4,
+                        help="序列长度 (默认: 4)")
+    parser.add_argument("--action_chunk", type=int, default=8,
+                        help="动作块大小 (默认: 8)")
+    parser.add_argument("--lr", type=float, default=1e-4,
+                        help="学习率 (默认: 1e-4)")
+    parser.add_argument("--num_epochs", type=int, default=500,
+                        help="训练轮数 (默认: 500)")
+    parser.add_argument("--start_epoch", type=int, default=0,
+                        help="起始轮数，用于断点续训 (默认: 0)")
+    parser.add_argument("--lambda_joints", type=float, default=10.0,
+                        help="关节损失权重 (默认: 10.0)")
+    parser.add_argument("--lambda_grip", type=float, default=5.0,
+                        help="夹爪损失权重 (默认: 5.0)")
+    parser.add_argument("--lambda_success", type=float, default=2.0,
+                        help="成功标志损失权重 (默认: 2.0)")
+
+    parser.add_argument("--log_path", type=str, default=None,
+                        help="日志文件路径 (可选)")
+    parser.add_argument("--ckpt_dir", type=str, default=None,
+                        help="检查点保存目录 (可选)")
+    parser.add_argument("--success_mode", type=str, default="within_horizon",
+                        choices=["within_horizon", "terminal_only"],
+                        help="成功标志模式 (默认: within_horizon)")
+    parser.add_argument("--report_url", type=str, default=None,
+                        help="训练过程上报回调地址 (可选)")
+
+    parser.add_argument("--monitor", action="store_true",
+                        help="启动训练后监控任务状态")
+    parser.add_argument("--monitor_interval", type=int, default=10,
+                        help="监控查询间隔（秒）(默认: 10)")
+
+    parser.add_argument("--status", action="store_true",
+                        help="查询任务状态（不启动新训练）")
+    parser.add_argument("--stop", action="store_true",
+                        help="停止指定任务的训练")
+
+    args = parser.parse_args()
+    client = TrainClient(base_url=args.url)
+
+    if args.status:
+        if not args.task_id or not args.ability_id or not args.model_id:
+            print("查询状态需要提供 --task-id、--ability-id 和 --model-id")
+            return 1
+        try:
+            status = client.get_training_status(args.task_id, args.ability_id, args.model_id)
+            print(json.dumps(status, indent=2, ensure_ascii=False))
+        except Exception as e:
+            logger.error(f"查询失败: {e}")
+            print(f"查询失败: {e}")
+            return 1
+        return 0
+
+    if args.stop:
+        if not args.task_id or not args.ability_id or not args.model_id:
+            print("停止训练需要提供 --task-id、--ability-id 和 --model-id")
+            return 1
+        try:
+            result = client.stop_training(args.task_id, args.ability_id, args.model_id)
+            print(json.dumps(result, indent=2, ensure_ascii=False))
+        except Exception as e:
+            logger.error(f"停止失败: {e}")
+            print(f"停止失败: {e}")
+            return 1
+        return 0
+
+    if not args.task_id or not args.ability_id or not args.dataset_id or not args.model_id:
+        print("启动训练需要提供 --task-id、--ability-id、--dataset-id 和 --model-id")
+        return 1
+
+    try:
+        result = client.start_training(
+            algo_type=args.algo_type,
+            action_type=args.action_type,
+            task_id=args.task_id,
+            ability_id=args.ability_id,
+            dataset_id=args.dataset_id,
+            model_id=args.model_id,
+            batch_size=args.batch_size,
+            seq_len=args.seq_len,
+            action_chunk=args.action_chunk,
+            lr=args.lr,
+            num_epochs=args.num_epochs,
+            start_epoch=args.start_epoch,
+            lambda_joints=args.lambda_joints,
+            lambda_grip=args.lambda_grip,
+            lambda_success=args.lambda_success,
+            log_path=args.log_path,
+            ckpt_dir=args.ckpt_dir,
+            success_mode=args.success_mode,
+            report_url=args.report_url,
+        )
+
+        if args.monitor:
+            client.monitor_training(args.task_id, args.ability_id, args.model_id, interval=args.monitor_interval)
+        else:
+            print("\n提示: 使用以下命令查询训练状态:")
+            print(f"  python train_client.py --status --task-id {args.task_id} --ability-id {args.ability_id} --model-id {args.model_id} --url {args.url}")
+
+    except Exception as e:
+        logger.error(f"启动训练失败: {e}")
+        print(f"启动训练失败: {e}")
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

zyworkflow/example/train_client_example.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+from zyworkflow.example.train_client import TrainClient
+
+
+ALGO_TYPE = "bnn"
+ACTION_TYPE = "pick"
+TASK_ID = "t1"
+ABILITY_ID = "a1"
+DATASET_ID = "d1"
+MODEL_ID = "m1"
+API_URL = "http://localhost:8001"
+
+
+def main():
+    client = TrainClient(base_url=API_URL)
+    result = client.start_training(
+        algo_type=ALGO_TYPE,
+        action_type=ACTION_TYPE,
+        task_id=TASK_ID,
+        ability_id=ABILITY_ID,
+        dataset_id=DATASET_ID,
+        model_id=MODEL_ID,
+        batch_size=48,
+        seq_len=4,
+        action_chunk=8,
+        lr=1e-4,
+        num_epochs=500,
+        start_epoch=0,
+        lambda_joints=10.0,
+        lambda_grip=5.0,
+        lambda_success=2.0,
+        success_mode="within_horizon",
+    )
+
+    task_name = result.get("task_name")
+    if task_name:
+        print(f"训练任务已启动，任务名称: {task_name}")
+    else:
+        print("未能获取任务名称")
+
+
+if __name__ == "__main__":
+    main()