Step 3: Implement the Flow to manage user requests¶
In step 1 and 2, we have implemented the GithubRepoRunner
and PyTorchLightningGithubRepoRunner
components.
Now, we are going to develop a component to dynamically handle user requests. Let’s dive in on how to develop the component with the following code:
class TensorflowGithubRepoRunner(GithubRepoRunner):
"""Left to the users to implement."""
GITHUB_REPO_RUNNERS = {
"PyTorch Lightning": PyTorchLightningGithubRepoRunner,
"Keras": KerasGithubRepoRunner,
"Tensorflow": TensorflowGithubRepoRunner,
}
class Flow(LightningFlow):
def __init__(self):
super().__init__()
# 1: Keep track of the requests within the state
self.requests = []
# 2: Create a dictionary of components.
self.ws = structures.Dict()
def run(self):
# Iterate continuously over all requests
for request_id, request in enumerate(self.requests):
self._handle_request(request_id, deepcopy(request))
def _handle_request(self, request_id: int, request: Dict):
# 1: Create a name and find selected framework
name = f"w_{request_id}"
ml_framework = request["train"].pop("ml_framework")
# 2: If the component hasn't been created yet, create it.
if name not in self.ws:
work_cls = GITHUB_REPO_RUNNERS[ml_framework]
work = work_cls(id=request["id"], **request["train"])
self.ws[name] = work
# 3: Run the component
self.ws[name].run()
# 4: Once the component has finished,
# add metadata to the original request for the UI.
if self.ws[name].best_model_path:
request = self.requests[request_id]
request["best_model_score"] = self.ws[name].best_model_score
request["best_model_path"] = self.ws[name].best_model_path
def configure_layout(self):
# Create a StreamLit UI for the user to run his Github Repo.
return StreamlitFrontend(render_fn=render_fn)