diff --git a/src/handler.py b/src/handler.py index 51c045d..2d539e6 100644 --- a/src/handler.py +++ b/src/handler.py @@ -45,17 +45,24 @@ async def async_handler(job): max_concurrency = int(os.getenv("MAX_CONCURRENCY", 100)) print(f"MAX_CONCURRENCY {max_concurrency}") -async def main(): +# Use a synchronous handler that wraps the async handler +def handler(job): + return asyncio.get_event_loop().run_until_complete(async_handler(job)) + +async def init(): await engine.wait_for_server() - await runpod.serverless.start({ - "handler": async_handler, + +if __name__ == "__main__": + # Initialize the server + asyncio.get_event_loop().run_until_complete(init()) + + # Start the serverless function with a synchronous handler + runpod.serverless.start({ + "handler": handler, "concurrency_modifier": lambda x: max_concurrency, "return_aggregate_stream": True }) -if __name__ == "__main__": - asyncio.run(main()) - # Ensure the server is shut down when the serverless function is terminated import atexit atexit.register(engine.shutdown) \ No newline at end of file