v0.3.15-beta (#660)

beta version

fix some bugs

new cli: task stop

swanlab/cli/commands/task/__init__.py

@@ -11,6 +11,7 @@
 from .launch import launch
 from .list import list
 from .search import search
+from .stop import stop
 import click
 
 __all__ = ["task"]
@@ -30,3 +31,5 @@ def task():
 task.add_command(list)
 # noinspection PyTypeChecker
 task.add_command(search)
+# noinspection PyTypeChecker
+task.add_command(stop)

swanlab/cli/commands/task/launch.py

@@ -11,6 +11,7 @@
 from .utils import login_init_sid, UseTaskHttp
 # noinspection PyPackageRequirements
 from qcloud_cos import CosConfig, CosS3Client
+from swanlab.error import ApiError
 from swanlab.log import swanlog
 from swankit.log import FONT
 import zipfile
@@ -57,13 +58,27 @@
     ),
     help="The entry file of the task, default by main.py",
 )
+@click.option(
+    "-y",
+    is_flag=True,
+    help="Skip the confirmation prompt and proceed with the task launch",
+)
 @click.option(
     "--python",
     default="python3.10",
     nargs=1,
     type=click.Choice(["python3.8", "python3.9", "python3.10"]),
     help="The python version of the task, default by python3.10",
 )
+@click.option(
+    "--combo",
+    "-c",
+    default=None,
+    nargs=1,
+    type=str,
+    help="The plan of the task. Swanlab will use the default plan if not specified. "
+         "You can check the plans in the official documentation.",
+)
 @click.option(
     "--name",
     "-n",
@@ -72,26 +87,31 @@
     type=str,
     help="The name of the task, default by Task_{current_time}",
 )
-def launch(path: str, entry: str, python: str, name: str):
+def launch(path: str, entry: str, python: str, name: str, combo: str, y: bool):
+    """
+    Launch a task!
+    """
     if not entry.startswith(path):
         raise ValueError(f"Error: Entry file '{entry}' must be in directory '{path}'")
     entry = os.path.relpath(entry, path)
     # 获取访问凭证，生成http会话对象
     login_info = login_init_sid()
     print(FONT.swanlab("Login successfully. Hi, " + FONT.bold(FONT.default(login_info.username))) + "!")
-    # 上传文件
-    text = f"The target folder {FONT.yellow(path)} will be packaged and uploaded, "
-    text += f"and you have specified {FONT.yellow(entry)} as the task entry point. "
-    swanlog.info(text)
-    ok = click.confirm(FONT.swanlab("Do you wish to proceed?"), abort=False)
-    if not ok:
-        return
+    # 确认
+    if not y:
+        swanlog.info("Please confirm the following information:")
+        swanlog.info(f"The target folder {FONT.yellow(path)} will be packaged and uploaded")
+        swanlog.info(f"You have specified {FONT.yellow(entry)} as the task entry point. ")
+        combo and swanlog.info(f"The task will use the combo {FONT.yellow(combo)}")
+        ok = click.confirm(FONT.swanlab("Do you wish to proceed?"), abort=False)
+        if not ok:
+            return
     # 压缩文件夹
     memory_file = zip_folder(path)
     # 上传文件
     src = upload_memory_file(memory_file)
     # 发布任务
-    ctm = CreateTaskModel(login_info.username, src, login_info.api_key, python, name, entry)
+    ctm = CreateTaskModel(login_info.username, src, login_info.api_key, python, name, entry, combo)
     ctm.create()
     swanlog.info(f"Task launched successfully. You can use {FONT.yellow('swanlab task list')} to view the task.")
 
@@ -211,23 +231,35 @@ def upload_memory_file(memory_file: io.BytesIO) -> str:
 
 
 class CreateTaskModel:
-    def __init__(self, username, src, key, python, name, index):
+    def __init__(self, username, src, key, python, name, index, combo):
         """
         :param username: 用户username
         :param key: 用户的api_key
         :param src: 任务zip文件路径
         :param python: 任务的python版本
         :param name: 任务名称
         :param index: 任务入口文件
+        :param combo: 任务的套餐类型
         """
         self.username = username
         self.src = src
         self.key = key
         self.python = python
         self.name = name
         self.index = index
+        self.combo = combo
 
     def __dict__(self):
+        if self.combo is not None:
+            return {
+                "username": self.username,
+                "src": self.src,
+                "index": self.index,
+                "python": self.python,
+                "conf": {"key": self.key},
+                "combo": self.combo,
+                "name": self.name
+            }
         return {
             "username": self.username,
             "src": self.src,
@@ -241,5 +273,12 @@ def create(self):
         """
         创建任务
         """
-        with UseTaskHttp() as http:
-            http.post("/task", self.__dict__())
+        try:
+            with UseTaskHttp() as http:
+                http.post("/task", self.__dict__())
+        except ApiError as e:
+            if e.resp.status_code == 406:
+                raise click.exceptions.UsageError("You have reached the maximum number of tasks")
+            elif e.resp.status_code == 400:
+                raise click.exceptions.UsageError("Incorrect combo name, please check it")
+            raise e

swanlab/cli/commands/task/list.py

@@ -28,14 +28,42 @@
     help="The maximum number of tasks to display, default by 10, maximum by 100",
 )
 def list(max_num: int):  # noqa
+    """
+    List tasks
+    """
     # 获取访问凭证，生成http会话对象
     login_info = login_init_sid()
     # 获取任务列表
     ltm = ListTasksModel(num=max_num, username=login_info.username)
-    layout = ListTaskLayout(ltm)
+    aqm = AskQueueModel()
+    layout = ListTaskLayout(ltm, aqm)
     layout.show()
 
 
+class AskQueueModel:
+    def __init__(self):
+        self.num = None
+
+    def ask(self):
+        with UseTaskHttp() as http:
+            queue_info = http.get("/task/queuing")
+            self.num = queue_info["sum"]
+
+    def table(self):
+        qi = Table(
+            expand=True,
+            show_header=False,
+            header_style="bold",
+            title="[blue][b]Now Global Queue[/b]",
+            highlight=True,
+            border_style="blue",
+        )
+        qi.add_column("Queue Info", "Queue Info")
+        self.ask()
+        qi.add_row(f"📦[b]Task Queuing count: {self.num}[/b]")
+        return qi
+
+
 class ListTasksModel:
     def __init__(self, num: int, username: str):
         """
@@ -59,6 +87,7 @@ def table(self):
             title="[magenta][b]Now Task[/b]",
             highlight=True,
             border_style="magenta",
+            show_lines=True,
         )
         st.add_column("Task ID", justify="right")
         st.add_column("Task Name", justify="center")
@@ -105,7 +134,7 @@ class ListTaskLayout:
     任务列表展示
     """
 
-    def __init__(self, ltm: ListTasksModel):
+    def __init__(self, ltm: ListTasksModel, aqm: AskQueueModel):
         self.event = []
         self.add_event(f"👏Welcome, [b]{ltm.username}[/b].")
         self.add_event("⌛️Task board is loading...")
@@ -116,11 +145,17 @@ def __init__(self, ltm: ListTasksModel):
         )
         self.layout["main"].split_row(
             Layout(name="task_table", ratio=16),
+            Layout(name="info_side", ratio=5)
+        )
+        self.layout["info_side"].split_column(
+            Layout(name="queue_info", ratio=1),
             Layout(name="term_output", ratio=5)
         )
         self.layout["header"].update(ListTaskHeader())
         self.layout["task_table"].update(Panel(ltm.table(), border_style="magenta"))
+        self.layout["queue_info"].update(Panel(aqm.table(), border_style="blue"))
         self.ltm = ltm
+        self.aqm = aqm
         self.add_event("🍺Task board is loaded.")
         self.redraw_term_output()
 
@@ -142,27 +177,36 @@ def term_output(self):
         )
         return to
 
-    def redraw_term_output(self, ):
+    def redraw_term_output(self):
         term_output = self.term_output
         for row in self.event:
             term_output.add_row(row)
         self.layout["term_output"].update(Panel(term_output, border_style="blue"))
 
-    def add_event(self, info: str, max_length=15):
+    def redraw_queue_info(self):
+        pass
+
+    def add_event(self, info: str, max_length=10):
         # 事件格式：yyyy-mm-dd hh:mm:ss - info
         self.event.append(f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} - {info}")
         while len(self.event) > max_length:
             self.event.pop(0)
 
     def show(self):
         with Live(self.layout, refresh_per_second=10, screen=True) as live:
-            now = time.time()
+            search_now = time.time()
+            queue_now = time.time()
             while True:
                 time.sleep(1)
                 self.layout["header"].update(ListTaskHeader())
-                if time.time() - now > 5:
-                    now = time.time()
+                if time.time() - search_now > 5:
+                    search_now = time.time()
                     self.add_event("🔍Searching for new tasks...")
                     self.layout["task_table"].update(Panel(self.ltm.table(), border_style="magenta"))
                     self.redraw_term_output()
+                if time.time() - queue_now > 3:
+                    queue_now = time.time()
+                    self.add_event("📦Asking queue info...")
+                    self.layout["queue_info"].update(Panel(self.aqm.table(), border_style="blue"))
+                    self.redraw_term_output()
                 live.refresh()

swanlab/cli/commands/task/search.py

@@ -10,6 +10,7 @@
 import click
 from .utils import TaskModel, login_init_sid, UseTaskHttp
 from rich.syntax import Console, Syntax
+from swanlab.error import ApiError
 
 
 def validate_six_char_string(_, __, value):
@@ -30,7 +31,11 @@ def search(cuid):
     """
     login_info = login_init_sid()
     with UseTaskHttp() as http:
-        data = http.get(f"/task/{cuid}")
+        try:
+            data = http.get(f"/task/{cuid}")
+        except ApiError as e:
+            if e.resp.status_code == 404:
+                raise click.BadParameter("Task not found")
     tm = TaskModel(login_info.username, data)
     """
     任务名称，python版本，入口文件，任务状态，URL，创建时间，执行时间，结束时间，错误信息
@@ -44,9 +49,12 @@ def search(cuid):
     icon = '✅'
     if tm.status == 'CRASHED':
         icon = '❌'
+    elif tm.status == 'STOPPED':
+        icon = '🛑'
     elif tm.status != 'COMPLETED':
         icon = '🏃'
     console.print(f"[bold]Status:[/bold] {icon} {tm.status}")
+    console.print(f"[bold]Combo:[/bold] [white]{tm.combo}[/white]")
     tm.url is not None and console.print(f"[bold]SwanLab URL:[/bold] {tm.url}")
     console.print(f"[bold]Created At:[/bold] {tm.created_at}")
     tm.started_at is not None and console.print(f"[bold]Started At:[/bold] {tm.started_at}")

swanlab/cli/commands/task/stop.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+r"""
+@DATE: 2024/7/30 16:13
+@File: stop.py
+@IDE: pycharm
+@Description:
+    停止任务
+"""
+import click
+from .utils import login_init_sid, UseTaskHttp, validate_six_char_string
+from swanlab.error import ApiError
+
+
+@click.command()
+@click.argument("cuid", type=str, callback=validate_six_char_string)
+def stop(cuid):
+    """
+    Stop a task by cuid
+    """
+    login_init_sid()
+    with UseTaskHttp() as http:
+        try:
+            http.patch(f"/task/status", {"cuid": cuid, "status": "STOPPED", "msg": "User stopped by sdk"})
+        except ApiError as e:
+            if e.resp.status_code == 404:
+                raise click.BadParameter("Task not found")
+    click.echo("Task stopped successfully, there may be a few minutes of delay online.")

swanlab/cli/commands/task/utils.py

@@ -14,6 +14,17 @@
 from typing import Optional
 from swanlab.log import swanlog
 import sys
+import click
+
+
+def validate_six_char_string(_, __, value):
+    if value is None:
+        raise click.BadParameter('Parameter is required')
+    if not isinstance(value, str):
+        raise click.BadParameter('Value must be a string')
+    if len(value) != 6:
+        raise click.BadParameter('String must be exactly 6 characters long')
+    return value
 
 
 def login_init_sid() -> LoginInfo:
@@ -32,7 +43,7 @@ class TaskModel:
     获取到的任务列表模型
     """
 
-    def __init__(self, username: str, task: dict, ):
+    def __init__(self, username: str, task: dict):
         self.cuid = task["cuid"]
         self.username = username
         self.name = task["name"]
@@ -60,6 +71,7 @@ def __init__(self, username: str, task: dict, ):
         self.finished_at = self.fmt_time(task.get("finishedAt", None))
         self.status = task["status"]
         self.msg = task.get("msg", None)
+        self.combo = task["combo"]
 
     @property
     def url(self):

swanlab/package.json

@@ -1,6 +1,6 @@
 {
   "name": "swanlab",
-  "version": "0.3.15-alpha.4",
+  "version": "0.3.15-beta",
   "description": "",
   "python": "true"
 }