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

flyte/app/extras/_auth_middleware.py

@@ -0,0 +1,217 @@
+"""
+FastAPI middleware for automatic Flyte authentication passthrough.
+
+This module provides middleware that automatically extracts authentication headers
+from incoming requests and sets them in the Flyte context, eliminating the need
+for manual auth_metadata() wrapping in every endpoint.
+
+Example:
+    Basic usage with default extractors (Authorization + Cookie headers)::
+
+        from fastapi import FastAPI
+        from flyte.app.extras import FastAPIAuthMiddleware
+
+        app = FastAPI()
+        app.add_middleware(FastAPIAuthMiddleware, excluded_paths={"/health"})
+
+        @app.get("/me")
+        async def get_current_user():
+            # Auth metadata automatically set from request headers
+            user = await remote.User.get.aio()
+            return {"subject": user.subject()}
+
+    Advanced usage with custom extractors and path exclusions::
+
+        from flyte.app.extras import FastAPIAuthMiddleware
+
+        app.add_middleware(
+            FastAPIAuthMiddleware,
+            header_extractors=[
+                FastAPIAuthMiddleware.extract_authorization_header,
+                FastAPIAuthMiddleware.extract_custom_header("x-api-key"),
+            ],
+            excluded_paths={"/health", "/metrics"},
+        )
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+    from fastapi import Request
+    from starlette.middleware.base import BaseHTTPMiddleware
+    from starlette.responses import Response
+else:
+    try:
+        from starlette.middleware.base import BaseHTTPMiddleware
+    except ImportError:
+
+        class BaseHTTPMiddleware:
+            pass
+
+
+logger = logging.getLogger(__name__)
+
+# Header extractor type: takes a Request, returns (key, value) tuple or None
+HeaderExtractor = Callable[["Request"], tuple[str, str] | None]
+
+
+class FastAPIPassthroughAuthMiddleware(BaseHTTPMiddleware):
+    """
+    FastAPI middleware that automatically sets Flyte auth metadata from request headers.
+
+    This middleware extracts authentication headers from incoming HTTP requests and
+    sets them in the Flyte context using the auth_metadata() context manager. This
+    eliminates the need to manually wrap endpoint handlers with auth_metadata().
+
+    The middleware is highly configurable:
+    - Custom header extractors can be provided
+    - Specific paths can be excluded from auth requirements
+    - Auth can be optional or required
+
+    Attributes:
+        app: The FastAPI application (this is a mandatory framework parameter)
+        header_extractors: List of functions to extract headers from requests
+        excluded_paths: Set of URL paths that bypass auth extraction
+
+    Thread Safety:
+        This middleware is async-safe and properly isolates auth metadata per request
+        using Python's contextvars. Multiple concurrent requests with different
+        authentication will not interfere with each other.
+    """
+
+    def __init__(
+        self,
+        app,
+        header_extractors: list[HeaderExtractor] | None = None,
+        excluded_paths: set[str] | None = None,
+    ):
+        """
+        Initialize the Flyte authentication middleware.
+
+        Args:
+            app: The FastAPI/Starlette application
+            header_extractors: List of functions to extract headers. Each function
+                takes a Request and returns (key, value) tuple or None.
+                Defaults to [extract_authorization_header, extract_cookie_header].
+            excluded_paths: Set of URL paths to exclude from auth extraction.
+                Requests to these paths proceed without setting auth context.
+        """
+        super().__init__(app)
+
+        if header_extractors is None:
+            self.header_extractors: list[HeaderExtractor] = [
+                self.extract_authorization_header,
+                self.extract_cookie_header,
+            ]
+        else:
+            self.header_extractors = header_extractors
+
+        self.excluded_paths = excluded_paths or set()
+
+    async def dispatch(self, request: "Request", call_next) -> "Response":
+        """
+        Process each request, extracting auth headers and setting Flyte context.
+
+        Args:
+            request: The incoming HTTP request
+            call_next: The next middleware or route handler to call
+
+        Returns:
+            The HTTP response from the handler
+        """
+        from starlette.responses import JSONResponse
+
+        # Skip auth extraction for excluded paths
+        if request.url.path in self.excluded_paths:
+            return await call_next(request)
+
+        # Extract auth headers using all configured extractors
+        auth_tuples = []
+        for extractor in self.header_extractors:
+            try:
+                result = extractor(request)
+                if result is not None:
+                    auth_tuples.append(result)
+            except Exception as e:
+                logger.warning(f"Header extractor {extractor.__name__} failed: {e}")
+
+        # Require auth headers
+        if not auth_tuples:
+            logger.info("No auth tuples found")
+            return JSONResponse(
+                status_code=401,
+                content={"detail": "Authentication credentials required"},
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+
+        # Set auth metadata in Flyte context for the duration of this request
+        from flyte.remote import auth_metadata
+
+        with auth_metadata(*auth_tuples):
+            return await call_next(request)
+
+    @staticmethod
+    def extract_authorization_header(request: "Request") -> tuple[str, str] | None:
+        """
+        Extract the Authorization header from the request.
+
+        Args:
+            request: The FastAPI/Starlette request object
+
+        Returns:
+            Tuple of ("authorization", header_value) if present, None otherwise
+        """
+        auth_header = request.headers.get("authorization")
+        if auth_header:
+            return "authorization", auth_header
+        return None
+
+    @staticmethod
+    def extract_cookie_header(request: "Request") -> tuple[str, str] | None:
+        """
+        Extract the Cookie header from the request.
+
+        Args:
+            request: The FastAPI/Starlette request object
+
+        Returns:
+            Tuple of ("cookie", header_value) if present, None otherwise
+        """
+        cookie_header = request.headers.get("cookie")
+        if cookie_header:
+            return "cookie", cookie_header
+        return None
+
+    @staticmethod
+    def extract_custom_header(header_name: str) -> HeaderExtractor:
+        """
+        Create a header extractor for a custom header name.
+
+        Args:
+            header_name: The name of the header to extract (case-insensitive)
+
+        Returns:
+            A header extractor function that extracts the specified header
+
+        Example::
+
+            # Create extractor for X-API-Key header
+            api_key_extractor = extract_custom_header("x-api-key")
+
+            app.add_middleware(
+                FastAPIAuthMiddleware,
+                header_extractors=[api_key_extractor],
+            )
+        """
+
+        def extractor(request: "Request") -> tuple[str, str] | None:
+            header_value = request.headers.get(header_name.lower())
+            if header_value:
+                return header_name.lower(), header_value
+            return None
+
+        extractor.__name__ = f"extract_{header_name.replace('-', '_')}_header"
+        return extractor

flyte/app/extras/_fastapi.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import rich.repr
+
+import flyte.app
+from flyte.models import SerializationContext
+
+if TYPE_CHECKING:
+    import fastapi
+    import uvicorn
+
+
+@rich.repr.auto
+@dataclass(kw_only=True, repr=True)
+class FastAPIAppEnvironment(flyte.app.AppEnvironment):
+    app: fastapi.FastAPI
+    type: str = "FastAPI"
+    uvicorn_config: uvicorn.Config | None = None
+    _caller_frame: inspect.FrameInfo | None = None
+
+    def __post_init__(self):
+        try:
+            import fastapi
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "fastapi is not installed. Please install 'fastapi' to use FastAPIAppEnvironment."
+            )
+
+        # starlette is a dependency of fastapi, so if fastapi is installed, starlette is also installed.
+        try:
+            from starlette.datastructures import State
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "starlette is not installed. Please install 'starlette' to use FastAPIAppEnvironment."
+            )
+
+        class PicklableState(State):
+            def __getstate__(self):
+                state = self.__dict__.copy()
+                # Replace the unpicklable State with an empty dict
+                state["_state"] = {}
+                return state
+
+            def __setstate__(self, state):
+                self.__dict__.update(state)
+                # Restore a fresh State object
+                self.state = State()
+
+        # NOTE: since FastAPI cannot be pickled (because starlette.datastructures.State cannot be pickled due to
+        # circular references), we need to patch the state object to make it picklable.
+        self.app.state = PicklableState()
+
+        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)}")
+
+        self.links = [flyte.app.Link(path="/docs", title="FastAPI OpenAPI Docs", is_relative=True), *self.links]
+        self._server = self._fastapi_app_server
+
+        # 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)
+
+    async def _fastapi_app_server(self):
+        try:
+            import uvicorn
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "uvicorn is not installed. Please install 'uvicorn' to use FastAPIAppEnvironment."
+            )
+
+        if self.uvicorn_config is None:
+            self.uvicorn_config = uvicorn.Config(self.app, port=self.port.port)
+        elif self.uvicorn_config is not None:
+            if self.uvicorn_config.port is None:
+                self.uvicorn_config.port = self.port.port
+
+        await uvicorn.Server(self.uvicorn_config).serve()
+
+    def container_command(self, serialization_context: SerializationContext) -> list[str]:
+        return []

flyte/app/extras/_model_loader/__init__.py

@@ -0,0 +1,3 @@
+from .loader import SafeTensorsStreamer, prefetch
+
+__all__ = ["SafeTensorsStreamer", "prefetch"]

flyte/app/extras/_model_loader/config.py

@@ -0,0 +1,7 @@
+import os
+
+REMOTE_MODEL_PATH = os.getenv("FLYTE_MODEL_LOADER_REMOTE_MODEL_PATH")
+LOCAL_MODEL_PATH = os.getenv("FLYTE_MODEL_LOADER_LOCAL_MODEL_PATH", "/srv/model")
+CHUNK_SIZE = int(os.getenv("FLYTE_MODEL_LOADER_CHUNK_SIZE", str(16 * 1024 * 1024)))
+MAX_CONCURRENCY = int(os.getenv("FLYTE_MODEL_LOADER_MAX_CONCURRENCY", str(32)))
+STREAM_SAFETENSORS = os.getenv("FLYTE_MODEL_LOADER_STREAM_SAFETENSORS", "false").lower() == "true"

flyte/app/extras/_model_loader/loader.py

@@ -0,0 +1,288 @@
+import asyncio
+import json
+import logging
+import pathlib
+import struct
+import time
+import typing
+from collections import defaultdict
+
+import obstore
+import pydantic
+from obstore.store import ObjectStore
+from typing_extensions import Annotated
+
+from flyte.storage._parallel_reader import (
+    BufferProtocol,
+    Chunk,
+    DownloadQueueEmpty,
+    DownloadTask,
+    ObstoreParallelReader,
+    Source,
+)
+from flyte.storage._storage import get_underlying_filesystem
+
+try:
+    import torch
+except ModuleNotFoundError:
+    raise ModuleNotFoundError("torch is not installed. Please install 'torch', to use the model loader.")
+
+
+from flyte.app.extras._model_loader.config import (
+    CHUNK_SIZE,
+    MAX_CONCURRENCY,
+)
+
+logger = logging.getLogger(__name__)
+
+LITTLE_ENDIAN_LONG_LONG_STRUCT_FORMAT = "<Q"
+
+SAFETENSORS_FORMAT_KEY = "format"
+SAFETENSORS_FORMAT_VALUE = "pt"
+SAFETENSORS_SUFFIX = ".safetensors"
+SAFETENSORS_DEFAULT_PATTERN = f"*{SAFETENSORS_SUFFIX}"
+SAFETENSORS_SHARDED_PATTERN = f"model-rank-{{rank}}-part-*{SAFETENSORS_SUFFIX}"
+SAFETENSORS_INTERNAL_METADATA_KEY = "__metadata__"
+SAFETENSORS_INDEX_PATH = "model.safetensors.index.json"
+SAFETENSORS_HEADER_BUFFER_SIZE = 8
+SAFETENSORS_TO_TORCH_DTYPE = {
+    "F64": torch.float64,
+    "F32": torch.float32,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+    "I64": torch.int64,
+    "I32": torch.int32,
+    "I16": torch.int16,
+    "I8": torch.int8,
+    "U8": torch.uint8,
+    "BOOL": torch.bool,
+    "F8_E5M2": torch.float8_e5m2,
+    "F8_E4M3": torch.float8_e4m3fn,
+}
+
+
+async def prefetch(remote_model_path, local_model_path, exclude_safetensors=True):
+    from flyte.storage._storage import _get_obstore_bypass
+
+    logger.info(f"Pre-fetching model artifacts from {remote_model_path} to {local_model_path}...")
+    if exclude_safetensors:
+        logger.info(f"Deferring download of safetensor files from {remote_model_path}")
+    start = time.perf_counter()
+
+    try:
+        # Exclude safetensors if model streaming is enabled, which will be handled by the flyte vllm model loader
+        await _get_obstore_bypass(
+            remote_model_path,
+            local_model_path,
+            recursive=True,
+            exclude=[SAFETENSORS_DEFAULT_PATTERN] if exclude_safetensors else None,
+        )
+    except DownloadQueueEmpty:
+        logger.warning("No model artifacts found to pre-fetch.")
+    else:
+        logger.info(f"Pre-fetched model artifacts in {time.perf_counter() - start:.2f}s")
+
+
+def _dtype_to_torch_dtype(dtype: str) -> torch.dtype:
+    try:
+        return SAFETENSORS_TO_TORCH_DTYPE[dtype]
+    except KeyError:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+class TensorMetadata(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    name: str
+    shape: list[int]
+    dtype: Annotated[torch.dtype, pydantic.BeforeValidator(_dtype_to_torch_dtype)]
+    data_offsets: tuple[int, int]
+
+    @pydantic.computed_field  # type: ignore[prop-decorator]
+    @property
+    def size(self) -> int:
+        start, end = self.data_offsets
+        return end - start
+
+    @pydantic.computed_field  # type: ignore[prop-decorator]
+    @property
+    def length(self) -> int:
+        count = 1
+        for dim in self.shape:
+            count *= dim
+        return count
+
+    def __len__(self):
+        return self.length
+
+
+class SafeTensorsMetadata(pydantic.BaseModel):
+    path: str
+    data_start: int
+    tensors: list[TensorMetadata]
+
+
+class SafeTensorsStreamer:
+    def __init__(
+        self,
+        remote_path,
+        local_path,
+        chunk_size=CHUNK_SIZE,
+        max_concurrency=MAX_CONCURRENCY,
+        rank=0,
+        tensor_parallel_size=1,
+        store_kwargs=None,
+    ):
+        fs = get_underlying_filesystem(path=remote_path)
+        bucket, prefix = fs._split_path(remote_path)  # pylint: disable=W0212
+
+        self._store: ObjectStore = fs._construct_store(bucket)
+        self._bucket = bucket
+        self._prefix: pathlib.Path = pathlib.Path(prefix)
+        self._local_path = pathlib.Path(local_path)
+        self._reader = ObstoreParallelReader(self._store, chunk_size=chunk_size, max_concurrency=max_concurrency)
+        self._rank = rank
+        self._tensor_parallel_size = tensor_parallel_size
+
+    async def _parse_safetensors_metadata(self, path):
+        header_len = await obstore.get_range_async(self._store, str(path), start=0, end=SAFETENSORS_HEADER_BUFFER_SIZE)
+        header_size = struct.unpack(
+            LITTLE_ENDIAN_LONG_LONG_STRUCT_FORMAT,
+            header_len,
+        )[0]
+        header_data = json.loads(
+            (
+                await obstore.get_range_async(
+                    self._store,
+                    str(path),
+                    start=SAFETENSORS_HEADER_BUFFER_SIZE,
+                    end=SAFETENSORS_HEADER_BUFFER_SIZE + header_size,
+                )
+            ).to_bytes()
+        )
+        if (
+            format := header_data.pop(SAFETENSORS_INTERNAL_METADATA_KEY, {}).get(SAFETENSORS_FORMAT_KEY)
+        ) and format != SAFETENSORS_FORMAT_VALUE:
+            raise ValueError(f"Unsupported format: {format}")
+        return SafeTensorsMetadata(
+            path=str(path),
+            data_start=SAFETENSORS_HEADER_BUFFER_SIZE + header_size,
+            tensors=[TensorMetadata.model_validate({"name": k, **v}) for k, v in header_data.items()],
+        )
+
+    async def _list_safetensors_files_with_index(self):
+        # Get index of expected tensors if it exists
+        weight_map_resp = await obstore.get_async(self._store, str(self._prefix / SAFETENSORS_INDEX_PATH))
+        weight_map_bytes = bytes(await weight_map_resp.bytes_async())
+        tensor_to_path_map = json.loads(weight_map_bytes)["weight_map"]
+
+        # Create index for path -> tensors
+        index = defaultdict(set)
+        for tensor, path in tensor_to_path_map.items():
+            index[path].add(tensor)
+
+        return index.items()
+
+    async def _load_safetensors_metadata_from_index(self):
+        for path, expected in await self._list_safetensors_files_with_index():
+            stm = await self._parse_safetensors_metadata(self._prefix / path)
+            # Keep only the tensors we expect (should already be deduplicated)
+            keep = {tm.name: tm for tm in filter(lambda tm: tm.name in expected, stm.tensors)}
+            # We have missing tensors at the path. Bail out!
+            if missing := expected - keep.keys():
+                raise ValueError(f"Missing {len(missing)} tensors at {path!r}: {' '.join(missing)}")
+            stm.tensors = list(keep.values())
+            yield stm
+
+    async def _list_safetensors_files_with_pattern(self, pattern):
+        paths = set()
+        list_result = await obstore.list_with_delimiter_async(self._store, prefix=str(self._prefix))
+        for obj in list_result["objects"]:
+            path = pathlib.Path(obj["path"])
+            if path.match(pattern):
+                paths.add(path)
+        if not paths:
+            raise ValueError(f"No files found matching pattern: {pattern}")
+        return paths
+
+    async def _load_safetensors_metadata_with_pattern(self, pattern):
+        seen = set()
+        stms = await asyncio.gather(
+            *(
+                self._parse_safetensors_metadata(path)
+                for path in await self._list_safetensors_files_with_pattern(pattern)
+            )
+        )
+        for stm in stms:
+            stm.tensors = list[TensorMetadata](
+                filter(
+                    lambda tm: tm.name not in seen and not seen.add(tm.name),
+                    stm.tensors,
+                )
+            )
+            yield stm
+
+    async def _load_safetensors_metadata(self):
+        # When using tensor parallelism, we can't rely on the index. Fallback to using a pattern.
+        if self._tensor_parallel_size > 1:
+            async for stm in self._load_safetensors_metadata_with_pattern(
+                SAFETENSORS_SHARDED_PATTERN.format(rank=self._rank)
+            ):
+                yield stm
+            return
+
+        # No tensor parallelism. Try to use the index first, then fallback to a pattern.
+        try:
+            async for stm in self._load_safetensors_metadata_from_index():
+                yield stm
+        except (
+            json.decoder.JSONDecodeError,
+            FileNotFoundError,
+            KeyError,
+        ):
+            async for stm in self._load_safetensors_metadata_with_pattern(SAFETENSORS_DEFAULT_PATTERN):
+                yield stm
+
+    async def _get_tensors_async(self) -> typing.AsyncGenerator[tuple[str, torch.Tensor], None]:
+        async def _to_tensor(buf: BufferProtocol, source: Source) -> torch.Tensor:
+            assert isinstance(source.metadata, TensorMetadata)
+            return torch.frombuffer(
+                await buf.read(),
+                dtype=source.metadata.dtype,
+                count=len(source.metadata),
+                offset=0,
+            ).view(source.metadata.shape)
+
+        async def _gen() -> typing.AsyncGenerator[DownloadTask, None]:
+            async for stm in self._load_safetensors_metadata():
+                for tm in stm.tensors:
+                    source = Source(
+                        id=tm.name,
+                        path=stm.path,
+                        length=tm.size,
+                        offset=stm.data_start + tm.data_offsets[0],
+                        metadata=tm,
+                    )
+                    for offset, length in self._reader._chunks(tm.size):
+                        yield DownloadTask(
+                            source=source,
+                            chunk=Chunk(offset, length),
+                        )
+
+        # Yield tensors as they are downloaded
+        async for result in self._reader._as_completed(_gen(), transformer=_to_tensor):
+            yield result
+
+    def get_tensors(self) -> typing.Generator[tuple[str, torch.Tensor], None, None]:
+        logger.info("Streaming tensors...")
+        start = time.perf_counter()
+        counter = 0
+        gen = self._get_tensors_async()
+        with asyncio.Runner() as runner:
+            try:
+                while True:
+                    yield runner.run(gen.__anext__())
+                    counter += 1
+            except StopAsyncIteration:
+                pass
+        logger.info(f"Streamed {counter} tensors in {time.perf_counter() - start:.2f}s")

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.")