File size: 3,060 Bytes
da6e788 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import os
import pickle
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from io import BytesIO
import subprocess
app = FastAPI()
@app.post("/container_task")
async def container_task(file: UploadFile = File(...)):
# Save the uploaded file to disk
input_filepath = "input.pkl"
with open(input_filepath, "wb") as f:
f.write(await file.read())
# Process the unpickle_param from the file
try:
with open(input_filepath, 'rb') as f:
unpickle_param = pickle.load(f)
except Exception as e:
raise HTTPException(status_code=400, detail=f"Failed to unpickle the input file: {str(e)}")
# Execute the Docker command
command = ["docker", "run", "--rm", "bbdown", str(unpickle_param)]
process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
# Stream the output of the command
stdout_stream = BytesIO()
stderr_stream = BytesIO()
# Create a generator to stream output
def stream_output():
while True:
output = process.stdout.readline()
if output == b"" and process.poll() is not None:
break
if output:
stdout_stream.write(output)
stderr_output = process.stderr.read()
stderr_stream.write(stderr_output)
yield ""
# Return the StreamingResponse for the current output
async def response_stream():
for _ in stream_output():
yield stdout_stream.getvalue()
stdout_stream.seek(0) # Rewind for next read
stdout_stream.truncate() # Clear for next fill
# Run the process and wait for completion
process.wait()
# Check for errors
if process.returncode != 0:
raise HTTPException(status_code=500, detail=f"Docker command failed with error: {stderr_stream.getvalue().decode()}")
# Create a new pickle file as output
output_filepath = "output.pkl"
with open(output_filepath, 'wb') as f:
f.write(b"Your output data here.") # Replace this with actual output data
# Return the output file
return StreamingResponse(open(output_filepath, "rb"), media_type='application/octet-stream',
headers={"Content-Disposition": f"attachment; filename={os.path.basename(output_filepath)}"})
# To run the application, use: uvicorn your_file_name:app --reload
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import time
import asyncio
app = FastAPI()
async def stream_generator():
for i in range(10):
yield f"Data chunk {i}\n"
await asyncio.sleep(1) # Simulating some delay
@app.get("/stream")
async def stream_response():
return StreamingResponse(stream_generator(), media_type="text/plain")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="127.0.0.1", port=8000)
def client_call(*args, **kwargs):
result = execute(*args, **kwargs)
result.text
result.file_manifest
result.files
|