Vscode decorator to turn any python (AutoContainer) task to a vscode server

[troychiu]

Introduction

Why:
When

1. tasks fail and users need to debug them.
2. users write tasks from scratch and want to develop in a container in a Kubernetes (k8s) cluster, likely because only the cluster side has GPU resources.

Flyte users might currently use a workaround of using time.sleep, SSHing into the cluster, and then developing/debugging. However, using the raw terminal for this purpose can be cumbersome and unintuitive. For instance, in a multi-level directory structure or a complex code base that requires heavy navigation, users may find the terminal environment frustrating.

Solution:
Therefore, @ByronHsu and I propose the addition of a decorator as a Flyte plugin, enabling users to turn any task into a Visual Studio Code (VSCode) server with just one line of code.

Additionally, we can auto-generate a launch.json file, providing a "one-click to run" feature.

Code sample

The @vscode decorator provides a simple way for users to run a task as an interactive VSCode server within any container image. @vscode is a decorator that can be used within the @task and user functions, like this:

@task
@vscode
def train():
    ...

By doing so, the task will install VSCode dependencies and run a VSCode server instead of the user-defined function. It also starts an exit handler that continuously polls the status of the heartbeat file and terminates if the server is idle for a set duration. This is crucial because the pod is most likely using several GPUs, and we want to free them up for other users.

Heartbeat file can monitor the active connection to the code-server

As long as there is an active browser connection, code-server touches ~/.local/share/code-server/heartbeat once a minute. If you want to shutdown code-server if there hasn't been an active connection after a predetermined amount of time, you can do so by checking continuously for the last modified time on the heartbeat file. If it is older than X minutes (or whatever amount of time you'd like), you can kill code-server.

def vscode(
    _task_function: Optional[Callable] =  None,
    max_idle_seconds: Optional[int] = 3600 * 24 * 2, # 2 days
    port: Optional[int] = 8001,
    enable: Optional[bool] = True,
    code_server_remote_path: str = "XXX",
    # The untarred directory name may be different from the tarball name
    code_server_dir_name: str = "YYY",
    plugins_remote_paths: List[str] = [
        "ZZZ",
        "KKK"
    ],
    pre_execute: Optional[Callable] = None,
):
    """
    vscode decorator modifies a container to run a VSCode server:
    1. Overrides the user function with a VSCode setup function.
    2. Download vscode server and plugins from remote to local.
    4. Launches and monitors the VSCode server.
    5. Terminates if the server is idle for a set duration.

    Parameters:
    - _task_function (function, optional): The user function to be decorated. Defaults to None.
    - max_idle_seconds (int, optional): The duration in seconds to live after no activity detected. 
    - port (int, optional): The port to be used by the VSCode server. Defaults to 8080.
    - enable (bool, optional): Whether to enable the VSCode decorator. Defaults to True.
    - code_server_remote_path (str, optional): The URL of the code-server tarball.
    - code_server_dir_name (str, optional): The name of the code-server directory.
    - plugins_remote_paths (List[str], optional): The URLs of the VSCode plugins.
    - pre_execute (function, optional): The function to be executed before the vscode setup function.
    """

Vscode example screenshot:

[troychiu]

Since it will be a plugin that doesn’t require changes in flytekit, no RFC required (discussed in the meetup on Nov. 9)