flyte 2.0.0b32__py3-none-any.whl

flyte/_logging.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+from datetime import datetime
+from typing import Literal, Optional
+
+import flyte
+
+from ._tools import ipython_check
+
+LogFormat = Literal["console", "json"]
+
+DEFAULT_LOG_LEVEL = logging.WARNING
+
+
+def make_hyperlink(label: str, url: str):
+    """
+    Create a hyperlink in the terminal output.
+    """
+    BLUE = "\033[94m"
+    RESET = "\033[0m"
+    OSC8_BEGIN = f"\033]8;;{url}\033\\"
+    OSC8_END = "\033]8;;\033\\"
+    return f"{BLUE}{OSC8_BEGIN}{label}{RESET}{OSC8_END}"
+
+
+def is_rich_logging_disabled() -> bool:
+    """
+    Check if rich logging is enabled
+    """
+    return os.environ.get("DISABLE_RICH_LOGGING") is not None
+
+
+def get_env_log_level() -> int:
+    return int(os.environ.get("LOG_LEVEL", DEFAULT_LOG_LEVEL))
+
+
+def log_format_from_env() -> LogFormat:
+    """
+    Get the log format from the environment variable.
+    """
+    format_str = os.environ.get("LOG_FORMAT", "console")
+    if format_str not in ("console", "json"):
+        return "console"
+    return format_str  # type: ignore[return-value]
+
+
+def _get_console():
+    """
+    Get the console.
+    """
+    from rich.console import Console
+
+    try:
+        width = os.get_terminal_size().columns
+    except Exception as e:
+        logger.debug(f"Failed to get terminal size: {e}")
+        width = 160
+
+    return Console(width=width)
+
+
+def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
+    """
+    Upgrades the global loggers to use Rich logging.
+    """
+    ctx = flyte.ctx()
+    if ctx and ctx.is_in_cluster():
+        return None
+    if not ipython_check() and is_rich_logging_disabled():
+        return None
+
+    import click
+    from rich.highlighter import NullHighlighter
+    from rich.logging import RichHandler
+
+    handler = RichHandler(
+        tracebacks_suppress=[click],
+        rich_tracebacks=False,
+        omit_repeated_times=False,
+        show_path=False,
+        log_time_format="%H:%M:%S.%f",
+        console=_get_console(),
+        level=log_level,
+        highlighter=NullHighlighter(),
+        markup=True,
+    )
+
+    formatter = logging.Formatter(fmt="%(filename)s:%(lineno)d - %(message)s")
+    handler.setFormatter(formatter)
+    return handler
+
+
+class JSONFormatter(logging.Formatter):
+    """
+    Formatter that outputs JSON strings for each log record.
+    """
+
+    def format(self, record: logging.LogRecord) -> str:
+        log_data = {
+            "timestamp": datetime.fromtimestamp(record.created).isoformat(),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+            "filename": record.filename,
+            "lineno": record.lineno,
+            "funcName": record.funcName,
+        }
+
+        # Add context fields if present
+        if getattr(record, "run_name", None):
+            log_data["run_name"] = record.run_name  # type: ignore[attr-defined]
+        if getattr(record, "action_name", None):
+            log_data["action_name"] = record.action_name  # type: ignore[attr-defined]
+        if getattr(record, "is_flyte_internal", False):
+            log_data["is_flyte_internal"] = True
+
+        # Add exception info if present
+        if record.exc_info:
+            log_data["exc_info"] = self.formatException(record.exc_info)
+
+        return json.dumps(log_data)
+
+
+def initialize_logger(log_level: int | None = None, log_format: LogFormat | None = None, enable_rich: bool = False):
+    """
+    Initializes the global loggers to the default configuration.
+    When enable_rich=True, upgrades to Rich handler for local CLI usage.
+    """
+    global logger  # noqa: PLW0603
+
+    if log_level is None:
+        log_level = get_env_log_level()
+    if log_format is None:
+        log_format = log_format_from_env()
+
+    # Clear existing handlers to reconfigure
+    root = logging.getLogger()
+    root.handlers.clear()
+
+    flyte_logger = logging.getLogger("flyte")
+    flyte_logger.handlers.clear()
+
+    # Determine log format (JSON takes precedence over Rich)
+    use_json = log_format == "json"
+    use_rich = enable_rich and not use_json
+
+    # Set up root logger handler
+    root_handler: logging.Handler | None = None
+    if use_json:
+        root_handler = logging.StreamHandler()
+        root_handler.setFormatter(JSONFormatter())
+    elif use_rich:
+        root_handler = get_rich_handler(log_level)
+
+    if root_handler is None:
+        root_handler = logging.StreamHandler()
+
+    # Add context filter to root handler for all logging
+    root_handler.addFilter(ContextFilter())
+    root_handler.setLevel(logging.DEBUG)
+    root.addHandler(root_handler)
+
+    # Set up Flyte logger handler
+    flyte_handler: logging.Handler | None = None
+    if use_json:
+        flyte_handler = logging.StreamHandler()
+        flyte_handler.setLevel(log_level)
+        flyte_handler.setFormatter(JSONFormatter())
+    elif use_rich:
+        flyte_handler = get_rich_handler(log_level)
+
+    if flyte_handler is None:
+        flyte_handler = logging.StreamHandler()
+        flyte_handler.setLevel(log_level)
+        formatter = logging.Formatter(fmt="%(message)s")
+        flyte_handler.setFormatter(formatter)
+
+    # Add both filters to Flyte handler
+    flyte_handler.addFilter(FlyteInternalFilter())
+    flyte_handler.addFilter(ContextFilter())
+
+    flyte_logger.addHandler(flyte_handler)
+    flyte_logger.setLevel(log_level)
+    flyte_logger.propagate = False  # Prevent double logging
+
+    logger = flyte_logger
+
+
+def log(fn=None, *, level=logging.DEBUG, entry=True, exit=True):
+    """
+    Decorator to log function calls.
+    """
+
+    def decorator(func):
+        if logger.isEnabledFor(level):
+
+            def wrapper(*args, **kwargs):
+                if entry:
+                    logger.log(level, f"[{func.__name__}] with args: {args} and kwargs: {kwargs}")
+                try:
+                    return func(*args, **kwargs)
+                finally:
+                    if exit:
+                        logger.log(level, f"[{func.__name__}] completed")
+
+            return wrapper
+        return func
+
+    if fn is None:
+        return decorator
+    return decorator(fn)
+
+
+class ContextFilter(logging.Filter):
+    """
+    A logging filter that adds the current action's run name and name to all log records.
+    Applied globally to capture context for both user and Flyte internal logging.
+    """
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        from flyte._context import ctx
+
+        c = ctx()
+        if c:
+            action = c.action
+            # Add as attributes for structured logging (JSON)
+            record.run_name = action.run_name
+            record.action_name = action.name
+            # Also modify message for console/Rich output
+            record.msg = f"[{action.run_name}][{action.name}] {record.msg}"
+        else:
+            record.run_name = None
+            record.action_name = None
+        return True
+
+
+class FlyteInternalFilter(logging.Filter):
+    """
+    A logging filter that adds [flyte] prefix to internal Flyte logging only.
+    """
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        is_internal = record.name.startswith("flyte")
+        # Add as attribute for structured logging (JSON)
+        record.is_flyte_internal = is_internal
+        # Also modify message for console/Rich output
+        if is_internal:
+            record.msg = f"[flyte] {record.msg}"
+        return True
+
+
+def _setup_root_logger():
+    """
+    Configure the root logger to capture all logging with context information.
+    This ensures both user code and Flyte internal logging get the context.
+    """
+    root = logging.getLogger()
+    root.handlers.clear()  # Remove any existing handlers to prevent double logging
+
+    # Create a basic handler for the root logger
+    handler = logging.StreamHandler()
+    # Add context filter to ALL logging
+    handler.addFilter(ContextFilter())
+    handler.setLevel(logging.DEBUG)
+
+    # Simple formatter since filters handle prefixes
+    root.addHandler(handler)
+
+
+def _create_flyte_logger() -> logging.Logger:
+    """
+    Create the internal Flyte logger with [flyte] prefix.
+    """
+    flyte_logger = logging.getLogger("flyte")
+    flyte_logger.setLevel(get_env_log_level())
+
+    # Add a handler specifically for flyte logging with the prefix filter
+    handler = logging.StreamHandler()
+    handler.setLevel(get_env_log_level())
+    handler.addFilter(FlyteInternalFilter())
+    handler.addFilter(ContextFilter())
+
+    formatter = logging.Formatter(fmt="%(message)s")
+    handler.setFormatter(formatter)
+
+    # Prevent propagation to root to avoid double logging
+    flyte_logger.propagate = False
+    flyte_logger.addHandler(handler)
+
+    return flyte_logger
+
+
+# Initialize root logger for global context
+_setup_root_logger()
+
+# Create the Flyte internal logger
+logger = _create_flyte_logger()

flyte/_map.py

@@ -0,0 +1,312 @@
+import asyncio
+import functools
+import logging
+from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast, overload
+
+from flyte.syncify import syncify
+
+from ._group import group
+from ._logging import logger
+from ._task import AsyncFunctionTaskTemplate, F, P, R
+
+
+class MapAsyncIterator(Generic[P, R]):
+    """AsyncIterator implementation for the map function results"""
+
+    def __init__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        args: tuple,
+        name: str,
+        concurrency: int,
+        return_exceptions: bool,
+    ):
+        self.func = func
+        self.args = args
+        self.name = name
+        self.concurrency = concurrency
+        self.return_exceptions = return_exceptions
+        self._tasks: List[asyncio.Task] = []
+        self._current_index = 0
+        self._completed_count = 0
+        self._exception_count = 0
+        self._task_count = 0
+        self._initialized = False
+
+    def __aiter__(self) -> AsyncIterator[Union[R, Exception]]:
+        """Return self as the async iterator"""
+        return self
+
+    async def __anext__(self) -> Union[R, Exception]:
+        """Get the next result"""
+        # Initialize on first call
+        if not self._initialized:
+            await self._initialize()
+
+        # Check if we've exhausted all tasks
+        if self._current_index >= self._task_count:
+            raise StopAsyncIteration
+
+        # Get the next task result
+        task = self._tasks[self._current_index]
+        self._current_index += 1
+
+        try:
+            result = await task
+            self._completed_count += 1
+            logger.debug(f"Task {self._current_index - 1} completed successfully")
+            return result
+        except Exception as e:
+            self._exception_count += 1
+            logger.debug(
+                f"Task {self._current_index - 1} failed with exception: {e}, return_exceptions={self.return_exceptions}"
+            )
+            if self.return_exceptions:
+                return e
+            else:
+                # Cancel remaining tasks
+                for remaining_task in self._tasks[self._current_index + 1 :]:
+                    remaining_task.cancel()
+                logger.warning("Exception raising is `ON`, raising exception and cancelling remaining tasks")
+                raise e
+
+    async def _initialize(self):
+        """Initialize the tasks - called lazily on first iteration"""
+        # Create all tasks at once
+        tasks = []
+        task_count = 0
+
+        if isinstance(self.func, functools.partial):
+            # Handle partial functions by merging bound args/kwargs with mapped args
+            base_func = cast(AsyncFunctionTaskTemplate, self.func.func)
+            bound_args = self.func.args
+            bound_kwargs = self.func.keywords or {}
+
+            for arg_tuple in zip(*self.args):
+                # Merge bound positional args with mapped args
+                merged_args = bound_args + arg_tuple
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(f"Running {base_func.name} with args: {merged_args} and kwargs: {bound_kwargs}")
+                task = asyncio.create_task(base_func.aio(*merged_args, **bound_kwargs))
+                tasks.append(task)
+                task_count += 1
+        else:
+            # Handle regular TaskTemplate functions
+            for arg_tuple in zip(*self.args):
+                task = asyncio.create_task(self.func.aio(*arg_tuple))
+                tasks.append(task)
+                task_count += 1
+
+        if task_count == 0:
+            logger.info(f"Group '{self.name}' has no tasks to process")
+            self._tasks = []
+            self._task_count = 0
+        else:
+            logger.info(f"Starting {task_count} tasks in group '{self.name}' with unlimited concurrency")
+            self._tasks = tasks
+            self._task_count = task_count
+
+        self._initialized = True
+
+    async def collect(self) -> List[Union[R, Exception]]:
+        """Convenience method to collect all results into a list"""
+        results = []
+        async for result in self:
+            results.append(result)
+        return results
+
+    def __repr__(self):
+        return f"MapAsyncIterator(group_name='{self.name}', concurrency={self.concurrency})"
+
+
+class _Mapper(Generic[P, R]):
+    """
+    Internal mapper class to handle the mapping logic
+
+    NOTE: The reason why we do not use the `@syncify` decorator here is because, in `syncify` we cannot use
+    context managers like `group` directly in the function body. This is because the `__exit__` method of the
+    context manager is called after the function returns. An for `_context` the `__exit__` method releases the
+    token (for contextvar), which was created in a separate thread. This leads to an exception like:
+
+    """
+
+    @classmethod
+    def _get_name(cls, task_name: str, group_name: str | None) -> str:
+        """Get the name of the group, defaulting to 'map' if not provided."""
+        return f"{task_name}_{group_name or 'map'}"
+
+    @staticmethod
+    def validate_partial(func: functools.partial[R]):
+        """
+        This method validates that the provided partial function is valid for mapping, i.e. only the one argument
+        is left for mapping and the rest are provided as keywords or args.
+
+        :param func: partial function to validate
+        :raises TypeError: if the partial function is not valid for mapping
+        """
+        f = cast(AsyncFunctionTaskTemplate, func.func)
+        inputs = f.native_interface.inputs
+        params = list(inputs.keys())
+        total_params = len(params)
+        provided_args = len(func.args)
+        provided_kwargs = len(func.keywords or {})
+
+        # Calculate how many parameters are left unspecified
+        unspecified_count = total_params - provided_args - provided_kwargs
+
+        # Exactly one parameter should be left for mapping
+        if unspecified_count != 1:
+            raise TypeError(
+                f"Partial function must leave exactly one parameter unspecified for mapping. "
+                f"Found {unspecified_count} unspecified parameters in {f.name}, "
+                f"params: {inputs.keys()}"
+            )
+
+        # Validate that no parameter is both in args and keywords
+        if func.keywords:
+            param_names = list(inputs.keys())
+            for i, arg_name in enumerate(param_names[: provided_args + 1]):
+                if arg_name in func.keywords:
+                    raise TypeError(
+                        f"Parameter '{arg_name}' is provided both as positional argument and keyword argument "
+                        f"in partial function {f.name}."
+                    )
+
+    @overload
+    def __call__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+    ) -> Iterator[R]: ...
+
+    @overload
+    def __call__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+        return_exceptions: bool = True,
+    ) -> Iterator[Union[R, Exception]]: ...
+
+    def __call__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+        return_exceptions: bool = True,
+    ) -> Iterator[Union[R, Exception]]:
+        """
+        Map a function over the provided arguments with concurrent execution.
+
+        :param func: The async function to map.
+        :param args: Positional arguments to pass to the function (iterables that will be zipped).
+        :param group_name: The name of the group for the mapped tasks.
+        :param concurrency: The maximum number of concurrent tasks to run. If 0, run all tasks concurrently.
+        :param return_exceptions: If True, yield exceptions instead of raising them.
+        :return: AsyncIterator yielding results in order.
+        """
+        if not args:
+            return
+
+        if isinstance(func, functools.partial):
+            f = cast(AsyncFunctionTaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(AsyncFunctionTaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
+        logger.debug(f"Blocking Map for {name}")
+        with group(name):
+            import flyte
+
+            tctx = flyte.ctx()
+            if tctx is None or tctx.mode == "local":
+                logger.warning("Running map in local mode, which will run every task sequentially.")
+                for v in zip(*args):
+                    try:
+                        yield func(*v)  # type: ignore
+                    except Exception as e:
+                        if return_exceptions:
+                            yield e
+                        else:
+                            raise e
+                return
+
+            i = 0
+            for x in cast(
+                Iterator[R],
+                _map(
+                    func,
+                    *args,
+                    name=name,
+                    concurrency=concurrency,
+                    return_exceptions=return_exceptions,
+                ),
+            ):
+                logger.debug(f"Mapped {x}, task {i}")
+                i += 1
+                yield x
+
+    async def aio(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+        return_exceptions: bool = True,
+    ) -> AsyncGenerator[Union[R, Exception], None]:
+        if not args:
+            return
+
+        if isinstance(func, functools.partial):
+            f = cast(AsyncFunctionTaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(AsyncFunctionTaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
+        with group(name):
+            import flyte
+
+            tctx = flyte.ctx()
+            if tctx is None or tctx.mode == "local":
+                logger.warning("Running map in local mode, which will run every task sequentially.")
+                for v in zip(*args):
+                    try:
+                        yield func(*v)  # type: ignore
+                    except Exception as e:
+                        if return_exceptions:
+                            yield e
+                        else:
+                            raise e
+                return
+            async for x in _map.aio(
+                func,
+                *args,
+                name=name,
+                concurrency=concurrency,
+                return_exceptions=return_exceptions,
+            ):
+                yield cast(Union[R, Exception], x)
+
+
+@syncify
+async def _map(
+    func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+    *args: Iterable[Any],
+    name: str = "map",
+    concurrency: int = 0,
+    return_exceptions: bool = True,
+) -> AsyncIterator[Union[R, Exception]]:
+    iter = MapAsyncIterator(
+        func=func, args=args, name=name, concurrency=concurrency, return_exceptions=return_exceptions
+    )
+    async for result in iter:
+        yield result
+
+
+map: _Mapper = _Mapper()

flyte/_module.py

@@ -0,0 +1,72 @@
+import inspect
+import os
+import pathlib
+import sys
+from types import ModuleType
+from typing import Tuple
+
+
+def extract_obj_module(obj: object, /, source_dir: pathlib.Path | None = None) -> Tuple[str, ModuleType]:
+    """
+    Extract the module from the given object. If source_dir is provided, the module will be relative to the source_dir.
+
+    Args:
+        obj: The object to extract the module from.
+        source_dir: The source directory to use for relative paths.
+
+    Returns:
+        The module name as a string.
+    """
+    if source_dir is None:
+        raise ValueError("extract_obj_module: source_dir cannot be None - specify root-dir")
+    # Get the module containing the object
+    entity_module = inspect.getmodule(obj)
+    if entity_module is None:
+        obj_name = getattr(obj, "__name__", str(obj))
+        raise ValueError(f"Object {obj_name} has no module.")
+
+    fp = entity_module.__file__
+    if fp is None:
+        obj_name = getattr(obj, "__name__", str(obj))
+        raise ValueError(f"Object {obj_name} has no module.")
+
+    file_path = pathlib.Path(fp)
+    try:
+        # Get the relative path to the current directory
+        # Will raise ValueError if the file is not in the source directory
+        relative_path = file_path.relative_to(str(pathlib.Path(source_dir).absolute()))
+
+        if relative_path == pathlib.Path("_internal/resolvers"):
+            entity_module_name = entity_module.__name__
+        else:
+            # Replace file separators with dots and remove the '.py' extension
+            dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
+            entity_module_name = dotted_path
+    except ValueError:
+        # If source_dir is not provided or file is not in source_dir, fallback to module name
+        # File is not relative to source_dir - check if it's an installed package
+        file_path_str = str(file_path)
+        if "site-packages" in file_path_str or "dist-packages" in file_path_str:
+            # It's an installed package - use the module's __name__ directly
+            # This will be importable via importlib.import_module()
+            entity_module_name = entity_module.__name__
+        else:
+            # File is not in source_dir and not in site-packages - re-raise the error
+            obj_name = getattr(obj, "__name__", str(obj))
+            raise ValueError(
+                f"Object {obj_name} module file {file_path} is not relative to "
+                f"source directory {source_dir} and is not an installed package."
+            )
+
+    if entity_module_name == "__main__":
+        """
+        This case is for the case in which the object is run from the main module.
+        """
+        fp = sys.modules["__main__"].__file__
+        if fp is None:
+            obj_name = getattr(obj, "__name__", str(obj))
+            raise ValueError(f"Object {obj_name} has no module.")
+        main_path = pathlib.Path(fp)
+        entity_module_name = main_path.stem
+
+    return entity_module_name, entity_module