diff --git a/pyproject.toml b/pyproject.toml index c80baab..f8a5e9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,3 +92,8 @@ packages = ["src/BowongModalFunctions", "src/Douyin_TikTok_Download_API"] sources = ["src"] only-packages = true require-runtime-dependencies = true + +[tool.uv.workspace] +members = [ + "cmd/wsl", +] diff --git a/src/BowongModalFunctions/utils/HTTPUtils.py b/src/BowongModalFunctions/utils/HTTPUtils.py new file mode 100644 index 0000000..4ea07a3 --- /dev/null +++ b/src/BowongModalFunctions/utils/HTTPUtils.py @@ -0,0 +1,100 @@ +from typing import Optional, Any, Union +from functools import wraps +import backoff +import httpx +import asyncio +from loguru import logger +import aiofiles +from pathlib import Path + + +class HTTPDownloadUtils: + # 创建一个类级别的 Semaphore 来控制并发 + _semaphore = asyncio.Semaphore(5) # 限制最大并发数为5 + + @staticmethod + @backoff.on_exception( + backoff.expo, + (httpx.RequestError, httpx.HTTPStatusError, Exception), + max_tries=10, + max_time=300, # 最大重试时间5分钟 + giveup=lambda e: isinstance(e, httpx.HTTPStatusError) and e.response.status_code == 404 + ) + async def async_download_file( + url: str, + output_path: Union[str, Path], + timeout: float = 60.0, + verify: bool = True + ) -> Union[str, None]: + """ + 异步下载文件或获取内容 + + :param url: 要下载的URL + :param output_path: 输出文件路径，如果为None则返回内容 + :param timeout: 请求超时时间（秒） + :param verify: 是否验证SSL证书 + + :return: 如果output_path为None，返回下载的内容；否则返回保存的文件路径 + + :exception httpx.RequestError: 请求错误 + :exception httpx.HTTPStatusError: HTTP状态错误 + :exception Exception: 其他错误 + """ + async with HTTPDownloadUtils._semaphore: # 使用信号量控制并发 + try: + logger.info(f"Starting download from {url}") + + async with httpx.AsyncClient( + timeout=timeout, + verify=verify, + follow_redirects=True + ) as client: + response = await client.get(url) + response.raise_for_status() # 检查HTTP状态码 + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + async with aiofiles.open(output_path, 'wb') as f: + await f.write(response.content) + logger.success(f"Successfully downloaded to {output_path}") + return str(output_path) + + + except httpx.HTTPStatusError as e: + logger.error(f"HTTP error occurred: {e.response.status_code} - {url}") + raise + except httpx.RequestError as e: + logger.error(f"Request error occurred: {str(e)} - {url}") + raise + except Exception as e: + logger.error(f"Unexpected error occurred: {str(e)} - {url}") + raise + + @staticmethod + async def batch_download( + urls: list[str], + output_paths: list[Union[str, Path]] = None, + max_concurrent: int = 5 + ) -> list[Union[str, None]]: + """ + 批量下载多个文件 + + :param urls: URL列表 + :param output_paths: 输出路径列表 + :param max_concurrent: 最大并发数 + + :return: 下载结果列表 + """ + # 更新信号量的值 + HTTPDownloadUtils._semaphore = asyncio.Semaphore(max_concurrent) + + tasks = [] + for i, url in enumerate(urls): + output_path = output_paths[i] + task = asyncio.create_task( + HTTPDownloadUtils.async_download_file(url, output_path) + ) + tasks.append(task) + + return await asyncio.gather(*tasks, return_exceptions=True) diff --git a/src/BowongModalFunctions/utils/PathUtils.py b/src/BowongModalFunctions/utils/PathUtils.py index f71bbfa..348f9d1 100644 --- a/src/BowongModalFunctions/utils/PathUtils.py +++ b/src/BowongModalFunctions/utils/PathUtils.py @@ -6,6 +6,13 @@ from typing import List class FileUtils: @staticmethod def file_path_extend(media_path: str, extend: str) -> str: + """ + 基于现有文件路径添加后缀, 例如 extend = "def" : 123/abc.txt -> 123/abc_def.txt + :param media_path: 现有文件路径 + :param extend: 后缀名 + + :return: 处理过的新文件路径 + """ media_filename = os.path.basename(media_path) media_dir = os.path.dirname(media_path) + '/' filenames = media_filename.split('.') @@ -14,11 +21,29 @@ class FileUtils: return os.path.join(media_dir, extend_filename) @staticmethod - def file_path_replace_root_prefix(media_path: str, prefix: str, depth: int = 1) -> str: + def replace_root_by_depth(media_path: str, root: str, depth: int = 1) -> str: + """ + 使用根目录名替换路径起始位置, 例如: + + prefix="pre",depth=0 : ./abc/def.txt -> pre\\.\\abc\\def.txt + + prefix="pre",depth=1 : ./abc/def.txt -> pre\\abc\\def.txt + + prefix="pre",depth=2 : ./abc/def.txt -> pre\\def.txt + + prefix="pre",depth=3 : IndexError("Depth is out of range") + + :param media_path: 现有文件路径 + :param root: 替换用的根目录名 + :param depth: root所占的路径深度层级 + + :return: 处理后的文件路径 + :exception IndexError: 根目录深度超过实际路径深度 + """ media_dirs = media_path.split('/') if depth >= len(media_dirs): raise IndexError("Depth is out of range") - media_dir_prefix = prefix + '/'.join(media_dirs[depth:]) + media_dir_prefix = os.path.join(root, *media_dirs[depth:]) return media_dir_prefix @staticmethod @@ -32,6 +57,14 @@ class FileUtils: @staticmethod def get_folder_size(folder_path: str) -> int: + """ + + Args: + folder_path: + + Returns: + + """ total_size = 0 for path in Path(folder_path).rglob('*'): if path.is_file(): @@ -48,3 +81,4 @@ class FileUtils: for file in files: total_size += os.path.getsize(file) return total_size + diff --git a/src/BowongModalFunctions/utils/VideoUtils.py b/src/BowongModalFunctions/utils/VideoUtils.py index 31de4ca..28b3fe3 100644 --- a/src/BowongModalFunctions/utils/VideoUtils.py +++ b/src/BowongModalFunctions/utils/VideoUtils.py @@ -1,5 +1,6 @@ import asyncio import re +import shutil import tempfile from datetime import datetime, timedelta @@ -703,14 +704,10 @@ class VideoUtils: @staticmethod def purge_temp_ts_dir(temp_dir: str) -> None: # 6. 删除临时文件和目录 - for file in os.listdir(temp_dir): - file_path = os.path.join(temp_dir, file) - try: - if os.path.isfile(file_path): - os.unlink(file_path) - except Exception as e: - logger.exception(f"删除文件失败 {file_path}: {e}") - os.rmdir(temp_dir) + try: + shutil.rmtree(temp_dir) + except Exception as e: + logger.exception(e) @staticmethod async def ffmpeg_convert_stream_media_multithread(media_stream_url: str, diff --git a/src/cluster/ffmpeg_app.py b/src/cluster/ffmpeg_app.py index 32a886b..ef8674e 100644 --- a/src/cluster/ffmpeg_app.py +++ b/src/cluster/ffmpeg_app.py @@ -1,8 +1,14 @@ +import asyncio +from pathlib import Path + +import botocore import httpx import modal from dotenv import dotenv_values from watchdog.events import DirMovedEvent +from BowongModalFunctions.utils.ModalUtils import ModalUtils + ffmpeg_worker_image = ( modal.Image.debian_slim(python_version="3.11") .apt_install('ffmpeg') @@ -19,7 +25,7 @@ app = modal.App( with ffmpeg_worker_image.imports(): import shutil, os, backoff, sentry_sdk - from typing import List, Optional, Tuple, Dict, Any, Union + from typing import List, Optional, Tuple, Dict, Any, Union, Set from loguru import logger from modal import current_function_call_id from ffmpeg.asyncio import FFmpeg @@ -158,15 +164,10 @@ with ffmpeg_worker_image.imports(): outputs = await ffmpeg_hls_slice_process(media_source=media, media_markers=markers, fn_id=fn_id) - case MediaProtocol.s3: # todo : 临时强行分支出来的逻辑，有空再重构一下 - # if media.urn.endswith(".m3u8"): - # outputs = await ffmpeg_hls_slice_process(media_source=media, - # media_markers=markers, - # fn_id=fn_id) - # else: - outputs = await ffmpeg_slice_process(media_source=media, - media_markers=markers, - fn_id=fn_id) + # case MediaProtocol.s3: + # outputs = await ffmpeg_slice_process(media_source=media, + # media_markers=markers, + # fn_id=fn_id) case _: # webhook不会报错，需要确认 raise NotImplementedError("暂不支持的协议") @@ -562,6 +563,90 @@ with ffmpeg_worker_image.imports(): hls_recording_volume = modal.Volume.from_name("stream_records", create_if_missing=True) hls_recording_mount_point = "/mnt/stream_records" + from watchdog.events import FileMovedEvent, FileCreatedEvent, FileSystemEventHandler + from watchdog.observers import Observer + import boto3 + from botocore.client import BaseClient, Config + + + class PlaylistEventHandler(FileSystemEventHandler): + update_counter: int = 0 + fn_id: str + webhook: Optional[WebhookNotify] = None + s3_mount_output_dir: str + s3_root_output_dir: str + boto3_client: BaseClient + + def __init__(self, fn_id: str, s3_mount_output_dir: str, webhook: Optional[WebhookNotify] = None, *args, + **kwargs): + super().__init__(*args, **kwargs) + self.webhook = webhook + self.fn_id = fn_id + self.s3_mount_output_dir = s3_mount_output_dir + self.boto3_client = boto3.client("s3", + aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"), + aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"), + region_name=config.S3_region, + endpoint_url="https://s3-accelerate.amazonaws.com", + config=Config( + s3={'addressing_style': 'virtual'}, + signature_version='s3v4', ) + ) + self.s3_root_output_dir = s3_mount_output_dir.replace(s3_mount, "") + + + def on_created(self, event: FileCreatedEvent) -> None: + logger.info(f"[created] {event.src_path}") + if event.src_path.endswith(".ts"): + filename = os.path.basename(event.src_path) + mount_path = f"{self.s3_mount_output_dir}/{filename}" + # 将Volume内的ts文件复制到S3挂载点 + shutil.copy(event.src_path, mount_path) + logger.info(f"[copy] {event.src_path} -> {mount_path}") + else: + return + + def on_moved(self, event: FileMovedEvent) -> None: + logger.info(f"[moved] {event.src_path} -> {event.dest_path}") + if not event.dest_path.endswith(".m3u8"): + return + filename = os.path.basename(event.dest_path) + filename.replace('.tmp', '') + mount_path = f"{self.s3_mount_output_dir}/{filename}" + # 将Volume内的playlist.m3u8.tmp复制到S3挂载点的playlist.m3u8文件 + try: + shutil.copy(event.dest_path, mount_path) + except Exception as e: + logger.exception(e) + filename = os.path.basename(event.dest_path) + self.boto3_client.upload_file(event.dest_path, f"{self.s3_root_output_dir}/{filename}") + logger.info(f"[copy] {event.dest_path} -> {mount_path}") + self.update_counter += 1 + if self.webhook and self.update_counter == 1: + logger.info("[Start] webhook trigger") + try: + self.webhook_on_start() + except Exception as e: + logger.exception(e) + + @backoff.on_exception(exception=Exception, wait_gen=backoff.constant, + max_time=5, max_tries=5, raise_on_giveup=True) + def webhook_on_start(self): + """ + 开始录制第一次更新时回调 + """ + webhook = self.webhook + if webhook.method is not WebhookMethodEnum.POST: + logger.warning(f"webhook method {webhook.method.value} not supported") + body = BaseFFMPEGTaskStatusResponse(taskId=self.fn_id, + task_type="ffmpeg_stream_record_as_hls", + status=TaskStatus.running).model_dump() + response = httpx.post(url=webhook.endpoint.__str__(), + json=body, + headers=webhook.headers) + response.raise_for_status() + logger.info(f"[Start] webhook {response.status_code} {response.text}") + @app.function(timeout=43200 + 300, # 最长处理12h的录制任务 + 5分钟的清理缓存工作 cloud="aws", @@ -587,65 +672,6 @@ with ffmpeg_worker_image.imports(): volume_output_dir = f"{hls_recording_mount_point}/{output_dir}" os.makedirs(volume_output_dir, exist_ok=True) logger.info(f"manifest = {volume_output_dir}/playlist.m3u8") - from watchdog.events import FileMovedEvent, FileCreatedEvent, FileSystemEventHandler - from watchdog.observers import Observer - - class PlaylistEventHandler(FileSystemEventHandler): - update_counter: int = 0 - fn_id: str - webhook: Optional[WebhookNotify] = None - - def __init__(self, fn_id: str, webhook: Optional[WebhookNotify] = None, *args, **kwargs): - super().__init__(*args, **kwargs) - self.webhook = webhook - self.fn_id = fn_id - - def on_created(self, event: FileCreatedEvent) -> None: - logger.info(f"[created] {event.src_path}") - if event.src_path.endswith(".ts"): - filename = os.path.basename(event.src_path) - mount_path = f"{s3_mount_output_dir}/{filename}" - # 将Volume内的ts文件复制到S3挂载点 - shutil.copy(event.src_path, mount_path) - logger.info(f"[copy] {event.src_path} -> {mount_path}") - else: - return - - def on_moved(self, event: FileMovedEvent) -> None: - logger.info(f"[moved] {event.src_path} -> {event.dest_path}") - if not event.dest_path.endswith(".m3u8"): - return - filename = os.path.basename(event.dest_path) - filename.replace('.tmp', '') - mount_path = f"{s3_mount_output_dir}/{filename}" - # 将Volume内的playlist.m3u8.tmp复制到S3挂载点的playlist.m3u8文件 - shutil.copy(event.dest_path, mount_path) - logger.info(f"[copy] {event.dest_path} -> {mount_path}") - self.update_counter += 1 - if self.webhook and self.update_counter == 1: - logger.info("[Start] webhook trigger") - try: - self.webhook_on_start() - except Exception as e: - logger.exception(e) - - @backoff.on_exception(exception=Exception, wait_gen=backoff.constant, - max_time=5, max_tries=5, raise_on_giveup=True) - def webhook_on_start(self): - """ - 开始录制第一次更新时回调 - """ - webhook = self.webhook - if webhook.method is not WebhookMethodEnum.POST: - logger.warning(f"webhook method {webhook.method.value} not supported") - body = BaseFFMPEGTaskStatusResponse(taskId=self.fn_id, - task_type="ffmpeg_stream_record_as_hls", - status=TaskStatus.running).model_dump() - response = httpx.post(url=webhook.endpoint.__str__(), - json=body, - headers=webhook.headers) - response.raise_for_status() - logger.info(f"[Start] webhook {response.status_code} {response.text}") @backoff.on_exception(exception=Exception, wait_gen=backoff.constant, max_time=5, max_tries=5, raise_on_giveup=False) @@ -675,7 +701,7 @@ with ffmpeg_worker_image.imports(): headers=webhook.headers) logger.info(f"[End] webhook {response.status_code} {response.text}") - playlist_handler = PlaylistEventHandler(webhook=webhook, fn_id=fn_id) + playlist_handler = PlaylistEventHandler(webhook=webhook, fn_id=fn_id, s3_mount_output_dir=s3_mount_output_dir) playlist_observer = Observer() os.makedirs(volume_output_dir, exist_ok=True) # 监控本地Volume下录制缓存目录 @@ -708,3 +734,89 @@ with ffmpeg_worker_image.imports(): x_baggage=sentry_sdk.get_baggage()) shutil.rmtree(volume_output_dir) return result, sentry_trace + + + @app.function(timeout=43200 + 300, # 最长处理12h的录制任务 + 5分钟的清理缓存工作 + cloud="aws", + volumes={ + s3_mount: modal.CloudBucketMount( + bucket_name=config.S3_bucket_name, + secret=modal.Secret.from_name("aws-s3-secret", + environment_name=config.modal_environment), + ), + hls_recording_mount_point: hls_recording_volume + }, ) + @modal.concurrent(max_inputs=1) + async def ffmpeg_stream_record_restore(fn_id: str): + + def get_files_set(directory: str) -> Set[str]: + """ + 获取目录下所有文件的相对路径集合 + + Args: + directory (str): 目录路径 + + Returns: + Set[str]: 文件相对路径的集合 + """ + directory_path = Path(directory) + return {str(f.relative_to(directory_path)) for f in directory_path.rglob("*") if f.is_file()} + + def copy_directory_contents(src_dir: str, dst_dir: str): + """ + 使用集合操作优化文件复制，跳过已存在的文件 + + Args: + src_dir (str): 源目录路径 + dst_dir (str): 目标目录路径 + """ + src_path = Path(src_dir) + dst_path = Path(dst_dir) + + # 确保目标目录存在 + dst_path.mkdir(parents=True, exist_ok=True) + + # 获取源目录和目标目录的文件集合 + src_files = get_files_set(src_dir) + dst_files = get_files_set(dst_dir) if dst_path.exists() else set() + + # 计算需要复制的文件（源目录有但目标目录没有的文件） + files_to_copy = src_files - dst_files + + # 复制文件 + for rel_path in files_to_copy: + src_file = src_path / rel_path + dst_file = dst_path / rel_path + + # 确保目标文件的父目录存在 + dst_file.parent.mkdir(parents=True, exist_ok=True) + + # 复制文件 + shutil.copy(src_file, dst_file) + logger.info(f"已复制: {rel_path}") + + logger.info(f"\n复制完成:") + logger.info(f"源目录文件数: {len(src_files)}") + logger.info(f"目标目录文件数: {len(dst_files)}") + logger.info(f"本次复制文件数: {len(files_to_copy)}") + + output_dir = f"{config.modal_environment}/records/hls/{fn_id}" + s3_mount_output_dir = f"{s3_mount}/{output_dir}" + volume_output_dir = f"{hls_recording_mount_point}/{output_dir}" + + copy_directory_contents(volume_output_dir, s3_mount_output_dir) + shutil.copy(f"{volume_output_dir}/playlist.m3u8", f"{s3_mount_output_dir}/playlist.m3u8") + + playlist_handler = PlaylistEventHandler(webhook=None, fn_id=fn_id, + s3_mount_output_dir=s3_mount_output_dir) + playlist_observer = Observer() + # 监控本地Volume下录制缓存目录 + playlist_observer.schedule(playlist_handler, path=volume_output_dir, recursive=False) + status = await ModalUtils.get_modal_task_status(fn_id) + if status.status == TaskStatus.running: + playlist_observer.start() + while (status.status == TaskStatus.running): + await asyncio.sleep(5) + status = await ModalUtils.get_modal_task_status(fn_id) + playlist_observer.stop() + logger.info("Stream restore end")