General-AI-Platform-Backend/website/handlers/alg_model_hub/handler.py

# -*- coding: utf-8 -*-
import logging
import os
import requests

from sqlalchemy import text

from website import consts
from website import db_mysql
from website import errors
from website import settings
from website.handler import APIHandler, authenticated


class ListHandler(APIHandler):
    """
    - 描述： 模型运行库列表
    - 请求方式：post
    - 请求参数：
    > - pageNo, int
    > - pageSize, int
    > - name, string, 名称
    - 返回值：
    ```
        {
            "count": 123,
            "data": [
                {
                    "id": 123,
                    "name": "xx",
                    "create_time": "xxx",
                    "update_time": "xxx"
                },
                ...
            ]
        }
    ```
    """

    @authenticated
    def post(self):
        pageNo = self.get_int_argument("pageNo", 1)
        pageSize = self.get_int_argument("pageSize", consts.PAGE_SIZE)
        name = self.get_escaped_argument("name", "")

        result = []
        count = 0

        with self.app_mysql.connect() as conn:
            sql = "select id, name, path, create_time, update_time from model_hub where 1=1"
            param = {}

            sql_count = "select count(id) from model_hub where 1=1"
            param_count = {}

            if name:
                sql += " and name like :name"
                param["name"] = "%{}%".format(name)

                sql_count += " and name like :name"
                param_count["name"] = "%{}%".format(name)

            sql += " order by id desc limit :pageSize offset :offset"
            param["pageSize"] = pageSize
            param["offset"] = (pageNo - 1) * pageSize

            cur = conn.execute(text(sql), param)
            result = db_mysql.to_json_list(cur)

            count = conn.execute(text(sql_count), param_count).fetchone()[0]

        data = []
        for item in result:
            data.append(
                {
                    "id": item["id"],
                    "name": item["name"],
                    "path": item["path"],
                    "create_time": item["create_time"].strftime("%Y-%m-%d %H:%M:%S"),
                    "update_time": item["update_time"].strftime("%Y-%m-%d %H:%M:%S"),
                }
            )

        self.finish({"count": count, "data": data})


class SyncHandler(APIHandler):
    """
        - 描述： 查询docker registry中的镜像
        - 请求方式：post
        - 请求参数：
        > - host, string, ip地址
        > - port, int, 端口
        - 返回值：
        ```
            {
                "data": [
                    "xxx",    # docker registry中docker images的地址
                    "xxx",
                    ...
                ]
            }
    ```
    """

    @authenticated
    def post(self):
        host = self.get_escaped_argument("host", "")
        port = self.get_int_argument("port")
        if not host or not port:
            raise errors.HTTPAPIError(
                errors.ERROR_BAD_REQUEST, "host and port must be provided."
            )

        images = []
        # 查询docker registry中的镜像
        repositories = requests.get(
            "http://{}:{}/v2/_catalog".format(host, port)
        ).json()["repositories"]
        for repository in repositories:
            # 查询docker registry中的镜像的tag
            tags = requests.get(
                "http://{}:{}/v2/{}/tags/list".format(host, port, repository)
            ).json()["tags"]
            for tag in tags:
                image_name = "{}:{}/{}:{}".format(host, port, repository, tag)
                images.append(image_name)
        self.finish({"data": images})


class AddHandler(APIHandler):
    """
    - 描述： 新建模型运行库
    - 请求方式：post
    - 请求参数：
    > - name, string, 名称
    > - host, string,
    > - port, int
    > - path, string, 镜像路径
    > - comment, string, 备注
    - 返回值：无
    """

    @authenticated
    def post(self):
        name = self.get_escaped_argument("name", "")
        host = self.get_escaped_argument("host", "")
        port = self.get_int_argument("port")
        path = self.get_escaped_argument("path", "")
        comment = self.get_escaped_argument("comment", "")
        if not name or not host or not port:
            raise errors.HTTPAPIError(
                errors.ERROR_BAD_REQUEST, "name and host and port must be provided."
            )

        with self.app_mysql.connect() as conn:
            conn.execute(
                text(
                    """insert into model_hub (name, host, port, path, comment, create_time, update_time) 
                              values (:name, :host, :port, :path, :comment, NOW(), NOW())"""
                ),
                {
                    "name": name,
                    "host": host,
                    "port": port,
                    "path": path,
                    "comment": comment,
                },
            )

            conn.commit()

        self.finish()


class EditHandler(APIHandler):
    """
    - 描述： 编辑模型运行库
    - 请求方式：post
    - 请求参数：
    > - id, int
    > - name, string, 名称
    > - host, string,
    > - port, int
    > - path, string, 镜像路径
    > - comment, string, 备注
    - 返回值：无
    """

    @authenticated
    def post(self):
        id = self.get_int_argument("id")
        name = self.get_escaped_argument("name", "")
        host = self.get_escaped_argument("host", "")
        port = self.get_int_argument("port")
        path = self.get_escaped_argument("path", "")
        comment = self.get_escaped_argument("comment", "")
        if not id or not name or not host or not port or path:
            raise errors.HTTPAPIError(errors.ERROR_BAD_REQUEST, "parameter error")
        with self.app_mysql.connect() as conn:
            conn.execute(
                text(
                    """update model_hub set name=:name, host=:host, port=:port, path=:path, comment=:comment, update_time=NOW() 
                              where id=:id"""
                ),
                {
                    "id": id,
                    "name": name,
                    "host": host,
                    "port": port,
                    "path": path,
                    "comment": comment,
                },
            )
            conn.commit()
        self.finish()


class InfoHandler(APIHandler):
    """
    - 描述： 模型运行库信息
    - 请求方式：post
    - 请求参数：
    > - id, int
    - 返回值：
    ```
    {
        "name": "xxx",
        "host": "xxx",
        "port": 123,
        "path": "xxx",
        "comment": "xxx",
    }
    ```
    """

    @authenticated
    def post(self):
        hid = self.get_int_argument("id")
        if not id:
            raise errors.HTTPAPIError(errors.ERROR_BAD_REQUEST, "parameter error")

        result = {}
        with self.app_mysql.connect() as conn:
            cur = conn.execute(
                text(
                    """select name, host, port, path, comment from model_hub where id=:id"""
                ),
                {"id": hid},
            )
            result = db_mysql.to_json(cur)
            if not result:
                raise errors.HTTPAPIError(
                    errors.ERROR_BAD_REQUEST, "model hub not found"
                )
        self.finish(result)


class DeleteHandler(APIHandler):
    """
    - 描述： 删除模型运行库
    - 请求方式：post
    - 请求参数：
    > - id, int
    - 返回值：无
    """

    @authenticated
    def post(self):
        hid = self.get_int_argument("id")
        if not id:
            raise errors.HTTPAPIError(errors.ERROR_BAD_REQUEST, "parameter error")
        with self.app_mysql.connect() as conn:
            conn.execute(text("""delete from model_hub where id=:id"""), {"id": hid})
            conn.commit()
        self.finish()