53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
from fastapi import FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from lightning import CloudCompute, LightningApp, LightningFlow, LightningWork
|
|
from uvicorn import run
|
|
|
|
|
|
class Work(LightningWork):
|
|
def __init__(self, **kwargs):
|
|
super().__init__(parallel=True, **kwargs)
|
|
|
|
def run(self):
|
|
fastapi_service = FastAPI()
|
|
|
|
fastapi_service.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
@fastapi_service.get("/")
|
|
def get_root():
|
|
return {"Hello Word!"}
|
|
|
|
run(fastapi_service, host=self.host, port=self.port)
|
|
|
|
|
|
class Flow(LightningFlow):
|
|
def __init__(self):
|
|
super().__init__()
|
|
# In the Cloud: All the works defined without passing explicitly a CloudCompute object
|
|
# are running on the default machine.
|
|
# This would apply to `work_a`, `work_b` and the dynamically created `work_d`.
|
|
|
|
self.work_a = Work()
|
|
self.work_b = Work()
|
|
|
|
self.work_c = Work(cloud_compute=CloudCompute(name="cpu-small"))
|
|
|
|
def run(self):
|
|
if not hasattr(self, "work_d"):
|
|
self.work_d = Work()
|
|
|
|
for work in self.works():
|
|
work.run()
|
|
|
|
def configure_layout(self):
|
|
return [{"name": w.name, "content": w} for i, w in enumerate(self.works())]
|
|
|
|
|
|
app = LightningApp(Flow(), log_level="debug")
|