From b0e26f6d5e56c9196e17baa1e18b93a0812ff56e Mon Sep 17 00:00:00 2001 From: KAAANG <79990647+SAKURA-CAT@users.noreply.github.com> Date: Tue, 30 Jul 2024 16:48:29 +0800 Subject: [PATCH] v0.3.15-beta (#660) beta version fix some bugs new cli: task stop --- swanlab/cli/commands/task/__init__.py | 3 ++ swanlab/cli/commands/task/launch.py | 63 ++++++++++++++++++++++----- swanlab/cli/commands/task/list.py | 58 +++++++++++++++++++++--- swanlab/cli/commands/task/search.py | 10 ++++- swanlab/cli/commands/task/stop.py | 28 ++++++++++++ swanlab/cli/commands/task/utils.py | 14 +++++- swanlab/package.json | 2 +- 7 files changed, 156 insertions(+), 22 deletions(-) create mode 100644 swanlab/cli/commands/task/stop.py diff --git a/swanlab/cli/commands/task/__init__.py b/swanlab/cli/commands/task/__init__.py index 3dbcdc520..b2bd3bc28 100644 --- a/swanlab/cli/commands/task/__init__.py +++ b/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) diff --git a/swanlab/cli/commands/task/launch.py b/swanlab/cli/commands/task/launch.py index c7854738d..e1e394c6d 100644 --- a/swanlab/cli/commands/task/launch.py +++ b/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,6 +58,11 @@ ), 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", @@ -64,6 +70,15 @@ 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,7 +231,7 @@ 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 @@ -219,6 +239,7 @@ def __init__(self, username, src, key, python, name, index): :param python: 任务的python版本 :param name: 任务名称 :param index: 任务入口文件 + :param combo: 任务的套餐类型 """ self.username = username self.src = src @@ -226,8 +247,19 @@ def __init__(self, username, src, key, python, name, index): 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 diff --git a/swanlab/cli/commands/task/list.py b/swanlab/cli/commands/task/list.py index 7c1d538b0..bb190f789 100644 --- a/swanlab/cli/commands/task/list.py +++ b/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,13 +177,16 @@ 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: @@ -156,13 +194,19 @@ def add_event(self, info: str, max_length=15): 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() diff --git a/swanlab/cli/commands/task/search.py b/swanlab/cli/commands/task/search.py index e1dac4f12..e6a52f3c6 100644 --- a/swanlab/cli/commands/task/search.py +++ b/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}") diff --git a/swanlab/cli/commands/task/stop.py b/swanlab/cli/commands/task/stop.py new file mode 100644 index 000000000..f3cde6abe --- /dev/null +++ b/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.") diff --git a/swanlab/cli/commands/task/utils.py b/swanlab/cli/commands/task/utils.py index c52cf9547..50dc10e2d 100644 --- a/swanlab/cli/commands/task/utils.py +++ b/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): diff --git a/swanlab/package.json b/swanlab/package.json index b5056a1bf..5b5b884f6 100644 --- a/swanlab/package.json +++ b/swanlab/package.json @@ -1,6 +1,6 @@ { "name": "swanlab", - "version": "0.3.15-alpha.4", + "version": "0.3.15-beta", "description": "", "python": "true" }