flyte 2.0.0b13__py3-none-any.whl → 2.0.0b15__py3-none-any.whl

flyte/_bin/debug.py

@@ -0,0 +1,38 @@
+import click
+
+
+@click.group()
+def _debug():
+    """Debug commands for Flyte."""
+
+
+@_debug.command("resume")
+@click.option("--pid", "-m", type=int, required=True, help="PID of the vscode server.")
+def resume(pid):
+    """
+    Resume a Flyte task for debugging purposes.
+
+    Args:
+        pid (int): PID of the vscode server.
+    """
+    import os
+    import signal
+
+    print("Terminating server and resuming task.")
+    answer = (
+        input(
+            "This operation will kill the server. All unsaved data will be lost,"
+            " and you will no longer be able to connect to it. Do you really want to terminate? (Y/N): "
+        )
+        .strip()
+        .upper()
+    )
+    if answer == "Y":
+        os.kill(pid, signal.SIGTERM)
+        print("The server has been terminated and the task has been resumed.")
+    else:
+        print("Operation canceled.")
+
+
+if __name__ == "__main__":
+    _debug()

flyte/_bin/runtime.py

@@ -26,6 +26,7 @@ DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
 ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
+FLYTE_ENABLE_VSCODE_KEY = "_F_E_VS"
 
 # TODO: Remove this after proper auth is implemented
 _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
@@ -49,6 +50,8 @@ def _pass_through():
 @click.option("--project", envvar=PROJECT_NAME, required=False)
 @click.option("--domain", envvar=DOMAIN_NAME, required=False)
 @click.option("--org", envvar=ORG_NAME, required=False)
+@click.option("--debug", envvar=FLYTE_ENABLE_VSCODE_KEY, type=click.BOOL, required=False)
+@click.option("--interactive-mode", type=click.BOOL, required=False)
 @click.option("--image-cache", required=False)
 @click.option("--tgz", required=False)
 @click.option("--pkl", required=False)
@@ -59,12 +62,16 @@ def _pass_through():
     type=click.UNPROCESSED,
     nargs=-1,
 )
+@click.pass_context
 def main(
+    ctx: click.Context,
     run_name: str,
     name: str,
     project: str,
     domain: str,
     org: str,
+    debug: bool,
+    interactive_mode: bool,
     image_cache: str,
     version: str,
     inputs: str,
@@ -109,6 +116,11 @@ def main(
     if name.startswith("{{"):
         name = os.getenv("ACTION_NAME", "")
 
+    if debug and name == "a0":
+        from flyte._debug.vscode import _start_vscode_server
+
+        asyncio.run(_start_vscode_server(ctx))
+
     # Figure out how to connect
     # This detection of api key is a hack for now.
     controller_kwargs: dict[str, Any] = {"insecure": False}
@@ -143,6 +155,7 @@ def main(
         version=version,
         controller=controller,
         image_cache=ic,
+        interactive_mode=interactive_mode or debug,
     )
     # Create a coroutine to watch for errors
     controller_failure = controller.watch_for_errors()

flyte/_code_bundle/_utils.py

@@ -240,6 +240,8 @@ def list_imported_modules_as_files(source_path: str, modules: List[ModuleType])
 
         if not _file_is_in_directory(mod_file, source_path):
             # Only upload files where the module file in the source directory
+            # print log line for files that have common ancestor with source_path, but not in it.
+            logger.debug(f"{mod_file} is not in {source_path}")
             continue
 
         files.append(mod_file)

flyte/_code_bundle/bundle.py

@@ -178,9 +178,10 @@ async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
     # TODO make storage apis better to accept pathlib.Path
     if bundle.tgz:
         downloaded_bundle = dest / os.path.basename(bundle.tgz)
+        if downloaded_bundle.exists():
+            return downloaded_bundle.absolute()
         # Download the tgz file
-        path = await storage.get(bundle.tgz, str(downloaded_bundle.absolute()))
-        downloaded_bundle = pathlib.Path(path)
+        await storage.get(bundle.tgz, str(downloaded_bundle.absolute()))
         # NOTE the os.path.join(destination, ''). This is to ensure that the given path is in fact a directory and all
         # downloaded data should be copied into this directory. We do this to account for a difference in behavior in
         # fsspec, which requires a trailing slash in case of pre-existing directory.
@@ -204,8 +205,7 @@ async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
 
         downloaded_bundle = dest / os.path.basename(bundle.pkl)
         # Download the tgz file
-        path = await storage.get(bundle.pkl, str(downloaded_bundle.absolute()))
-        downloaded_bundle = pathlib.Path(path)
+        await storage.get(bundle.pkl, str(downloaded_bundle.absolute()))
         return downloaded_bundle.absolute()
     else:
         raise ValueError("Code bundle should be either tgz or pkl, found neither.")

flyte/_debug/constants.py

@@ -0,0 +1,39 @@
+import os
+from pathlib import Path
+
+# Where the code-server tar and plugins are downloaded to
+EXECUTABLE_NAME = "code-server"
+DOWNLOAD_DIR = Path.home() / ".code-server"
+HOURS_TO_SECONDS = 60 * 60
+DEFAULT_UP_SECONDS = 10 * HOURS_TO_SECONDS  # 10 hours
+DEFAULT_CODE_SERVER_REMOTE_PATHS = {
+    "amd64": "https://github.com/coder/code-server/releases/download/v4.18.0/code-server-4.18.0-linux-amd64.tar.gz",
+    "arm64": "https://github.com/coder/code-server/releases/download/v4.18.0/code-server-4.18.0-linux-arm64.tar.gz",
+}
+DEFAULT_CODE_SERVER_EXTENSIONS = [
+    "https://raw.githubusercontent.com/flyteorg/flytetools/master/flytekitplugins/flyin/ms-python.python-2023.20.0.vsix",
+    "https://raw.githubusercontent.com/flyteorg/flytetools/master/flytekitplugins/flyin/ms-toolsai.jupyter-2023.9.100.vsix",
+]
+
+# Duration to pause the checking of the heartbeat file until the next one
+HEARTBEAT_CHECK_SECONDS = 60
+MAX_IDLE_SECONDS = 180
+
+# The path is hardcoded by code-server
+# https://coder.com/docs/code-server/latest/FAQ#what-is-the-heartbeat-file
+HEARTBEAT_PATH = os.path.expanduser("~/.local/share/code-server/heartbeat")
+
+INTERACTIVE_DEBUGGING_FILE_NAME = "flyteinteractive_interactive_entrypoint.py"
+RESUME_TASK_FILE_NAME = "flyteinteractive_resume_task.py"
+# Config keys to store in task template
+VSCODE_TYPE_KEY = "flyteinteractive_type"
+VSCODE_PORT_KEY = "flyteinteractive_port"
+
+TASK_FUNCTION_SOURCE_PATH = "TASK_FUNCTION_SOURCE_PATH"
+
+# Default max idle seconds to terminate the flyteinteractive server
+HOURS_TO_SECONDS = 60 * 60
+MAX_IDLE_SECONDS = 10 * HOURS_TO_SECONDS  # 10 hours
+
+# Subprocess constants
+EXIT_CODE_SUCCESS = 0

flyte/_debug/utils.py

@@ -0,0 +1,17 @@
+import asyncio
+
+from flyte._debug.constants import EXIT_CODE_SUCCESS
+from flyte._logging import logger
+
+
+async def execute_command(cmd: str):
+    """
+    Execute a command in the shell.
+    """
+    process = await asyncio.create_subprocess_shell(cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
+    logger.info(f"cmd: {cmd}")
+    stdout, stderr = await process.communicate()
+    if process.returncode != EXIT_CODE_SUCCESS:
+        raise RuntimeError(f"Command {cmd} failed with error: {stderr!r}")
+    logger.info(f"stdout: {stdout!r}")
+    logger.info(f"stderr: {stderr!r}")

flyte/_debug/vscode.py

@@ -0,0 +1,300 @@
+import asyncio
+import json
+import multiprocessing
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import tarfile
+import time
+from typing import List
+
+import aiofiles
+import click
+import httpx
+
+from flyte import storage
+from flyte._debug.constants import (
+    DEFAULT_CODE_SERVER_EXTENSIONS,
+    DEFAULT_CODE_SERVER_REMOTE_PATHS,
+    DOWNLOAD_DIR,
+    EXECUTABLE_NAME,
+    EXIT_CODE_SUCCESS,
+    HEARTBEAT_PATH,
+    MAX_IDLE_SECONDS,
+)
+from flyte._debug.utils import (
+    execute_command,
+)
+from flyte._internal.runtime.rusty import download_tgz
+from flyte._logging import logger
+
+
+async def download_file(url: str, target_dir: str) -> str:
+    """
+    Downloads a file from a given URL using HTTPX and saves it locally.
+
+    Args:
+        url (str): The URL of the file to download.
+        target_dir (str): The directory where the file should be saved. Defaults to current directory.
+    """
+    try:
+        filename = os.path.join(target_dir, os.path.basename(url))
+        if url.startswith("http"):
+            response = httpx.get(url, follow_redirects=True)
+            response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
+            async with aiofiles.open(filename, "wb") as f:
+                await f.write(response.content)
+        else:
+            await storage.get(url, filename)
+        logger.info(f"File '{filename}' downloaded successfully from '{url}'.")
+        return filename
+
+    except httpx.RequestError as e:
+        raise RuntimeError(f"An error occurred while requesting '{url}': {e}")
+    except httpx.HTTPStatusError as e:
+        raise RuntimeError(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
+    except Exception as e:
+        raise RuntimeError(f"An unexpected error occurred: {e}")
+
+
+def get_default_extensions() -> List[str]:
+    extensions = os.getenv("_F_CS_E")
+    if extensions is not None:
+        return extensions.split(",")
+    return DEFAULT_CODE_SERVER_EXTENSIONS
+
+
+def get_code_server_info() -> str:
+    """
+    Returns the code server information based on the system's architecture.
+
+    This function checks the system's architecture and returns the corresponding
+    code server information from the provided dictionary. The function currently
+    supports AMD64 and ARM64 architectures.
+
+    Returns:
+        str: The code server information corresponding to the system's architecture.
+
+    Raises:
+        ValueError: If the system's architecture is not AMD64 or ARM64.
+    """
+    code_server_path = os.getenv("_F_CS_RP")
+    if code_server_path is not None:
+        return code_server_path
+
+    machine_info = platform.machine()
+    logger.info(f"machine type: {machine_info}")
+    code_server_info_dict = DEFAULT_CODE_SERVER_REMOTE_PATHS
+
+    if "aarch64" == machine_info:
+        return code_server_info_dict["arm64"]
+    elif "x86_64" == machine_info:
+        return code_server_info_dict["amd64"]
+    else:
+        raise ValueError(
+            "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."
+        )
+
+
+def get_installed_extensions() -> List[str]:
+    """
+    Get the list of installed extensions.
+
+    Returns:
+        List[str]: The list of installed extensions.
+    """
+    installed_extensions = subprocess.run(
+        ["code-server", "--list-extensions"], check=False, 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)
+
+
+async def download_vscode():
+    """
+    Download vscode server and extension from remote to local and add the directory of binary executable to $PATH.
+    """
+    # If the code server already exists in the container, skip downloading
+    executable_path = shutil.which(EXECUTABLE_NAME)
+    if executable_path is not None or os.path.exists(DOWNLOAD_DIR):
+        logger.info(f"Code server binary already exists at {executable_path}")
+        logger.info("Skipping downloading code server...")
+    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)
+
+        logger.info(f"Start downloading files to {DOWNLOAD_DIR}")
+        # Download remote file to local
+        code_server_remote_path = get_code_server_info()
+        code_server_tar_path = await download_file(code_server_remote_path, str(DOWNLOAD_DIR))
+
+        # Extract the tarball
+        with tarfile.open(code_server_tar_path, "r:gz") as tar:
+            tar.extractall(path=DOWNLOAD_DIR)
+
+    if os.path.exists(DOWNLOAD_DIR):
+        code_server_dir_name = os.path.basename(get_code_server_info()).removesuffix(".tar.gz")
+        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 the extension already exists in the container, skip downloading
+    installed_extensions = get_installed_extensions()
+    coros = []
+
+    for extension in get_default_extensions():
+        if not is_extension_installed(extension, installed_extensions):
+            coros.append(download_file(extension, str(DOWNLOAD_DIR)))
+    extension_paths = await asyncio.gather(*coros)
+
+    coros = []
+    for p in extension_paths:
+        logger.info(f"Execute extension installation command to install extension {p}")
+        coros.append(execute_command(f"code-server --install-extension {p}"))
+
+    await asyncio.gather(*coros)
+
+
+def prepare_launch_json(ctx: click.Context, pid: int):
+    """
+    Generate the launch.json and settings.json for users to easily launch interactive debugging and task resumption.
+    """
+
+    virtual_venv = os.getenv("VIRTUAL_ENV")
+    if virtual_venv is None:
+        raise RuntimeError("VIRTUAL_ENV is not found in environment variables.")
+
+    run_name = ctx.params["run_name"]
+    name = ctx.params["name"]
+    # TODO: Executor should pass correct name.
+    if run_name.startswith("{{"):
+        run_name = os.getenv("RUN_NAME", "")
+    if name.startswith("{{"):
+        name = os.getenv("ACTION_NAME", "")
+
+    launch_json = {
+        "version": "0.2.0",
+        "configurations": [
+            {
+                "name": "Interactive Debugging",
+                "type": "python",
+                "request": "launch",
+                "program": f"{virtual_venv}/bin/runtime.py",
+                "console": "integratedTerminal",
+                "justMyCode": True,
+                "args": [
+                    "a0",
+                    "--inputs",
+                    ctx.params["inputs"],
+                    "--outputs-path",
+                    ctx.params["outputs_path"],
+                    "--version",
+                    ctx.params["version"],
+                    "--run-base-dir",
+                    ctx.params["run_base_dir"],
+                    "--name",
+                    name,
+                    "--run-name",
+                    run_name,
+                    "--project",
+                    ctx.params["project"],
+                    "--domain",
+                    ctx.params["domain"],
+                    "--org",
+                    ctx.params["org"],
+                    "--image-cache",
+                    ctx.params["image_cache"],
+                    "--debug",
+                    "False",
+                    "--interactive-mode",
+                    "True",
+                    "--tgz",
+                    ctx.params["tgz"],
+                    "--dest",
+                    ctx.params["dest"],
+                    "--resolver",
+                    ctx.params["resolver"],
+                    *ctx.params["resolver_args"],
+                ],
+            },
+            {
+                "name": "Resume Task",
+                "type": "python",
+                "request": "launch",
+                "program": f"{virtual_venv}/bin/debug.py",
+                "console": "integratedTerminal",
+                "justMyCode": True,
+                "args": ["resume", "--pid", str(pid)],
+            },
+        ],
+    }
+
+    vscode_directory = os.path.join(os.getcwd(), ".vscode")
+    if not os.path.exists(vscode_directory):
+        os.makedirs(vscode_directory)
+
+    with open(os.path.join(vscode_directory, "launch.json"), "w") as file:
+        json.dump(launch_json, file, indent=4)
+
+    settings_json = {"python.defaultInterpreterPath": sys.executable}
+    with open(os.path.join(vscode_directory, "settings.json"), "w") as file:
+        json.dump(settings_json, file, indent=4)
+
+
+async def _start_vscode_server(ctx: click.Context):
+    await asyncio.gather(download_tgz(ctx.params["dest"], ctx.params["version"], ctx.params["tgz"]), download_vscode())
+    child_process = multiprocessing.Process(
+        target=lambda cmd: asyncio.run(asyncio.run(execute_command(cmd))),
+        kwargs={"cmd": f"code-server --bind-addr 0.0.0.0:8080 --disable-workspace-trust --auth none {os.getcwd()}"},
+    )
+    child_process.start()
+    if child_process.pid is None:
+        raise RuntimeError("Failed to start vscode server.")
+    prepare_launch_json(ctx, child_process.pid)
+
+    start_time = time.time()
+    check_interval = 60  # Interval for heartbeat checking in seconds
+    last_heartbeat_check = time.time() - check_interval
+
+    def terminate_process():
+        if child_process.is_alive():
+            child_process.terminate()
+        child_process.join()
+
+    logger.info("waiting for task to resume...")
+    while child_process.is_alive():
+        current_time = time.time()
+        if current_time - last_heartbeat_check >= check_interval:
+            last_heartbeat_check = current_time
+            if not os.path.exists(HEARTBEAT_PATH):
+                delta = current_time - start_time
+                logger.info(f"Code server has not been connected since {delta} seconds ago.")
+                logger.info("Please open the browser to connect to the running server.")
+            else:
+                delta = current_time - os.path.getmtime(HEARTBEAT_PATH)
+                logger.info(f"The latest activity on code server is {delta} seconds ago.")
+
+            # If the time from last connection is longer than max idle seconds, terminate the vscode server.
+            if delta > MAX_IDLE_SECONDS:
+                logger.info(f"VSCode server is idle for more than {MAX_IDLE_SECONDS} seconds. Terminating...")
+                terminate_process()
+                sys.exit()
+
+        await asyncio.sleep(1)
+
+    logger.info("User has resumed the task.")
+    terminate_process()
+    return

flyte/_image.py

@@ -252,9 +252,15 @@ class AptPackages(Layer):
 @dataclass(frozen=True, repr=True)
 class Commands(Layer):
     commands: Tuple[str, ...]
+    secret_mounts: Optional[Tuple[str | Secret, ...]] = None
 
     def update_hash(self, hasher: hashlib._Hash):
-        hasher.update("".join(self.commands).encode("utf-8"))
+        hash_input = "".join(self.commands)
+
+        if self.secret_mounts:
+            for secret_mount in self.secret_mounts:
+                hash_input += str(secret_mount)
+        hasher.update(hash_input.encode("utf-8"))
 
 
 @rich.repr.auto
@@ -782,9 +788,24 @@ class Image:
 
         Example:
         ```python
-        @flyte.task(image=(flyte.Image
-                        .ubuntu_python()
-                        .with_pip_packages("requests", "numpy")))
+        @flyte.task(image=(flyte.Image.from_debian_base().with_pip_packages("requests", "numpy")))
+        def my_task(x: int) -> int:
+            import numpy as np
+            return np.sum([x, 1])
+        ```
+
+        To mount secrets during the build process to download private packages, you can use the `secret_mounts`.
+        In the below example, "GITHUB_PAT" will be mounted as env var "GITHUB_PAT",
+         and "apt-secret" will be mounted at /etc/apt/apt-secret.
+        Example:
+        ```python
+        private_package = "git+https://$GITHUB_PAT@github.com/flyteorg/flytex.git@2e20a2acebfc3877d84af643fdd768edea41d533"
+        @flyte.task(
+            image=(
+                flyte.Image.from_debian_base()
+                .with_pip_packages("private_package", secret_mounts=[Secret(key="GITHUB_PAT")])
+                .with_apt_packages("git", secret_mounts=[Secret(key="apt-secret", mount="/etc/apt/apt-secret")])
+        )
         def my_task(x: int) -> int:
             import numpy as np
             return np.sum([x, 1])
@@ -909,16 +930,21 @@ class Image:
         )
         return new_image
 
-    def with_commands(self, commands: List[str]) -> Image:
+    def with_commands(self, commands: List[str], secret_mounts: Optional[SecretRequest] = None) -> Image:
         """
         Use this method to create a new image with the specified commands layered on top of the current image
         Be sure not to use RUN in your command.
 
         :param commands: list of commands to run
+        :param secret_mounts: list of secret mounts to use for the build process.
         :return: Image
         """
         new_commands: Tuple = _ensure_tuple(commands)
-        new_image = self.clone(addl_layer=Commands(commands=new_commands))
+        new_image = self.clone(
+            addl_layer=Commands(
+                commands=new_commands, secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None
+            )
+        )
         return new_image
 
     def with_local_v2(self) -> Image:

flyte/_initialize.py

@@ -205,6 +205,14 @@ async def init(
                 http_proxy_url=http_proxy_url,
             )
 
+        if not root_dir:
+            editable_root = get_cwd_editable_install()
+            if editable_root:
+                logger.info(f"Using editable install as root directory: {editable_root}")
+                root_dir = editable_root
+            else:
+                logger.info("No editable install found, using current working directory as root directory.")
+                root_dir = Path.cwd()
         root_dir = root_dir or get_cwd_editable_install() or Path.cwd()
         _init_config = _InitConfig(
             root_dir=root_dir,
@@ -242,17 +250,19 @@ async def init_from_config(
     cfg: config.Config
     if path_or_config is None or isinstance(path_or_config, str):
         # If a string is passed, treat it as a path to the config file
-        if path_or_config:
+        if root_dir and path_or_config:
+            cfg = config.auto(str(root_dir / path_or_config))
+        elif path_or_config:
             if not Path(path_or_config).exists():
                 raise InitializationError(
                     "ConfigFileNotFoundError",
                     "user",
                     f"Configuration file '{path_or_config}' does not exist., current working directory is {Path.cwd()}",
                 )
-        if root_dir and path_or_config:
-            cfg = config.auto(str(root_dir / path_or_config))
-        else:
             cfg = config.auto(path_or_config)
+        else:
+            # If no path is provided, use the default config file
+            cfg = config.auto()
     else:
         # If a Config object is passed, use it directly
         cfg = path_or_config
@@ -374,30 +384,6 @@ def ensure_client():
         )
 
 
-def requires_client(func: T) -> T:
-    """
-    Decorator that checks if the client has been initialized before executing the function.
-    Raises InitializationError if the client is not initialized.
-
-    :param func: Function to decorate
-    :return: Decorated function that checks for initialization
-    """
-
-    @functools.wraps(func)
-    async def wrapper(*args, **kwargs) -> T:
-        init_config = _get_init_config()
-        if init_config is None or init_config.client is None:
-            raise InitializationError(
-                "ClientNotInitializedError",
-                "user",
-                f"Function '{func.__name__}' requires client to be initialized. "
-                f"Call flyte.init() with a valid endpoint or api-key before using this function.",
-            )
-        return func(*args, **kwargs)
-
-    return typing.cast(T, wrapper)
-
-
 def requires_storage(func: T) -> T:
     """
     Decorator that checks if the storage has been initialized before executing the function.

flyte/_internal/controllers/remote/_action.py

@@ -184,7 +184,7 @@ class Action:
         et.nanos = int((end_time % 1) * 1e9)
 
         spec = (
-            task_definition_pb2.TaskSpec(task_template=tasks_pb2.TaskTemplate(interface=typed_interface))
+            task_definition_pb2.TraceSpec(interface=typed_interface)
             if typed_interface
             else None
         )