flyte 0.2.0b1__py3-none-any.whl → 0.2.0b3__py3-none-any.whl

flyte/_utils/async_cache.py

@@ -0,0 +1,139 @@
+import asyncio
+import time
+from collections import OrderedDict
+from typing import Awaitable, Callable, Dict, Generic, Optional, TypeVar
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class AsyncLRUCache(Generic[K, V]):
+    """
+    A high-performance async-compatible LRU cache.
+
+    Examples:
+    ```python
+    # Create a cache instance
+    cache = AsyncLRUCache[str, dict](maxsize=100)
+
+    async def fetch_data(user_id: str) -> dict:
+        # Define the expensive operation as a local function
+        async def get_user_data():
+            await asyncio.sleep(1)  # Simulating network/DB delay
+            return {"id": user_id, "name": f"User {user_id}"}
+
+    # Use the cache
+    return await cache.get(f"user:{user_id}", get_user_data)
+    ```
+    This cache can be used from async coroutines and handles concurrent access safely.
+    """
+
+    def __init__(self, maxsize: int = 128, ttl: Optional[float] = None):
+        """
+        Initialize the async LRU cache.
+
+        Args:
+            maxsize: Maximum number of items to keep in the cache
+            ttl: Time-to-live for cache entries in seconds, or None for no expiration
+        """
+        self._cache: OrderedDict[K, tuple[V, float]] = OrderedDict()
+        self._maxsize = maxsize
+        self._ttl = ttl
+        self._locks: Dict[K, asyncio.Lock] = {}
+        self._access_lock = asyncio.Lock()
+
+    async def get(self, key: K, value_func: Callable[[], V | Awaitable[V]]) -> V:
+        """
+        Get a value from the cache, computing it if necessary.
+
+        Args:
+            key: The cache key
+            value_func: Function or coroutine to compute the value if not cached
+
+        Returns:
+            The cached or computed value
+        """
+        # Fast path: check if key exists and is not expired
+        if key in self._cache:
+            value, timestamp = self._cache[key]
+            if self._ttl is None or time.time() - timestamp < self._ttl:
+                # Move the accessed item to the end (most recently used)
+                async with self._access_lock:
+                    self._cache.move_to_end(key)
+                return value
+
+        # Slow path: compute the value
+        # Get or create a lock for this key to prevent redundant computation
+        async with self._access_lock:
+            lock = self._locks.get(key)
+            if lock is None:
+                lock = asyncio.Lock()
+                self._locks[key] = lock
+
+        async with lock:
+            # Check again in case another coroutine computed the value while we waited
+            if key in self._cache:
+                value, timestamp = self._cache[key]
+                if self._ttl is None or time.time() - timestamp < self._ttl:
+                    async with self._access_lock:
+                        self._cache.move_to_end(key)
+                    return value
+
+            # Compute the value
+            if asyncio.iscoroutinefunction(value_func):
+                value = await value_func()
+            else:
+                value = value_func()  # type: ignore
+
+            # Store in cache
+            async with self._access_lock:
+                self._cache[key] = (value, time.time())
+                # Evict least recently used items if needed
+                while len(self._cache) > self._maxsize:
+                    self._cache.popitem(last=False)
+                # Clean up the lock
+                self._locks.pop(key, None)
+
+            return value
+
+    async def set(self, key: K, value: V) -> None:
+        """
+        Explicitly set a value in the cache.
+
+        Args:
+            key: The cache key
+            value: The value to cache
+        """
+        async with self._access_lock:
+            self._cache[key] = (value, time.time())
+            # Evict least recently used items if needed
+            while len(self._cache) > self._maxsize:
+                self._cache.popitem(last=False)
+
+    async def invalidate(self, key: K) -> None:
+        """Remove a specific key from the cache."""
+        async with self._access_lock:
+            self._cache.pop(key, None)
+
+    async def clear(self) -> None:
+        """Clear the entire cache."""
+        async with self._access_lock:
+            self._cache.clear()
+            self._locks.clear()
+
+    async def contains(self, key: K) -> bool:
+        """Check if a key exists in the cache and is not expired."""
+        if key not in self._cache:
+            return False
+
+        if self._ttl is None:
+            return True
+
+        _, timestamp = self._cache[key]
+        return time.time() - timestamp < self._ttl
+
+
+# Example usage:
+"""
+
+"""

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0b1'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b1')
+__version__ = version = '0.2.0b3'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b3')

flyte/config/__init__.py

@@ -1,8 +1,13 @@
+from __future__ import annotations
+
 import os
+import pathlib
 import typing
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
+import rich.repr
+
 from flyte._logging import logger
 from flyte.config import _internal
 from flyte.config._config import ConfigFile, get_config_file, read_file_if_exists
@@ -13,6 +18,7 @@ if TYPE_CHECKING:
     from flyte.remote._client.auth import AuthType
 
 
+@rich.repr.auto
 @dataclass(init=True, repr=True, eq=True, frozen=True)
 class PlatformConfig(object):
     """
@@ -57,7 +63,6 @@ class PlatformConfig(object):
         :param config_file:
         :return:
         """
-        from .._initialize import set_if_exists
 
         config_file = get_config_file(config_file)
         kwargs: typing.Dict[str, typing.Any] = {}
@@ -106,6 +111,7 @@ class PlatformConfig(object):
         return PlatformConfig(endpoint=endpoint, insecure=insecure)
 
 
+@rich.repr.auto
 @dataclass(init=True, repr=True, eq=True, frozen=True)
 class TaskConfig(object):
     org: str | None = None
@@ -119,8 +125,6 @@ class TaskConfig(object):
         :param config_file:
         :return:
         """
-        from flyte._initialize import set_if_exists
-
         config_file = get_config_file(config_file)
         kwargs: typing.Dict[str, typing.Any] = {}
         kwargs = set_if_exists(kwargs, "org", _internal.Task.ORG.read(config_file))
@@ -129,6 +133,7 @@ class TaskConfig(object):
         return TaskConfig(**kwargs)
 
 
+@rich.repr.auto
 @dataclass(init=True, repr=True, eq=True, frozen=True)
 class Config(object):
     """
@@ -142,6 +147,7 @@ class Config(object):
 
     platform: PlatformConfig = field(default=PlatformConfig())
     task: TaskConfig = field(default=TaskConfig())
+    source: pathlib.Path | None = None
 
     def with_params(
         self,
@@ -165,4 +171,20 @@ class Config(object):
         :return: Config
         """
         config_file = get_config_file(config_file)
-        return Config(platform=PlatformConfig.auto(config_file), task=TaskConfig.auto(config_file))
+        if config_file is None:
+            logger.debug("No config file found, using default values")
+            return Config()
+        return Config(
+            platform=PlatformConfig.auto(config_file), task=TaskConfig.auto(config_file), source=config_file.path
+        )
+
+
+def set_if_exists(d: dict, k: str, val: typing.Any) -> dict:
+    """
+    Given a dict ``d`` sets the key ``k`` with value of config ``v``, if the config value ``v`` is set
+    and return the updated dictionary.
+    """
+    exists = isinstance(val, bool) or bool(val is not None and val)
+    if exists:
+        d[k] = val
+    return d

flyte/config/_config.py

@@ -99,6 +99,14 @@ class ConfigFile(object):
         self._location = location
         self._yaml_config = self._read_yaml_config(location)
 
+    @property
+    def path(self) -> pathlib.Path:
+        """
+        Returns the path to the config file.
+        :return: Path to the config file
+        """
+        return pathlib.Path(self._location)
+
     @staticmethod
     def _read_yaml_config(location: str) -> typing.Optional[typing.Dict[str, typing.Any]]:
         with open(location, "r") as fh:
@@ -136,6 +144,11 @@ def resolve_config_path() -> pathlib.Path | None:
     4. ~/.flyte/config.yaml if it exists
     5. ./config.yaml if it exists
     """
+    current_location_config = Path("config.yaml")
+    if current_location_config.exists():
+        return current_location_config
+    logger.debug("No ./config.yaml found, returning None")
+
     uctl_path_from_env = getenv(UCTL_CONFIG_ENV_VAR, None)
     if uctl_path_from_env:
         return pathlib.Path(uctl_path_from_env)
@@ -156,10 +169,6 @@ def resolve_config_path() -> pathlib.Path | None:
         return home_dir_flytectl_config
     logger.debug("No ~/.flyte/config.yaml found, checking current directory")
 
-    current_location_config = Path("config.yaml")
-    if current_location_config.exists():
-        return current_location_config
-    logger.debug("No ./config.yaml found, returning None")
     return None
 
 

flyte/extras/_container.py

@@ -5,9 +5,9 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
 from flyteidl.core import tasks_pb2
 
 from flyte import Image, storage
-from flyte._datastructures import NativeInterface, SerializationContext
 from flyte._logging import logger
 from flyte._task import TaskTemplate
+from flyte.models import NativeInterface, SerializationContext
 
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
 
@@ -258,8 +258,8 @@ class ContainerTask(TaskTemplate):
         }
 
         return tasks_pb2.DataLoadingConfig(
-            input_path=self._input_data_dir,
-            output_path=self._output_data_dir,
+            input_path=str(self._input_data_dir),
+            output_path=str(self._output_data_dir),
             enabled=True,
             format=literal_to_protobuf.get(self._metadata_format, "JSON"),
         )

flyte/{_datastructures.py → models.py}

@@ -56,11 +56,12 @@ class ActionID:
             name = generate_random_name()
         return replace(self, name=name)
 
-    def new_sub_action_from(self, task_name: str, input_hash: str, group: str | None) -> ActionID:
+    def new_sub_action_from(self, task_call_seq: int, task_hash: str, input_hash: str, group: str | None) -> ActionID:
         """Make a deterministic name"""
         import hashlib
 
-        components = f"{self.run_name}-{self.name}-{input_hash}-{task_name}" + (f"-{group}" if group else "")
+        components = f"{self.name}-{input_hash}-{task_hash}-{task_call_seq}" + (f"-{group}" if group else "")
+        logger.debug(f"----- Generating sub-action ID from components: {components}")
         # has the components into something deterministic
         bytes_digest = hashlib.md5(components.encode()).digest()
         new_name = base36_encode(bytes_digest)

flyte/remote/_client/auth/_auth_utils.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import base64
+from typing import Literal
+
+
+def decode_api_key(encoded_str: str) -> tuple[str, str, str, str | Literal["None"]]:
+    """Decode encoded base64 string into app credentials. endpoint, client_id, client_secret, org"""
+    endpoint, client_id, client_secret, org = base64.b64decode(encoded_str.encode("utf-8")).decode("utf-8").split(":")
+    # For consistency, let's make sure org is always a non-empty string
+    if not org:
+        org = "None"
+
+    return endpoint, client_id, client_secret, org

flyte/remote/_client/auth/_channel.py

@@ -44,7 +44,9 @@ def bootstrap_ssl_from_server(endpoint: str) -> grpc.ChannelCredentials:
 
 
 async def create_channel(
-    endpoint: str,
+    endpoint: str | None,
+    api_key: str | None = None,
+    /,
     insecure: typing.Optional[bool] = None,
     insecure_skip_verify: typing.Optional[bool] = False,
     ca_cert_file_path: typing.Optional[str] = None,
@@ -66,6 +68,7 @@ async def create_channel(
     and create authentication interceptors that perform async operations.
 
     :param endpoint: The endpoint URL for the gRPC channel
+    :param api_key: API key for authentication; if provided, it will be used to detect the endpoint and credentials.
     :param insecure: Whether to use an insecure channel (no SSL)
     :param insecure_skip_verify: Whether to skip SSL certificate verification
     :param ca_cert_file_path: Path to CA certificate file for SSL verification
@@ -104,6 +107,18 @@ async def create_channel(
             - refresh_access_token_params: Parameters to add when refreshing access token
     :return: grpc.aio.Channel with authentication interceptors configured
     """
+    assert endpoint or api_key, "Either endpoint or api_key must be specified"
+
+    if api_key:
+        from flyte.remote._client.auth._auth_utils import decode_api_key
+
+        endpoint, client_id, client_secret, org = decode_api_key(api_key)
+        kwargs["auth_type"] = "ClientSecret"
+        kwargs["client_id"] = client_id
+        kwargs["client_secret"] = client_secret
+        kwargs["client_credentials_secret"] = client_secret
+
+    assert endpoint, "Endpoint must be specified by this point"
 
     if not ssl_credentials:
         if insecure_skip_verify:
@@ -119,7 +134,12 @@ async def create_channel(
 
     # Create an unauthenticated channel first to use to get the server metadata
     if insecure:
-        unauthenticated_channel = grpc.aio.insecure_channel(endpoint, **kwargs)
+        insecure_kwargs = {}
+        if kw_opts := kwargs.get("options"):
+            insecure_kwargs["options"] = kw_opts
+        if compression:
+            insecure_kwargs["compression"] = compression
+        unauthenticated_channel = grpc.aio.insecure_channel(endpoint, **insecure_kwargs)
     else:
         unauthenticated_channel = grpc.aio.secure_channel(
             target=endpoint,
@@ -173,7 +193,12 @@ async def create_channel(
     interceptors.extend(auth_interceptors)
 
     if insecure:
-        return grpc.aio.insecure_channel(endpoint, interceptors=interceptors, **kwargs)
+        insecure_kwargs = {}
+        if kw_opts := kwargs.get("options"):
+            insecure_kwargs["options"] = kw_opts
+        if compression:
+            insecure_kwargs["compression"] = compression
+        return grpc.aio.insecure_channel(endpoint, interceptors=interceptors, **insecure_kwargs)
 
     return grpc.aio.secure_channel(
         target=endpoint,

flyte/remote/_client/auth/_token_client.py

@@ -94,7 +94,7 @@ async def get_token(
     http_proxy_url: typing.Optional[str] = None,
     verify: typing.Optional[typing.Union[bool, str]] = None,
     refresh_token: typing.Optional[str] = None,
-) -> typing.Tuple[str, str, int]:
+) -> typing.Tuple[str, str | None, int]:
     """
     Retrieves an access token from the specified token endpoint.
 
@@ -165,7 +165,7 @@ async def get_token(
     if "refresh_token" in j:
         new_refresh_token = j["refresh_token"]
     else:
-        raise AuthenticationError("Token not yet available, try again in some time")
+        logger.info("No refresh token received, this is expected for client credentials flow")
 
     return j["access_token"], new_refresh_token, j["expires_in"]
 
@@ -213,7 +213,7 @@ async def poll_token_endpoint(
     scopes: typing.Optional[typing.List[str]] = None,
     http_proxy_url: typing.Optional[str] = None,
     verify: typing.Optional[typing.Union[bool, str]] = None,
-) -> typing.Tuple[str, str, int]:
+) -> typing.Tuple[str, str | None, int]:
     """
     Polls the token endpoint until authentication is complete or times out.
 

flyte/remote/_client/controlplane.py

@@ -39,21 +39,21 @@ class ClientSet:
 
     @classmethod
     async def for_endpoint(cls, endpoint: str, *, insecure: bool = False, **kwargs) -> ClientSet:
-        if insecure:
-            del kwargs["api_key"]
-            del kwargs["auth_type"]
-            del kwargs["headless"]
-            del kwargs["command"]
-            del kwargs["client_id"]
-            del kwargs["client_credentials_secret"]
-            del kwargs["client_config"]
-            del kwargs["rpc_retries"]
-            del kwargs["http_proxy_url"]
-        return cls(await create_channel(endpoint, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs)
+        return cls(
+            await create_channel(endpoint, None, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs
+        )
 
     @classmethod
-    async def for_api_key(cls, api_key: str, **kwargs) -> ClientSet:
-        raise NotImplementedError
+    async def for_api_key(cls, api_key: str, *, insecure: bool = False, **kwargs) -> ClientSet:
+        from flyte.remote._client.auth._auth_utils import decode_api_key
+
+        # Parsing the API key is done in create_channel, but cleaner to redo it here rather than getting create_channel
+        # to return the endpoint
+        endpoint, _, _, _ = decode_api_key(api_key)
+
+        return cls(
+            await create_channel(None, api_key, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs
+        )
 
     @classmethod
     async def for_serverless(cls) -> ClientSet:

flyte/remote/_logs.py

@@ -111,6 +111,6 @@ class Logs:
             log_source=cls.tail.aio(cls, action_id=action_id, attempt=attempt),
             max_lines=max_lines,
             show_ts=show_ts,
-            name=f"{action_id.run.name}:{action_id.name}",
+            name=f"{action_id.run.name}:{action_id.name} ({attempt})",
         )
         await viewer.run()

flyte/remote/_run.py

@@ -243,14 +243,10 @@ class Run:
                 return
             yield ad
 
-    async def show_logs(self, attempt: int = 1, max_lines: int = 100, show_ts: bool = False, raw: bool = False):
-        return await Logs.create_viewer(
-            action_id=self.action.action_id,
-            attempt=attempt,
-            max_lines=max_lines,
-            show_ts=show_ts,
-            raw=raw,
-        )
+    async def show_logs(
+        self, attempt: int | None = None, max_lines: int = 100, show_ts: bool = False, raw: bool = False
+    ):
+        await self.action.show_logs(attempt, max_lines, show_ts, raw)
 
     async def details(self) -> RunDetails:
         """

flyte/remote/_task.py

@@ -11,9 +11,9 @@ import flyte
 import flyte.errors
 from flyte._api_commons import syncer
 from flyte._context import internal_ctx
-from flyte._datastructures import NativeInterface
 from flyte._initialize import get_client, get_common_config
 from flyte._protos.workflow import task_definition_pb2, task_service_pb2
+from flyte.models import NativeInterface
 
 
 class LazyEntity:
@@ -187,7 +187,7 @@ class Task:
             # We will also check if we are not initialized, It is not expected to be not initialized
             from flyte._internal.controllers import get_controller
 
-            controller = await get_controller()
+            controller = get_controller()
             if controller:
                 return await controller.submit_task_ref(self.pb2, *args, **kwargs)
         raise flyte.errors

flyte/storage/__init__.py

@@ -1,4 +1,8 @@
 __all__ = [
+    "ABFS",
+    "GCS",
+    "S3",
+    "Storage",
     "get",
     "get_random_local_directory",
     "get_random_local_path",
@@ -11,6 +15,7 @@ __all__ = [
     "put_stream",
 ]
 
+from ._config import ABFS, GCS, S3, Storage
 from ._storage import (
     get,
     get_random_local_directory,