arkruntime 0.1.0__py3-none-any.whl

arkruntime/_utils/__init__.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from .._compat import (
+    parse_date as parse_date,
+)
+from .._compat import (
+    parse_datetime as parse_datetime,
+)
+from ._logs import setup_logging as setup_logging
+from ._streams import (
+    consume_async_iterator as consume_async_iterator,
+)
+from ._streams import (
+    consume_sync_iterator as consume_sync_iterator,
+)
+from ._transform import (
+    PropertyInfo as PropertyInfo,
+)
+from ._transform import (
+    async_maybe_transform as async_maybe_transform,
+)
+from ._transform import (
+    async_transform as async_transform,
+)
+from ._transform import (
+    maybe_transform as maybe_transform,
+)
+from ._transform import (
+    transform as transform,
+)
+from ._typing import (
+    extract_type_arg as extract_type_arg,
+)
+from ._typing import (
+    extract_type_var_from_base as extract_type_var_from_base,
+)
+from ._typing import (
+    is_annotated_type as is_annotated_type,
+)
+from ._typing import (
+    is_iterable_type as is_iterable_type,
+)
+from ._typing import (
+    is_list_type as is_list_type,
+)
+from ._typing import (
+    is_required_type as is_required_type,
+)
+from ._typing import (
+    is_type_alias_type as is_type_alias_type,
+)
+from ._typing import (
+    is_union_type as is_union_type,
+)
+from ._typing import (
+    strip_annotated_type as strip_annotated_type,
+)
+from ._utils import (
+    async_with_sts_token,
+    with_sts_token,
+)
+from ._utils import (
+    coerce_boolean as coerce_boolean,
+)
+from ._utils import (
+    coerce_float as coerce_float,
+)
+from ._utils import (
+    coerce_integer as coerce_integer,
+)
+from ._utils import (
+    deepcopy_minimal as deepcopy_minimal,
+)
+from ._utils import (
+    extract_files as extract_files,
+)
+from ._utils import (
+    file_from_path as file_from_path,
+)
+from ._utils import (
+    flatten as flatten,
+)
+from ._utils import (
+    get_async_library as get_async_library,
+)
+from ._utils import (
+    is_dict as is_dict,
+)
+from ._utils import (
+    is_given as is_given,
+)
+from ._utils import (
+    is_iterable as is_iterable,
+)
+from ._utils import (
+    is_list as is_list,
+)
+from ._utils import (
+    is_mapping as is_mapping,
+)
+from ._utils import (
+    is_mapping_t as is_mapping_t,
+)
+from ._utils import (
+    is_sequence as is_sequence,
+)
+from ._utils import (
+    is_sequence_t as is_sequence_t,
+)
+from ._utils import (
+    is_tuple as is_tuple,
+)
+from ._utils import (
+    is_tuple_t as is_tuple_t,
+)
+from ._utils import (
+    json_safe as json_safe,
+)
+from ._utils import (
+    lru_cache as lru_cache,
+)
+from ._utils import (
+    maybe_coerce_boolean as maybe_coerce_boolean,
+)
+from ._utils import (
+    maybe_coerce_float as maybe_coerce_float,
+)
+from ._utils import (
+    maybe_coerce_integer as maybe_coerce_integer,
+)
+from ._utils import (
+    removeprefix as removeprefix,
+)
+from ._utils import (
+    removesuffix as removesuffix,
+)
+from ._utils import (
+    required_args as required_args,
+)
+from ._utils import (
+    strip_not_given as strip_not_given,
+)

arkruntime/_utils/_key_agreement.py

@@ -0,0 +1,217 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from __future__ import annotations
+
+import base64
+import re
+import time
+from typing import Tuple
+
+
+def get_cert_info(cert_pem: str) -> Tuple[str, str, float]:
+    from cryptography import x509
+    from cryptography.hazmat.backends import default_backend
+
+    cert = x509.load_pem_x509_certificate(cert_pem.encode(), default_backend())
+    try:
+        dns = cert.extensions.get_extension_for_class(x509.SubjectAlternativeName).value.get_values_for_type(
+            x509.DNSName
+        )
+        if dns and len(dns) > 1 and re.match(r"^ring\..*$", dns[0]) and re.match(r"^key\..*$", dns[1]):
+            return dns[0][5:], dns[1][4:], cert.not_valid_after_utc.timestamp()
+    except Exception:
+        pass
+    return "", "", cert.not_valid_after_utc.timestamp()
+
+
+def aes_gcm_encrypt_bytes(key: bytes, iv: bytes, plain_bytes: bytes, associated_data: bytes = b"") -> bytes:
+    # aes_gcm_encrypt_bytes encrypt message using AES-GCM
+    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+    encryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv),
+    ).encryptor()
+    # associated_data will be authenticated but not encrypted,
+    # it must also be passed in on decryption.
+    encryptor.authenticate_additional_data(associated_data)
+    # Encrypt the plaintext and get the associated ciphertext.
+    # GCM does not require padding.
+    ciphertext = encryptor.update(plain_bytes) + encryptor.finalize()
+    return ciphertext + encryptor.tag
+
+
+def aes_gcm_encrypt_base64_string(key: bytes, nonce: bytes, plaintext: str) -> str:
+    """aes_gcm_encrypt_base64_string Encrypt message from base64 string to string using AES-GCM"""
+    plain_bytes = plaintext.encode()
+    # Encrypt message to string using AES-GCM
+    c = aes_gcm_encrypt_bytes(key, nonce, plain_bytes)
+    return base64.b64encode(c).decode()
+
+
+def aes_gcm_decrypt_bytes(key: bytes, iv: bytes, cipher_bytes: bytes, associated_data: bytes = b"") -> bytes:
+    """aes_gcm_decrypt_bytes Decrypt message from bytes to bytes using AES-GCM"""
+    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+    tag_length = 16  # default aes gcm tag length
+    cipher = cipher_bytes[:-tag_length]
+    tag = cipher_bytes[-tag_length:]
+    # Construct a Cipher object, with the key, iv, and additionally the
+    # GCM tag used for authenticating the message.
+    decryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv, tag),
+    ).decryptor()
+    # We put associated_data back in or the tag will fail to verify
+    # when we finalize the decryptor.
+    decryptor.authenticate_additional_data(associated_data)
+    # Decryption gets us the authenticated plaintext.
+    # If the tag does not match an InvalidTag exception will be raised.
+    return decryptor.update(cipher) + decryptor.finalize()
+
+
+def aes_gcm_decrypt_base64_string(key: bytes, nonce: bytes, ciphertext: str) -> str:
+    # Decrypt message(base64.std.string) using AES-GCM
+    cipher_bytes = base64.decodebytes(ciphertext.encode())
+    return aes_gcm_decrypt_bytes(key, nonce, cipher_bytes).decode()
+
+
+base64_pattern = r"(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=|[A-Za-z0-9+/]{4})"
+
+
+def decrypt_corner_case(key: bytes, nonce: bytes, data: str) -> str:
+    """decrypt_corner_case Decrypt corner case data"""
+    if len(data) < 24:
+        return ""
+    for i in range(20, len(data), 4):
+        try:
+            decrypted = aes_gcm_decrypt_base64_string(key, nonce, data[: i + 4])
+            if i + 4 == len(data):
+                return decrypted
+            return decrypted + decrypt_corner_case(key, nonce, data[i + 4 :])
+        except Exception:
+            pass
+
+
+def aes_gcm_decrypt_base64_list(key: bytes, nonce: bytes, ciphertext: str) -> str:
+    # Decrypt
+    base64_array = re.findall(base64_pattern, ciphertext)
+    result = []
+    for b64 in base64_array:
+        try:
+            result.append(aes_gcm_decrypt_base64_string(key, nonce, b64))
+        except Exception:
+            result.append(decrypt_corner_case(key, nonce, b64))
+    return "".join(result)
+
+
+def decrypt_validate(ciphertext: str) -> bool:
+    cipher_bytes = ciphertext.encode()
+    cipher_b64_bytes = base64.decodebytes(cipher_bytes)
+    return len(cipher_bytes) / 4 >= len(cipher_b64_bytes) / 3 >= len(cipher_bytes) / 4 - 1
+
+
+def marshal_cryptography_pub_key(key) -> bytes:
+    # python version of crypto/elliptic/elliptic.go Marshal
+    # without point on curve check
+    return bytes([4]) + key.x.to_bytes(32, "big") + key.y.to_bytes(32, "big")
+
+
+class key_agreement_client:
+    def __init__(self, certificate_pem_string: str) -> None:
+        """Load cert and extract public key"""
+        __fixed_version__ = "42.0.0"  # version check
+        from cryptography import __version__
+
+        if __version__ < __fixed_version__:
+            raise Exception(
+                "The cryptography package of Ark SDK only supports versions after {}, "
+                'please install the cryptography package by using pip install "cryptography>={}"'.format(
+                    __fixed_version__, __fixed_version__
+                )
+            )
+        from cryptography import x509
+        from cryptography.hazmat.primitives.asymmetric import ec
+
+        pem_data = certificate_pem_string.encode()
+        self._cert = x509.load_pem_x509_certificate(pem_data)
+        cert_pub = self._cert.public_key().public_numbers()
+        self._curve = ec._CURVE_TYPES[self._cert.public_key().curve.name]
+        self._public_key = ec.EllipticCurvePublicNumbers(cert_pub.x, cert_pub.y, self._curve).public_key()
+        self._not_valid_after_utc = self._cert.not_valid_after_utc.timestamp()
+        self._reload_time = time.time() + 60 * 60 * 12 * 14  # 14 days
+        self.init_cert_ring_key_id()
+
+    def encrypt_string(self, plaintext: str) -> Tuple[bytes, bytes, str, str]:
+        """encrypt_string encrypt plaintext with ECIES DH protocol"""
+        key, nonce, token = self.generate_ecies_key_pair()
+        # Encrypt message using AES-GCM
+        ciphertext = aes_gcm_encrypt_base64_string(key, nonce, plaintext)
+        return key, nonce, token, ciphertext
+
+    def encrypt_string_with_key(self, key: bytes, nonce: bytes, plaintext: str) -> str:
+        """encrypt_string_with_key encrypt plaintext with ECIES DH protocol"""
+        # Encrypt message using AES-GCM
+        ciphertext = aes_gcm_encrypt_base64_string(key, nonce, plaintext)
+        return ciphertext
+
+    def decrypt_string_with_key(self, key: bytes, nonce: bytes, ciphertext: str) -> str:
+        """decrypt_string_with_key decrypt ciphertext with ECIES DH protocol"""
+        # Decrypt message using AES-GCM
+        return aes_gcm_decrypt_base64_string(key, nonce, ciphertext)
+
+    def generate_ecies_key_pair(self) -> Tuple[bytes, bytes, str]:
+        """generate_ecies_key_pair generate ECIES key pair"""
+        from cryptography.hazmat.primitives import hashes
+        from cryptography.hazmat.primitives.asymmetric import ec
+        from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+
+        # Generate an ephemeral elliptic curve scalar and point
+        peer_private_key = ec.generate_private_key(self._curve)
+        dh = peer_private_key.exchange(ec.ECDH(), self._public_key)
+        R = peer_private_key.public_key().public_numbers()
+
+        # Derive symmetric key and nonce via HKDF
+        length = 32 + 12
+        buf = HKDF(
+            algorithm=hashes.SHA256(),
+            length=length,
+            salt=None,
+            info=None,
+        ).derive(dh)
+        key = buf[:32]
+        nonce = buf[32:length]
+
+        token = marshal_cryptography_pub_key(R)
+        return key, nonce, base64.b64encode(token).decode()
+
+    def need_reload(self) -> bool:
+        """need_reload check if the cert need reload"""
+        return time.time() > self._reload_time
+
+    def init_cert_ring_key_id(self) -> None:
+        """init_cert_ring_key_id init ring id and key id from cert"""
+        try:
+            from cryptography import x509
+
+            dns = self._cert.extensions.get_extension_for_class(x509.SubjectAlternativeName).value.get_values_for_type(
+                x509.DNSName
+            )
+            if dns and len(dns) > 1 and re.match(r"^ring\..*$", dns[0]) and re.match(r"^key\..*$", dns[1]):
+                self._ring_id = dns[0][5:]
+                self._key_id = dns[1][4:]
+            else:
+                self._ring_id = ""
+                self._key_id = ""
+        except Exception:
+            pass
+
+    def get_cert_ring_key_id(self) -> Tuple[str, str]:
+        """get_cert_ring_key_id get ring id and key id from cert"""
+        return self._ring_id, self._key_id
+
+    def get_cert_expiration_time(self) -> float:
+        """get_cert_expiration_time get cert expiration time"""
+        return self._not_valid_after_utc

arkruntime/_utils/_logs.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+import logging
+import os
+
+logger: logging.Logger = logging.getLogger("ark")
+httpx_logger: logging.Logger = logging.getLogger("httpx")
+
+
+def _basic_config() -> None:
+    logging.basicConfig(
+        format="%(asctime)s %(levelname)s %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+
+def setup_logging() -> None:
+    env = os.environ.get("ARK_LOG")
+    if env == "debug":
+        _basic_config()
+        logger.setLevel(logging.DEBUG)
+        httpx_logger.setLevel(logging.DEBUG)
+    elif env == "info":
+        _basic_config()
+        logger.setLevel(logging.INFO)
+        httpx_logger.setLevel(logging.INFO)

arkruntime/_utils/_model_breaker.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+import asyncio
+import threading
+import time
+import uuid
+
+
+class _QuerySet(object):
+    def __init__(self):
+        self._items = list()
+        self._index = dict()
+        self._lock = threading.Lock()
+
+    def add(self, item: int) -> None:
+        with self._lock:
+            if item in self._index:
+                return
+
+            self._items.append(item)
+            self._index[item] = len(self._items) - 1
+
+    def remove(self, item: int) -> None:
+        with self._lock:
+            if item not in self._index:
+                return
+
+            index = self._index[item]
+            self._items[index] = self._items[-1]
+            self._index[self._items[-1]] = index
+            self._items.pop()
+            del self._index[item]
+
+    def query(self, item: int) -> int:
+        with self._lock:
+            return self._index[item]
+
+
+class ModelBreaker(object):
+    def __init__(self):
+        # 初始化 allow_time 为当前时间
+        self._allow_time = time.perf_counter()
+        self._waiters = _QuerySet()
+
+    def _allow(self, id: int) -> bool:
+        cur = time.perf_counter()
+        # 如果当前时间小于等于 allow_time，不允许通过
+        if cur <= self._allow_time:
+            return 0
+        # 如果当前时间与 allow_time 的差值大于 10，允许通过
+        if cur - self._allow_time > 10:
+            return True
+        # 如果当前时间与 allow_time 的差值小于等于 10，慢启动通过
+        return self._waiters.query(id) < 2 ** (cur - self._allow_time)
+
+    def _get_allowed_duration(self) -> float:
+        # 计算当前时间与 allow_time 之间的持续时间
+        allow_duration = self._allow_time - time.perf_counter()
+
+        # 至少有 1 秒的等待时间
+        return max(allow_duration, 1)
+
+    def _acquire(self) -> int:
+        id = uuid.uuid4().int
+        self._waiters.add(id)
+        return id
+
+    def _release(self, id: int) -> None:
+        self._waiters.remove(id)
+
+    def reset(self, duration: float) -> None:
+        # 将 allow_time 重置为当前时间加上指定的持续时间
+        self._allow_time = time.perf_counter() + duration
+
+    def wait(self) -> None:
+        id = self._acquire()
+        while not self._allow(id):
+            time.sleep(self._get_allowed_duration())
+        self._release(id)
+
+    async def asyncwait(self) -> None:
+        id = self._acquire()
+        while not self._allow(id):
+            await asyncio.sleep(self._get_allowed_duration())
+        self._release(id)

arkruntime/_utils/_streams.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+# License text: https://github.com/volcengine/ark-runtime-python/blob/main/LICENSE
+
+from typing import Any
+
+from typing_extensions import AsyncIterator, Iterator
+
+
+def consume_sync_iterator(iterator: Iterator[Any]) -> None:
+    for _ in iterator:
+        ...
+
+
+async def consume_async_iterator(iterator: AsyncIterator[Any]) -> None:
+    async for _ in iterator:
+        ...