-
Notifications
You must be signed in to change notification settings - Fork 245
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add vscode plugin to enable interactive debugging #1922
Changes from all commits
fbcde94
699a44d
05885fe
be5ecc0
677f9a6
bf1dbd8
78c6366
64e3968
e4d9a3d
0c0c98c
ba73cd6
1fff564
402d450
7df1ffa
933e586
c870612
b0d3997
9435d4e
268bb1a
6eef2e1
f5d4fb3
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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}" | ||
``` |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,13 @@ | ||
""" | ||
.. currentmodule:: flytekitplugins.vscode | ||
|
||
This package contains things that are useful when extending Flytekit. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is this sentence from the plugin template? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This sentence is in every plugin so I assume it's a kind of a template. |
||
|
||
.. autosummary:: | ||
:template: custom.rst | ||
:toctree: generated/ | ||
|
||
vscode | ||
""" | ||
|
||
from .decorator import vscode |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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" |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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"] | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. is There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do you mean os.path.join()? In this case, I think we cannot use os.path.join() because we are going to concatenate $PATH with new path by ":" (os.pathsep) instead of combining two path to one path. For example, if $PATH is /usr/bin:/local/bin and code_server_bin_dir is /tmp/bin, then the result will be /tmp/bin:/usr/bin:/local/bin. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. good point! |
||
|
||
|
||
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: | ||
troychiu marked this conversation as resolved.
Show resolved
Hide resolved
|
||
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 | ||
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
. | ||
-e file:.#egg=flytekitplugins-vscode |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
did you experiment at all with
@task + @vscode
?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do you mean the end to end test of
@task + @vscode
. If yes, then I have done it. You can check the End to end test section of the PR description.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
no... i mean the syntax... that was byron's first idea way back in the day.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What does that mean? Could you give an example