flyte 2.0.0b32__py3-none-any.whl

flyte/app/extras/_fastapi.py

@@ -0,0 +1,151 @@
+import importlib.util
+import inspect
+import pathlib
+import sys
+from dataclasses import dataclass
+from typing import Tuple
+
+import rich.repr
+
+import flyte.app
+from flyte.models import SerializationContext
+
+try:
+    import fastapi
+except ModuleNotFoundError:
+    raise ModuleNotFoundError("fastapi is not installed. Please install the 'fastapi', to use FastAPI apps.")
+
+
+def _extract_fastapi_app_module_and_var(
+    app: fastapi.FastAPI, caller_frame: inspect.FrameInfo | None, serialization_context: SerializationContext
+) -> Tuple[str, str]:
+    """
+    Extract the module name and variable name for a FastAPI app instance.
+
+    This function solves the challenge that `inspect.getmodule(app)` returns the
+    `fastapi.applications` module (where the FastAPI class is defined) rather than
+    the user's module (where the `app` variable is defined). Since FastAPI apps are
+    instances rather than classes or functions, we need special handling to locate
+    the correct module and variable name.
+
+    The function uses the caller frame (captured when FastAPIAppEnvironment was
+    instantiated) to determine which file contains the app definition, then inspects
+    that module's globals to find FastAPI instances.
+
+    Args:
+        app: The FastAPI application instance to locate.
+        caller_frame: Frame information from where FastAPIAppEnvironment was instantiated.
+                     If None, falls back to extract_obj_module (which may not work correctly).
+        serialization_context: Context containing the root directory for calculating
+                              relative module paths.
+
+    Returns:
+        A tuple of (module_name, variable_name) where:
+        - module_name: Dotted module path (e.g., "examples.apps.single_script_fastapi")
+        - variable_name: The name of the variable holding the FastAPI app (e.g., "app")
+
+    Raises:
+        RuntimeError: If the module cannot be loaded or the app variable cannot be found.
+
+    Example:
+        >>> frame = inspect.getframeinfo(inspect.currentframe().f_back)
+        >>> module_name, var_name = _extract_fastapi_app_module_and_var(
+        ...     app, frame, serialization_context
+        ... )
+        >>> # Returns: ("examples.apps.my_app", "app")
+        >>> # Can be used as: uvicorn examples.apps.my_app:app
+    """
+    if caller_frame is None:
+        raise RuntimeError("Caller frame cannot be None")
+
+    # Get the file path where the app was defined
+    file_path = pathlib.Path(caller_frame.filename)
+
+    # Calculate module name relative to source_dir
+    try:
+        relative_path = file_path.relative_to(serialization_context.root_dir or pathlib.Path("."))
+        module_name = pathlib.Path(relative_path).with_suffix("").as_posix().replace("/", ".")
+    except ValueError:
+        # File is not relative to source_dir, use the stem
+        module_name = file_path.stem
+
+    # Instead of reloading the module, inspect the caller frame's local variables
+    # The app variable should be in the frame's globals
+    caller_globals = None
+
+    # Try to get globals from the main module if it matches our file
+    if hasattr(sys.modules.get("__main__"), "__file__"):
+        main_file = pathlib.Path(sys.modules["__main__"].__file__ or ".").resolve()
+        if main_file == file_path.resolve():
+            caller_globals = sys.modules["__main__"].__dict__
+
+    if caller_globals is None:
+        # Load the module to inspect it
+        spec = importlib.util.spec_from_file_location(file_path.stem, file_path)
+        if spec is None or spec.loader is None:
+            raise RuntimeError(f"Could not load module from {file_path}")
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        caller_globals = module.__dict__
+
+    # Extract variable name from module - look for FastAPI instances
+    app_var_name = None
+    for var_name, obj in caller_globals.items():
+        if isinstance(obj, fastapi.FastAPI):
+            # Found a FastAPI app - this is likely the one we want
+            # Store the first one we find
+            if app_var_name is None:
+                app_var_name = var_name
+            # If the objects match by identity, use this one
+            if obj is app:
+                app_var_name = var_name
+                break
+
+    if app_var_name is None:
+        raise RuntimeError("Could not find variable name for FastAPI app in module")
+
+    return module_name, app_var_name
+
+
+@rich.repr.auto
+@dataclass(kw_only=True, repr=True)
+class FastAPIAppEnvironment(flyte.app.AppEnvironment):
+    app: fastapi.FastAPI
+    type: str = "FastAPI"
+    _caller_frame: inspect.FrameInfo | None = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.app is None:
+            raise ValueError("app cannot be None for FastAPIAppEnvironment")
+        if not isinstance(self.app, fastapi.FastAPI):
+            raise TypeError(f"app must be of type fastapi.FastAPI, got {type(self.app)}")
+
+        # Capture the frame where this environment was instantiated
+        # This helps us find the module where the app variable is defined
+        frame = inspect.currentframe()
+        if frame and frame.f_back:
+            # Go up the call stack to find the user's module
+            # Skip the dataclass __init__ frame
+            caller_frame = frame.f_back
+            if caller_frame and caller_frame.f_back:
+                self._caller_frame = inspect.getframeinfo(caller_frame.f_back)
+
+    def container_args(self, serialization_context: SerializationContext) -> list[str]:
+        """
+        Generate the container arguments for running the FastAPI app with uvicorn.
+
+        Returns:
+            A list of command arguments in the format:
+            ["uvicorn", "<module_name>:<app_var_name>", "--port", "<port>"]
+        """
+        module_name, app_var_name = _extract_fastapi_app_module_and_var(
+            self.app, self._caller_frame, serialization_context
+        )
+
+        p = self.port
+        assert isinstance(p, flyte.app.Port)
+        return ["uvicorn", f"{module_name}:{app_var_name}", "--port", str(p.port)]
+
+    def container_command(self, serialization_context: SerializationContext) -> list[str]:
+        return []

flyte/cli/__init__.py

@@ -0,0 +1,12 @@
+import os
+
+from flyte.cli.main import main
+
+__all__ = ["main"]
+
+
+# Set GRPC_VERBOSITY to NONE if not already set to silence unwanted output
+# This addresses the issue with grpcio >=1.68.0 causing unwanted output
+# https://github.com/flyteorg/flyte/issues/6082
+if "GRPC_VERBOSITY" not in os.environ:
+    os.environ["GRPC_VERBOSITY"] = "NONE"

flyte/cli/_abort.py

@@ -0,0 +1,28 @@
+import rich_click as click
+
+from flyte.cli import _common as common
+
+
+@click.group(name="abort")
+def abort():
+    """
+    Abort an ongoing process.
+    """
+
+
+@abort.command(cls=common.CommandBase)
+@click.argument("run-name", type=str, required=True)
+@click.pass_obj
+def run(cfg: common.CLIConfig, run_name: str, project: str | None = None, domain: str | None = None):
+    """
+    Abort a run.
+    """
+    from flyte.remote import Run
+
+    cfg.init(project=project, domain=domain)
+    r = Run.get(name=run_name)
+    if r:
+        console = common.get_console()
+        with console.status(f"Aborting run '{run_name}'...", spinner="dots"):
+            r.abort()
+        console.print(f"Run '{run_name}' has been aborted.")

flyte/cli/_build.py

@@ -0,0 +1,114 @@
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Dict, List, cast
+
+import rich_click as click
+
+import flyte
+
+from . import _common as common
+from ._common import CLIConfig
+
+
+@dataclass
+class BuildArguments:
+    noop: bool = field(
+        default=False,
+        metadata={
+            "click.option": click.Option(
+                ["--noop"],
+                type=bool,
+                help="Dummy parameter, placeholder for future use. Does not affect the build process.",
+            )
+        },
+    )
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "BuildArguments":
+        return cls(**d)
+
+    @classmethod
+    def options(cls) -> List[click.Option]:
+        """
+        Return the set of base parameters added to every flyte run workflow subcommand.
+        """
+        return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
+
+
+class BuildEnvCommand(click.Command):
+    def __init__(self, obj_name: str, obj: Any, build_args: BuildArguments, *args, **kwargs):
+        self.obj_name = obj_name
+        self.obj = obj
+        self.build_args = build_args
+        super().__init__(*args, **kwargs)
+
+    def invoke(self, ctx: click.Context):
+        console = common.get_console()
+        console.print(f"Building Environment: {self.obj_name}")
+        obj: CLIConfig = ctx.obj
+        obj.init()
+        with console.status("Building...", spinner="dots"):
+            image_cache = flyte.build_images(self.obj)
+
+        console.print(common.format("Images", image_cache.repr(), obj.output_format))
+
+
+class EnvPerFileGroup(common.ObjectsPerFileGroup):
+    """
+    Group that creates a command for each task in the current directory that is not `__init__.py`.
+    """
+
+    def __init__(self, filename: Path, build_args: BuildArguments, *args, **kwargs):
+        args = (filename, *args)
+        super().__init__(*args, **kwargs)
+        self.build_args = build_args
+
+    def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
+        return {k: v for k, v in module.__dict__.items() if isinstance(v, flyte.Environment)}
+
+    def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
+        obj = cast(flyte.Environment, obj)
+        return BuildEnvCommand(
+            name=obj_name,
+            obj_name=obj_name,
+            obj=obj,
+            help=f"{obj.name}" + (f": {obj.description}" if obj.description else ""),
+            build_args=self.build_args,
+        )
+
+
+class EnvFiles(common.FileGroup):
+    """
+    Group that creates a command for each file in the current directory that is not `__init__.py`.
+    """
+
+    common_options_enabled = False
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        if "params" not in kwargs:
+            kwargs["params"] = []
+        kwargs["params"].extend(BuildArguments.options())
+        super().__init__(*args, **kwargs)
+
+    def get_command(self, ctx, filename):
+        build_args = BuildArguments.from_dict(ctx.params)
+        return EnvPerFileGroup(
+            filename=Path(filename),
+            build_args=build_args,
+            name=filename,
+            help=f"Run, functions decorated `env.task` or instances of Tasks in {filename}",
+        )
+
+
+build = EnvFiles(
+    name="build",
+    help="""
+    Build the environments defined in a python file or directory. This will build the images associated with the
+    environments.
+    """,
+)