flyte 2.0.0b22__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/_internal/runtime/trigger_serde.py

@@ -0,0 +1,160 @@
+import asyncio
+from typing import Union
+
+from flyteidl2.core import interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2, run_pb2, task_definition_pb2
+from google.protobuf import timestamp_pb2, wrappers_pb2
+
+import flyte.types
+from flyte import Cron, FixedRate, Trigger, TriggerTime
+
+
+def _to_schedule(m: Union[Cron, FixedRate], kickoff_arg_name: str | None = None) -> common_pb2.Schedule:
+    if isinstance(m, Cron):
+        return common_pb2.Schedule(
+            cron=common_pb2.Cron(
+                expression=m.expression,
+                timezone=m.timezone,
+            ),
+            kickoff_time_input_arg=kickoff_arg_name,
+        )
+    elif isinstance(m, FixedRate):
+        start_time = None
+        if m.start_time is not None:
+            start_time = timestamp_pb2.Timestamp()
+            start_time.FromDatetime(m.start_time)
+
+        return common_pb2.Schedule(
+            rate=common_pb2.FixedRate(
+                value=m.interval_minutes,
+                unit=common_pb2.FixedRateUnit.FIXED_RATE_UNIT_MINUTE,
+                start_time=start_time,
+            ),
+            kickoff_time_input_arg=kickoff_arg_name,
+        )
+
+
+async def process_default_inputs(
+    default_inputs: dict,
+    task_name: str,
+    task_inputs: interface_pb2.VariableMap,
+    task_default_inputs: list[common_pb2.NamedParameter],
+) -> list[common_pb2.NamedLiteral]:
+    """
+    Process default inputs and convert them to NamedLiteral objects.
+
+    Args:
+        default_inputs: Dictionary of default input values
+        task_name: Name of the task for error messages
+        task_inputs: Task input variable map
+        task_default_inputs: List of default parameters from task
+
+    Returns:
+        List of NamedLiteral objects
+    """
+    keys = []
+    literal_coros = []
+    for k, v in default_inputs.items():
+        if k not in task_inputs.variables:
+            raise ValueError(
+                f"Trigger default input '{k}' must be an input to the task, but not found in task {task_name}. "
+                f"Available inputs: {list(task_inputs.variables.keys())}"
+            )
+        else:
+            literal_coros.append(flyte.types.TypeEngine.to_literal(v, type(v), task_inputs.variables[k].type))
+            keys.append(k)
+
+    final_literals: list[literals_pb2.Literal] = await asyncio.gather(*literal_coros)
+
+    for p in task_default_inputs or []:
+        if p.name not in keys:
+            keys.append(p.name)
+            final_literals.append(p.parameter.default)
+
+    literals: list[common_pb2.NamedLiteral] = []
+    for k, lit in zip(keys, final_literals):
+        literals.append(
+            common_pb2.NamedLiteral(
+                name=k,
+                value=lit,
+            )
+        )
+
+    return literals
+
+
+async def to_task_trigger(
+    t: Trigger,
+    task_name: str,
+    task_inputs: interface_pb2.VariableMap,
+    task_default_inputs: list[common_pb2.NamedParameter],
+) -> task_definition_pb2.TaskTrigger:
+    """
+    Converts a Trigger object to a TaskTrigger protobuf object.
+    Args:
+        t:
+        task_name:
+        task_inputs:
+        task_default_inputs:
+    Returns:
+
+    """
+    env = None
+    if t.env_vars:
+        env = run_pb2.Envs()
+        for k, v in t.env_vars.items():
+            env.values.append(literals_pb2.KeyValuePair(key=k, value=v))
+
+    labels = run_pb2.Labels(values=t.labels) if t.labels else None
+
+    annotations = run_pb2.Annotations(values=t.annotations) if t.annotations else None
+
+    run_spec = run_pb2.RunSpec(
+        overwrite_cache=t.overwrite_cache,
+        envs=env,
+        interruptible=wrappers_pb2.BoolValue(value=t.interruptible) if t.interruptible is not None else None,
+        cluster=t.queue,
+        labels=labels,
+        annotations=annotations,
+    )
+
+    kickoff_arg_name = None
+    default_inputs = {}
+    if t.inputs:
+        for k, v in t.inputs.items():
+            if v is TriggerTime:
+                if k == "trigger_time" and k not in task_inputs.variables:
+                    # the 'trigger_time' input name that by default Triggers look for so it's always added.
+                    # Remove it here by skipping if it's not actually an input to the task
+                    continue
+                kickoff_arg_name = k
+            else:
+                default_inputs[k] = v
+
+    # assert that default_inputs and the kickoff_arg_name are infact in the task inputs
+    if kickoff_arg_name is not None and kickoff_arg_name not in task_inputs.variables:
+        raise ValueError(
+            f"For a scheduled trigger, the TriggerTime input '{kickoff_arg_name}' "
+            f"must be an input to the task, but not found in task {task_name}. "
+            f"Available inputs: {list(task_inputs.variables.keys())}"
+        )
+
+    literals = await process_default_inputs(default_inputs, task_name, task_inputs, task_default_inputs)
+
+    automation = _to_schedule(
+        t.automation,
+        kickoff_arg_name=kickoff_arg_name,
+    )
+
+    return task_definition_pb2.TaskTrigger(
+        name=t.name,
+        spec=task_definition_pb2.TaskTriggerSpec(
+            active=t.auto_activate,
+            run_spec=run_spec,
+            inputs=common_pb2.Inputs(literals=literals),
+        ),
+        automation_spec=common_pb2.TriggerAutomationSpec(
+            type=common_pb2.TriggerAutomationSpecType.TYPE_SCHEDULE,
+            schedule=automation,
+        ),
+    )

flyte/_internal/runtime/types_serde.py

@@ -1,6 +1,6 @@
 from typing import Dict, Optional, TypeVar
 
-from flyteidl.core import interface_pb2
+from flyteidl2.core import interface_pb2
 
 from flyte.models import NativeInterface
 from flyte.types._type_engine import TypeEngine

flyte/_keyring/file.py

@@ -10,9 +10,27 @@ _FLYTE_KEYRING_PATH: Path = Path.home() / ".flyte" / "keyring.cfg"
 
 
 class SimplePlainTextKeyring(KeyringBackend):
-    """Simple plain text keyring"""
+    """
+    Simple plain text keyring for remote notebook environments.
 
-    priority = 0.5
+    This backend is only active when running in IPython/Jupyter notebooks.
+    For local development, the system keyring is used instead.
+    """
+
+    @property
+    def priority(self):
+        """
+        Return priority based on whether we're in a notebook environment.
+        Negative priority means this backend will be skipped by keyring.
+        """
+        from flyte._tools import ipython_check
+
+        if ipython_check():
+            # In IPython/Jupyter - use this backend
+            return 0.5
+        else:
+            # Not in IPython - skip this backend, let system keyring handle it
+            return -1
 
     def get_password(self, service: str, username: str) -> Optional[str]:
         """Get password."""
@@ -72,13 +90,25 @@ class SimplePlainTextKeyring(KeyringBackend):
 
     @property
     def file_path(self) -> Path:
-        from flyte._initialize import get_common_config
-
-        config_path = get_common_config().source_config_path
-        if config_path and str(config_path.parent) == ".flyte":
-            # if the config is in a .flyte directory, use that as the path
-            return config_path.parent / "keyring.cfg"
-        # otherwise use the default path
+        from flyte._initialize import get_init_config, is_initialized
+        from flyte._logging import logger
+
+        # Only try to use source_config_path if flyte.init() has been called
+        if is_initialized():
+            try:
+                config = get_init_config()
+                config_path = config.source_config_path
+                if config_path and str(config_path.parent.name) == ".flyte":
+                    # if the config is in a .flyte directory, use that as the path
+                    return config_path.parent / "keyring.cfg"
+            except Exception as e:
+                # If anything fails, fall back to default path
+                logger.debug(f"Skipping config-based keyring path due to error: {e}")
+        else:
+            # flyte.init() hasn't been called, use default path
+            logger.debug("flyte.init() not called, using default keyring path")
+
+        # Default path
         return _FLYTE_KEYRING_PATH
 
     def __repr__(self):

flyte/_logging.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
+import json
 import logging
 import os
-from typing import Optional
+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
 
 
@@ -33,11 +37,14 @@ def get_env_log_level() -> int:
     return int(os.environ.get("LOG_LEVEL", DEFAULT_LOG_LEVEL))
 
 
-def log_format_from_env() -> str:
+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():
@@ -71,7 +78,7 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
 
     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",
@@ -86,13 +93,49 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
     return handler
 
 
-def initialize_logger(log_level: int = get_env_log_level(), enable_rich: bool = False):
+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()
@@ -100,9 +143,16 @@ def initialize_logger(log_level: int = get_env_log_level(), enable_rich: bool =
     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 = None
-    if enable_rich:
+    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:
@@ -110,11 +160,16 @@ def initialize_logger(log_level: int = get_env_log_level(), enable_rich: bool =
 
     # 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 = None
-    if enable_rich:
+    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:
@@ -165,13 +220,20 @@ class ContextFilter(logging.Filter):
     Applied globally to capture context for both user and Flyte internal logging.
     """
 
-    def filter(self, record):
+    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
 
 
@@ -180,8 +242,12 @@ class FlyteInternalFilter(logging.Filter):
     A logging filter that adds [flyte] prefix to internal Flyte logging only.
     """
 
-    def filter(self, record):
-        if record.name.startswith("flyte"):
+    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
 
@@ -198,6 +264,7 @@ def _setup_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)

flyte/_map.py

@@ -1,13 +1,13 @@
 import asyncio
 import functools
 import logging
-from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast
+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 P, R, TaskTemplate
+from ._task import AsyncFunctionTaskTemplate, F, P, R
 
 
 class MapAsyncIterator(Generic[P, R]):
@@ -15,7 +15,7 @@ class MapAsyncIterator(Generic[P, R]):
 
     def __init__(
         self,
-        func: TaskTemplate[P, R] | functools.partial[R],
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
         args: tuple,
         name: str,
         concurrency: int,
@@ -78,7 +78,7 @@ class MapAsyncIterator(Generic[P, R]):
 
         if isinstance(self.func, functools.partial):
             # Handle partial functions by merging bound args/kwargs with mapped args
-            base_func = cast(TaskTemplate, self.func.func)
+            base_func = cast(AsyncFunctionTaskTemplate, self.func.func)
             bound_args = self.func.args
             bound_kwargs = self.func.keywords or {}
 
@@ -144,7 +144,7 @@ class _Mapper(Generic[P, R]):
         :param func: partial function to validate
         :raises TypeError: if the partial function is not valid for mapping
         """
-        f = cast(TaskTemplate, func.func)
+        f = cast(AsyncFunctionTaskTemplate, func.func)
         inputs = f.native_interface.inputs
         params = list(inputs.keys())
         total_params = len(params)
@@ -172,9 +172,28 @@ class _Mapper(Generic[P, R]):
                         f"in partial function {f.name}."
                     )
 
+    @overload
     def __call__(
         self,
-        func: TaskTemplate[P, R] | functools.partial[R],
+        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,
@@ -194,10 +213,10 @@ class _Mapper(Generic[P, R]):
             return
 
         if isinstance(func, functools.partial):
-            f = cast(TaskTemplate, func.func)
+            f = cast(AsyncFunctionTaskTemplate, func.func)
             self.validate_partial(func)
         else:
-            f = cast(TaskTemplate, func)
+            f = cast(AsyncFunctionTaskTemplate, func)
 
         name = self._get_name(f.name, group_name)
         logger.debug(f"Blocking Map for {name}")
@@ -234,7 +253,7 @@ class _Mapper(Generic[P, R]):
 
     async def aio(
         self,
-        func: TaskTemplate[P, R] | functools.partial[R],
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -244,10 +263,10 @@ class _Mapper(Generic[P, R]):
             return
 
         if isinstance(func, functools.partial):
-            f = cast(TaskTemplate, func.func)
+            f = cast(AsyncFunctionTaskTemplate, func.func)
             self.validate_partial(func)
         else:
-            f = cast(TaskTemplate, func)
+            f = cast(AsyncFunctionTaskTemplate, func)
 
         name = self._get_name(f.name, group_name)
         with group(name):
@@ -277,7 +296,7 @@ class _Mapper(Generic[P, R]):
 
 @syncify
 async def _map(
-    func: TaskTemplate[P, R] | functools.partial[R],
+    func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
     *args: Iterable[Any],
     name: str = "map",
     concurrency: int = 0,

flyte/_module.py

@@ -0,0 +1,70 @@
+import inspect
+import os
+import pathlib
+import sys
+
+
+def extract_obj_module(obj: object, /, source_dir: pathlib.Path) -> str:
+    """
+    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

flyte/_pod.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, Optional
 
 if TYPE_CHECKING:
-    from flyteidl.core.tasks_pb2 import K8sPod
+    from flyteidl2.core.tasks_pb2 import K8sPod
     from kubernetes.client import V1PodSpec
 
 
@@ -20,7 +20,7 @@ class PodTemplate(object):
     annotations: Optional[Dict[str, str]] = None
 
     def to_k8s_pod(self) -> "K8sPod":
-        from flyteidl.core.tasks_pb2 import K8sObjectMetadata, K8sPod
+        from flyteidl2.core.tasks_pb2 import K8sObjectMetadata, K8sPod
         from kubernetes.client import ApiClient
 
         return K8sPod(