flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_custom_context.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+from flyte._context import ctx
+
+from ._context import internal_ctx
+
+
+def get_custom_context() -> dict[str, str]:
+    """
+    Get the current input context. This can be used within a task to retrieve
+    context metadata that was passed to the action.
+
+    Context will automatically propagate to sub-actions.
+
+    Example:
+    ```python
+    import flyte
+
+    env = flyte.TaskEnvironment(name="...")
+
+    @env.task
+    def t1():
+        # context can be retrieved with `get_custom_context`
+        ctx = flyte.get_custom_context()
+        print(ctx)  # {'project': '...', 'entity': '...'}
+    ```
+
+    :return: Dictionary of context key-value pairs
+    """
+    tctx = ctx()
+    if tctx is None or tctx.custom_context is None:
+        return {}
+    return tctx.custom_context
+
+
+@contextmanager
+def custom_context(**context: str):
+    """
+    Synchronous context manager to set input context for tasks spawned within this block.
+
+    Example:
+    ```python
+    import flyte
+
+    env = flyte.TaskEnvironment(name="...")
+
+    @env.task
+    def t1():
+        ctx = flyte.get_custom_context()
+        print(ctx)
+
+    @env.task
+    def main():
+        # context can be passed via a context manager
+        with flyte.custom_context(project="my-project"):
+            t1()  # will have {'project': 'my-project'} as context
+    ```
+
+    :param context: Key-value pairs to set as input context
+    """
+    ctx = internal_ctx()
+    if ctx.data.task_context is None:
+        yield
+        return
+
+    tctx = ctx.data.task_context
+    new_tctx = tctx.replace(custom_context={**tctx.custom_context, **context})
+
+    with ctx.replace_task_context(new_tctx):
+        yield
+        # Exit the context and restore the previous context

flyte/_debug/constants.py

@@ -0,0 +1,37 @@
+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.106.3/code-server-4.106.3-linux-amd64.tar.gz",
+    "arm64": "https://github.com/coder/code-server/releases/download/v4.106.3/code-server-4.106.3-linux-arm64.tar.gz",
+}
+DEFAULT_CODE_SERVER_EXTENSIONS = [
+    "https://raw.githubusercontent.com/flyteorg/flytetools/master/flytekitplugins/flyin/ms-python.python-2023.20.0.vsix",
+]
+
+# Duration to pause the checking of the heartbeat file until the next one
+HEARTBEAT_CHECK_SECONDS = 60
+
+# 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,315 @@
+import asyncio
+import json
+import multiprocessing
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import tarfile
+import time
+from pathlib import Path
+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", str(Path(sys.executable).parent.parent))
+    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,
+        "remote.autoForwardPorts": False,
+        "remote.autoForwardPortsFallback": 0,
+    }
+    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):
+    if ctx.params["tgz"] is None:
+        await download_vscode()
+    else:
+        await asyncio.gather(
+            download_tgz(ctx.params["dest"], ctx.params["version"], ctx.params["tgz"]), download_vscode()
+        )
+    code_server_idle_timeout_seconds = os.getenv("CODE_SERVER_IDLE_TIMEOUT_SECONDS", str(MAX_IDLE_SECONDS))
+    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:6060 --idle-timeout-seconds {code_server_idle_timeout_seconds}"
+            f" --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