modalDeploy/src/cluster/video_apps/cache_submit.py

import modal

from ..video import downloader_image, app, config, batch_update_cloudflare_kv, modal_kv_cache

with downloader_image.imports():
    import os, httpx, crcmod
    import sentry_sdk
    from tqdm import tqdm
    from typing import Tuple
    from loguru import logger
    from datetime import datetime, UTC, timedelta
    from modal import current_function_call_id

    from tencentcloud.common.credential import Credential
    from tencentcloud.vod.v20180717.vod_client import VodClient
    from tencentcloud.vod.v20180717 import models as vod_request_models

    from BowongModalFunctions.models.media_model import MediaSource, MediaCacheStatus, MediaProtocol
    from BowongModalFunctions.models.web_model import SentryTransactionInfo


    @app.function(cpu=1, timeout=1800,
                  cloud="aws",
                  max_containers=config.video_downloader_concurrency,
                  volumes={
                      "/mntS3": modal.CloudBucketMount(
                          bucket_name=config.S3_bucket_name,
                          secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.modal_environment),
                      ),
                  },
                  secrets=[modal.Secret.from_name("tencent-cloud-secret", environment_name=config.modal_environment)])
    @modal.concurrent(max_inputs=10)
    async def cache_submit(media: MediaSource, sentry_trace: SentryTransactionInfo) -> MediaSource:
        def vod_init():
            tencent_secret_id = os.environ["VOD_SECRET_ID"]
            tencent_secret_key = os.environ["VOD_SECRET_KEY"]
            cred = Credential(secret_id=tencent_secret_id, secret_key=tencent_secret_key)
            return VodClient(credential=cred, region='ap-shanghai')

        def vod_info(media: MediaSource) -> Tuple[str, str, str]:
            logger.info(f"Downloading media {media}")
            request = vod_request_models.DescribeMediaInfosRequest()
            request.SubAppId = int(media.bucket)
            # 兼容fileId带文件类型的格式和不带文件类型的格式
            request.FileIds = [media.path.split('.')[0] if '.' in media.path else media.path]
            response = vod_client.DescribeMediaInfos(request)
            if len(response.MediaInfoSet) > 0:
                media_info = response.MediaInfoSet[0].BasicInfo
                logger.info(f"VOD info = {media_info}")
                file_extension = media_info.Type
                cache_dir = f"{config.S3_mount_dir}/{media.protocol.value}/{media.endpoint}/{media.bucket}"
                cache_file = media.path if '.' in media.path else f"{media.path}.{file_extension}"
                return cache_dir, cache_file, media_info.MediaUrl
            else:
                raise FileNotFoundError(
                    f"FileId : {media.path} not found in SubAppId: {media.bucket} at {media.endpoint}")

        def vod_download(media: MediaSource, on_progress_update: callable(float) = None) -> str:
            cache_dir, cache_file, url = vod_info(media)
            local_cache_filepath = os.path.join(cache_dir, cache_file)
            download_large_file(url=url, output_path=local_cache_filepath,
                                on_progress_callback=on_progress_update)
            return local_cache_filepath

        def download_large_file(url: str, output_path: str, protocol: MediaProtocol = MediaProtocol.vod,
                                on_progress_callback: callable(float) = None) -> None:
            # 配置日志
            logger.info(f"Starting download from {url}")
            try:
                # 使用 httpx 发送 HEAD 请求获取文件大小
                # 设置请求头，支持断点续传
                headers = {'Range': 'bytes=0-'}
                with httpx.Client() as client:
                    match protocol:
                        case MediaProtocol.vod:
                            head_response = client.head(url)
                            file_size = int(head_response.headers.get('content-length', 0))
                            remote_crc64 = int(head_response.headers.get('X-Cos-Hash-Crc64ecma', 0))
                            logger.info(f"File size: {file_size / (1024 * 1024 * 1024):.2f} GB")

                            if os.path.exists(output_path):
                                local_file_size = os.path.getsize(output_path)
                                logger.info(f"File size match check {local_file_size} = {file_size}")
                                if local_file_size == file_size:
                                    logger.info(f"Check file CRC64...")
                                    # CRC64使用ECMA-182标准校验 ref: https://cloud.tencent.com/document/product/436/40334#python-sdk
                                    c64 = crcmod.mkCrcFun(0x142F0E1EBA9EA3693, initCrc=0, xorOut=0xffffffffffffffff,
                                                          rev=True)
                                    with open(output_path, "rb") as local_file:
                                        local_crc64 = c64(local_file.read())
                                    logger.info(f"File crc64 check {local_crc64} = {remote_crc64}")
                                    if local_crc64 == remote_crc64:
                                        logger.success("File size verification passed!")
                                    return
                            os.makedirs(os.path.dirname(output_path), exist_ok=True)

                    logger.info(f"Downloading {url}...")
                    # 发起流式请求
                    with client.stream('GET', url, headers=headers) as response:
                        response.raise_for_status()
                        file_size = int(response.headers.get('content-length', 0))
                        # 设置进度条
                        progress_bar = tqdm(
                            total=file_size,
                            unit='iB',
                            unit_scale=True,
                            desc='Downloading'
                        )
                        # 以二进制写模式打开文件
                        with open(output_path, 'wb') as file:
                            # 分块下载，每次读取1MB
                            chunk_size = 1024 * 1024  # 1MB
                            downloaded_size = 0
                            for chunk in response.iter_bytes(chunk_size=chunk_size):
                                if chunk:
                                    file.write(chunk)
                                    downloaded_size += len(chunk)
                                    if on_progress_callback:
                                        on_progress_callback(downloaded_size / file_size)
                                    progress_bar.update(len(chunk))
                                    # 每下载100MB记录一次日志
                                    if downloaded_size % (100 * 1024 * 1024) == 0:
                                        logger.info(
                                            f"Downloaded: {downloaded_size / (1024 * 1024 * 1024):.2f} GB")

                        progress_bar.close()

                        # 验证下载是否完成
                        if os.path.exists(output_path):
                            final_size = os.path.getsize(output_path)
                        logger.info(f"Download completed successfully!")
                        logger.info(f"Final file size: {final_size / (1024 * 1024 * 1024):.2f} GB")
                        logger.info(f"File saved to: {os.path.abspath(output_path)}")
                        # 验证文件大小是否匹配
                        if final_size == file_size:
                            logger.info("File size verification passed!")
                        else:
                            logger.warning(f"File size mismatch! Expected: {file_size}, Got: {final_size}")
            except httpx.RequestError as e:
                logger.error(f"An error occurred while requesting {url}: {str(e)}")
                raise e
            except httpx.HTTPStatusError as e:
                logger.error(f"HTTP error occurred: {str(e)}")
                raise e
            except Exception as e:
                logger.error(f"An unexpected error occurred: {str(e)}")
                raise e
            finally:
                if 'progress_bar' in locals():
                    progress_bar.close()

        vod_client = vod_init()
        modal_kv = modal_kv_cache
        fn_id = current_function_call_id()

        with sentry_sdk.continue_trace(environ_or_headers={"sentry-trace": sentry_trace.x_trace_id,
                                                           "baggage": sentry_trace.x_baggage, }) as transaction:
            transaction.set_context("runtime_environment", {
                "MODAL_CLOUD_PROVIDER": os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown'),
                "MODAL_ENVIRONMENT": config.modal_environment,
                "MODAL_IMAGE_ID": os.environ.get('MODAL_IMAGE_ID', 'unknown'),
                "MODAL_IS_REMOTE": os.environ.get('MODAL_IS_REMOTE', 'unknown'),
                "MODAL_REGION": os.environ.get('MODAL_REGION', 'unknown'),
                "MODAL_TASK_ID": os.environ.get('MODAL_TASK_ID', 'unknown'),
                "MODAL_IDENTITY_TOKEN": os.environ.get('MODAL_IDENTITY_TOKEN', 'unknown'),
            })
            with transaction.start_child(name="收到缓存视频任务", op="queue.receive") as receive_span:
                receive_span.set_data("messaging.message.id", fn_id)
                receive_span.set_data("messaging.destination.name", "video-downloader.cache_submit")
                receive_span.set_data("messaging.message.retry.count", 0)
                receive_span.set_data("cache.key", media.urn)
                with receive_span.start_child(name="处理缓存视频任务", op="queue.process") as process_span:
                    process_span.set_data("messaging.message.id", fn_id)
                    process_span.set_data("messaging.destination.name", "video-downloader.cache_submit")
                    process_span.set_data("messaging.message.retry.count", 0)
                    process_span.set_data("cache.key", media.urn)
                    volume_cache_path = None
                    match media.protocol:
                        case MediaProtocol.vod:
                            try:
                                volume_cache_path = vod_download(media)
                                process_span.set_status("success")
                            except Exception as e:
                                logger.exception(e)
                                media.status = MediaCacheStatus.failed
                                modal_kv.set_cache(media)
                                batch_update_cloudflare_kv([media])
                                process_span.set_status("failed")
                                raise e
                        case MediaProtocol.http:
                            try:
                                cache_filepath = f"{config.S3_mount_dir}/{media.protocol.value}/{media.cache_filepath}"
                                os.makedirs(os.path.dirname(cache_filepath), exist_ok=True)
                                download_large_file(url=media.__str__(), output_path=cache_filepath)
                            except Exception as e:
                                logger.exception(e)
                                media.status = MediaCacheStatus.failed
                                modal_kv.set_cache(media)
                                batch_update_cloudflare_kv([media])
                                process_span.set_status("failed")
                                raise e
                        case MediaProtocol.s3:
                            # 本地挂载缓存
                            if media.protocol == MediaProtocol.s3 and media.endpoint == config.S3_region and media.bucket == config.S3_bucket_name:
                                volume_cache_path = f"{config.S3_mount_dir}/{media.cache_filepath}"
                            else:
                                logger.error("protocol not yet supported")
                        case _:
                            process_span.set_status("failed")
                            raise NotImplementedError("protocol not yet supported")
                    media.downloader_id = fn_id
                    media.status = MediaCacheStatus.ready if volume_cache_path else MediaCacheStatus.failed
                    media.progress = 1 if volume_cache_path else 0
                    media.expired_at = datetime.now(UTC) + timedelta(days=7) if volume_cache_path else None
                    modal_kv.set_cache(media)
                    batch_update_cloudflare_kv([media])
        return media