diff --git a/src/cluster/config.py b/src/cluster/config.py index 4efe7bf..6859381 100644 --- a/src/cluster/config.py +++ b/src/cluster/config.py @@ -11,6 +11,7 @@ class WorkerConfig(BaseSettings): modal_kv_name: Optional[str] = Field(default='media-cache', description="Modal视频缓存KV库") s3_region: Optional[str] = Field(default='ap-northeast-2', description="S3挂载桶的地域") s3_bucket_name: Optional[str] = Field(default='modal-media-cache', description="集群挂载的S3存储桶") + s3_mount_dir: str = Field(default='/mntS3', description="集群S3存储桶挂载在本地的根目录") cdn_endpoint: Optional[str] = Field(default="https://d2nj71io21vkj2.cloudfront.net", description="集群挂载S3存储桶的对应AWS Cloudfront CDN") environment: Optional[str] = Field(default="dev", description="Modal worker运行环境") @@ -19,6 +20,7 @@ class WorkerConfig(BaseSettings): config = WorkerConfig( + app_name="bowong-ai-video", video_downloader_concurrency=10, ffmpeg_worker_concurrency=10, modal_kv_name="media-cache", diff --git a/src/cluster/ffmpeg_worker/Utils/SentryUtils.py b/src/cluster/ffmpeg_worker/Utils/SentryUtils.py index 987b50d..5256de3 100644 --- a/src/cluster/ffmpeg_worker/Utils/SentryUtils.py +++ b/src/cluster/ffmpeg_worker/Utils/SentryUtils.py @@ -37,3 +37,4 @@ class SentryUtils: return wrapper return decorator + diff --git a/src/cluster/ffmpeg_worker/Utils/VideoUtils.py b/src/cluster/ffmpeg_worker/Utils/VideoUtils.py index ac1685b..7df0350 100644 --- a/src/cluster/ffmpeg_worker/Utils/VideoUtils.py +++ b/src/cluster/ffmpeg_worker/Utils/VideoUtils.py @@ -1,7 +1,26 @@ -from typing import Union, List, Tuple +from typing import Union, List, Tuple, Optional + +import numpy as np from pydantic import BaseModel, ConfigDict, computed_field from ffmpeg import FFmpeg from ffmpeg.asyncio import FFmpeg as AsyncFFmpeg +import soundfile as sf +import pyloudnorm as pyln +import noisereduce as nr +from pedalboard import ( + Pedalboard, + Compressor, + Limiter, + HighpassFilter, + LowpassFilter, + Gain, + Reverb, + Chorus, + Distortion, +) +from pedalboard.io import AudioFile + +from src.cluster.ffmpeg_worker.Utils.PathUtils import FileUtils class MediaStream(BaseModel): @@ -84,3 +103,79 @@ class VideoUtils: ) video_metadata = VideoMetadata.model_validate_json(await ffprobe.execute()) return video_metadata.streams[0].width, video_metadata.streams[0].height + + @staticmethod + def noise_reduce(media_path: str, noise_sample_path: Optional[str] = None) -> str: + samplerate = 44100 + with AudioFile(media_path).resampled_to(float(samplerate)) as f: + audio = f.read(f.frames) + + if noise_sample_path: + with AudioFile(noise_sample_path).resampled_to(float(samplerate)) as f: + noise_sample = f.read(f.frames) + else: + # 获取前2秒作为噪声样本 + noise_sample_length = min(int(2 * samplerate), audio.shape[0]) + noise_sample = audio[:noise_sample_length] + + reduced_noise = nr.reduce_noise(y=audio, y_noise=noise_sample, sr=samplerate, + stationary=True, prop_decrease=0.75, n_std_thresh_stationary=1.5, + n_fft=2048, win_length=1024, hop_length=512, n_jobs=1) + + board = Pedalboard( + [ + HighpassFilter(cutoff_frequency_hz=150), + LowpassFilter(cutoff_frequency_hz=8000), + Reverb(room_size=0.08, damping=0.7, wet_level=0.08, + dry_level=0.92, width=0.4), + Chorus(rate_hz=0.7, depth=0.12, centre_delay_ms=3.0, mix=0.10), + Distortion(drive_db=3.0), + Compressor(threshold_db=-30, ratio=1.8, attack_ms=20, release_ms=200), + Compressor(threshold_db=-24, ratio=2.2, attack_ms=15, release_ms=180), + Compressor(threshold_db=-18, ratio=1.5, attack_ms=10, release_ms=150), + Gain(gain_db=4), + Limiter(threshold_db=-6, release_ms=200), + ] + ) + # Convert to float32 if not already + reduced_noise = reduced_noise.astype(np.float32) + # Ensure audio is in the correct range (-1.0 to 1.0) + if np.abs(reduced_noise).max() > 1.0: + reduced_noise = reduced_noise / np.abs(reduced_noise).max() + + processed_audio = board(reduced_noise, samplerate) + # 格式处理 + if len(processed_audio.shape) == 1: + processed_audio = processed_audio.reshape(-1, 1) + elif len(processed_audio.shape) == 2: + if processed_audio.shape[0] < processed_audio.shape[1]: + processed_audio = processed_audio.T + if processed_audio.shape[1] > 2: + processed_audio = processed_audio[:, :2] + # 响度标准化 + meter = pyln.Meter(samplerate) + min_samples = int(0.4 * samplerate) + + if processed_audio.shape[0] < min_samples: + normalized_audio = processed_audio + else: + loudness = meter.integrated_loudness(processed_audio) + safety_factor = 0.7 + processed_audio = processed_audio * safety_factor + normalized_audio = pyln.normalize.loudness( + processed_audio, loudness, -16.0 + ) + + max_peak = np.max(np.abs(normalized_audio)) + if max_peak > 0.85: + additional_safety_factor = 0.85 / max_peak + normalized_audio = normalized_audio * additional_safety_factor + processed_audio_path = FileUtils.file_path_extend(media_path, "nr") + sf.write( + processed_audio_path, + normalized_audio, + samplerate, + format="WAV", + subtype="PCM_16", + ) + return processed_audio_path diff --git a/src/cluster/ffmpeg_worker/worker.py b/src/cluster/ffmpeg_worker/worker.py index b4352e2..01bef9a 100644 --- a/src/cluster/ffmpeg_worker/worker.py +++ b/src/cluster/ffmpeg_worker/worker.py @@ -1,4 +1,5 @@ import modal +from modal import current_function_call_id ffmpeg_worker_image = ( modal.Image.debian_slim(python_version="3.11") @@ -7,9 +8,7 @@ ffmpeg_worker_image = ( 'pedalboard==0.9.2', 'soundfile', 'pyloudnorm', 'noisereduce') .add_local_python_source("src.cluster.config", copy=True) .add_local_python_source("src.cluster.ffmpeg_worker", copy=True) - # .add_local_python_source("ffmpeg_worker", copy=True) .add_local_python_source("src.cluster.video_downloader.model", copy=True) - # .env({"PYTHONPATH": "/root/src"}) ) app = modal.App(image=ffmpeg_worker_image, ) @@ -25,27 +24,9 @@ with ffmpeg_worker_image.imports(): from src.cluster.ffmpeg_worker.model import FFMpegSliceSegment from src.cluster.video_downloader.model import (MediaSources, MediaSource, - MediaCache, - MediaProtocol, - MediaCacheStatus) + MediaProtocol) from src.cluster.web.model import SentryTransactionInfo import shutil, psutil, os, backoff, json, sentry_sdk - from modal import current_function_call_id - import soundfile as sf - import pyloudnorm as pyln - import noisereduce as nr - from pedalboard import ( - Pedalboard, - Compressor, - Limiter, - HighpassFilter, - LowpassFilter, - Gain, - Reverb, - Chorus, - Distortion, - ) - from pedalboard.io import AudioFile sentry_sdk.init( dsn="https://75cca970bfcc3d45d24361e7c0f1833c@sentry.bowongai.com/4", @@ -59,28 +40,10 @@ with ffmpeg_worker_image.imports(): modal_kv = modal.Dict.from_name(config.modal_kv_name, environment_name=config.environment, create_if_missing=True) - s3_mount = "/mntS3" + s3_mount = config.s3_mount_dir output_path_prefix = "/mnt/outputs" - @sentry_sdk.trace - def get_cache_filepath(media: MediaSource) -> MediaCache: - """ - 获取缓存路径 - :param media: 媒体文件 - :return: - """ - # 本地挂载缓存 - if media.protocol == MediaProtocol.s3 and media.endpoint == config.s3_region and media.bucket == config.s3_bucket_name: - cache_data = MediaCache(status=MediaCacheStatus.ready, cache_filepath=f"{s3_mount}/{media.get_cdn_url()}") - else: - cache_data_json = modal_kv.get(media.cache_key) - if not cache_data_json: - raise FileNotFoundError(f"{media.cache_key} cache not found") - cache_data = MediaCache.model_validate_json(cache_data_json) - return cache_data - - @sentry_sdk.trace def capture_hardware_info() -> Tuple[Any, int, Any]: cpu_freq = psutil.cpu_freq() @@ -145,7 +108,7 @@ with ffmpeg_worker_image.imports(): # region='ap-northeast', max_containers=config.ffmpeg_worker_concurrency, volumes={ - s3_mount: modal.CloudBucketMount( + f"/mntS3": modal.CloudBucketMount( bucket_name=config.s3_bucket_name, secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.environment), @@ -156,6 +119,8 @@ with ffmpeg_worker_image.imports(): sentry_trace: Optional[SentryTransactionInfo] = None) -> Tuple[ str, Optional[SentryTransactionInfo]]: + fn_id = current_function_call_id() + @sentry_sdk.trace async def print_media_streams_info(media_paths: List[str]) -> List[Dict[str, Any]]: """ @@ -175,16 +140,17 @@ with ffmpeg_worker_image.imports(): continue return res - @sentry_sdk.trace - async def ffmpeg_process(media_sources: MediaSources, output_filepath: str) -> None: - ffmpeg_cmd = FFmpeg().option('y').option('hide_banner') - + @SentryUtils.sentry_tracker(name="视频合并任务", op="ffmpeg.concat", fn_id=fn_id, + sentry_trace_id=sentry_trace.x_trace_id if sentry_trace else None, + sentry_baggage=sentry_trace.x_baggage if sentry_trace else None) + async def ffmpeg_process(media_sources: MediaSources, output_filepath: str) -> str: + ffmpeg_cmd = ffmpeg_init() input_videos = [] for media_source in media_sources.inputs: - cache = get_cache_filepath(media_source) - logger.info(cache.cache_filepath) - ffmpeg_cmd.input(cache.cache_filepath) - input_videos.append(cache.cache_filepath) + cache_filepath = f"{s3_mount}/{media_source.cache_filepath}" + logger.info(cache_filepath) + ffmpeg_cmd.input(cache_filepath) + input_videos.append(cache_filepath) target_width = 1080 target_height = 1920 # video_info = await print_media_streams_info(media_paths=input_videos) @@ -232,67 +198,14 @@ with ffmpeg_worker_image.imports(): }, ) - @ffmpeg_cmd.on("start") - def on_start(arguments: list[str]): - try: - filter_index = arguments.index("-filter_complex") - filter_content = arguments[filter_index + 1] - arguments[filter_index + 1] = f'"{filter_content}"' - args = " ".join(arguments) - logger.info(f"FFmpeg command:{args}") - arguments[filter_index + 1] = filter_content - except ValueError: - args = " ".join(arguments) - logger.info(f"FFmpeg command:{args}") - - @ffmpeg_cmd.on("progress") - def on_progress(progress): - logger.info(f"处理进度: {progress}") - - @ffmpeg_cmd.on("completed") - def on_completed(): - logger.info(f"FFMpeg task completed.") - - @ffmpeg_cmd.on("stderr") - def on_stderr(line: str): - if line.startswith('Error'): - logger.error(line) - raise RuntimeError(line) - else: - logger.warning(line) - await ffmpeg_cmd.execute() + s3_outputs = local_copy_to_s3([output_filepath]) + return s3_outputs[0] - fn_id = current_function_call_id() - cpu_freq, cpu_count, mem = capture_hardware_info() - total_mb = mem.total >> 20 - available_mb = mem.available >> 20 - logger.info((f"[{fn_id}] Current CPUs {cpu_count} * {cpu_freq.current / 1000:.2f} GHz")) - logger.info(f"Current Memories {available_mb} Mb / {total_mb} Mb | {mem.percent}% used") - - if sentry_trace: - transaction = sentry_sdk.continue_trace(environ_or_headers={ - "sentry-trace": sentry_trace.x_trace_id, - "baggage": sentry_trace.x_baggage, - }) - else: - transaction = sentry_sdk.start_transaction(op='modal.function', name="Modal Function直接调用") - sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(), - x_baggage=sentry_sdk.get_baggage()) - - with transaction: - with transaction.start_child(op='ffmpeg.concat', name='视频合并任务') as span: - sentry_sdk.set_tag('fn.id', fn_id) - span.set_data('fn.id', fn_id) - span.set_data('cpu.count', cpu_count) - span.set_data('cpu.frequency', f"{cpu_freq.current / 1000:.2f} GHz") - span.set_data('memory.available', f"{total_mb} Mb") - output_path = f"{output_path_prefix}/concat/outputs/{fn_id}/output.mp4" - os.makedirs(os.path.dirname(output_path), exist_ok=True) - await ffmpeg_process(medias, output_filepath=output_path) - s3_outputs = local_copy_to_s3([output_path]) - - return s3_outputs[0], sentry_trace + output_path = f"{output_path_prefix}/concat/outputs/{fn_id}/output.mp4" + os.makedirs(os.path.dirname(output_path), exist_ok=True) + s3_output = await ffmpeg_process(medias, output_filepath=output_path) + return s3_output, sentry_trace @app.function( @@ -300,7 +213,7 @@ with ffmpeg_worker_image.imports(): cloud="aws", max_containers=config.ffmpeg_worker_concurrency, volumes={ - s3_mount: modal.CloudBucketMount( + f"/mntS3": modal.CloudBucketMount( bucket_name=config.s3_bucket_name, secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.environment), @@ -313,14 +226,18 @@ with ffmpeg_worker_image.imports(): webhook: Optional[str], # todo : handle webhook callback sentry_trace: Optional[SentryTransactionInfo] = None) -> Tuple[ List[str], Optional[SentryTransactionInfo]]: + fn_id = current_function_call_id() - @sentry_sdk.trace + @SentryUtils.sentry_tracker(name="视频切割任务", op="ffmpeg.slice", fn_id=fn_id, + sentry_trace_id=sentry_trace.x_trace_id if sentry_trace else None, + sentry_baggage=sentry_trace.x_baggage if sentry_trace else None) async def ffmpeg_slice_process(media_source: MediaSource, media_markers: List[FFMpegSliceSegment], fn_id: str) -> List[str]: - cache = get_cache_filepath(media_source) - logger.info(cache.cache_filepath) + cache_filepath = f"{s3_mount}/{media_source.cache_filepath}" + logger.info(cache_filepath) outputs: List[str] = [] - ffmpeg_cmd = FFmpeg().option('y').option('hide_banner').input(cache.cache_filepath) + ffmpeg_cmd = ffmpeg_init() + ffmpeg_cmd.input(cache_filepath) filter_complex: List[str] = [] for index, marker in enumerate(media_markers): filter_complex.extend( @@ -332,7 +249,7 @@ with ffmpeg_worker_image.imports(): ffmpeg_cmd.option('filter_complex', ';'.join(filter_complex)) for i, marker in enumerate(media_markers): - filename = FileUtils.file_path_extend(os.path.basename(cache.cache_filepath), str(i)) + filename = FileUtils.file_path_extend(os.path.basename(cache_filepath), str(i)) output_filepath = f"{output_path_prefix}/slice/outputs/{fn_id}/{filename}" output_filepath = output_filepath.replace('//', '/') workdir = os.path.dirname(output_filepath) @@ -349,52 +266,26 @@ with ffmpeg_worker_image.imports(): crf=16, r=30, ) - @ffmpeg_cmd.on("start") - def on_start(arguments: list[str]): - try: - filter_index = arguments.index("-filter_complex") - filter_content = arguments[filter_index + 1] - arguments[filter_index + 1] = f'"{filter_content}"' - args = " ".join(arguments) - logger.info(f"FFmpeg command:{args}") - arguments[filter_index + 1] = filter_content - except ValueError: - args = " ".join(arguments) - logger.info(f"FFmpeg command:{args}") - - @ffmpeg_cmd.on("progress") - def on_progress(progress): - logger.info(f"处理进度: {progress}") - - @ffmpeg_cmd.on("completed") - def on_completed(): - logger.info(f"FFMpeg task completed.") - - @ffmpeg_cmd.on("stderr") - def on_stderr(line: str): - if line.startswith('Error'): - logger.error(line) - raise RuntimeError(line) - else: - logger.warning(line) - await ffmpeg_cmd.execute() s3_outputs = local_copy_to_s3(outputs) return s3_outputs - @sentry_sdk.trace + @SentryUtils.sentry_tracker(name="视频切割任务", op="ffmpeg.slice", fn_id=fn_id, + sentry_trace_id=sentry_trace.x_trace_id if sentry_trace else None, + sentry_baggage=sentry_trace.x_baggage if sentry_trace else None) async def ffmpeg_hls_slice_process(media_source: MediaSource, media_markers: List[FFMpegSliceSegment], fn_id: str) -> List[str]: - hls_m3u8_url = media_source.url + hls_m3u8_url = media_source.path logger.info(hls_m3u8_url) outputs: List[str] = [] - ffmpeg_cmd = FFmpeg().option('y').option('hide_banner').input(hls_m3u8_url, - protocol_whitelist="file,http,https,tcp,tls", - reconnect="1", # 自动重连 - reconnect_streamed="1", - reconnect_delay_max="5") + ffmpeg_cmd = ffmpeg_init() + ffmpeg_cmd.input(hls_m3u8_url, + protocol_whitelist="file,http,https,tcp,tls", + reconnect="1", # 自动重连 + reconnect_streamed="1", + reconnect_delay_max="5") filter_complex: List[str] = [] for index, marker in enumerate(media_markers): filter_complex.extend( @@ -425,67 +316,17 @@ with ffmpeg_worker_image.imports(): crf=16, r=30, ) - @ffmpeg_cmd.on("start") - def on_start(arguments: list[str]): - try: - filter_index = arguments.index("-filter_complex") - filter_content = arguments[filter_index + 1] - arguments[filter_index + 1] = f'"{filter_content}"' - args = " ".join(arguments) - logger.info(f"FFmpeg command:{args}") - arguments[filter_index + 1] = filter_content - except ValueError: - args = " ".join(arguments) - logger.info(f"FFmpeg command:{args}") - - @ffmpeg_cmd.on("progress") - def on_progress(progress): - logger.info(f"处理进度: {progress}") - - @ffmpeg_cmd.on("completed") - def on_completed(): - logger.info(f"FFMpeg task completed.") - - @ffmpeg_cmd.on("stderr") - def on_stderr(line: str): - if line.startswith('Error'): - logger.error(line) - raise RuntimeError(line) - else: - logger.warning(line) - await ffmpeg_cmd.execute() s3_outputs = local_copy_to_s3(outputs) return s3_outputs - fn_id = current_function_call_id() - cpu_freq, cpu_count, mem = capture_hardware_info() - total_mb = mem.total >> 20 - available_mb = mem.available >> 20 - logger.info((f"[{fn_id}] Current CPUs {cpu_count} * {cpu_freq.current / 1000:.2f} GHz")) - logger.info(f"Current Memories {available_mb} Mb / {total_mb} Mb | {mem.percent}% used") - - if sentry_trace: - transaction = sentry_sdk.continue_trace(environ_or_headers={"sentry-trace": sentry_trace.x_trace_id, - "baggage": sentry_trace.x_baggage, }) - else: - transaction = sentry_sdk.start_transaction(op='modal.function', name="Modal Function直接调用") - sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(), - x_baggage=sentry_sdk.get_baggage()) - with transaction: - with transaction.start_child(op='ffmpeg.slice', name='视频切割任务') as span: - sentry_sdk.set_tag('fn.id', fn_id) - span.set_data('fn.id', fn_id) - span.set_data('cpu.count', cpu_count) - span.set_data('cpu.frequency', f"{cpu_freq.current / 1000:.2f} GHz") - span.set_data('memory.available', f"{total_mb} Mb") - match media.protocol: - case MediaProtocol.hls: - outputs = await ffmpeg_hls_slice_process(media_source=media, media_markers=markers, fn_id=fn_id) - case _: - outputs = await ffmpeg_slice_process(media_source=media, media_markers=markers, fn_id=fn_id) + match media.protocol: + case MediaProtocol.hls: + outputs = await ffmpeg_hls_slice_process(media_source=media, media_markers=markers, fn_id=fn_id) + case _: + outputs = await ffmpeg_slice_process(media_source=media, media_markers=markers, fn_id=fn_id) return outputs, sentry_trace @@ -495,7 +336,7 @@ with ffmpeg_worker_image.imports(): cloud="aws", max_containers=config.ffmpeg_worker_concurrency, volumes={ - s3_mount: modal.CloudBucketMount( + f"/mntS3": modal.CloudBucketMount( bucket_name=config.s3_bucket_name, secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.environment), @@ -509,16 +350,13 @@ with ffmpeg_worker_image.imports(): Optional[str], Optional[SentryTransactionInfo]]: fn_id = current_function_call_id() - cpu_freq, cpu_count, mem = capture_hardware_info() - total_mb = mem.total >> 20 - available_mb = mem.available >> 20 - logger.info((f"[{fn_id}] Current CPUs {cpu_count} * {cpu_freq.current / 1000:.2f} GHz")) - logger.info(f"Current Memories {available_mb} Mb / {total_mb} Mb | {mem.percent}% used") - @sentry_sdk.trace + @SentryUtils.sentry_tracker(name="视频切割任务", op="ffmpeg.extract.audio", fn_id=fn_id, + sentry_trace_id=sentry_trace.x_trace_id if sentry_trace else None, + sentry_baggage=sentry_trace.x_baggage if sentry_trace else None) async def ffmpeg_process(media: MediaSource, fn_id: str) -> str: - cache = get_cache_filepath(media) - ffprobe_cmd = FFmpeg('ffprobe').input(cache.cache_filepath, + cache_filepath = f"{s3_mount}/{media.cache_filepath}" + ffprobe_cmd = FFmpeg('ffprobe').input(cache_filepath, v="quiet", print_format="json", select_streams="a", @@ -561,36 +399,22 @@ with ffmpeg_worker_image.imports(): output_path = f"{output_path_prefix}/extract_audio/outputs/{fn_id}/output.wav" os.makedirs(os.path.dirname(output_path), exist_ok=True) ffmpeg_cmd = ffmpeg_init() - ffmpeg_cmd.input(cache.cache_filepath).output(output_path, - map="0:a", - acodec="pcm_s16le", - ar=44100, - ac=1) + ffmpeg_cmd.input(cache_filepath).output(output_path, + map="0:a", + acodec="pcm_s16le", + ar=44100, + ac=1) await ffmpeg_cmd.execute() s3_outputs = local_copy_to_s3([output_path]) return s3_outputs[0] - if sentry_trace: - transaction = sentry_sdk.continue_trace(environ_or_headers={"sentry-trace": sentry_trace.x_trace_id, - "baggage": sentry_trace.x_baggage, }) - else: - transaction = sentry_sdk.start_transaction(op='modal.function', name="Modal Function直接调用") - sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(), - x_baggage=sentry_sdk.get_baggage()) - with transaction: - with transaction.start_child(op='ffmpeg.extract.audio', name='视频切割任务') as span: - sentry_sdk.set_tag('fn.id', fn_id) - span.set_data('fn.id', fn_id) - span.set_data('cpu.count', cpu_count) - span.set_data('cpu.frequency', f"{cpu_freq.current / 1000:.2f} GHz") - span.set_data('memory.available', f"{total_mb} Mb") - match media_source.protocol: - case MediaProtocol.hls: - return None, sentry_trace - case _: - output = await ffmpeg_process(media_source, fn_id=fn_id) - return output, sentry_trace + match media_source.protocol: + case MediaProtocol.hls: + return None, sentry_trace + case _: + output = await ffmpeg_process(media_source, fn_id=fn_id) + return output, sentry_trace async def ffmpeg_overlay_gif(media: MediaSource, gif: MediaSource): @@ -602,7 +426,7 @@ with ffmpeg_worker_image.imports(): cloud="aws", max_containers=config.ffmpeg_worker_concurrency, volumes={ - s3_mount: modal.CloudBucketMount( + f"/mntS3": modal.CloudBucketMount( bucket_name=config.s3_bucket_name, secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.environment), @@ -622,8 +446,7 @@ with ffmpeg_worker_image.imports(): async def ffmpeg_process(media: MediaSource, func_id: str, mirror_scale_down_size: int = 6, mirror_from_right: bool = True, mirror_position: tuple[float, float] = (40, 40)) -> str: - cache = get_cache_filepath(media) - media_filepath = cache.cache_filepath + media_filepath = f"{s3_mount}/{media.cache_filepath}" video_metadata = VideoUtils.ffprobe_video_format(media_filepath) mirror_x = ( f"main_w-overlay_w-{mirror_position[0]}" @@ -669,7 +492,12 @@ with ffmpeg_worker_image.imports(): raise NotImplementedError - async def ffmpeg_bgm_nosie_reduce(media: MediaSource, bgm: MediaSource, ): + async def ffmpeg_bgm_nosie_reduce(media: MediaSource, bgm: MediaSource, noise_sample: Optional[MediaSource] = None): + fn_id = current_function_call_id() + + async def ffmpeg_process(video: MediaSource, bgm: MediaSource, noise_sample: Optional[MediaSource] = None): + raise NotImplementedError + raise NotImplementedError diff --git a/src/cluster/video_downloader/model.py b/src/cluster/video_downloader/model.py index c5296d8..60005c8 100644 --- a/src/cluster/video_downloader/model.py +++ b/src/cluster/video_downloader/model.py @@ -2,8 +2,11 @@ import os from datetime import datetime from enum import Enum from typing import List, Union, Optional, Any, Dict -from pydantic import BaseModel, Field, field_validator, ValidationError, field_serializer, SerializationInfo +from urllib.parse import urlparse +from pydantic import BaseModel, Field, field_validator, ValidationError, field_serializer, SerializationInfo, \ + computed_field from pydantic.json_schema import JsonSchemaValue +from src.cluster.config import config class MediaProtocol(str, Enum): @@ -20,57 +23,68 @@ class MediaCacheStatus(str, Enum): ready = "ready" deleted = "deleted" missing = "missing" + unknown = "unknown" class MediaSource(BaseModel): - url: str = Field() - protocol: MediaProtocol = Field() - endpoint: Optional[str] = Field() - bucket: Optional[str] = Field() - cache_key: Optional[str] = Field() + path: str = Field(description="媒体源的路径") + protocol: MediaProtocol = Field(description="媒体源的来源协议") + endpoint: Optional[str] = Field(description="媒体源来源的终端地址,根据使用协议的不同,会有没有endpoint的情况") + bucket: Optional[str] = Field(description="媒体源所使用的存储桶(s3)/SubAppId(vod)") + urn: Optional[str] = Field(description="媒体源的唯一指定标识") + status: MediaCacheStatus = Field(default=MediaCacheStatus.unknown, description="媒体源在Modal集群挂载的缓存状态") + expired_at: Optional[datetime] = Field(description="缓存过期时间点", default=None) + downloader_id: Optional[str] = Field(description="正在处理下载的Downloader ID", default=None) + progress: int = Field(description="缓存进度", default=0) + + # cache_filepath: Optional[str] = Field(description="s3挂载路径下的缓存相对路径", default=None) @classmethod - def from_str(cls, mediaUrl: str) -> 'MediaSource': - if mediaUrl.startswith('http://') or mediaUrl.startswith('https://'): - return MediaSource(url=mediaUrl, + def from_str(cls, media_url: str) -> 'MediaSource': + if media_url.startswith('http://') or media_url.startswith('https://'): + parsed_url = urlparse(media_url) + path_str = f"{parsed_url.path}?{parsed_url.query}" if parsed_url.query else parsed_url.path + return MediaSource(path=path_str, protocol=MediaProtocol.http, - endpoint=None, + endpoint=parsed_url.netloc, # domain of http url bucket=None, - cache_key=mediaUrl) - elif mediaUrl.startswith('s3://'): # s3://{endpoint}/{bucket}/{url} - paths = mediaUrl[5:].split('/') - return MediaSource(url='/'.join(paths[2:]), + urn=media_url) + elif media_url.startswith('s3://'): # s3://{endpoint}/{bucket}/{url} + paths = media_url[5:].split('/') + return MediaSource(path='/'.join(paths[2:]), protocol=MediaProtocol.s3, endpoint=paths[0], bucket=paths[1], - cache_key=mediaUrl) - elif mediaUrl.startswith('vod://'): # vod://{endpoint}/{subAppId}/{fileId} - paths = mediaUrl[6:].split('/') + urn=media_url) + elif media_url.startswith('vod://'): # vod://{endpoint}/{subAppId}/{fileId} + paths = media_url[6:].split('/') # 兼容有文件类型后缀和没有文件类型后缀的格式 url = paths[2] if '.' in os.path.basename(paths[2]) else paths[2] + ".mp4" - return MediaSource(url=url, + return MediaSource(path=url, protocol=MediaProtocol.vod, bucket=paths[1], endpoint=paths[0], - cache_key=mediaUrl) - elif mediaUrl.startswith('cos://'): # cos://{endpoint}/{bucket}/{url} - paths = mediaUrl[6:].split('/') - return MediaSource(url='/'.join(paths[2:]), + urn=media_url) + elif media_url.startswith('cos://'): # cos://{endpoint}/{bucket}/{url} + paths = media_url[6:].split('/') + return MediaSource(path='/'.join(paths[2:]), protocol=MediaProtocol.cos, endpoint=paths[0], bucket=paths[1], - cache_key=mediaUrl) - elif mediaUrl.startswith('hls://'): + urn=media_url) + elif media_url.startswith('hls://'): # hls://merge-local-1324682537.cos.ap-shanghai.myqcloud.com/streams/1264/30322.m3u8 - paths = mediaUrl[6:] - return MediaSource(url=f"https://{paths}", + paths = media_url[6:] + return MediaSource(path=f"https://{paths}", protocol=MediaProtocol.hls, endpoint=None, bucket=None, - cache_key=mediaUrl + urn=media_url ) else: - raise ValidationError("mediaUrl必须以http[s]、s3或vod协议开头") + available_schemas = [member.value for member in MediaProtocol] + available_schemas_str = ','.join(available_schemas) + raise ValidationError(f"mediaUrl必须以{available_schemas_str}协议开头") @classmethod def __get_pydantic_json_schema__(cls, core_schema: Any, handler: Any) -> JsonSchemaValue: @@ -85,16 +99,44 @@ class MediaSource(BaseModel): ] } + @computed_field(description="s3挂载路径下的缓存相对路径") + @property + def cache_filepath(self) -> str: + match self.protocol: + case MediaProtocol.s3: + # 本地挂载缓存 + if self.protocol == MediaProtocol.s3 and self.endpoint == config.s3_region and self.bucket == config.s3_bucket_name: + return f"{self.path}" + else: + return f"{self.protocol.value}/{self.endpoint}/{self.bucket}/{self.path}" + case MediaProtocol.http: + clean_path = self.path.split('?')[0] if '?' in self.path else self.path + return f"{self.protocol.value}/{self.endpoint}/{clean_path}" + case _: + return f"{self.protocol.value}/{self.endpoint}/{self.bucket}/{self.path}" + + @field_serializer('expired_at') + def serialize_datetime(self, value: Optional[datetime], info: SerializationInfo) -> Optional[str]: + if value: + return value.isoformat() + else: + return None + def get_cdn_url(self) -> str: if self.protocol == MediaProtocol.s3: - return f"{self.url}" + return f"{self.path}" elif self.protocol == MediaProtocol.http: - return f"{self.protocol.value}/{self.url}" - return f"{self.protocol.value}/{self.endpoint}/{self.bucket}/{self.url}" + return f"{self.protocol.value}/{self.path}" + return f"{self.protocol.value}/{self.endpoint}/{self.bucket}/{self.path}" - -class CachedMediaSource(MediaSource): - cached_s3_path: str = Field(description="s3上的相对缓存路径") + def __str__(self): + match self.protocol: + case MediaProtocol.http: + return f"{self.protocol.value}://{self.path}" + case MediaProtocol.hls: + return f"{self.protocol.value}://{self.path[8:]}" # strip "https://" from url + case _: + return f"{self.protocol.value}://{self.endpoint}/{self.bucket}/{self.path}" class MediaSources(BaseModel): @@ -126,23 +168,8 @@ class MediaSources(BaseModel): } -class MediaCache(BaseModel): - status: MediaCacheStatus = Field(description="缓存被处理的状态") - progress: float = Field(description="缓存下载进度0-1", default=0) - expired_at: Optional[datetime] = Field(description="缓存过期时间点", default=None) - downloader_id: Optional[str] = Field(description="正在处理下载的Downloader ID", default=None) - cache_filepath: Optional[str] = Field(description="缓存的文件地址", default=None) - - @field_serializer('expired_at') - def serialize_datetime(self, value: Optional[datetime], info: SerializationInfo) -> Optional[str]: - if value: - return value.isoformat() - else: - return None - - class CacheResult(BaseModel): - caches: Dict[str, MediaCache] = Field(description="Cache ID") + caches: Dict[str, MediaSource] = Field(description="Cache ID") class DownloadResult(BaseModel): diff --git a/src/cluster/video_downloader/worker.py b/src/cluster/video_downloader/worker.py index b9272f5..1b7c2b1 100644 --- a/src/cluster/video_downloader/worker.py +++ b/src/cluster/video_downloader/worker.py @@ -14,11 +14,12 @@ downloader_image = ( with downloader_image.imports(): import os, httpx, crcmod from tqdm import tqdm - from typing import Tuple, Dict + from typing import Tuple, Dict, List from loguru import logger from modal import current_function_call_id from src.cluster.config import config - from src.cluster.video_downloader.model import MediaSource, MediaCacheStatus, MediaCache, MediaProtocol + from src.cluster.web.Utils.KVCache import KVCache + from src.cluster.video_downloader.model import MediaSource, MediaCacheStatus, MediaProtocol from src.cluster.web.model import SentryTransactionInfo from datetime import datetime, UTC, timedelta from tencentcloud.common.credential import Credential @@ -45,9 +46,11 @@ with downloader_image.imports(): cf_kv_api_token = os.environ.get("CF_KV_API_TOKEN") cf_kv_namespace_id = os.environ.get("CF_KV_NAMESPACE_ID") + modal_kv_cache = KVCache(kv_name=config.modal_kv_name, environment=config.environment) + @sentry_sdk.trace - def batch_update_cloudflare_kv(caches: Dict[str, MediaCache]): + def batch_update_cloudflare_kv(caches: List[MediaSource]): with httpx.Client() as client: try: response = client.put( @@ -56,10 +59,10 @@ with downloader_image.imports(): json=[ { "based64": False, - "key": mediaKey, - "value": cache_data.model_dump_json(), + "key": cache.urn, + "value": cache.model_dump_json(), } - for (mediaKey, cache_data) in caches.items() + for cache in caches ] ) response.raise_for_status() @@ -75,13 +78,13 @@ with downloader_image.imports(): @sentry_sdk.trace - def batch_remove_cloudflare_kv(caches: Dict[str, MediaCache]): + def batch_remove_cloudflare_kv(caches: List[MediaSource]): with httpx.Client() as client: try: response = client.post( f"https://api.cloudflare.com/client/v4/accounts/{cf_account_id}/storage/kv/namespaces/{cf_kv_namespace_id}/bulk/delete", headers={"Authorization": f"Bearer {cf_kv_api_token}"}, - json=[mediaKey for (mediaKey, cache_data) in caches.items()] + json=[cache.urn for cache in caches] ) response.raise_for_status() except httpx.RequestError as e: @@ -107,7 +110,7 @@ with downloader_image.imports(): }, secrets=[modal.Secret.from_name("tencent-cloud-secret", environment_name=config.environment)]) @modal.concurrent(max_inputs=10) - async def cache_submit(media: MediaSource, sentry_trace: SentryTransactionInfo) -> MediaCache: + 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"] @@ -119,18 +122,18 @@ with downloader_image.imports(): request = vod_request_models.DescribeMediaInfosRequest() request.SubAppId = int(media.bucket) # 兼容fileId带文件类型的格式和不带文件类型的格式 - request.FileIds = [media.url.split('.')[0] if '.' in media.url else media.url] + 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"/mntS3/{media.protocol.value}/{media.endpoint}/{media.bucket}" - cache_file = f"{media.url}.{file_extension}" + 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.url} not found in SubAppId: {media.bucket} at {media.endpoint}") + 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) @@ -226,8 +229,7 @@ with downloader_image.imports(): progress_bar.close() vod_client = vod_init() - modal_kv = modal.Dict.from_name(config.modal_kv_name, environment_name=config.environment, - create_if_missing=True) + 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, @@ -245,56 +247,50 @@ with downloader_image.imports(): 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.cache_key) + 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.cache_key) + process_span.set_data("cache.key", media.urn) volume_cache_path = None - media_cache_downloading = MediaCache(status=MediaCacheStatus.downloading, - downloader_id=fn_id) match media.protocol: case MediaProtocol.vod: - # def on_progress_callback(progress: float): - # media_cache_downloading.progress = progress - # caches_data = {media.cache_key: media_cache_downloading} - # modal_kv.put(media.cache_key, media_cache_downloading.model_dump_json()) - # batch_update_cloudflare_kv(caches_data) - try: volume_cache_path = vod_download(media) process_span.set_status("success") except Exception as e: logger.exception(e) - media_cache_downloading.status = MediaCacheStatus.failed - caches = {f"{media.cache_key}": media_cache_downloading} - modal_kv.put(media.cache_key, media_cache_downloading) - batch_update_cloudflare_kv(caches) + media.status = MediaCacheStatus.failed + modal_kv.set_cache(media) + batch_update_cloudflare_kv([media]) process_span.set_status("failed") case MediaProtocol.http: try: - - download_large_file(url=url, output_path=local_cache_filepath) + cache_filepath = f"{config.s3_mount_dir}/{media.cache_filepath}" + download_large_file(url=media.__str__(), output_path=cache_filepath) except Exception as e: logger.exception(e) - media_cache_downloading.status = MediaCacheStatus.failed - caches = {f"{media.cache_key}": media_cache_downloading} - modal_kv.put(media.cache_key, media_cache_downloading) - batch_update_cloudflare_kv(caches) + media.status = MediaCacheStatus.failed + modal_kv.set_cache(media) + batch_update_cloudflare_kv([media]) process_span.set_status("failed") + 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_cache_ready = MediaCache( - downloader_id=fn_id, - cache_filepath=volume_cache_path, - status=MediaCacheStatus.ready if volume_cache_path else MediaCacheStatus.failed, - progress=1 if volume_cache_path else 0, - expired_at=datetime.now(UTC) + timedelta(days=7) if volume_cache_path else None, ) - modal_kv.put(media.cache_key, media_cache_ready.model_dump_json()) - batch_update_cloudflare_kv({media.cache_key: media_cache_ready}) - return media_cache_ready + logger.error(f"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 @worker_app.function(cpu=1, timeout=300, @@ -306,16 +302,7 @@ with downloader_image.imports(): ), }) @modal.concurrent(max_inputs=10) - async def cache_delete(cache: MediaCache) -> MediaCache: - sentry_sdk.add_breadcrumb({ - "MODAL_CLOUD_PROVIDER": os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown'), - "MODAL_ENVIRONMENT": os.environ.get('MODAL_ENVIRONMENT', 'unknown'), - "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'), - }) + async def cache_delete(cache: MediaSource) -> MediaSource: if os.path.exists(cache.cache_filepath): os.remove(cache.cache_filepath) cache.status = MediaCacheStatus.deleted diff --git a/src/cluster/web/Utils/KVCache.py b/src/cluster/web/Utils/KVCache.py new file mode 100644 index 0000000..43d6569 --- /dev/null +++ b/src/cluster/web/Utils/KVCache.py @@ -0,0 +1,33 @@ +from typing import Optional + +import modal + +from src.cluster.video_downloader.model import MediaSource + + +class KVCache: + kv: modal.Dict + + def __init__(self, kv_name: str, environment: str): + self.kv = modal.Dict.from_name(kv_name, environment_name=environment, create_if_missing=True) + + def get_cache(self, urn: str) -> Optional[MediaSource]: + cache_json = self.kv.get(urn) + if not cache_json: + return None + return MediaSource.model_validate_json(cache_json) + + def set_cache(self, media: MediaSource): + cache_json = media.model_dump_json() + self.kv.put(media.urn, cache_json) + + def clear(self): + self.kv.clear() + + def pop(self, urn: str, raise_exception: bool = True) -> Optional[MediaSource]: + cache_json = self.kv.pop(urn) + if not cache_json: + if raise_exception: + raise KeyError("URN错误,资源不存在") + return None + return MediaSource.model_validate_json(cache_json) \ No newline at end of file diff --git a/src/cluster/web/Utils/__init__.py b/src/cluster/web/Utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/cluster/web/model.py b/src/cluster/web/model.py index 2e31d86..8e9a055 100644 --- a/src/cluster/web/model.py +++ b/src/cluster/web/model.py @@ -124,5 +124,6 @@ class FFMPEGCornerMirrorRequest(BaseFFMPEGTaskRequest): else: raise TypeError(v) + class FFMPEGCornerMirrorTaskStatusResponse(BaseFFMPEGTaskStatusResponse): result: Optional[str] = Field(default=None, description="任务运行结果") diff --git a/src/cluster/web/worker.py b/src/cluster/web/worker.py index 897639c..b902ba6 100644 --- a/src/cluster/web/worker.py +++ b/src/cluster/web/worker.py @@ -1,64 +1,74 @@ import modal +from src.cluster.config import config fastapi_image = ( modal.Image .debian_slim(python_version="3.11") .pip_install("fastapi[standard]", "sentry-sdk[fastapi]", - 'loguru', 'pydantic', 'pydantic_settings', 'scalar-fastapi') - .add_local_python_source("src.cluster.video_downloader.model", copy=True) - .add_local_python_source("src.cluster.web.model", copy=True) + 'loguru', 'pydantic', 'pydantic_settings', 'scalar-fastapi', 'psutil') .add_local_python_source("src.cluster.config", copy=True) + .add_local_python_source("src.cluster.web", copy=True) + .add_local_python_source("src.cluster.video_downloader", copy=True) + .add_local_python_source("src.cluster.ffmpeg_worker", copy=True) ) -with fastapi_image.imports(): - import os - import httpx - from typing import Dict, List, Annotated, Tuple, Any, Optional - from modal import current_function_call_id - from modal.call_graph import InputStatus - from loguru import logger - from fastapi import FastAPI, Response, Depends, Header - from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials - from fastapi.responses import JSONResponse, RedirectResponse - from fastapi.exceptions import HTTPException - from starlette import status - from scalar_fastapi import get_scalar_api_reference - import sentry_sdk - from sentry_sdk.integrations.loguru import LoguruIntegration, LoggingLevels - from sentry_sdk.integrations.fastapi import FastApiIntegration - from src.cluster.config import config - from src.cluster.video_downloader.model import (MediaSource, MediaSources, MediaCacheStatus, MediaCache, - CacheResult, DownloadResult) - from src.cluster.web.model import (TaskStatus, ErrorCode, - SentryTransactionInfo, - SentryTransactionHeader, - ModalTaskResponse, - FFMPEGSliceRequest, - FFMPEGConcatRequest, - FFMPEGExtractAudioRequest, - FFMPEGCornerMirrorRequest, - FFMPEGSliceTaskStatusResponse, - FFMPEGConcatTaskStatusResponse, - FFMPEGExtractAudioTaskStatusResponse, - FFMPEGCornerMirrorTaskStatusResponse) - - fastapi_app = modal.App(image=fastapi_image) +fastapi_app = modal.App(image=fastapi_image, ) - @fastapi_app.function(scaledown_window=60, - secrets=[ - modal.Secret.from_name("cf-kv-secret", environment_name=config.environment), - ]) - @modal.concurrent(max_inputs=100) - @modal.asgi_app() - def fastapi_webapp(): +@fastapi_app.function(scaledown_window=60, + secrets=[ + modal.Secret.from_name("cf-kv-secret", environment_name=config.environment), + ]) +@modal.concurrent(max_inputs=100) +@modal.asgi_app(custom_domains=["modal-dev.bowong.cc"]) +def fastapi_webapp(): + with fastapi_image.imports(): + import os + import httpx + import asyncio + from typing import List, Annotated, Tuple, Any, Optional + from modal import current_function_call_id + from modal.call_graph import InputStatus + from loguru import logger + from fastapi import FastAPI, Response, Depends, Header + from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials + from fastapi.responses import JSONResponse, RedirectResponse + from fastapi.exceptions import HTTPException + from starlette import status + from scalar_fastapi import get_scalar_api_reference + import sentry_sdk + from sentry_sdk.integrations.loguru import LoguruIntegration, LoggingLevels + from sentry_sdk.integrations.fastapi import FastApiIntegration + from src.cluster.web.Utils.KVCache import KVCache + from src.cluster.ffmpeg_worker.Utils.SentryUtils import SentryUtils + from src.cluster.video_downloader.model import (MediaSource, + MediaSources, + MediaCacheStatus, + CacheResult, + DownloadResult) + from src.cluster.web.model import (TaskStatus, ErrorCode, + SentryTransactionInfo, + SentryTransactionHeader, + ModalTaskResponse, + FFMPEGSliceRequest, + FFMPEGConcatRequest, + FFMPEGExtractAudioRequest, + FFMPEGCornerMirrorRequest, + FFMPEGSliceTaskStatusResponse, + FFMPEGConcatTaskStatusResponse, + FFMPEGExtractAudioTaskStatusResponse, + FFMPEGCornerMirrorTaskStatusResponse) + bearer_scheme = HTTPBearer() web_app = FastAPI(title="Modal worker API", summary="Modal Worker的API, 包括缓存视频, 发起生产任务等", servers=[ - {'url': 'https://bowongai-dev--video-downloader-fastapi-webapp.modal.run', - 'description': 'modal dev环境测试服务'}]) + # {'url': 'https://bowongai-dev--video-downloader-fastapi-webapp.modal.run', + # 'description': 'modal dev环境测试服务'}, + {'url': 'https://modal-dev.bowong.cc', + 'description': 'modal dev环境测试服务'}, + ]) async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)): token = credentials.credentials @@ -87,8 +97,7 @@ with fastapi_image.imports(): FastApiIntegration() ] ) - media_cache_kv = modal.Dict.from_name('media-cache', environment_name=config.environment, - create_if_missing=True) + modal_kv_cache = KVCache(kv_name=config.modal_kv_name, environment=config.environment) cf_account_id = os.environ.get("CF_ACCOUNT_ID") cf_kv_api_token = os.environ.get("CF_KV_API_TOKEN") cf_kv_namespace_id = os.environ.get("CF_KV_NAMESPACE_ID") @@ -109,7 +118,7 @@ with fastapi_image.imports(): } @sentry_sdk.trace - def batch_update_cloudflare_kv(caches: Dict[str, MediaCache]): + def batch_update_cloudflare_kv(caches: List[MediaSource]): with httpx.Client() as client: try: response = client.put( @@ -118,10 +127,10 @@ with fastapi_image.imports(): json=[ { "based64": False, - "key": mediaKey, - "value": cache_data.model_dump_json(), + "key": cache.urn, + "value": cache.model_dump_json(), } - for (mediaKey, cache_data) in caches.items() + for cache in caches ] ) response.raise_for_status() @@ -215,62 +224,74 @@ with fastapi_image.imports(): @sentry_sdk.trace async def cache(medias: MediaSources) -> CacheResult: fn_id = current_function_call_id() - caches: Dict[str, MediaCache] = {} - parent = sentry_sdk.get_current_span() + caches: MediaSources sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(), x_baggage=sentry_sdk.get_baggage()) - for media in medias.inputs: - with parent.start_child(name="同步视频缓存", op="cache.get") as cache_span: - cache_span.set_data("runner_id", fn_id) - cache_span.set_data("cache.key", [media.cache_key]) - video_cache_status_json = media_cache_kv.get(media.cache_key) - video_cache: MediaCache - cache_hit: bool = False - if not video_cache_status_json: + @SentryUtils.sentry_tracker(name="同步视频缓存", op="cache.get", fn_id=fn_id, + sentry_trace_id=None, sentry_baggage=None) + async def cache_handler(media: MediaSource): + cache_span = sentry_sdk.get_current_span() + cache_span.set_data("runner_id", fn_id) + cache_span.set_data("cache.key", [media.urn]) + cached_media = modal_kv_cache.get_cache(media.urn) + cache_hit: bool = False + if not cached_media: + # start new download task + with cache_span.start_child(name="视频缓存任务入队", + op="queue.publish") as queue_publish_span: + fn = modal.Function.from_name(config.app_name, 'cache_submit') + fn_task = fn.spawn(media, sentry_trace) + queue_publish_span.set_data("cache.key", media.urn) + queue_publish_span.set_data("messaging.message.id", fn_task.object_id) + queue_publish_span.set_data("messaging.destination.name", + "video-downloader.cache_submit") + queue_publish_span.set_data("messaging.message.body.size", 0) + # video_cache = MediaCache(status=MediaCacheStatus.downloading, + # downloader_id=fn_task.object_id) + media.status = MediaCacheStatus.downloading + media.downloader_id = fn_task.object_id + # video_cache_status_json = video_cache.model_dump_json() + modal_kv_cache.set_cache(media) + else: + media = cached_media + # video_cache = MediaCache.model_validate_json(video_cache_status_json) + match media.status: + case MediaCacheStatus.ready: + cache_hit = True + case MediaCacheStatus.downloading: # 下载任务已经在进行 + cache_hit = True + case _: # start new download task with cache_span.start_child(name="视频缓存任务入队", op="queue.publish") as queue_publish_span: fn = modal.Function.from_name(config.app_name, 'cache_submit') fn_task = fn.spawn(media, sentry_trace) - queue_publish_span.set_data("cache.key", media.cache_key) + queue_publish_span.set_data("cache.key", media.urn) queue_publish_span.set_data("messaging.message.id", fn_task.object_id) queue_publish_span.set_data("messaging.destination.name", "video-downloader.cache_submit") queue_publish_span.set_data("messaging.message.body.size", 0) - video_cache = MediaCache(status=MediaCacheStatus.downloading, - downloader_id=fn_task.object_id) - video_cache_status_json = video_cache.model_dump_json() - media_cache_kv.put(media.cache_key, video_cache_status_json) + media.status = MediaCacheStatus.downloading + media.downloader_id = fn_task.object_id + # video_cache = MediaCache(status=MediaCacheStatus.downloading, + # downloader_id=fn_task.object_id) + # video_cache_status_json = video_cache.model_dump_json() + modal_kv_cache.set_cache(media) + cache_hit = False - video_cache = MediaCache.model_validate_json(video_cache_status_json) - match video_cache.status: - case MediaCacheStatus.ready: - cache_hit = True - case MediaCacheStatus.downloading: # 下载任务已经在进行 - cache_hit = True - case _: - # start new download task - with cache_span.start_child(name="视频缓存任务入队", - op="queue.publish") as queue_publish_span: - fn = modal.Function.from_name(config.app_name, 'cache_submit') - fn_task = fn.spawn(media, sentry_trace) - queue_publish_span.set_data("cache.key", media.cache_key) - queue_publish_span.set_data("messaging.message.id", fn_task.object_id) - queue_publish_span.set_data("messaging.destination.name", - "video-downloader.cache_submit") - queue_publish_span.set_data("messaging.message.body.size", 0) - video_cache = MediaCache(status=MediaCacheStatus.downloading, - downloader_id=fn_task.object_id) - video_cache_status_json = video_cache.model_dump_json() - media_cache_kv.put(media.cache_key, video_cache_status_json) + # caches[media.urn] = video_cache + # logger.info(f"Media cache hit ? {cache_hit}") + cache_span.set_data("cache.hit", cache_hit) + return media - caches[media.cache_key] = video_cache - logger.info(f"Media cache hit ? {cache_hit}") - cache_span.set_data("cache.hit", cache_hit) - batch_update_cloudflare_kv(caches) - return CacheResult(caches=caches) - # return JSONResponse(content={"caches": jsonable_encoder(caches)}) + async with asyncio.TaskGroup() as group: + tasks = [group.create_task(cache_handler(media)) for media in medias.inputs] + + cache_task_result = [task.result() for task in tasks] + + batch_update_cloudflare_kv(cache_task_result) + return CacheResult(caches={media.urn: media for media in cache_task_result}) @web_app.post("/cache/download", tags=["缓存"], @@ -303,7 +324,7 @@ with fastapi_image.imports(): async def purge_kv_all(): parent = sentry_sdk.get_current_span() span = parent.start_child(name="清除缓存KV", op="purge.flush") - media_cache_kv.clear() + modal_kv_cache.clear() span.set_data("cache.success", True) span.finish() return JSONResponse(content={"success": True}) @@ -316,8 +337,8 @@ with fastapi_image.imports(): async def purge_kv(medias: MediaSources): try: for media in medias.inputs: - media_cache_kv.pop(media.cache_key) - keys = [media.cache_key for media in medias.inputs] + modal_kv_cache.pop(media.urn) + keys = [media.urn for media in medias.inputs] batch_remove_cloudflare_kv(keys) return JSONResponse(content={"success": True, "keys": keys}) except Exception as e: @@ -329,20 +350,22 @@ with fastapi_image.imports(): description="清除指定的所有缓存(包括KV记录和S3存储文件)", dependencies=[Depends(verify_token)]) async def purge_media(medias: MediaSources): - caches: Dict[str, MediaCache] = {} - for media in medias.inputs: - try: - cache_data_json = media_cache_kv.pop(media.cache_key) - cache_data = MediaCache.model_validate_json(cache_data_json) - fn = modal.Function.from_name(config.app_name, "cache_delete", - environment_name=config.environment) - deleted_cache = await fn.remote.aio(cache_data) - caches[media.cache_key] = deleted_cache - except KeyError: - logger.warning("cache key not found") - caches[media.cache_key] = MediaCache(status=MediaCacheStatus.missing) - continue - keys = [key for key in caches.keys()] + fn_id = current_function_call_id() + fn = modal.Function.from_name(config.app_name, "cache_delete", environment_name=config.environment) + + @SentryUtils.sentry_tracker(name="清除媒体源缓存", op="cache.purge", fn_id=fn_id, + sentry_trace_id=None, sentry_baggage=None) + async def purge_handle(media: MediaSource): + cache_media = modal_kv_cache.pop(media.urn) + if cache_media: + deleted_cache: MediaSource = await fn.remote.aio(cache_media) + return deleted_cache.urn + return None + + async with asyncio.TaskGroup() as group: + tasks = [group.create_task(purge_handle(media)) for media in medias.inputs] + + keys = [task.result() for task in tasks] batch_remove_cloudflare_kv(keys) return JSONResponse(content={"success": True, "keys": keys}) diff --git a/src/deploy.py b/src/deploy.py index 75a9da1..aa255e5 100644 --- a/src/deploy.py +++ b/src/deploy.py @@ -2,4 +2,6 @@ import modal.cli.run from cluster.config import config if __name__ == '__main__': - modal.cli.run.deploy(app_ref='.\\src\\cluster\\app.py', env=config.environment, use_module_mode=False) + # modal.cli.run.deploy(app_ref='src.cluster.app', env=config.environment, use_module_mode=True) + # todo : 通过python命令部署不成功,换https://modal.com/docs/guide/continuous-deployment推荐的方式通过命令行部署 + modal.cli.run.deploy(app_ref='src\\cluster\\app.py', env=config.environment, use_module_mode=False)