flyte 2.0.0b32__py3-none-any.whl

flyte/__init__.py

@@ -0,0 +1,108 @@
+"""
+Flyte SDK for authoring compound AI applications, services and workflows.
+"""
+
+from __future__ import annotations
+
+import sys
+
+from ._build import build
+from ._cache import Cache, CachePolicy, CacheRequest
+from ._context import ctx
+from ._custom_context import custom_context, get_custom_context
+from ._deploy import build_images, deploy
+from ._environment import Environment
+from ._excepthook import custom_excepthook
+from ._group import group
+from ._image import Image
+from ._initialize import current_domain, init, init_from_config, init_in_cluster
+from ._logging import logger
+from ._map import map
+from ._pod import PodTemplate
+from ._resources import AMD_GPU, GPU, HABANA_GAUDI, TPU, Device, DeviceClass, Neuron, Resources
+from ._retry import RetryStrategy
+from ._reusable_environment import ReusePolicy
+from ._run import run, with_runcontext
+from ._secret import Secret, SecretRequest
+from ._task_environment import TaskEnvironment
+from ._timeout import Timeout, TimeoutType
+from ._trace import trace
+from ._trigger import Cron, FixedRate, Trigger, TriggerTime
+from ._version import __version__
+
+sys.excepthook = custom_excepthook
+
+
+def _silence_grpc_warnings():
+    """
+    Silences gRPC warnings that can clutter the output.
+    """
+    import os
+
+    # Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
+    # before importing grpc
+    if "GRPC_VERBOSITY" not in os.environ:
+        os.environ["GRPC_VERBOSITY"] = "ERROR"
+        os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
+        # Disable fork support (stops "skipping fork() handlers")
+        os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "0"
+        # Reduce absl/glog verbosity
+        os.environ["GLOG_minloglevel"] = "2"
+        os.environ["ABSL_LOG"] = "0"
+
+
+_silence_grpc_warnings()
+
+
+def version() -> str:
+    """
+    Returns the version of the Flyte SDK.
+    """
+    return __version__
+
+
+__all__ = [
+    "AMD_GPU",
+    "GPU",
+    "HABANA_GAUDI",
+    "TPU",
+    "Cache",
+    "CachePolicy",
+    "CacheRequest",
+    "Cron",
+    "Device",
+    "DeviceClass",
+    "Environment",
+    "FixedRate",
+    "Image",
+    "Neuron",
+    "PodTemplate",
+    "Resources",
+    "RetryStrategy",
+    "ReusePolicy",
+    "Secret",
+    "SecretRequest",
+    "TaskEnvironment",
+    "Timeout",
+    "TimeoutType",
+    "Trigger",
+    "TriggerTime",
+    "__version__",
+    "build",
+    "build_images",
+    "ctx",
+    "current_domain",
+    "custom_context",
+    "deploy",
+    "get_custom_context",
+    "group",
+    "init",
+    "init_from_config",
+    "init_in_cluster",
+    "logger",
+    "map",
+    "run",
+    "trace",
+    "version",
+    "with_runcontext",
+]

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

@@ -0,0 +1,195 @@
+"""
+Flyte runtime module, this is the entrypoint script for the Flyte runtime.
+
+Caution: Startup time for this module is very important, as it is the entrypoint for the Flyte runtime.
+Refrain from importing any modules here. If you need to import any modules, do it inside the main function.
+"""
+
+import asyncio
+import os
+import sys
+from typing import List
+
+import click
+
+from flyte.models import PathRewrite
+
+# Todo: work with pvditt to make these the names
+# ACTION_NAME = "_U_ACTION_NAME"
+# RUN_NAME = "_U_RUN_NAME"
+# PROJECT_NAME = "_U_PROJECT_NAME"
+# DOMAIN_NAME = "_U_DOMAIN_NAME"
+# ORG_NAME = "_U_ORG_NAME"
+
+ACTION_NAME = "ACTION_NAME"
+RUN_NAME = "RUN_NAME"
+PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
+ORG_NAME = "_U_ORG_NAME"
+ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+INSECURE_SKIP_VERIFY_OVERRIDE = "_U_INSECURE_SKIP_VERIFY"
+RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
+FLYTE_ENABLE_VSCODE_KEY = "_F_E_VS"
+
+_UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+_F_PATH_REWRITE = "_F_PATH_REWRITE"
+
+
+@click.group()
+def _pass_through():
+    pass
+
+
+@_pass_through.command("a0")
+@click.option("--inputs", "-i", required=True)
+@click.option("--outputs-path", "-o", required=True)
+@click.option("--version", "-v", required=True)
+@click.option("--run-base-dir", envvar=RUN_OUTPUT_BASE_DIR, required=True)
+@click.option("--raw-data-path", "-r", required=False)
+@click.option("--checkpoint-path", "-c", required=False)
+@click.option("--prev-checkpoint", "-p", required=False)
+@click.option("--name", envvar=ACTION_NAME, required=False)
+@click.option("--run-name", envvar=RUN_NAME, required=False)
+@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)
+@click.option("--dest", required=False)
+@click.option("--resolver", required=False)
+@click.argument(
+    "resolver-args",
+    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,
+    run_base_dir: str,
+    outputs_path: str,
+    raw_data_path: str,
+    checkpoint_path: str,
+    prev_checkpoint: str,
+    tgz: str,
+    pkl: str,
+    dest: str,
+    resolver: str,
+    resolver_args: List[str],
+):
+    sys.path.insert(0, ".")
+
+    import faulthandler
+    import signal
+
+    import flyte
+    import flyte._utils as utils
+    import flyte.errors
+    import flyte.storage as storage
+    from flyte._initialize import init_in_cluster
+    from flyte._internal.controllers import create_controller
+    from flyte._internal.imagebuild.image_builder import ImageCache
+    from flyte._internal.runtime.entrypoints import load_and_run_task
+    from flyte._logging import logger
+    from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+
+    logger.info("Registering faulthandler for SIGUSR1")
+    faulthandler.register(signal.SIGUSR1)
+
+    logger.info(f"Initializing flyte runtime - version {flyte.__version__}")
+    assert org, "Org is required for now"
+    assert project, "Project is required"
+    assert domain, "Domain is required"
+    assert run_name, f"Run name is required {run_name}"
+    assert name, f"Action name is required {name}"
+
+    if run_name.startswith("{{"):
+        run_name = os.getenv("RUN_NAME", "")
+    if name.startswith("{{"):
+        name = os.getenv("ACTION_NAME", "")
+
+    logger.warning(f"Flyte runtime started for action {name} with run name {run_name}")
+
+    if debug and name == "a0":
+        from flyte._debug.vscode import _start_vscode_server
+
+        asyncio.run(_start_vscode_server(ctx))
+
+    controller_kwargs = init_in_cluster(org=org, project=project, domain=domain)
+    bundle = None
+    if tgz or pkl:
+        bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+    # Controller is created with the same kwargs as init, so that it can be used to run tasks
+    controller = create_controller(ct="remote", **controller_kwargs)
+
+    ic = ImageCache.from_transport(image_cache) if image_cache else None
+
+    path_rewrite_cfg = os.getenv(_F_PATH_REWRITE, None)
+    path_rewrite = None
+    if path_rewrite_cfg:
+        potential_path_rewrite = PathRewrite.from_str(path_rewrite_cfg)
+        if storage.exists_sync(potential_path_rewrite.new_prefix):
+            path_rewrite = potential_path_rewrite
+            logger.info(f"Path rewrite configured for {path_rewrite.new_prefix}")
+        else:
+            logger.error(
+                f"Path rewrite failed for path {potential_path_rewrite.new_prefix}, "
+                f"not found, reverting to original path {potential_path_rewrite.old_prefix}"
+            )
+
+    # Create a coroutine to load the task and run it
+    task_coroutine = load_and_run_task(
+        resolver=resolver,
+        resolver_args=resolver_args,
+        action=ActionID(name=name, run_name=run_name, project=project, domain=domain, org=org),
+        raw_data_path=RawDataPath(path=raw_data_path, path_rewrite=path_rewrite),
+        checkpoints=Checkpoints(checkpoint_path, prev_checkpoint),
+        code_bundle=bundle,
+        input_path=inputs,
+        output_path=outputs_path,
+        run_base_dir=run_base_dir,
+        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()
+
+    # Run both coroutines concurrently and wait for first to finish and cancel the other
+    async def _run_and_stop():
+        loop = asyncio.get_event_loop()
+        loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
+        try:
+            await utils.run_coros(controller_failure, task_coroutine)
+            await controller.stop()
+        except flyte.errors.RuntimeSystemError as e:
+            logger.error(f"Runtime system error: {e}")
+            from flyte._internal.runtime.convert import convert_from_native_to_error
+            from flyte._internal.runtime.io import upload_error
+
+            logger.error(f"Flyte runtime failed for action {name} with run name {run_name}, error: {e}")
+            err = convert_from_native_to_error(e)
+            path = await upload_error(err.err, outputs_path)
+            logger.error(f"Run {run_name} Action {name} failed with error: {err}. Uploaded error to {path}")
+            await controller.stop()
+            raise
+
+    asyncio.run(_run_and_stop())
+    logger.warning(f"Flyte runtime completed for action {name} with run name {run_name}")
+
+
+if __name__ == "__main__":
+    _pass_through()

flyte/_bin/serve.py

@@ -0,0 +1,178 @@
+"""
+Flyte runtime serve module. This is used to serve Apps/serving.
+"""
+
+import asyncio
+import logging
+from typing import Tuple
+
+import click
+
+from flyte.models import CodeBundle
+
+logger = logging.getLogger(__name__)
+
+PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
+ORG_NAME = "_U_ORG_NAME"
+_UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+_F_PATH_REWRITE = "_F_PATH_REWRITE"
+ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+
+
+async def sync_inputs(serialized_inputs: str, dest: str) -> Tuple[dict, dict]:
+    """
+    Converts inputs into simple dict of name to value, downloading any files/directories as needed.
+
+    Args:
+        serialized_inputs (str): The serialized inputs string.
+        dest: Destination to download inputs to
+
+    Returns:
+        Tuple[dict, dict]: A tuple containing the output dictionary and the environment variables dictionary.
+        The output dictionary maps input names to their values.
+        The environment variables dictionary maps environment variable names to their values.
+    """
+    import flyte.storage as storage
+    from flyte.app._input import SerializableInputCollection
+
+    user_inputs = SerializableInputCollection.from_transport(serialized_inputs)
+
+    output = {}
+    env_vars = {}
+
+    for input in user_inputs.inputs:
+        if input.download:
+            user_dest = input.dest or dest
+            if input.type == "file":
+                value = await storage.get(input.value, user_dest)
+            elif input.type == "directory":
+                value = await storage.get(input.value, user_dest, recursive=True)
+            else:
+                raise ValueError("Can only download files or directories")
+        else:
+            value = input.value
+
+        output[input.name] = value
+
+        if input.env_var:
+            env_vars[input.env_var] = value
+
+    return output, env_vars
+
+
+async def download_code_inputs(
+    serialized_inputs: str, tgz: str, pkl: str, dest: str, version: str
+) -> Tuple[dict, dict, CodeBundle | None]:
+    from flyte._internal.runtime.entrypoints import download_code_bundle
+
+    user_inputs: dict[str, str] = {}
+    env_vars: dict[str, str] = {}
+    if serialized_inputs and len(serialized_inputs) > 0:
+        user_inputs, env_vars = await sync_inputs(serialized_inputs, dest)
+    code_bundle: CodeBundle | None = None
+    if tgz or pkl:
+        bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+        code_bundle = await download_code_bundle(bundle)
+
+    return user_inputs, env_vars, code_bundle
+
+
+@click.command()
+@click.option("--inputs", "-i", required=False)
+@click.option("--version", required=True)
+@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)
+@click.option("--dest", required=False)
+@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.argument("command", nargs=-1, type=click.UNPROCESSED)
+def main(
+    inputs: str | None,
+    version: str,
+    interactive_mode: bool,
+    image_cache: str,
+    tgz: str,
+    pkl: str,
+    dest: str,
+    command: Tuple[str, ...] | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+    org: str | None = None,
+):
+    import json
+    import os
+    import signal
+    from subprocess import Popen
+
+    from flyte.app._input import RUNTIME_INPUTS_FILE
+
+    logger.info("Starting flyte-serve")
+    # TODO Do we need to init here?
+    # from flyte._initialize import init
+    # remote_kwargs: dict[str, Any] = {"insecure": False}
+    # if api_key := os.getenv(_UNION_EAGER_API_KEY_ENV_VAR):
+    #     logger.info("Using api key from environment")
+    #     remote_kwargs["api_key"] = api_key
+    # else:
+    #     ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+    #     remote_kwargs["endpoint"] = ep
+    #     if "localhost" in ep or "docker" in ep:
+    #         remote_kwargs["insecure"] = True
+    #     logger.debug(f"Using controller endpoint: {ep} with kwargs: {remote_kwargs}")
+    # init(org=org, project=project, domain=domain, image_builder="remote")  # , **remote_kwargs)
+
+    materialized_inputs, env_vars, _code_bundle = asyncio.run(
+        download_code_inputs(
+            serialized_inputs=inputs or "",
+            tgz=tgz or "",
+            pkl=pkl or "",
+            dest=dest or os.getcwd(),
+            version=version,
+        )
+    )
+
+    for key, value in env_vars.items():
+        # set environment variables defined in the AppEnvironment Inputs
+        logger.info(f"Setting environment variable {key}='{value}'")
+        os.environ[key] = value
+
+    inputs_file = os.path.join(os.getcwd(), RUNTIME_INPUTS_FILE)
+    with open(inputs_file, "w") as f:
+        json.dump(materialized_inputs, f)
+
+    os.environ[RUNTIME_INPUTS_FILE] = inputs_file
+
+    if command is None or len(command) == 0:
+        raise ValueError("No command provided to execute")
+
+    command_list = []
+    for arg in command:
+        logger.info(f"Processing arg: {arg}")
+        if arg.startswith("$"):
+            # expand environment variables in the user-defined command
+            val = os.getenv(arg[1:])
+            if val is None:
+                raise ValueError(f"Environment variable {arg[1:]} not found")
+            logger.info(f"Found env var {arg}.")
+            command_list.append(val)
+        else:
+            command_list.append(arg)
+
+    command_joined = " ".join(command_list)
+    logger.info(f"Serving command: {command_joined}")
+    p = Popen(command_joined, env=os.environ, shell=True)
+
+    def handle_sigterm(signum, frame):
+        p.send_signal(signum)
+
+    signal.signal(signal.SIGTERM, handle_sigterm)
+    returncode = p.wait()
+    exit(returncode)
+
+
+if __name__ == "__main__":
+    main()

flyte/_build.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from flyte.syncify import syncify
+
+from ._image import Image
+
+
+@syncify
+async def build(image: Image) -> str:
+    """
+    Build an image. The existing async context will be used.
+
+    Example:
+    ```
+    import flyte
+    image = flyte.Image("example_image")
+    if __name__ == "__main__":
+        asyncio.run(flyte.build.aio(image))
+    ```
+
+    :param image: The image(s) to build.
+    :return: The image URI.
+    """
+    from flyte._internal.imagebuild.image_builder import ImageBuildEngine
+
+    return await ImageBuildEngine.build(image)

flyte/_cache/__init__.py

@@ -0,0 +1,12 @@
+from .cache import Cache, CacheBehavior, CachePolicy, CacheRequest
+from .defaults import get_default_policies
+from .policy_function_body import FunctionBodyPolicy
+
+__all__ = [
+    "Cache",
+    "CacheBehavior",
+    "CachePolicy",
+    "CacheRequest",
+    "FunctionBodyPolicy",
+    "get_default_policies",
+]