Add an RPC mechanism to Lightning App

[tchaton]

🚀 Feature

As a user, I would like to build an app that can act as a platform. Once the app is running in the cloud, I would like my users to be able to connect to the app and execute commands (RPC) calls to launch dynamically jobs.

Ideally, I would like to be able to execute some python code on the user client side, so more complex RPC calls can be implemented.

Example

Imagine you want to build a simplified application version of Grid.ai.

Using Grid, users can log in with grid login.

Once logged in, a user can execute their PyTorch Lightning scripts in the cloud with the following command:

grid run train.py --trainer.max_epochs=uniform(10, 100)

In order to create the same functionality with an app, we need a RPC API. Let's assume it exposes a sweep command.

lightning run app simplified_grid_app.py

Once running, users of the app could run the following command.

lightning connect app_id/url/app_name
lightning sweep script.py --trainer.max_epochs=uniform(10, 100) --compute="fast-gpu"

Here are the steps behind the sweep command:

• [user machine] Enter the sweep command
• [user machine] Upload the local source code to s3, send a request to the App RestAPI with all information to know where the code is, etc...
• [Rest API] Receives the request and passes the payload to the flow.
• [Flow] Execute the associated method with the payload.

Here is an example of the Sweep Method:

def sweep(self, request):
    self.flows.append(
        HPOCodeRunner(**request)
    )

Motivation

Pitch

Alternatives

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @Borda @tchaton @justusschock @awaelchli

[justusschock]

Is the idea here to add the sweep method (and others)? Or to add possibilities to call virtually any script ( I'd really like that!)?

I'd suggest to do this and then use the entrypoints from #13602 to optionally do stuff like sweeps.

I don't think sweeps would be something that need to be available on every app

[ethanwharris]

I like this! One thing that could make UI development much cleaner IMO is if we also had a mechanism to trigger these commands from e.g. the react UI by sending an HTTP request or similar. That would avoid some tricks we currently use like setting a flag to trigger an action then unsetting the flag once the action starts (which works fine but gets cluttered when you are triggering several things in this way).

[tchaton]

Hey @justusschock @ethanwharris.

Actually, the API in #13602 is too limited as there isn't any code execution on the client side.

I chatted with @carmocca to brainstorm on another API and we thought about this one.

from lightning import LightningFlow
from pydantic import BaseModel
from typing import List

class SweepConfig(BaseModel):
    github_repo: str
    script_path: str
    requirements: List[str]


class HPOFlow(LightningFlow):


    def sweep(self, sweep: SweepConfig):
        ...

    def configure_commands(self):
        return [
            {
                "name": "sweep",
                "server": self.sweep,
                "client": "sweep.py"    # optional
            },
            ...
        ]

#sweep.py

def sweep(....):
     ....
     return payload

When running `lightning command sweep, this would download locally the self-contained sweep.py, extract the sweep method and execute it over the user arguments

In this use case, this would enable to upload of the source code to s3 and then launch a sweep with the associated information.

What are your thoughts ?

[tchaton]

After chatting with @justusschock and @lantiga.

Here is the proposed API

The commands need to be isolated from the rest of the codebase, so we can easily package them (versioned on PyPI or s3) and enable the users to download the associated code in a trusted manner.

Here is a Sweep Command

# commands/sweep.py

from lightning.app.commands import ClientCommand

class SweepClientCommand(ClientCommand):

    def __init__(self):
           super().__init__(requirements=["boto3", ...])

    def run(script_path, *args, **kwargs):
        remote_url = upload_local_dir_to_s3(script_path)
        response = self.invoke_handler(SweepConfig(remote_url, *args, **kwargs))  # This triggers the method on the CI and waits the response.

Here is a Predict Command

# commands/predict.py

class PredictClientCommand(ClientCommand):

    def run(image_folder):
        for file in os.listdir(image_folder):
            endpoint_url = self.invoke_handler()
            predictions = requests.post(endpoint_url, load_data(file))
       
         self.state.num_predictions += 1

from commands import PredictClientCommand, SweepClientCommand, SweepConfig

class HPOFlow(LightningFlow):

    def run(self):
        for flow in self.flows:
            flow.run()

    def command_without_client(self, name: str):
        ...

    def sweep(self, sweep: SweepConfig):
        self.flows.append(HPOCodeRunner(sweep))

    def fetch_endpoint(self):
        return self.current_server.url

    def configure_commands(self):
        return [
            {"command_without_client": self.command_without_client}, # This doesn't trigger code execution on the CLI
            {"sweep": SweepClientCommand(self.sweep)},
            {"predict_from_folders": PredictClientCommand(self.fetch_endpoint)}
        ]

Here is the App Builder Workflow

Once he decides to publish his app in the cloud, the commands code will be packaged and versioned to avoid malicious tampering.

Here is the User Workflow

Once the user triggers this command the first time, he will be prompted several of instructions.

lightning command connect --app_id X

This application contains a client package, do you agree for this to be installed? 
Here are the downloaded files:
.....sweep.py -> Users/..../sweep.py
....
...

[Yes, No]
The file is downloaded, here are the available commands:
command_without_client
sweep
predict_from_folders

The client code now lives on your machine and matches the one publicly versionned. Run a command if the code is trusted ... (to be re-phrased)

lightning command sweep script.py --lr=uniform(0, 1)

[lantiga]

Looks good. One extra thing we aligned about is that we may need to pass more handlers to the command (in case one of the handlers need to be called conditionally or sequentially):

    def configure_commands(self):
        return [
            {"sweep": SweepClientCommand(sweep=self.sweep, single_run=self.single_run)},
        ]

In that case the idea is that invoke_handler needs to take an extra argument with the name of the handler (”sweep” or ”single_run” in the example above)

[manskx]

Looks great!

The ClientCommand itself can be considered as a client package per app (perhaps with its own commands) and gets installed when a user does lightning connect --app_id x.

I imagine the workflow as follows:

1- lightning connect --app_id x
This installs a client package

2- the grid-client package has commands e.g. predict, sweep, and the implementation of the predict for example can be

# grid-client 

def predict(input_folder, output_folder):
        for file in os.listdir(input_folder):
            endpoint_url = self.app.run_command.fetch_endpoint()        # command_without_client equivalent to `lightning command run fetch_endpoint`
            predictions = requests.post(endpoint_url, load_data(file))
            # write predictions to output_folder

3- lightning command predict --input_folder=<> output_folder==<> which is mapped to grid-client predict --input_folder=<> output_folder==<>

does this make sense ?