Fix Flyin Plugin VSCode download bug

plugins/flytekit-flyin/Dockerfile

@@ -27,6 +27,7 @@ RUN apt-get update \
     && apt-get autoremove --yes \
     && rm -rf /var/lib/{apt,dpkg,cache,log}/ \
     && useradd -u 1000 flytekit \
+    && chown -R flytekit:flytekit /tmp/code-server \
     && chown flytekit: /root \
     && chown flytekit: /home \
     && :

plugins/flytekit-flyin/flytekitplugins/flyin/vscode_lib/constants.py

@@ -36,3 +36,6 @@
 
 # Context attribute name of the task function's source file path
 TASK_FUNCTION_SOURCE_PATH = "TASK_FUNCTION_SOURCE_PATH"
+
+# Subprocess constants
+EXIT_CODE_SUCCESS = 0

plugins/flytekit-flyin/flytekitplugins/flyin/vscode_lib/decorator.py

@@ -10,7 +10,7 @@
 import tarfile
 import time
 from threading import Event
-from typing import Callable, Optional
+from typing import Callable, List, Optional
 
 import fsspec
 from flytekitplugins.flyin.utils import load_module_from_path
@@ -23,6 +23,7 @@
 from .constants import (
     DOWNLOAD_DIR,
     EXECUTABLE_NAME,
+    EXIT_CODE_SUCCESS,
     HEARTBEAT_CHECK_SECONDS,
     HEARTBEAT_PATH,
     INTERACTIVE_DEBUGGING_FILE_NAME,
@@ -42,7 +43,7 @@
     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:
+    if process.returncode != EXIT_CODE_SUCCESS:
         raise RuntimeError(f"Command {cmd} failed with error: {stderr}")
     logger.info(f"stdout: {stdout}")
     logger.info(f"stderr: {stderr}")
@@ -176,6 +177,27 @@
         )
 
 
+def get_installed_extensions() -> List[str]:
+    """
+    Get the list of installed extensions.
+
+    Returns:
+        List[str]: The list of installed extensions.
+    """
+    logger = flytekit.current_context().logging
+
+    installed_extensions = subprocess.run(["code-server", "--list-extensions"], capture_output=True, text=True)
+    if installed_extensions.returncode != EXIT_CODE_SUCCESS:
+        logger.info(f"Command code-server --list-extensions failed with error: {installed_extensions.stderr}")
+        return []
+
+    return installed_extensions.stdout.splitlines()
+
+
+def is_extension_installed(extension: str, installed_extensions: List[str]) -> bool:
+    return any(installed_extension in extension for installed_extension in installed_extensions)
+
+
 def download_vscode(config: VscodeConfig):
     """
     Download vscode server and extension from remote to local and add the directory of binary executable to $PATH.
@@ -190,34 +212,34 @@
     if executable_path is not None:
         logger.info(f"Code server binary already exists at {executable_path}")
         logger.info("Skipping downloading code server...")
-        return
+    else:
+        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("Code server is not in $PATH, start downloading code server...")
+        logger.info(f"Start downloading files to {DOWNLOAD_DIR}")
+        # Download remote file to local
+        code_server_remote_path = get_code_server_info(config.code_server_remote_paths)
+        code_server_tar_path = download_file(code_server_remote_path, DOWNLOAD_DIR)
 
-    # Create DOWNLOAD_DIR if not exist
-    logger.info(f"DOWNLOAD_DIR: {DOWNLOAD_DIR}")
-    os.makedirs(DOWNLOAD_DIR, exist_ok=True)
+        # Extract the tarball
+        with tarfile.open(code_server_tar_path, "r:gz") as tar:
+            tar.extractall(path=DOWNLOAD_DIR)
 
-    logger.info(f"Start downloading files to {DOWNLOAD_DIR}")
+        code_server_dir_name = get_code_server_info(config.code_server_dir_names)
+        code_server_bin_dir = os.path.join(DOWNLOAD_DIR, code_server_dir_name, "bin")
 
-    # Download remote file to local
-    code_server_remote_path = get_code_server_info(config.code_server_remote_paths)
-    code_server_tar_path = download_file(code_server_remote_path, DOWNLOAD_DIR)
+        # Add the directory of code-server binary to $PATH
+        os.environ["PATH"] = code_server_bin_dir + os.pathsep + os.environ["PATH"]
 
+    # If the extension already exists in the container, skip downloading
+    installed_extensions = get_installed_extensions()
     extension_paths = []
     for extension in config.extension_remote_paths:
-        file_path = download_file(extension, DOWNLOAD_DIR)
-        extension_paths.append(file_path)
-
-    # Extract the tarball
-    with tarfile.open(code_server_tar_path, "r:gz") as tar:
-        tar.extractall(path=DOWNLOAD_DIR)
-
-    code_server_dir_name = get_code_server_info(config.code_server_dir_names)
-    code_server_bin_dir = os.path.join(DOWNLOAD_DIR, code_server_dir_name, "bin")
-
-    # Add the directory of code-server binary to $PATH
-    os.environ["PATH"] = code_server_bin_dir + os.pathsep + os.environ["PATH"]
+        if not is_extension_installed(extension, installed_extensions):
+            file_path = download_file(extension, DOWNLOAD_DIR)
+            extension_paths.append(file_path)
 
     for p in extension_paths:
         logger.info(f"Execute extension installation command to install extension {p}")

plugins/flytekit-flyin/tests/test_flyin_plugin.py

@@ -14,7 +14,12 @@
     jupyter,
     vscode,
 )
-from flytekitplugins.flyin.vscode_lib.decorator import get_code_server_info
+from flytekitplugins.flyin.vscode_lib.constants import EXIT_CODE_SUCCESS
+from flytekitplugins.flyin.vscode_lib.decorator import (
+    get_code_server_info,
+    get_installed_extensions,
+    is_extension_installed,
+)
 
 from flytekit import task, workflow
 from flytekit.configuration import Image, ImageConfig, SerializationSettings
@@ -216,6 +221,20 @@ def wf():
     mock_exit.assert_called_once()
 
 
+def test_is_extension_installed():
+    installed_extensions = [
+        "ms-python.python",
+        "ms-toolsai.jupyter",
+        "ms-toolsai.jupyter-keymap",
+        "ms-toolsai.jupyter-renderers",
+        "ms-toolsai.vscode-jupyter-cell-tags",
+        "ms-toolsai.vscode-jupyter-slideshow",
+    ]
+    config = VscodeConfig()
+    for extension in config.extension_remote_paths:
+        assert is_extension_installed(extension, installed_extensions)
+
+
 def test_vscode_config():
     config = VscodeConfig()
     assert config.code_server_remote_paths == DEFAULT_CODE_SERVER_REMOTE_PATHS
@@ -345,3 +364,58 @@ def test_platform_unsupported(mock_machine, mock_code_server_info_dict):
         match="Automatic download is only supported on AMD64 and ARM64 architectures. If you are using a different architecture, please visit the code-server official website to manually download the appropriate version for your image.",
     ):
         get_code_server_info(mock_code_server_info_dict)
+
+
+@mock.patch("subprocess.run")
+def test_get_installed_extensions_succeeded(mock_run):
+    # Set up the mock process
+    mock_process = mock.Mock()
+    mock_process.returncode = EXIT_CODE_SUCCESS
+    mock_process.stdout = (
+        "ms-python.python\n"
+        "ms-toolsai.jupyter\n"
+        "ms-toolsai.jupyter-keymap\n"
+        "ms-toolsai.jupyter-renderers\n"
+        "ms-toolsai.vscode-jupyter-cell-tags\n"
+        "ms-toolsai.vscode-jupyter-slideshow\n"
+    )
+    mock_run.return_value = mock_process
+
+    installed_extensions = get_installed_extensions()
+
+    # Verify the correct command was called
+    mock_run.assert_called_once_with(["code-server", "--list-extensions"], capture_output=True, text=True)
+
+    # Assert that the output matches the expected list of extensions
+    expected_extensions = [
+        "ms-python.python",
+        "ms-toolsai.jupyter",
+        "ms-toolsai.jupyter-keymap",
+        "ms-toolsai.jupyter-renderers",
+        "ms-toolsai.vscode-jupyter-cell-tags",
+        "ms-toolsai.vscode-jupyter-slideshow",
+    ]
+    assert installed_extensions == expected_extensions
+
+
+@mock.patch("subprocess.run")
+def test_get_installed_extensions_failed(mock_run):
+    # Set up the mock process
+    mock_process = mock.Mock()
+    mock_process.returncode = 1
+    mock_process.stdout = (
+        "ms-python.python\n"
+        "ms-toolsai.jupyter\n"
+        "ms-toolsai.jupyter-keymap\n"
+        "ms-toolsai.jupyter-renderers\n"
+        "ms-toolsai.vscode-jupyter-cell-tags\n"
+        "ms-toolsai.vscode-jupyter-slideshow\n"
+    )
+    mock_run.return_value = mock_process
+
+    installed_extensions = get_installed_extensions()
+
+    mock_run.assert_called_once_with(["code-server", "--list-extensions"], capture_output=True, text=True)
+
+    expected_extensions = []
+    assert installed_extensions == expected_extensions