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)