volcengine-python-sdk-aicc 0.0.41__py3-none-any.whl

bytedance/jeddak_secure_channel/utils.py

@@ -0,0 +1,469 @@
+__all__ = [
+    "RepeatTimer",
+    "WeakMethod",
+    "weak_repeat_timer",
+    "request_bytedance_gateway",
+]
+
+import hashlib
+import hmac
+import json
+import os
+import urllib.parse
+import weakref
+from datetime import datetime, timezone
+from threading import Timer
+from typing import Dict
+
+import requests
+import volcenginesdkcore
+import volcenginesdksts
+from typing_extensions import (
+    Any,
+    Callable,
+    Concatenate,
+    Generic,
+    Iterable,
+    Mapping,
+    Optional,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+)
+from volcenginesdkcore.rest import ApiException
+
+from . import error
+from .log import logger
+
+
+class RepeatTimer(Timer):
+    """
+    定时重复执行函数.
+    """
+
+    def run(self):
+        while not self.finished.wait(self.interval):
+            self.function(*self.args, **self.kwargs)
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+R = TypeVar("R")
+R_co = TypeVar("R_co", covariant=True)
+
+
+class MethodProtocol(Protocol[T, P, R]):
+    __self__: T
+
+    __func__: Callable[Concatenate[T, P], R]
+
+    __call__: Callable[P, R]
+
+
+class WeakMethod(Generic[T, P, R_co]):
+    """
+    对方法 (bound method) 的弱引用.
+    """
+
+    # instance: weakref.ReferenceType[T]
+    # func: weakref.ReferenceType[Callable[Concatenate[T, P], R_co]]
+
+    def __init__(self, method: Callable[P, R_co]):
+        # Type checkers do not differentiate methods from callables, so
+        # we have to ignore the casting error
+        _method: MethodProtocol[T, P, R_co] = method  # pyright: ignore
+
+        self.instance = weakref.ref(_method.__self__)
+        self.func = weakref.ref(_method.__func__)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Optional[R_co]:
+        instance = self.instance()
+        func = self.func()
+
+        if instance is None or func is None:
+            logger.info("Weak method called on freed object")
+            return
+
+        return func(instance, *args, **kwargs)
+
+
+def weak_repeat_timer(
+    interval: float,
+    method: Callable[P, None],
+    args: Optional[Iterable[Any]] = None,
+    kwargs: Optional[Mapping[str, Any]] = None,
+) -> RepeatTimer:
+    """
+    定时重复执行对象上的方法. 在对象被垃圾回收时自动结束线程.
+    """
+    # Due to limitations of 'ParamSpec', we cannot annotate 'args' and 'kwargs' precisely
+    timer = RepeatTimer(interval, WeakMethod(method), args, kwargs)
+    weakref.finalize(method.__self__, timer.cancel)  # pyright: ignore
+
+    return timer
+
+
+class TopRequester:
+    def __init__(
+        self,
+        ak,
+        sk,
+        token=None,
+        host="open.volcengineapi.com",
+        region="cn-beijing",
+        use_tls=True,
+        service="pcc",
+        version="2024-12-24",
+        debug=False,
+    ):
+        volc_config = volcenginesdkcore.Configuration()
+        volc_config.host = host
+        volc_config.ak = ak
+        volc_config.sk = sk
+        volc_config.session_token = token
+        volc_config.region = region
+        volc_config.debug = debug
+        if use_tls:
+            volc_config.scheme = "https"
+        else:
+            volc_config.scheme = "http"
+        self._api_client = volcenginesdkcore.ApiClient(configuration=volc_config)
+        self._service = service
+        self._version = version
+
+    def _construct_resource_path(
+        self,
+        action: str,
+        version: str,
+        service: str,
+        method: str = "post",
+        content_type: str = "application_json",
+    ) -> str:
+        return f"/{action}/{version}/{service}/{method.lower()}/{content_type}/"
+
+    def request(
+        self,
+        method,
+        query,
+        header,
+        param,
+        action,
+        body,
+        service=None,
+        version=None,
+        content_type="application/json",
+    ):
+        resp = self._api_client.call_api(
+            self._construct_resource_path(
+                action=action,
+                version=self._version if not version else version,
+                service=self._service if not service else service,
+                method=method,
+                content_type=content_type.replace("/", "_"),
+            ),
+            method,
+            body=body,
+            auth_settings=["volcengineSign"],
+            header_params={
+                "Accept": "application/json",
+                "Content-Type": content_type,
+                **header,
+            },
+            path_params=param,
+            query_params=query,
+            post_params=[],
+            files={},
+            response_type=bytes,
+            collection_formats={},
+            _return_http_data_only=False,
+            _preload_content=True,
+        )
+
+        status, headers, data = resp[1], resp[2], resp[0]
+        logger.info(f"Top request response,status={status}, headers: {headers}")
+
+        return data
+
+    def request_with_assume_role(
+        self,
+        method,
+        query,
+        header,
+        param,
+        service,
+        action,
+        body,
+        version,
+        target_uid,
+        target_resource=None,
+        region="cn-beijing",
+        content_type="application/json",
+        use_tls=True,
+        host="open.volcengineapi.com",
+    ):
+        data = {
+            "TargetAccount": target_uid,
+            "TargetResource": os.getenv("PCC_SERVICE", "pcc")
+            if not target_resource
+            else target_resource,
+        }
+        access_info = self.request("POST", [], {}, {}, "UpdateUserAccessInfo", data)
+        ak = access_info["Ak"]
+        sk = access_info["Sk"]
+        sts_token = access_info["StsToken"]
+
+        customer_requester = TopRequester(
+            ak=ak, sk=sk, token=sts_token, region=region, use_tls=use_tls, host=host
+        )
+        return customer_requester.request(
+            method, query, header, param, action, body, service, version, content_type
+        )
+
+
+def _hmac_sha256(key: bytes, content: str):
+    return hmac.new(key, content.encode(), hashlib.sha256).digest()
+
+
+def _hash_sha256(content: str):
+    return hashlib.sha256(content.encode()).hexdigest()
+
+
+def _format_query(params: Dict[str, str]) -> str:
+    return "&".join(
+        urllib.parse.quote(key, safe="-_.~")
+        + "="
+        + urllib.parse.quote(value, safe="-_.~")
+        for key, value in params.items()
+    )
+
+
+def top_request(
+    ra_url: str,
+    body: Any,
+    additional_headers: Optional[Dict[str, str]],
+    top_config: Dict[str, str],
+) -> requests.Response:
+    top_ak = top_config["ak"]
+    top_sk = top_config["sk"]
+    top_service = top_config["service"]
+    top_region = top_config.get("region", "cn-beijing")
+    top_method = top_config.get("method", "POST").upper()
+    top_action = top_config.get("action", "GetAttestationBackend")
+    top_version = top_config.get("version", "2024-12-24")
+    use_tls: bool = top_config.get("use_tls", True)  # type: ignore
+    sts_token = top_config.get("sts_token", None)
+    target_uid = top_config.get("target_uid", None)
+    target_resource = top_config.get("target_resource", "pcc")
+
+    if ra_url:
+        top_host = ra_url
+    else:
+        top_host = top_config.get("url", "open.volcengineapi.com")
+
+    requester = TopRequester(
+        ak=top_ak,
+        sk=top_sk,
+        token=sts_token,
+        region=top_region,
+        use_tls=use_tls,
+        host=top_host,
+        service=top_service,
+        version=top_version,
+    )
+
+    resp = None
+    if not target_uid:
+        try:
+            resp = requester.request(
+                top_method, [], additional_headers, {}, top_action, body
+            )
+        except ApiException as e:
+            logger.error(
+                "status: {}, reason: {}, body: {}, headers: {}".format(
+                    e.status, e.reason, e.body, e.headers
+                )
+            )
+            raise e
+        except Exception as e:
+            logger.error("request error: {}".format(e))
+            raise e
+    else:
+        try:
+            resp = requester.request_with_assume_role(
+                top_method,
+                [],
+                additional_headers,
+                {},
+                top_service,
+                top_action,
+                body,
+                top_version,
+                target_uid,
+                target_resource,
+                top_region,
+                use_tls=use_tls,
+                host=top_host
+            )
+        except ApiException as e:
+            logger.error(
+                "status: {}, reason: {}, body: {}, headers: {}".format(
+                    e.status, e.reason, e.body, e.headers
+                )
+            )
+            raise e
+        except Exception as e:
+            logger.error("request error: {}".format(e))
+            raise e
+
+    return resp
+
+
+def get_top_access_key(top_config: Dict[str, str]) -> tuple:
+    try:
+        configuration = volcenginesdkcore.Configuration()
+        configuration.ak = top_config.get("ak")
+        configuration.sk = top_config.get("sk")
+        configuration.region = top_config.get("region", "cn-beijing")
+
+        duration_seconds = int(top_config.get("duration_seconds", "3600"))
+        role_session_name = top_config.get("role_session_name", "temp_")
+        role_trn = top_config.get("aicc_saas_trn")
+
+        volcenginesdkcore.Configuration.set_default(configuration)
+
+        api_instance = volcenginesdksts.STSApi()
+        assume_role_request = volcenginesdksts.AssumeRoleRequest(
+            duration_seconds=duration_seconds,
+            role_session_name=role_session_name,
+            role_trn=role_trn,
+        )
+
+        ret = api_instance.assume_role(assume_role_request)
+        if ret and ret.credentials:
+            ak = ret.credentials.access_key_id
+            sk = ret.credentials.secret_access_key
+            session_token = ret.credentials.session_token
+
+            return ak, sk, session_token
+        else:
+            return None, None, None
+
+    except Exception as e:
+        logger.error("get_top_access_key: {}".format(e))
+        raise e
+
+
+# The first parameter was url, and is now unused
+def request_bytedance_gateway(
+    url: str,
+    body: Any,
+    additional_headers: Optional[Dict[str, str]],
+    top_config: Dict[str, str],
+) -> requests.Response:
+    top_ak = top_config["ak"]
+    top_sk = top_config["sk"]
+
+    if not additional_headers:
+        additional_headers = dict()
+
+    if top_config.get("aicc_saas_trn"):
+        ak, sk, session_token = get_top_access_key(top_config)
+        if ak and sk and session_token:
+            top_ak = ak
+            top_sk = sk
+            additional_headers["X-Security-Token"] = session_token
+            additional_headers["sts_token"] = session_token
+        if top_config.get("target_uid"):
+            additional_headers["UID"] = top_config.get("target_uid")
+
+    top_service = top_config["service"]
+    top_region = top_config.get("region", "cn-beijing")
+    top_method = top_config.get("method", "POST").upper()
+    top_action = top_config.get("action", "GetAttestationBackend")
+    top_version = top_config.get("version", "2024-12-24")
+    url_rewrite = top_config.get("url_rewrite", "")
+    http_scheme = top_config.get("http_scheme", "https")
+
+    top_url = top_config.get("url", "") if top_config.get("url", "") else url
+    if not top_url:
+        top_url = url_rewrite
+    if not top_url:
+        logger.error("ra request error url=None")
+        raise error.ParamError("top_url", "top_url is None")
+
+    if top_url.startswith("http"):
+        url_parsed = urllib.parse.urlparse(top_url)
+        top_host = url_parsed.hostname
+    else:
+        top_host = top_url
+
+    params = {"Action": top_action, "Version": top_version}
+    query_str = _format_query(params)
+
+    date_now = datetime.now(timezone.utc)
+    body_str = json.dumps(body, separators=(",", ":"))
+
+    x_date = date_now.strftime("%Y%m%dT%H%M%SZ")
+    short_x_date = x_date[:8]
+    x_content_sha256 = _hash_sha256(body_str)
+
+    headers = {
+        "Content-Type": "application/json",
+        "Host": top_host,
+        "X-Content-Sha256": x_content_sha256,
+        "X-Date": x_date,
+    }
+
+    signed_headers = ";".join(key.lower() for key in headers.keys())
+
+    canonical_request = "\n".join(
+        [
+            top_method,
+            "/",
+            query_str,
+            *[key.lower() + ":" + value for key, value in headers.items()],
+            "",
+            signed_headers,
+            x_content_sha256,
+        ]
+    )
+
+    canonical_request_hash = _hash_sha256(canonical_request)
+    credential_scope = "/".join([short_x_date, top_region, top_service, "request"])
+    string_to_sign = "\n".join(
+        ["HMAC-SHA256", x_date, credential_scope, canonical_request_hash]
+    )
+
+    logger.debug(
+        f"ra request={canonical_request}, hash={canonical_request_hash}, sign={string_to_sign}"
+    )
+
+    k_date = _hmac_sha256(top_sk.encode(), short_x_date)
+    k_region = _hmac_sha256(k_date, top_region)
+    k_service = _hmac_sha256(k_region, top_service)
+    k_signing = _hmac_sha256(k_service, "request")
+    signature = _hmac_sha256(k_signing, string_to_sign).hex()
+
+    headers["Authorization"] = (
+        f"HMAC-SHA256 Credential={top_ak}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}"
+    )
+
+    if url_rewrite:
+        url_parsed = urllib.parse.urlparse(url_rewrite)
+        headers["Host"] = url_parsed.hostname or ""
+        url_final = url_rewrite + "?" + query_str
+    else:
+        if top_url.startswith("http"):
+            url_final = top_url + "?" + query_str
+        else:
+            url_final = f"{http_scheme}://{top_url}/?{query_str}"
+
+    if additional_headers:
+        headers = {**headers, **additional_headers}
+
+    r = requests.request(
+        method=top_method, url=url_final, headers=headers, data=body_str, timeout=30
+    )
+    logger.debug(f"call ra: url={url_final}, header={headers}, body={body_str}, res={json.loads(r.text).get('Result')}")
+    logger.info(f"ra response headers: X-Tt-Logid={r.headers.get('X-Tt-Logid')}")
+    return r

bytedance/jeddak_secure_model/model_automation_tool.py

@@ -0,0 +1,181 @@
+from bytedance.jeddak_secure_model.model_encryption import \
+    EncryptionConfig, JeddakModelEncrypter
+from bytedance.jeddak_secure_model import pcc_utils
+import bytedance
+from bytedance.jeddak_secure_model.pcc_utils import \
+    set_tks_policy, \
+    publish_inference_model, \
+    deploy_inference_model, \
+    list_inference_model, \
+    test_inference_model, \
+    get_engine_conf
+import argparse
+import time
+import sys
+
+
+def parser():
+    parser = argparse.ArgumentParser(description="JeddakPCC Model Automation Tool")
+    parser.add_argument("--model_path", type=str, default="", required=True, help="本地模型的文件夹路径")
+    parser.add_argument("--bucket_name", type=str, default="", help="创建TOS的存储桶")
+    parser.add_argument("--AK", type=str, default="", help="your volc access key")
+    parser.add_argument("--SK", type=str, default="", help="your volc secret key")
+    parser.add_argument("--region", type=str, default="cn-beijing", help="存储桶的区域")
+    parser.add_argument("--endpoint", type=str, default="tos-cn-beijing.volces.com", help="tos endpoint")
+    parser.add_argument("--ring_id", type=str, default="", help="密钥环ID")
+    parser.add_argument("--ring_name", type=str, default="", help="密钥环名称")
+    parser.add_argument("--ring_desc", type=str, default="", help="密钥环描述")
+    parser.add_argument("--key_name", type=str, default="", help="密钥名称")
+    parser.add_argument("--key_desc", type=str, default="", help="密钥描述")
+    parser.add_argument("--app_id", type=str, default="", help="火山页面查看自己账号的app_id")
+    parser.add_argument("--service", type=str, default="pcc", help="pcc for online, pcc_test for ppe")
+    parser.add_argument("--encrypt_flag", '-e', action='store_true', help="是否加密模型，输入此参数表示加密，不输入此参数表示不加密")
+    parser.add_argument("--tks_addr", type=str, default="pcc.volcengineapi.com", help="tks top 地址")
+    parser.add_argument("--policy_id", type=str, default="", help="策略ID")
+    parser.add_argument("--model_name", type=str, default="", help="模型名称")
+    parser.add_argument("--model_version", type=str, default="", help="模型版本")
+    parser.add_argument("--model_description", type=str, default="", help="模型描述")
+    parser.add_argument("--cu_type", type=str, default="", help="CU类型，可选值为Basic, Advanced, Enterprise")
+    parser.add_argument("--replica_count", type=int, default=1, help="副本数量，默认1")
+    parser.add_argument("--inference_engine", type=str, default="vllm", help="推理引擎，可选值为vllm, sglang")
+    parser.add_argument("--silent", '-s', action='store_true', help="h是否打印http request日志")
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = parser()
+    config = EncryptionConfig("", "")
+    encrypter = JeddakModelEncrypter(config)
+
+    # 通过 sys.modules 修改 b 中的全局变量
+    sys.modules['bytedance.jeddak_secure_model.pcc_utils'].Host = args.tks_addr
+    sys.modules['bytedance.jeddak_secure_model.pcc_utils'].Region = args.region
+
+    print("Step1: Encrypting model and uploading to TOS...")
+    res = encrypter.encrypt_model_and_upload(
+        model_path=args.model_path,
+        bucket_name=args.bucket_name,
+        volc_ak=args.AK,
+        volc_sk=args.SK,
+        region=args.region,
+        endpoint=args.endpoint,
+        ring_id=args.ring_id,
+        ring_name=args.ring_name,
+        ring_desc=args.ring_desc,
+        key_name=args.key_name,
+        key_desc=args.key_desc,
+        app_id=args.app_id,
+        service=args.service,
+        tks_addr=args.tks_addr,
+        encrypt_flag=args.encrypt_flag
+    )
+    print(f"Step1: Finished with result={res}")
+
+    # Step2: Set TKS policy
+    # 注意：需要在JeddakPCC的控制台中创建密钥环和密钥，并获取密钥环ID和密钥ID
+    # 这里的key_id是上一步中返回的密钥ID
+    # policy_id是你在JeddakPCC控制台中创建的策略ID
+    # 你需要在JeddakPCC控制台中创建策略，并获取策略ID
+    step = 1
+    if args.encrypt_flag:
+        print("Step2: Setting TKS policy...")
+        set_tks_policy(
+            AK=args.AK,
+            SK=args.SK,
+            uid=args.app_id,
+            key_id=res["key_id"],
+            policy_id=args.policy_id,
+            silent=args.silent,
+            service=args.service
+        )
+        print("Step2: Finished setting TKS policy")
+        step = 2
+
+    print("Step{}: Publishing model to JeddakPCC MMS".format(1+step))
+    ring_id = res.get("ring_id")
+    key_id = res.get("key_id")
+    if ring_id is None:
+        ring_id = "None"
+    if key_id is None:
+        key_id = "None"
+    publish_res = publish_inference_model(
+        AK=args.AK,
+        SK=args.SK,
+        model_name=args.model_name,
+        model_version=args.model_version,
+        model_description=args.model_description,
+        model_baseline=res["baseline"],
+        tos_url=res["tos_url"],
+        ring_id=ring_id,
+        key_id=key_id,
+        encrypt_flag=args.encrypt_flag,
+        silent=args.silent,
+        service=args.service
+    )
+    print("Step{}: Published model to JeddakPCC MMS with response = {}".format(1+step, publish_res))
+
+    engine_conf_res = get_engine_conf(
+        AK=args.AK,
+        SK=args.SK,
+        cu_type=args.cu_type,
+        inference_engine=args.inference_engine,
+        model_id=publish_res["Result"]["ModelID"],
+        silent=args.silent,
+        service=args.service
+    )
+    engine_conf = engine_conf_res["Result"]
+
+    # Step3: Deploy model to JeddakPCC
+    print("Step{}: Deploying model to JeddakPCC".format(2+step))
+    deploy_inference_model(
+        AK=args.AK,
+        SK=args.SK,
+        model_id=publish_res["Result"]["ModelID"],
+        cu_type=args.cu_type,
+        replica_count=args.replica_count,
+        inference_engine=args.inference_engine,
+        engine_conf=engine_conf,
+        silent=args.silent,
+        service=args.service
+    )
+    print("Step{}: Deployed model to JeddakPCC".format(2+step))
+
+    # Step4: Query model deploy status in JeddakPCC using list_inference_model
+    print("Step{}: Querying model deploy status in JeddakPCC".format(3+step))
+    deploy_status = "wait"
+    while True:
+        if deploy_status == "running":
+            break
+        list_res = list_inference_model(
+            AK=args.AK,
+            SK=args.SK,
+            silent=args.silent,
+            service=args.service
+        )
+        model_list = list_res["Result"]["ModelList"]
+        if publish_res["Result"]["ModelID"] not in [model["ModelID"] for model in model_list]:
+            print("Step{}.1: Model not found in JeddakPCC".format(3+step))
+            raise Exception("Model not found in JeddakPCC, checkout MMS to see if the model is published or deleted")
+            break
+        for model in model_list:
+            if model["ModelID"] == publish_res["Result"]["ModelID"]:
+                print("    Step{}.1: Queried model deploy status in JeddakPCC with status={}".format(3+step, model["DeployStatus"]))
+                if model["DeployStatus"] == "running":
+                    deploy_status = "running"
+                    break
+        time.sleep(10)
+
+    print("Step{}: model deployment succeed".format(3+step))
+
+    # Step5: Test model in JeddakPCC
+    print("Step{}: Testing model in JeddakPCC".format(4+step))
+    test_res = test_inference_model(
+        AK=args.AK,
+        SK=args.SK,
+        model_id=publish_res["Result"]["ModelID"],
+        silent=args.silent,
+        service=args.service
+    )
+    print("Step{}: model test succeed with resp = {}".format(4+step, test_res))
+    print("Step{}: Finished".format(5+step))