Add vscode plugin to enable interactive debugging

.github/workflows/pythonbuild.yml

@@ -128,6 +128,7 @@ jobs:
           - flytekit-spark
           - flytekit-sqlalchemy
           - flytekit-vaex
+          - flytekit-vscode
           - flytekit-whylogs
         exclude:
           # flytekit-modin depends on ray which does not have a 3.11 wheel yet.

Dockerfile.dev

@@ -23,6 +23,7 @@ COPY . /flytekit
 RUN pip install -e /flytekit
 RUN pip install -e /flytekit/plugins/flytekit-k8s-pod
 RUN pip install -e /flytekit/plugins/flytekit-deck-standard
+RUN pip install -e /flytekit/plugins/flytekit-vscode
 RUN pip install scikit-learn
 
 ENV PYTHONPATH "/flytekit:/flytekit/plugins/flytekit-k8s-pod:/flytekit/plugins/flytekit-deck-standard:"

plugins/flytekit-vscode/README.md

@@ -0,0 +1,41 @@
+# Flytekit VSCode Plugin
+
+The Flytekit VSCode plugin offers an easy solution for users to run Python tasks within an interactive VSCode server, compatible with any image. `@vscode` is a decorator which users can put within @task and user function. With `@vscode`, the task will install vscode dependencies (skip if they already exist) and run a vscode server instead of the user defined functions.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-vscode
+```
+
+## Task Example
+```python
+from flytekit import task
+from flytekitplugins.vscode import vscode
+
+@task
+@vscode
+def train():
+    ...
+```
+
+## User Guide
+1. Build the image with Dockerfile.dev `docker build --push . -f Dockerfile.dev -t localhost:30000/flytekit:dev --build-arg PYTHON_VERSION=3.8`
+2. Run the decorated task on the remote. For example: `pyflyte run --remote --image localhost:30000/flytekit:dev [PYTHONFILE] [WORKFLOW|TASK] [ARGS]...`
+3. Once the code server is prepared, you can forward a local port to the pod. For example: `kubectl port-forward -n [NAMESPACE] [PODNAME] 8080:8080`.
+4. You can access the server by opening a web browser and navigating to `localhost:8080`.
+
+VSCode example screenshot:
+<img src="./docs/example.png">
+
+## Build Custom Image with VSCode Plugin
+If users want to skip the vscode downloading process at runtime, they have the option to create a custom image with vscode by including the following lines in their Dockerfile.
+```Dockerfile
+# Include this line if the image does not already have 'curl' installed.
++ RUN apt-get -y install curl
+# Download and extract the binary, and ensure it's added to the system's $PATH.
++ RUN mkdir /tmp/code-server
++ RUN curl -kfL -o /tmp/code-server/code-server-4.18.0-linux-amd64.tar.gz https://github.com/coder/code-server/releases/download/v4.18.0/code-server-4.18.0-linux-amd64.tar.gz
++ RUN tar -xzf /tmp/code-server/code-server-4.18.0-linux-amd64.tar.gz -C /tmp/code-server/
++ ENV PATH="/tmp/code-server/code-server-4.18.0-linux-amd64/bin:${PATH}"
+```

plugins/flytekit-vscode/flytekitplugins/vscode/__init__.py

@@ -0,0 +1,13 @@
+"""
+.. currentmodule:: flytekitplugins.vscode
+
+This package contains things that are useful when extending Flytekit.
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   vscode
+"""
+
+from .decorator import vscode

plugins/flytekit-vscode/flytekitplugins/vscode/constants.py

@@ -0,0 +1,9 @@
+# Where the code-server tar and plugins are downloaded to
+EXECUTABLE_NAME = "code-server"
+DOWNLOAD_DIR = "/tmp/code-server"
+HOURS_TO_SECONDS = 60 * 60
+DEFAULT_UP_SECONDS = 10 * HOURS_TO_SECONDS  # 10 hours
+DEFAULT_CODE_SERVER_REMOTE_PATH = (
+    "https://github.com/coder/code-server/releases/download/v4.18.0/code-server-4.18.0-linux-amd64.tar.gz"
+)
+DEFAULT_CODE_SERVER_DIR_NAME = "code-server-4.18.0-linux-amd64"

plugins/flytekit-vscode/flytekitplugins/vscode/decorator.py

@@ -0,0 +1,179 @@
+import multiprocessing
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+import time
+from functools import wraps
+from typing import Callable, Optional
+
+import fsspec
+
+from flytekit.loggers import logger
+
+from .constants import (
+    DEFAULT_CODE_SERVER_DIR_NAME,
+    DEFAULT_CODE_SERVER_REMOTE_PATH,
+    DEFAULT_UP_SECONDS,
+    DOWNLOAD_DIR,
+    EXECUTABLE_NAME,
+)
+
+
+def execute_command(cmd):
+    """
+    Execute a command in the shell.
+    """
+
+    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    logger.info(f"cmd: {cmd}")
+    stdout, stderr = process.communicate()
+    if process.returncode != 0:
+        raise RuntimeError(f"Command {cmd} failed with error: {stderr}")
+    logger.info(f"stdout: {stdout}")
+    logger.info(f"stderr: {stderr}")
+
+
+def download_file(url, target_dir="."):
+    """
+    Download a file from a given URL using fsspec.
+
+    Args:
+        url (str): The URL of the file to download.
+        target_dir (str, optional): The directory where the file should be saved. Defaults to current directory.
+
+    Returns:
+        str: The path to the downloaded file.
+    """
+
+    if not url.startswith("http"):
+        raise ValueError(f"URL {url} is not valid. Only http/https is supported.")
+
+    # Derive the local filename from the URL
+    local_file_name = os.path.join(target_dir, os.path.basename(url))
+
+    fs = fsspec.filesystem("http")
+
+    # Use fsspec to get the remote file and save it locally
+    logger.info(f"Downloading {url}... to {os.path.abspath(local_file_name)}")
+    fs.get(url, local_file_name)
+    logger.info("File downloaded successfully!")
+
+    return local_file_name
+
+
+def download_vscode(
+    code_server_remote_path: str,
+    code_server_dir_name: str,
+):
+    """
+    Download vscode server and plugins from remote to local and add the directory of binary executable to $PATH.
+
+    Args:
+        code_server_remote_path (str): The URL of the code-server tarball.
+        code_server_dir_name (str): The name of the code-server directory.
+    """
+
+    # If the code server already exists in the container, skip downloading
+    executable_path = shutil.which(EXECUTABLE_NAME)
+    if executable_path is not None:
+        logger.info(f"Code server binary already exists at {executable_path}")
+        logger.info("Skipping downloading code server...")
+        return
+
+    logger.info("Code server is not in $PATH, start downloading code server...")
+
+    # Create DOWNLOAD_DIR if not exist
+    logger.info(f"DOWNLOAD_DIR: {DOWNLOAD_DIR}")
+    os.makedirs(DOWNLOAD_DIR, exist_ok=True)
+
+    logger.info(f"Start downloading files to {DOWNLOAD_DIR}")
+
+    # Download remote file to local
+    code_server_tar_path = download_file(code_server_remote_path, DOWNLOAD_DIR)
+
+    # Extract the tarball
+    with tarfile.open(code_server_tar_path, "r:gz") as tar:
+        tar.extractall(path=DOWNLOAD_DIR)
+
+    code_server_dir_path = os.path.join(DOWNLOAD_DIR, code_server_dir_name)
+
+    code_server_bin_dir = os.path.join(code_server_dir_path, "bin")
+
+    # Add the directory of code-server binary to $PATH
+    os.environ["PATH"] = code_server_bin_dir + os.pathsep + os.environ["PATH"]
+
+
+def vscode(
+    _task_function: Optional[Callable] = None,
+    server_up_seconds: Optional[int] = DEFAULT_UP_SECONDS,
+    port: Optional[int] = 8080,
+    enable: Optional[bool] = True,
+    code_server_remote_path: Optional[str] = DEFAULT_CODE_SERVER_REMOTE_PATH,
+    # The untarred directory name may be different from the tarball name
+    code_server_dir_name: Optional[str] = DEFAULT_CODE_SERVER_DIR_NAME,
+    pre_execute: Optional[Callable] = None,
+    post_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 after server_up_seconds seconds.
+
+    Args:
+        _task_function (function, optional): The user function to be decorated. Defaults to None.
+        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.
+        pre_execute (function, optional): The function to be executed before the vscode setup function.
+        post_execute (function, optional): The function to be executed before the vscode is self-terminated.
+    """
+
+    def wrapper(fn):
+        if not enable:
+            return fn
+
+        @wraps(fn)
+        def inner_wrapper(*args, **kwargs):
+            # 0. Executes the pre_execute function if provided.
+            if pre_execute is not None:
+                pre_execute()
+                logger.info("Pre execute function executed successfully!")
+
+            # 1. Downloads the VSCode server from Internet to local.
+            download_vscode(
+                code_server_remote_path=code_server_remote_path,
+                code_server_dir_name=code_server_dir_name,
+            )
+
+            # 2. Launches and monitors the VSCode server.
+            # Run the function in the background
+            logger.info(f"Start the server for {server_up_seconds} seconds...")
+            child_process = multiprocessing.Process(
+                target=execute_command, kwargs={"cmd": f"code-server --bind-addr 0.0.0.0:{port} --auth none"}
+            )
+
+            child_process.start()
+            time.sleep(server_up_seconds)
+
+            # 3. Terminates the server after server_up_seconds
+            logger.info(f"{server_up_seconds} seconds passed. Terminating...")
+            if post_execute is not None:
+                post_execute()
+                logger.info("Post execute function executed successfully!")
+            child_process.terminate()
+            child_process.join()
+            sys.exit(0)
+
+        return inner_wrapper
+
+    # for the case when the decorator is used without arguments
+    if _task_function is not None:
+        return wrapper(_task_function)
+    # for the case when the decorator is used with arguments
+    else:
+        return wrapper

plugins/flytekit-vscode/requirements.in

@@ -0,0 +1,2 @@
+.
+-e file:.#egg=flytekitplugins-vscode