更新2个ffmpeg worker方法和web调用接口

src/cluster/ffmpeg_worker/worker.py

@@ -12,10 +12,10 @@ ffmpeg_worker_image = (
 app = modal.App(image=ffmpeg_worker_image, )
 
 with ffmpeg_worker_image.imports():
-    from src.cluster.config import config
     from typing import List, Optional, Tuple, Dict, Any
     from loguru import logger
     from ffmpeg.asyncio import FFmpeg
+    from src.cluster.config import config
     from src.cluster.ffmpeg_worker.Utils.PathUtils import FileUtils
     from src.cluster.video_downloader.model import (MediaSources,
                                                     MediaSource,
@@ -25,8 +25,6 @@ with ffmpeg_worker_image.imports():
     from src.cluster.web.model import SentryTransactionInfo
     from src.cluster.ffmpeg_worker.model import FFMpegSliceSegment
     import shutil, psutil, os, backoff, json, sentry_sdk
-    from psutil import svmem
-    from psutil._common import scpufreq
     from modal import current_function_call_id
 
     sentry_sdk.init(
@@ -44,6 +42,7 @@ with ffmpeg_worker_image.imports():
     s3_mount = "/mntS3"
     output_path_prefix = "/mnt/outputs"
 
+
     @sentry_sdk.trace
     def get_cache_filepath(media: MediaSource) -> MediaCache:
         """
@@ -52,7 +51,7 @@ with ffmpeg_worker_image.imports():
         :return:
         """
         # 本地挂载缓存
-        if media.protocol == MediaProtocol.s3 and media.endpoint == config.s3_region and media.bucket == config.bucket:
+        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)
@@ -63,7 +62,7 @@ with ffmpeg_worker_image.imports():
 
 
     @sentry_sdk.trace
-    def capture_hardware_info() -> Tuple[scpufreq, int, svmem]:
+    def capture_hardware_info() -> Tuple[Any, int, Any]:
         cpu_freq = psutil.cpu_freq()
         cpu_count = psutil.cpu_count()
         mem = psutil.virtual_memory()
@@ -71,73 +70,6 @@ with ffmpeg_worker_image.imports():
         return (cpu_freq, cpu_count, mem)
 
 
-    @sentry_sdk.trace
-    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)
-        outputs: List[str] = []
-        ffmpeg_cmd = FFmpeg().option('y').option('hide_banner').input(cache.cache_filepath)
-        filter_complex: List[str] = []
-        for index, marker in enumerate(media_markers):
-            segment_out_index = f"cut{index}"
-            filter_complex.append(
-                f"[0]trim=start={marker.start.total_seconds()}:end={marker.end.total_seconds()},setpts=PTS-STARTPTS[{segment_out_index}]")
-        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))
-            output_filepath = f"{output_path_prefix}/slice/outputs/{fn_id}/{filename}"
-            output_filepath = output_filepath.replace('//', '/')
-            workdir = os.path.dirname(output_filepath)
-            os.makedirs(workdir, exist_ok=True)
-            outputs.append(output_filepath)
-            ffmpeg_cmd.output(output_filepath,
-                              map=f"[cut{i}]",
-                              reset_timestamps="1",
-                              sc_threshold="0",
-                              g="1",
-                              force_key_frames="expr:gte(t,n_forced*1)",
-                              vcodec="libx264",
-                              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
-
-
     @backoff.on_exception(backoff.constant, exception=Exception, max_tries=5, max_time=60, raise_on_giveup=True)
     def local_copy_to_s3(local_outputs: List[str]) -> List[str]:
         s3_outputs = []
@@ -151,9 +83,10 @@ with ffmpeg_worker_image.imports():
                 out_s3.replace(s3_mount + '/', f"s3://{config.s3_region}/{config.s3_bucket_name}/"))
         return s3_outputs
 
+
     @app.function(
-        name="concat_medias",
+        cpu=12, timeout=1800,
-        cpu=12, timeout=1800, memory=(2048, 6144),
+        memory=6144,
         cloud="aws",
         # region='ap-northeast',
         max_containers=config.ffmpeg_worker_concurrency,
@@ -198,47 +131,53 @@ with ffmpeg_worker_image.imports():
                 logger.info(cache.cache_filepath)
                 ffmpeg_cmd.input(cache.cache_filepath)
                 input_videos.append(cache.cache_filepath)
+            target_width = 1080
+            target_height = 1920
+            # video_info = await print_media_streams_info(media_paths=input_videos)
 
-            video_info = print_media_streams_info(media_paths=input_videos)
             filter_complex = []
+            total_videos = len(input_videos)
 
-            for i in range(len(input_videos)):
+            # 2. 统一所有视频的格式、分辨率和帧率
-                info = video_info[i]
+            for i in range(total_videos):
-                streams = info.get("streams", [])
-                video_stream = streams[0]
-                # audio_stream = streams[1]
-                video_duration = video_stream.get("duration")
-                # audio_duration = audio_stream.get("duration")
-                '''
                 filter_complex.extend(
                     [
                         # 先缩放到统一分辨率，然后设置帧率和格式
                         f"[{i}:v]scale={target_width}:{target_height}:force_original_aspect_ratio=decrease,"
                         f"pad={target_width}:{target_height}:(ow-iw)/2:(oh-ih)/2,"
+                        f"setsar=1:1,"  # 新增强制设置SAR
                         f"fps=30,format=yuv420p[v{i}]",
                         f"[{i}:a]aformat=sample_fmts=fltp:sample_rates=44100:channel_layouts=stereo[a{i}]",
                     ]
                 )
-                '''
-                filter_complex.extend(
-                    [
-                        f"[{i}:v]trim=0:{video_duration},setpts=PTS-STARTPTS[v{i}]",
-                        f"[{i}:a]atrim=0:{video_duration},asetpts=PTS-STARTPTS[a{i}]",
-                    ]
-                )
 
             # 3. 准备处理后的视频流和音频流的连接字符串
-            video_streams = "".join(f"[v{i}]" for i in range(len(input_videos)))
+            video_streams = "".join(f"[v{i}]" for i in range(total_videos))
-            audio_streams = "".join(f"[a{i}]" for i in range(len(input_videos)))
+            audio_streams = "".join(f"[a{i}]" for i in range(total_videos))
 
             # 4. 使用concat过滤器合并视频和音频
             filter_complex.extend(
                 [
-                    f"{video_streams}concat=n={len(input_videos)}:v=1:a=0[vconcated]",
+                    f"{video_streams}concat=n={total_videos}:v=1:a=0[vconcated]",
-                    f"{audio_streams}concat=n={len(input_videos)}:v=0:a=1[aconcated]",
+                    f"{audio_streams}concat=n={total_videos}:v=0:a=1[aconcated]",
                 ]
             )
 
+            ffmpeg_cmd.output(
+                output_filepath,
+                {
+                    "filter_complex": ";".join(filter_complex),
+                    "map": ["[vconcated]", "[aconcated]"],
+                    "vcodec": "libx264",
+                    "crf": 16,
+                    "r": 30,
+                    "acodec": "aac",
+                    "ar": 44100,
+                    "ac": 2,
+                    "ab": "192k",
+                },
+            )
+
             @ffmpeg_cmd.on("start")
             def on_start(arguments: list[str]):
                 try:
@@ -268,21 +207,6 @@ with ffmpeg_worker_image.imports():
                 else:
                     logger.warning(line)
 
-            ffmpeg_cmd.output(
-                output_filepath,
-                {
-                    "filter_complex": ";".join(filter_complex),
-                    "map": ["[vconcated]", "[aconcated]"],
-                    "vcodec": "libx264",
-                    "crf": 16,
-                    "r": 30,
-                    "acodec": "aac",
-                    "ar": 44100,
-                    "ac": 2,
-                    "ab": "192k",
-                },
-            )
-
             await ffmpeg_cmd.execute()
 
         fn_id = current_function_call_id()
@@ -304,11 +228,13 @@ with ffmpeg_worker_image.imports():
 
         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])
 
@@ -316,8 +242,8 @@ with ffmpeg_worker_image.imports():
 
 
     @app.function(
-        name="slice_media",
+        cpu=12, timeout=900,
-        cpu=12, timeout=900, memory=(2048, 4096),
+        memory=6144,
         cloud="aws",
         max_containers=config.ffmpeg_worker_concurrency,
         volumes={
@@ -334,6 +260,73 @@ with ffmpeg_worker_image.imports():
                                  sentry_trace: Optional[SentryTransactionInfo] = None) -> Tuple[
         List[str], Optional[SentryTransactionInfo]]:
 
+        @sentry_sdk.trace
+        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)
+            outputs: List[str] = []
+            ffmpeg_cmd = FFmpeg().option('y').option('hide_banner').input(cache.cache_filepath)
+            filter_complex: List[str] = []
+            for index, marker in enumerate(media_markers):
+                segment_out_index = f"cut{index}"
+                filter_complex.append(
+                    f"[0]trim=start={marker.start.total_seconds()}:end={marker.end.total_seconds()},setpts=PTS-STARTPTS[{segment_out_index}]")
+            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))
+                output_filepath = f"{output_path_prefix}/slice/outputs/{fn_id}/{filename}"
+                output_filepath = output_filepath.replace('//', '/')
+                workdir = os.path.dirname(output_filepath)
+                os.makedirs(workdir, exist_ok=True)
+                outputs.append(output_filepath)
+                ffmpeg_cmd.output(output_filepath,
+                                  map=f"[cut{i}]",
+                                  reset_timestamps="1",
+                                  sc_threshold="0",
+                                  g="1",
+                                  force_key_frames="expr:gte(t,n_forced*1)",
+                                  vcodec="libx264",
+                                  acodec="aac",
+                                  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
@@ -350,6 +343,7 @@ with ffmpeg_worker_image.imports():
                                                  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")
@@ -357,3 +351,105 @@ with ffmpeg_worker_image.imports():
                 outputs = await ffmpeg_slice_process(media_source=media, media_markers=markers, fn_id=fn_id)
 
         return outputs, sentry_trace
+
+
+    @app.local_entrypoint()
+    async def test():
+        from ffmpeg import FFmpeg
+
+        async def ffmpeg_process(media_sources: List[str], output_filepath: str) -> None:
+            ffmpeg_cmd = FFmpeg().option('y').option('hide_banner')
+
+            input_videos = []
+            for media_source in media_sources:
+                # cache = get_cache_filepath(media_source)
+                ffmpeg_cmd.input(media_source)
+                input_videos.append(media_source)
+
+            target_width = 1080
+            target_height = 1920
+
+            # video_info = await print_media_streams_info(media_paths=input_videos)
+
+            filter_complex = []
+            total_videos = len(input_videos)
+
+            # 2. 统一所有视频的格式、分辨率和帧率
+            for i in range(total_videos):
+                filter_complex.extend(
+                    [
+                        # 先缩放到统一分辨率，然后设置帧率和格式
+                        f"[{i}:v]scale={target_width}:{target_height}:force_original_aspect_ratio=decrease,"
+                        f"pad={target_width}:{target_height}:(ow-iw)/2:(oh-ih)/2,"
+                        f"setsar=1:1,"  # 新增强制设置SAR
+                        f"fps=30,format=yuv420p[v{i}]",
+                        f"[{i}:a]aformat=sample_fmts=fltp:sample_rates=44100:channel_layouts=stereo[a{i}]",
+                    ]
+                )
+
+            # 3. 准备处理后的视频流和音频流的连接字符串
+            video_streams = "".join(f"[v{i}]" for i in range(total_videos))
+            audio_streams = "".join(f"[a{i}]" for i in range(total_videos))
+
+            # 4. 使用concat过滤器合并视频和音频
+            filter_complex.extend(
+                [
+                    f"{video_streams}concat=n={total_videos}:v=1:a=0[vconcated]",
+                    f"{audio_streams}concat=n={total_videos}:v=0:a=1[aconcated]",
+                ]
+            )
+
+            ffmpeg_cmd.output(
+                output_filepath,
+                {
+                    "filter_complex": ";".join(filter_complex),
+                    "map": ["[vconcated]", "[aconcated]"],
+                    "vcodec": "libx264",
+                    "crf": 16,
+                    "r": 30,
+                    "acodec": "aac",
+                    "ar": 44100,
+                    "ac": 2,
+                    "ab": "192k",
+                },
+            )
+
+            @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)
+
+            ffmpeg_cmd.execute()
+
+        medias = [
+            "./videos/1397757910405340824_0.mp4",
+            "./videos/1397757910405340824_1.mp4",
+            "./videos/1397757910405340824_2.mp4",
+            "./videos/1397757910405340824_3.mp4",
+            "./videos/1397757910405340824_4.mp4"
+        ]
+        await ffmpeg_process(medias, "./videos/output.mp4")

src/cluster/web/worker.py

@@ -35,6 +35,7 @@ with fastapi_image.imports():
 
     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),
@@ -153,8 +154,6 @@ with fastapi_image.imports():
             :param task_id:
             :return: (TaskStaus, errorReason, results, sentryTransactionInfo)
             """
-
-            try:
             fn_task = modal.FunctionCall.from_id(task_id)
             call_graph = fn_task.get_call_graph()
             task = call_graph[0]
@@ -170,8 +169,8 @@ with fastapi_image.imports():
                     except TimeoutError:
                         return TaskStatus.running, None, None, None
                     except Exception as e:
-                            logger.error(e)
+                        logger.exception(e)
-                            return TaskStatus.failed, None, None, None
+                        raise e
                 case _:
                     error = 'FAILURE'
                     match task.status:
@@ -182,9 +181,6 @@ with fastapi_image.imports():
                         case InputStatus.TIMEOUT:
                             error = 'TIMEOUT'
                     return TaskStatus.failed, error, None, None
-            except Exception as e:
-                logger.exception(e)
-                raise RuntimeError("Modal task status check failed")
 
         @web_app.get("/scalar", include_in_schema=False)
         async def scalar():
@@ -360,7 +356,7 @@ with fastapi_image.imports():
             if transaction:
                 response.headers["x-trace-id"] = transaction.x_trace_id
                 response.headers["x-baggage"] = transaction.x_baggage
-            return FFMPEGSliceTaskStatusResponse(taskId=task_id, status=TaskStatus.failed, error=reason)
+            return FFMPEGSliceTaskStatusResponse(taskId=task_id, status=task_status, error=reason, result=results)
 
         @web_app.post("/ffmpeg/concat",
                       tags=["发起任务"],