intellif-aihub 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

aihub/services/artifact.py

@@ -0,0 +1,332 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+"""制品管理服务模块
+
+该模块提供了制品管理相关的功能，包括创建制品、上传制品、获取制品信息等。
+制品可以是数据集、模型、指标、日志、检查点、图像、预测结果等类型。
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import List, Optional
+
+import httpx
+import minio
+from loguru import logger
+
+from ..exceptions import APIError
+from ..models.artifact import (
+    CreateArtifactsReq,
+    CreateArtifactsResponseData,
+    ArtifactResp,
+    InfinityPageSize,
+    ArtifactRespData,
+    ArtifactType,
+    StsResp,
+)
+from ..models.common import APIWrapper
+from ..utils.s3 import S3_path_to_info, upload_dir_to_s3, download_dir_from_s3
+
+# 制品管理API的基础路径
+_Base = "/artifact-management/api/v1"
+
+
+class ArtifactService:
+    """制品管理服务类，该类提供了制品管理相关的功能，包括创建制品、上传制品、获取制品信息等。
+
+
+    Methods:
+        get_by_run_id: 使用run_id获取制品
+        create_artifact: 创建一个制品文件
+        create_artifacts: 从目录创建一个制品
+        download_artifacts: 下载制品
+
+    """
+
+    def __init__(self, http: httpx.Client):
+        """初始化制品管理服务
+
+        Args:
+            http: HTTP客户端实例
+        """
+        self._http = http
+        self._Artifact = _Artifact(http)
+        self.sts = self._get_sts()
+        self.s3_client = minio.Minio(
+            self.sts.endpoint,
+            access_key=self.sts.access_key_id,
+            secret_key=self.sts.secret_access_key,
+            secure=False,
+            session_token=self.sts.session_token,
+        )
+
+    @property
+    def _artifact(self) -> _Artifact:
+        """获取内部制品管理实例
+
+        Returns:
+            _Artifact: 内部制品管理实例
+        """
+        return self._Artifact
+
+    def _create(self, payload: CreateArtifactsReq) -> CreateArtifactsResponseData:
+        """创建制品（内部方法）
+
+        Args:
+            payload: 创建制品请求参数
+
+        Returns:
+            CreateArtifactsResponseData: 创建制品响应数据
+        """
+        return self._artifact.create(payload)
+
+    def _upload_done(self, artifact_id: int) -> None:
+        """标记制品上传完成
+
+        Args:
+            artifact_id: 制品ID
+        """
+        return self._artifact.upload_done(artifact_id)
+
+    def _get_sts(self) -> StsResp:
+        """获取STS临时凭证
+
+        Returns:
+            StsResp: STS临时凭证信息
+        """
+        return self._artifact.get_sts()
+
+    def get_by_run_id(
+        self, run_id: str, artifact_path: Optional[str] = None
+    ) -> List[ArtifactResp]:
+        """根据运行ID获取制品列表
+
+        Args:
+            run_id: 运行ID
+            artifact_path: 制品路径，如果指定则只返回匹配的制品
+
+        Returns:
+            List[ArtifactResp]: 制品列表
+
+        Raises:
+            APIError: 当API调用失败时抛出
+        """
+        return self._artifact.get_by_run_id(run_id, artifact_path)
+
+    def create_artifact(
+        self,
+        local_path: str,
+        artifact_path: Optional[str] = None,
+        run_id: Optional[str] = None,
+        artifact_type: ArtifactType = ArtifactType.other,
+    ) -> None:
+        """创建单个文件制品并上传
+
+        该方法用于将本地文件上传为制品。过程包括：
+        1. 获取STS临时凭证
+        2. 创建S3客户端
+        3. 检查本地文件是否存在
+        4. 创建制品记录
+        5. 上传文件到S3
+        6. 标记上传完成
+
+        Args:
+            local_path (str): 本地文件路径
+            artifact_path (str): 制品路径，如果为None则使用本地文件名
+            run_id (str): 运行ID，关联制品与特定运行
+            artifact_type (ArtifactType): 制品类型，默认为other
+
+        Raises:
+            ValueError: 当本地文件不存在时抛出
+            APIError: 当API调用失败时抛出
+        """
+        logger.info(f"log artifact: {artifact_path},local path: {local_path} ")
+
+        # 检查文件是否存在
+        if not os.path.exists(local_path):
+            raise ValueError(f"File {local_path} does not exist")
+        req = CreateArtifactsReq(
+            entity_id=run_id,
+            entity_type=artifact_type,
+            src_path=artifact_path,
+            is_dir=False,
+        )
+        resp = self._create(req)
+        bucket, object_name = S3_path_to_info(resp.s3_path)
+
+        self.s3_client.fput_object(bucket, object_name, local_path)
+        self._upload_done(resp.id)
+        logger.info(f"log artifact done: {artifact_path}")
+        return
+
+    def create_artifacts(
+        self,
+        local_dir: str,
+        artifact_path: Optional[str] = None,
+        run_id: Optional[str] = None,
+        artifact_type: ArtifactType = ArtifactType.other,
+    ) -> None:
+        """创建目录制品并上传
+
+        该方法用于将本地目录上传为制品。过程包括：
+        1. 获取STS临时凭证
+        2. 创建S3客户端
+        3. 检查本地目录是否存在
+        4. 创建制品记录
+        5. 上传目录内容到S3
+        6. 标记上传完成
+
+        Args:
+            local_dir (str): 本地目录路径
+            artifact_path (str): 制品路径，如果为None则使用本地目录名
+            run_id (str): 运行ID，关联制品与特定运行
+            artifact_type (ArtifactType): 制品类型，默认为other
+
+        Raises:
+            ValueError: 当本地目录不存在时抛出
+            APIError: 当API调用失败时抛出
+        """
+
+        logger.info(f"log artifact: {artifact_path},local path: {local_dir} ")
+        if not os.path.exists(local_dir):
+            raise ValueError(f"File {local_dir} does not exist")
+        req = CreateArtifactsReq(
+            entity_id=run_id,
+            entity_type=artifact_type,
+            src_path=artifact_path,
+            is_dir=True,
+        )
+        resp = self._create(req)
+        bucket, object_name = S3_path_to_info(resp.s3_path)
+        upload_dir_to_s3(self.s3_client, local_dir, bucket, object_name)
+        self._upload_done(resp.id)
+        logger.info(f"log artifact done: {artifact_path}")
+        return
+
+    def download_artifacts(
+        self, run_id: str, artifact_path: Optional[str], local_dir: str
+    ) -> None:
+        """下载制品
+
+        Args:
+            run_id: 运行ID
+            artifact_path: 制品路径
+            local_dir: 本地目录路径
+
+        Raises:
+            APIError: 当API调用失败时抛出
+        """
+        artifacts = self.get_by_run_id(run_id, artifact_path)
+
+        for artifact_item in artifacts:
+            bucket, object_name = S3_path_to_info(artifact_item.s3_path)
+            if artifact_item.is_dir:
+                download_dir_from_s3(self.s3_client, bucket, object_name, local_dir)
+            else:
+                self.s3_client.fget_object(
+                    bucket, object_name, str(Path(local_dir) / artifact_item.src_path)
+                )
+
+        logger.info(f"download artifact done: {artifact_path}")
+        return
+
+
+class _Artifact:
+    """内部制品管理类
+
+    该类提供了与制品管理API交互的底层方法。
+    通常不直接使用该类，而是通过ArtifactService类访问其功能。
+    """
+
+    def __init__(self, http: httpx.Client):
+        """初始化内部制品管理类
+
+        Args:
+            http: HTTP客户端实例
+        """
+        self._http = http
+
+    def create(self, payload: CreateArtifactsReq) -> CreateArtifactsResponseData:
+        """创建制品记录
+
+        Args:
+            payload: 创建制品请求参数
+
+        Returns:
+            CreateArtifactsResponseData: 创建制品响应数据
+
+        Raises:
+            APIError: 当API调用失败时抛出
+        """
+        resp = self._http.post(f"{_Base}/artifacts", json=payload.model_dump())
+
+        wrapper = APIWrapper[CreateArtifactsResponseData].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def upload_done(self, artifact_id: int) -> None:
+        """标记制品上传完成
+
+        在制品文件上传到S3后，需要调用此方法标记上传完成。
+
+        Args:
+            artifact_id: 制品ID
+
+        Raises:
+            APIError: 当API调用失败时抛出
+        """
+        resp = self._http.post(f"{_Base}/artifacts/{artifact_id}/uploaded")
+        wrapper = APIWrapper.model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return
+
+    def get_by_run_id(
+        self, run_id: str, artifact_path: Optional[str]
+    ) -> List[ArtifactResp]:
+        """根据运行ID获取制品列表
+
+        Args:
+            run_id: 运行ID
+            artifact_path: 制品路径，如果指定则只返回匹配的制品
+
+        Returns:
+            List[ArtifactResp]: 制品列表
+
+        Raises:
+            APIError: 当API调用失败时抛出
+        """
+        resp = self._http.get(
+            f"{_Base}/artifacts?entity_id={run_id}&page_num=1&page_size={InfinityPageSize}"
+        )
+        wrapper = APIWrapper[ArtifactRespData].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        if artifact_path:
+            return [
+                artifact
+                for artifact in wrapper.data.data
+                if artifact.src_path == artifact_path
+            ]
+        else:
+            return wrapper.data.data
+
+    def get_sts(self) -> StsResp:
+        """获取STS临时凭证
+
+        获取用于访问S3存储的临时凭证。
+
+        Returns:
+            StsResp: STS临时凭证信息
+
+        Raises:
+            APIError: 当API调用失败时抛出
+        """
+        resp = self._http.get(f"{_Base}/artifacts/get-sts")
+        logger.info(f"get sts: {resp.text}")
+        wrapper = APIWrapper[StsResp].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data

aihub/services/dataset_management.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+import mimetypes
+import os
+import pathlib
+
+import httpx
+
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.dataset_management import *
+from ..utils.download import dataset_download
+
+_BASE = "/dataset-mng/api/v2"
+
+
+class DatasetManagementService:
+    """数据集管理服务，用于数据集的上传、下载
+
+    Methods:
+        create_dataset_and_version: 创建数据集版本
+        run_download: 下载
+
+
+    """
+
+    def __init__(self, http: httpx.Client):
+        self._dataset = _Dataset(http)
+        self._dataset_version = _DatasetVersion(http)
+        self._upload = _Upload(http)
+
+    # 直接把常用方法抛到一级，调用体验简单
+    def create_dataset(self, payload: CreateDatasetRequest) -> int:
+        return self._dataset.create(payload)
+
+    def get_dataset(self, dataset_id: int) -> DatasetDetail:
+        return self._dataset.get(dataset_id)
+
+    def create_dataset_version(self, payload: CreateDatasetVersionRequest) -> int:
+        return self._dataset_version.create(payload)
+
+    def upload_dataset_version(self, payload: UploadDatasetVersionRequest) -> int:
+        return self._dataset_version.upload(payload)
+
+    def get_dataset_version(self, version_id: int) -> DatasetVersionDetail:
+        return self._dataset_version.get(version_id)
+
+    def get_dataset_version_by_name(self, version_name: str) -> DatasetVersionDetail:
+        return self._dataset_version.get_by_name(version_name)
+
+    def upload_file(self, file_path: str) -> FileUploadData:
+        return self._upload.upload_file(file_path)
+
+    # 如果想要访问子对象，也保留属性
+    @property
+    def dataset(self) -> _Dataset:
+        return self._dataset
+
+    @property
+    def dataset_version(self) -> _DatasetVersion:
+        return self._dataset_version
+
+    def create_dataset_and_version(
+        self,
+        *,
+        dataset_name: str,
+        dataset_description: str = "",
+        is_local_upload: bool,
+        local_file_path: str | None = None,
+        server_file_path: str | None = None,
+        version_description: str = "",
+    ) -> tuple[int, int, str]:
+        """创建数据集及其版本。
+
+        根据参数创建数据集，并根据上传类型（本地或服务器路径）创建对应的数据集版本。
+
+        Args:
+            dataset_name: 数据集名称。
+            dataset_description: 数据集描述，默认为空。
+            is_local_upload: 是否为本地上传。若为 True，需提供 local_file_path；
+                             否则需提供 server_file_path。
+            local_file_path: 本地文件路径，当 is_local_upload=True 时必须提供。
+            server_file_path: 服务器已有文件路径，当 is_local_upload=False 时必须提供。
+            version_description: 版本描述，默认为空。
+
+        Returns:
+           tuple[int, int, str]: 一个三元组，包含：[数据集 ID,数据集版本 ID, 数据集版本标签（格式为 <dataset_name>/V<version_number>)]
+        """
+        if is_local_upload:
+            if not local_file_path:
+                raise ValueError("is_local_upload=True 时必须提供 local_file_path")
+            upload_type = 1
+        else:
+            if not server_file_path:
+                raise ValueError("is_local_upload=False 时必须提供 server_file_path")
+            upload_type = 3
+
+        dataset_id = self._dataset.create(
+            CreateDatasetRequest(
+                name=dataset_name,
+                description=dataset_description,
+                tags=[],
+            )
+        )
+
+        if is_local_upload:
+            upload_data = self._upload.upload_file(local_file_path)
+            upload_path = upload_data.path
+        else:
+            upload_path = server_file_path
+
+        version_id = self._dataset_version.upload(
+            UploadDatasetVersionRequest(
+                upload_path=upload_path,
+                upload_type=upload_type,
+                dataset_id=dataset_id,
+                description=version_description,
+            )
+        )
+
+        detail = self._dataset.get(dataset_id)
+        ver_num = next(
+            (v.version for v in detail.versions if v.id == version_id),
+            None,
+        )
+        if ver_num is None:
+            ver_num = 1
+
+        version_tag = f"{detail.name}/V{ver_num}"
+
+        return dataset_id, version_id, version_tag
+
+    def run_download(
+        self, *, dataset_version_name: str, local_dir: str, worker: int = 4
+    ) -> None:
+        """根据数据集版本名称下载对应的数据集文件。
+
+        Args:
+            dataset_version_name (str): 数据集版本名称。
+            local_dir (str): 下载文件保存的本地目录路径。
+            worker (int): 并发下载使用的线程数，默认为 4。
+
+        Raises:
+            APIError: 如果获取到的版本信息中没有 parquet_index_path，即无法进行下载时抛出异常。
+
+        Returns:
+            None
+        """
+        detail = self._dataset_version.get_by_name(dataset_version_name)
+        if not detail.parquet_index_path:
+            raise APIError("parquet_index_path 为空")
+        dataset_download(detail.parquet_index_path, local_dir, worker)
+
+
+class _Dataset:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def create(self, payload: CreateDatasetRequest) -> int:
+        resp = self._http.post(
+            f"{_BASE}/datasets",
+            json=payload.model_dump(by_alias=True, exclude_none=True),
+        )
+        wrapper = APIWrapper[CreateDatasetResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data.id
+
+    def get(self, dataset_id: int) -> DatasetDetail:
+        resp = self._http.get(f"{_BASE}/datasets/{dataset_id}")
+        wrapper = APIWrapper[DatasetDetail].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+
+class _DatasetVersion:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def create(self, payload: CreateDatasetVersionRequest) -> int:
+        resp = self._http.post(
+            f"{_BASE}/dataset-versions",
+            json=payload.model_dump(by_alias=True, exclude_none=True),
+        )
+        wrapper = APIWrapper[CreateDatasetVersionResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data.id
+
+    def upload(self, payload: UploadDatasetVersionRequest) -> int:
+        resp = self._http.post(
+            f"{_BASE}/dataset-versions-upload",
+            json=payload.model_dump(by_alias=True, exclude_none=True),
+        )
+        wrapper = APIWrapper[UploadDatasetVersionResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data.id
+
+    def get(self, version_id: int) -> DatasetVersionDetail:
+        resp = self._http.get(f"{_BASE}/dataset-versions/{version_id}")
+        wrapper = APIWrapper[DatasetVersionDetail].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+    def get_by_name(self, version_name: str) -> DatasetVersionDetail:
+        resp = self._http.get(
+            f"{_BASE}/dataset-versions-detail", params={"name": version_name}
+        )
+        wrapper = APIWrapper[DatasetVersionDetail].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data
+
+
+class _Upload:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def upload_file(self, file_path: str) -> FileUploadData:
+        if not os.path.isfile(file_path):
+            raise FileNotFoundError(file_path)
+
+        file_name = pathlib.Path(file_path).name
+        mime_type = mimetypes.guess_type(file_name)[0] or "application/octet-stream"
+
+        with open(file_path, "rb") as fp:
+            resp = self._http.post(
+                f"/dataset-mng/api/v1/uploads",
+                files={"file": (file_name, fp, mime_type)},
+                timeout=None,
+            )
+
+        wrapper = APIWrapper[FileUploadData].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+
+        return wrapper.data

aihub/services/document_center.py

@@ -0,0 +1,43 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+from __future__ import annotations
+
+import httpx
+
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.document_center import *
+
+_BASE = "/document-center/api/v1"
+
+
+class DocumentCenterService:
+    def __init__(self, http: httpx.Client):
+        self._document = _Document(http)
+
+    def get_documents(
+            self, page_size: int = 9999, page_num: int = 1, name: str = ""
+    ) -> List[Document]:
+        return self._document.get_documents(page_size, page_num, name)
+
+    @property
+    def document(self) -> _Document:
+        return self._document
+
+
+class _Document:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def get_documents(
+            self, page_size: int = 9999, page_num: int = 1, name: str = ""
+    ) -> List[Document]:
+        params = {"page_size": page_size, "page_num": page_num, "name": name}
+        resp = self._http.get(f"{_BASE}/documents", params=params)
+        if resp.status_code != 200:
+            raise APIError(f"backend code {resp.status_code}: {resp.text}")
+        res = resp.json()
+        wrapper = APIWrapper[GetDocumentsResponse].model_validate(res)
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data.data

aihub/services/labelfree.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import httpx
+
+from ..exceptions import APIError
+from ..models.common import APIWrapper
+from ..models.labelfree import *
+
+_BASE = "/labelfree/api/v2"
+
+
+class LabelfreeService:
+    def __init__(self, http: httpx.Client):
+        self._project = _Project(http)
+
+    def get_project_global_stats(self, project_name: str) -> GetGlobalStatsResponse:
+        """获取标注项目的进展信息
+
+        Args:
+            project_name: 标注项目名称
+
+        Returns:
+            GetGlobalStatsResponse
+
+        """
+        return self._project.get_global_stats(project_name)
+
+    @property
+    def project(self) -> _Project:
+        return self._project
+
+
+class _Project:
+    def __init__(self, http: httpx.Client):
+        self._http = http
+
+    def get_global_stats(self, project_name: str) -> GetGlobalStatsResponse:
+        resp = self._http.get(
+            f"{_BASE}/projects/global_stats", params={"project_name": project_name}
+        )
+        wrapper = APIWrapper[GetGlobalStatsResponse].model_validate(resp.json())
+        if wrapper.code != 0:
+            raise APIError(f"backend code {wrapper.code}: {wrapper.msg}")
+        return wrapper.data

aihub/services/quota_schedule_management.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+_ENV_KEY = "PRE_STOP_SENTINEL_FILE"
+_DEFAULT_SENTINEL = "/tmp/pre_stop_sentinel_file"
+_SENTINEL_PATH = Path(os.getenv(_ENV_KEY, _DEFAULT_SENTINEL))
+
+
+def is_pre_stopped() -> bool:
+    return _SENTINEL_PATH.exists()
+
+
+class PreStopService:
+    @staticmethod
+    def is_pre_stopped() -> bool:
+        return is_pre_stopped()

aihub/services/reporter.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import json
+import os
+
+
+class ReportService:
+    @staticmethod
+    def report_envs():
+        pod_namespace = os.environ.get('POD_NAMESPACE', '')
+        pod_name = os.environ.get('POD_NAME', '')
+        if not pod_namespace or not pod_name:
+            return
+        env_file_path = f'/var/mtp/log/{pod_namespace}/{pod_name}/env'
+        try:
+            os.makedirs(os.path.dirname(env_file_path), exist_ok=True)
+            with open(env_file_path, 'w', encoding='utf-8') as f:
+                f.write(json.dumps(dict(os.environ)))
+        except Exception as e:
+            print(f"[report_envs] 写入失败: {e}", flush=True)