FlyIn Flytekit Plugin

[ByronHsu]

Motivation: Why do you think this is important?

When

1. The tasks fail and users want to debug them
2. Users start writing the task from scratch and want to develop in a container in k8s cluster (likely because only the cluster side has GPU)

Flyte users might time.sleep the task, ssh inside, and develop/debug. However, using raw terminal to do so is cumbersome and unintuitive. For example, in multi-level directories structure or complex code base that need heavy navigation, users may feel pretty frustrated with the terminal environment.

I propose that we can add a decorator as a flyte plugin and enable users to turn any tasks to vscode server in just one line.

Goal: What should the final outcome look like, ideally?

@vscode provides an simple way for users to run the task as a interactive vscode server with any image. @vscode is a decorator which users can put within @task and user function.

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

By doing so, the task will install vscode dependencies and run a vscode server instead of the user defined functions. Also, it starts a exit handler which keeps polling the status of the heartbeat file and terminate if the server is idle for a set duration. This is very important because the pod is most likely using several GPU's, so we want to free them 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.
    3. Launches and monitors the VSCode server.
    4. 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:

Describe alternatives you've considered

Sleep the pod and ssh

Propose: Link/Inline OR Additional context

No response

Are you sure this issue hasn't been raised already?

• Yes

Have you read the Code of Conduct?

• Yes

Tasks

Beta Give feedback

No tasks being tracked yet.

[ByronHsu]

Note:

Got a very cool idea from flyte community sync regarding vscode plugin. The current drawback of vscode is that we cannot easily execute the task with task inputs unless we copy the command from pod yaml.

the proposal is that we can inject a launch.json to vscode container with all commands set up and users can easily click one button to iterate on the tasks with task inputs. Furthermore, we download the input protobuf to the container, and convert that to a json, so users can easily change the input themselves

# launch.json

{
  "command": "pyflyte run workflow.py t1 --inputs input.json",
}
}

# input.json

{
  "data_importer": {
     "hdfs_path": "abc" # user can change the path here
  },
}

the experience will be like, users enter the vscode, and click one button the execute the task with task inputs

[fg91]

# launch.json

{
  "command": "pyflyte run workflow.py t1 --inputs input.json",
}
}

Given this example ...

from flytekit import task, workflow


@task
def foo(inp: int = 1):
    print(f"bar {inp}")


@workflow
def wf():
    foo()

... we could automatically inject a launch.json containing the command pyflyte run test.py foo --inp 1 (already works today) which would allow modifying primitive task args directly in the command in the launch.json - without the additional json logic you proposed.

I don't see a good solution yet for providing a way to override task args that are offloaded to blob storage and which one cannot pass as CLI args yet. That being said, I'm not sure this even needs to be possible. I think it would be fair to only allow overriding primitive types via launch.json (the ones that pyflyte run supports as CLI args) and retrieve the remaining types such as e.g. torch.nn.Module from the input protobuf. Since users have control over the code, they could just instantiate e.g. a different torch.nn.Module in the task function to override what has been passed as a task argument from the input.pb.

[Future-Outlier]

I've added the sandbox email publisher 3 months ago.
I can help to add the notification integration with this feature.

Here are the advice from @ByronHsu

we can also see if we can add the regular notification by flyte notification.
e.g. sending a reminder email every X hours, or a termination email Y hours

PR Link: flyteorg/flyteadmin#595