lightning/docs/source-app/levels/intermediate/scripts/toy_payload.py

32 lines
959 B
Python

# app.py
import lightning as L
class EmbeddingProcessor(L.LightningWork):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.embeddings = None
def run(self):
print('PROCESSOR: Generating embeddings...')
fake_embeddings = [[1, 2, 3], [2, 3, 4]]
self.embeddings = L.storage.Payload(fake_embeddings)
class EmbeddingServer(L.LightningWork):
def run(self, payload):
print('SERVER: Using embeddings from processor', payload)
embeddings = payload.value
print('serving embeddings sent from EmbeddingProcessor: ', embeddings)
class WorkflowOrchestrator(L.LightningFlow):
def __init__(self) -> None:
super().__init__()
self.processor = EmbeddingProcessor()
self.server = EmbeddingServer()
def run(self):
self.processor.run()
self.server.run(self.processor.embeddings)
app = L.LightningApp(WorkflowOrchestrator())