""" DaaS (Docker as a Service) is a service that allows users to run docker commands on the server side. """ from fastapi import FastAPI from fastapi.responses import StreamingResponse from fastapi import FastAPI, File, UploadFile, HTTPException from pydantic import BaseModel, Field from typing import Optional, Dict import time import os import asyncio import subprocess import uuid import glob import threading import queue from shared_utils.docker_as_service_api import DockerServiceApiComModel app = FastAPI() def python_obj_to_pickle_file_bytes(obj): import pickle import io with io.BytesIO() as f: pickle.dump(obj, f) return f.getvalue() def yield_message(message): dsacm = DockerServiceApiComModel(server_message=message) return python_obj_to_pickle_file_bytes(dsacm) def read_output(stream, output_queue): while True: line_stdout = stream.readline() # print('recv') if line_stdout: output_queue.put(line_stdout) else: break async def stream_generator(request_obj): import tempfile # Create a temporary directory with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: # Construct the docker command download_folder = temp_dir # Get list of existing files before download existing_file_before_download = [] video_id = request_obj.client_command cmd = [ '/root/.dotnet/tools/BBDown', video_id, '--use-app-api', '--work-dir', f'{os.path.abspath(temp_dir)}' ] cmd_chmod = [] cmd = [ 'docker', 'run', '--rm', '-v', f'{os.path.abspath(temp_dir)}:/downloads', 'bbdown', video_id, '--use-app-api', '--work-dir', '/downloads' ] cmd_chmod = [ 'docker', 'run', '--rm', '-v', f'{os.path.abspath(temp_dir)}:/downloads', '--entrypoint=""', # override the entrypoint 'bbdown', # image name # chmod -R 777 /downloads 'chmod', '-R', '777', '/downloads' ] cmd = ' '.join(cmd) yield yield_message(cmd) process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, text=True) stdout_queue = queue.Queue() thread = threading.Thread(target=read_output, args=(process.stdout, stdout_queue)) thread.daemon = True thread.start() stderr_queue = queue.Queue() thread = threading.Thread(target=read_output, args=(process.stderr, stderr_queue)) thread.daemon = True thread.start() while True: print("looping") # Check if there is any output in the queue stdout_this_round = "" stderr_this_round = "" while True: try: output_stdout = stdout_queue.get_nowait() # Non-blocking get if output_stdout: stdout_this_round += output_stdout print(output_stdout) except queue.Empty: yield yield_message(stdout_this_round) break while True: try: output_stderr = stderr_queue.get_nowait() # Non-blocking get if output_stderr: stderr_this_round += output_stderr print(output_stderr) except queue.Empty: yield yield_message(stderr_this_round) break # Break the loop if the process has finished if process.poll() is not None: break await asyncio.sleep(0.5) # Get the return code return_code = process.returncode yield yield_message("(daas return code:) " + str(return_code)) # change files mod to 777 if cmd_chmod: docker_chmod_res = subprocess.call(' '.join(cmd_chmod), shell=True) # print(f"Successfully downloaded video {video_id}") existing_file_after_download = glob.glob(os.path.join(download_folder, '**', '*')) # existing_file_after_download = list(os.listdir(download_folder)) # get the difference downloaded_files = [ f for f in existing_file_after_download if f not in existing_file_before_download ] downloaded_files_path = [ os.path.join(download_folder, f) for f in existing_file_after_download if f not in existing_file_before_download ] # read file server_file_attach = {} for fp, fn in zip(downloaded_files_path, downloaded_files): if os.path.isdir(fp): continue with open(fp, "rb") as f: file_bytes = f.read() server_file_attach[fn] = file_bytes dsacm = DockerServiceApiComModel( server_message="complete", server_file_attach=server_file_attach, ) yield python_obj_to_pickle_file_bytes(dsacm) def simple_generator(return_obj): dsacm = DockerServiceApiComModel( server_message=return_obj, ) yield python_obj_to_pickle_file_bytes(dsacm) @app.post("/stream") async def stream_response(file: UploadFile = File(...)): # read the file in memory, treat it as pickle file, and unpickle it import pickle import io content = await file.read() with io.BytesIO(content) as f: request_obj = pickle.load(f) # process the request_obj return StreamingResponse(stream_generator(request_obj), media_type="application/octet-stream") @app.post("/search") async def stream_response(file: UploadFile = File(...)): # read the file in memory, treat it as pickle file, and unpickle it import pickle import io content = await file.read() with io.BytesIO(content) as f: request_obj = pickle.load(f) # process the request_obj keyword = request_obj.client_command from experimental_mods.get_search_kw_api_stop import search_videos # Default parameters for video search csrf_token = '40a227fcf12c380d7d3c81af2cd8c5e8' # Using default from main() search_type = 'default' max_pages = 1 output_path = 'search_results' config_path = 'experimental_mods/config.json' # Search for videos and return the first result videos = search_videos( keyword=keyword, csrf_token=csrf_token, search_type=search_type, max_pages=max_pages, output_path=output_path, config_path=config_path, early_stop=True ) return StreamingResponse(simple_generator(videos), media_type="application/octet-stream") @app.get("/") async def hi(): return "Hello, this is Docker as a Service (DaaS)! If you want to use this service, you must duplicate this space. " \ "您好，这里是Docker作为服务（DaaS）！如果您想使用此服务，您必须复制此空间。复制方法：点击https://huggingface.co/spaces/hamercity/bbdown页面右上角的三个点，然后选择“复制空间”。" \ "此外，在设置中，你还需要修改URL，例如：DAAS_SERVER_URL = \"https://你的用户名-你的空间名.hf.space/stream\"" if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=49000)