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

flyte/_logging.py

@@ -1,12 +1,36 @@
 from __future__ import annotations
 
+import json
 import logging
 import os
-from typing import Optional
+from datetime import datetime
+from typing import Literal, Optional
 
-from ._tools import ipython_check, is_in_cluster
+import flyte
 
-DEFAULT_LOG_LEVEL = logging.INFO
+from ._tools import ipython_check
+
+LogFormat = Literal["console", "json"]
+_LOG_LEVEL_MAP = {
+    "critical": logging.CRITICAL,  # 50
+    "error": logging.ERROR,  # 40
+    "warning": logging.WARNING,  # 30
+    "warn": logging.WARNING,  # 30
+    "info": logging.INFO,  # 20
+    "debug": logging.DEBUG,  # 10
+}
+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:
@@ -17,43 +41,69 @@ def is_rich_logging_disabled() -> bool:
 
 
 def get_env_log_level() -> int:
-    return int(os.environ.get("LOG_LEVEL", DEFAULT_LOG_LEVEL))
+    value = os.getenv("LOG_LEVEL")
+    if value is None:
+        return DEFAULT_LOG_LEVEL
+    # Case 1: numeric value ("10", "20", "5", etc.)
+    if value.isdigit():
+        return int(value)
 
+    # Case 2: named log level ("info", "debug", ...)
+    if value.lower() in _LOG_LEVEL_MAP:
+        return _LOG_LEVEL_MAP[value.lower()]
 
-def log_format_from_env() -> str:
+    return DEFAULT_LOG_LEVEL
+
+
+def log_format_from_env() -> LogFormat:
     """
     Get the log format from the environment variable.
     """
-    return os.environ.get("LOG_FORMAT", "json")
+    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.
     """
-    if is_in_cluster():
+    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.console import Console
+    from rich.highlighter import NullHighlighter
     from rich.logging import RichHandler
 
-    try:
-        width = os.get_terminal_size().columns
-    except Exception as e:
-        logger.debug(f"Failed to get terminal size: {e}")
-        width = 160
-
     handler = RichHandler(
         tracebacks_suppress=[click],
-        rich_tracebacks=True,
+        rich_tracebacks=False,
         omit_repeated_times=False,
         show_path=False,
         log_time_format="%H:%M:%S.%f",
-        console=Console(width=width),
+        console=_get_console(),
         level=log_level,
+        highlighter=NullHighlighter(),
+        markup=True,
     )
 
     formatter = logging.Formatter(fmt="%(filename)s:%(lineno)d - %(message)s")
@@ -61,39 +111,99 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
     return handler
 
 
-def get_default_handler(log_level: int) -> logging.Handler:
-    handler = logging.StreamHandler()
-    handler.setLevel(log_level)
-    formatter = logging.Formatter(fmt="[%(name)s] %(message)s")
-    if log_format_from_env() == "json":
-        pass
-        # formatter = jsonlogger.JsonFormatter(fmt="%(asctime)s %(name)s %(levelname)s %(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 metric fields if present
+        if getattr(record, "metric_type", None):
+            log_data["metric_type"] = record.metric_type  # type: ignore[attr-defined]
+            log_data["metric_name"] = record.metric_name  # type: ignore[attr-defined]
+            log_data["duration_seconds"] = record.duration_seconds  # type: ignore[attr-defined]
 
-def initialize_logger(log_level: int = DEFAULT_LOG_LEVEL, enable_rich: bool = False):
+        # 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,
+    reset_root_logger: 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
-    logger = _create_logger("flyte", log_level, enable_rich)
 
+    if log_level is None:
+        log_level = get_env_log_level()
+    if log_format is None:
+        log_format = log_format_from_env()
 
-def _create_logger(name: str, log_level: int = DEFAULT_LOG_LEVEL, enable_rich: bool = False) -> logging.Logger:
-    """
-    Creates a logger with the given name and log level.
-    """
-    logger = logging.getLogger(name)
-    logger.setLevel(log_level)
-    handler = None
-    logger.handlers = []
-    if enable_rich:
-        handler = get_rich_handler(log_level)
-    if handler is None:
-        handler = get_default_handler(log_level)
-    logger.addHandler(handler)
-    return logger
+    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
+
+    reset_root_logger = reset_root_logger or os.environ.get("FLYTE_RESET_ROOT_LOGGER") == "1"
+    if reset_root_logger:
+        _setup_root_logger(use_json=use_json, use_rich=use_rich, log_level=log_level)
+    else:
+        root_logger = logging.getLogger()
+        for h in root_logger.handlers:
+            h.addFilter(ContextFilter())
+
+    # 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):
@@ -121,4 +231,93 @@ def log(fn=None, *, level=logging.DEBUG, entry=True, exit=True):
     return decorator(fn)
 
 
-logger = _create_logger("flyte", get_env_log_level())
+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(use_json: bool, use_rich: bool, log_level: int):
+    """
+    Wipe all handlers from the root logger and reconfigure. This ensures
+    both user/library logging and Flyte internal logging get context information and look the same.
+    """
+    root = logging.getLogger()
+    root.handlers.clear()  # Remove any existing handlers to prevent double logging
+
+    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)
+
+    # get_rich_handler can return None in some environments
+    if not root_handler:
+        root_handler = logging.StreamHandler()
+
+    # Add context filter to ALL logging
+    root_handler.addFilter(ContextFilter())
+    root_handler.setLevel(log_level)
+
+    root.addHandler(root_handler)
+    root.setLevel(log_level)
+
+
+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
+
+
+# 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()