flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_initialize.py

@@ -1,260 +1,27 @@
 from __future__ import annotations
 
-import datetime
 import functools
-import os
 import threading
 import typing
 from dataclasses import dataclass, replace
-from datetime import timedelta
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, List, Literal, Optional, TypeVar
+from typing import TYPE_CHECKING, Callable, List, Literal, Optional, TypeVar
 
 from flyte.errors import InitializationError
+from flyte.syncify import syncify
 
-from ._api_commons import syncer
-from ._logging import initialize_logger
+from ._logging import initialize_logger, logger
 from ._tools import ipython_check
 
 if TYPE_CHECKING:
+    from flyte.config import Config
     from flyte.remote._client.auth import AuthType, ClientConfig
     from flyte.remote._client.controlplane import ClientSet
+    from flyte.storage import Storage
 
 Mode = Literal["local", "remote"]
 
 
-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
-
-
-@dataclass(init=True, repr=True, eq=True, frozen=True)
-class Storage(object):
-    """
-    Data storage configuration that applies across any provider.
-    """
-
-    retries: int = 3
-    backoff: datetime.timedelta = datetime.timedelta(seconds=5)
-    enable_debug: bool = False
-    attach_execution_metadata: bool = True
-
-    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
-        "enable_debug": "UNION_STORAGE_DEBUG",
-        "retries": "UNION_STORAGE_RETRIES",
-        "backoff": "UNION_STORAGE_BACKOFF_SECONDS",
-    }
-
-    def get_fsspec_kwargs(self, anonymous: bool = False, /, **kwargs) -> Dict[str, Any]:
-        """
-        Returns the configuration as kwargs for constructing an fsspec filesystem.
-        """
-        return {}
-
-    @classmethod
-    def _auto_as_kwargs(cls) -> Dict[str, Any]:
-        retries = os.getenv(cls._KEY_ENV_VAR_MAPPING["retries"])
-        backoff = os.getenv(cls._KEY_ENV_VAR_MAPPING["backoff"])
-        enable_debug = os.getenv(cls._KEY_ENV_VAR_MAPPING["enable_debug"])
-
-        kwargs: Dict[str, Any] = {}
-        kwargs = set_if_exists(kwargs, "enable_debug", enable_debug)
-        kwargs = set_if_exists(kwargs, "retries", retries)
-        kwargs = set_if_exists(kwargs, "backoff", backoff)
-        return kwargs
-
-    @classmethod
-    def auto(cls) -> Storage:
-        """
-        Construct the config object automatically from environment variables.
-        """
-        return cls(**cls._auto_as_kwargs())
-
-
-@dataclass(init=True, repr=True, eq=True, frozen=True)
-class S3(Storage):
-    """
-    S3 specific configuration
-    """
-
-    endpoint: typing.Optional[str] = None
-    access_key_id: typing.Optional[str] = None
-    secret_access_key: typing.Optional[str] = None
-
-    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
-        "endpoint": "FLYTE_AWS_ENDPOINT",
-        "access_key_id": "FLYTE_AWS_ACCESS_KEY_ID",
-        "secret_access_key": "FLYTE_AWS_SECRET_ACCESS_KEY",
-    } | Storage._KEY_ENV_VAR_MAPPING
-
-    # Refer to https://github.com/developmentseed/obstore/blob/33654fc37f19a657689eb93327b621e9f9e01494/obstore/python/obstore/store/_aws.pyi#L11
-    # for key and secret
-    _CONFIG_KEY_FSSPEC_S3_KEY_ID: ClassVar = "access_key_id"
-    _CONFIG_KEY_FSSPEC_S3_SECRET: ClassVar = "secret_access_key"
-    _CONFIG_KEY_ENDPOINT: ClassVar = "endpoint_url"
-    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
-
-    @classmethod
-    def auto(cls) -> S3:
-        """
-        :return: Config
-        """
-        endpoint = os.getenv(cls._KEY_ENV_VAR_MAPPING["endpoint"], None)
-        access_key_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["access_key_id"], None)
-        secret_access_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["secret_access_key"], None)
-
-        kwargs = super()._auto_as_kwargs()
-        kwargs = set_if_exists(kwargs, "endpoint", endpoint)
-        kwargs = set_if_exists(kwargs, "access_key_id", access_key_id)
-        kwargs = set_if_exists(kwargs, "secret_access_key", secret_access_key)
-
-        return S3(**kwargs)
-
-    @classmethod
-    def for_sandbox(cls) -> S3:
-        """
-        :return:
-        """
-        kwargs = super()._auto_as_kwargs()
-        final_kwargs = kwargs | {
-            "endpoint": "http://localhost:4566",
-            "access_key_id": "minio",
-            "secret_access_key": "miniostorage",
-        }
-        return S3(**final_kwargs)
-
-    def get_fsspec_kwargs(self, anonymous: bool = False, /, **kwargs) -> Dict[str, Any]:
-        # Construct the config object
-        config: Dict[str, Any] = {}
-        if self._CONFIG_KEY_FSSPEC_S3_KEY_ID in kwargs or self.access_key_id:
-            config[self._CONFIG_KEY_FSSPEC_S3_KEY_ID] = kwargs.pop(
-                self._CONFIG_KEY_FSSPEC_S3_KEY_ID, self.access_key_id
-            )
-        if self._CONFIG_KEY_FSSPEC_S3_SECRET in kwargs or self.secret_access_key:
-            config[self._CONFIG_KEY_FSSPEC_S3_SECRET] = kwargs.pop(
-                self._CONFIG_KEY_FSSPEC_S3_SECRET, self.secret_access_key
-            )
-        if self._CONFIG_KEY_ENDPOINT in kwargs or self.endpoint:
-            config["endpoint_url"] = kwargs.pop(self._CONFIG_KEY_ENDPOINT, self.endpoint)
-
-        retries = kwargs.pop("retries", self.retries)
-        backoff = kwargs.pop("backoff", self.backoff)
-
-        if anonymous:
-            config[self._KEY_SKIP_SIGNATURE] = True
-
-        retry_config = {
-            "max_retries": retries,
-            "backoff": {
-                "base": 2,
-                "init_backoff": backoff,
-                "max_backoff": timedelta(seconds=16),
-            },
-            "retry_timeout": timedelta(minutes=3),
-        }
-
-        client_options = {"timeout": "99999s", "allow_http": True}
-
-        if config:
-            kwargs["config"] = config
-        kwargs["client_options"] = client_options or None
-        kwargs["retry_config"] = retry_config or None
-
-        return kwargs
-
-
-@dataclass(init=True, repr=True, eq=True, frozen=True)
-class GCS(Storage):
-    """
-    Any GCS specific configuration.
-    """
-
-    gsutil_parallelism: bool = False
-
-    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
-        "gsutil_parallelism": "GCP_GSUTIL_PARALLELISM",
-    }
-
-    @classmethod
-    def auto(cls) -> GCS:
-        gsutil_parallelism = os.getenv(cls._KEY_ENV_VAR_MAPPING["gsutil_parallelism"], None)
-
-        kwargs: Dict[str, Any] = {}
-        kwargs = set_if_exists(kwargs, "gsutil_parallelism", gsutil_parallelism)
-        return GCS(**kwargs)
-
-    def get_fsspec_kwargs(self, anonymous: bool = False, /, **kwargs) -> Dict[str, Any]:
-        return kwargs
-
-
-@dataclass(init=True, repr=True, eq=True, frozen=True)
-class ABFS(Storage):
-    """
-    Any Azure Blob Storage specific configuration.
-    """
-
-    account_name: typing.Optional[str] = None
-    account_key: typing.Optional[str] = None
-    tenant_id: typing.Optional[str] = None
-    client_id: typing.Optional[str] = None
-    client_secret: typing.Optional[str] = None
-
-    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
-        "account_name": "AZURE_STORAGE_ACCOUNT_NAME",
-        "account_key": "AZURE_STORAGE_ACCOUNT_KEY",
-        "tenant_id": "AZURE_TENANT_ID",
-        "client_id": "AZURE_CLIENT_ID",
-        "client_secret": "AZURE_CLIENT_SECRET",
-    }
-    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
-
-    @classmethod
-    def auto(cls) -> ABFS:
-        account_name = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_name"], None)
-        account_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_key"], None)
-        tenant_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["tenant_id"], None)
-        client_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_id"], None)
-        client_secret = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_secret"], None)
-
-        kwargs: Dict[str, Any] = {}
-        kwargs = set_if_exists(kwargs, "account_name", account_name)
-        kwargs = set_if_exists(kwargs, "account_key", account_key)
-        kwargs = set_if_exists(kwargs, "tenant_id", tenant_id)
-        kwargs = set_if_exists(kwargs, "client_id", client_id)
-        kwargs = set_if_exists(kwargs, "client_secret", client_secret)
-        return ABFS(**kwargs)
-
-    def get_fsspec_kwargs(self, anonymous: bool = False, /, **kwargs) -> Dict[str, Any]:
-        config: Dict[str, Any] = {}
-        if "account_name" in kwargs or self.account_name:
-            config["account_name"] = kwargs.get("account_name", self.account_name)
-        if "account_key" in kwargs or self.account_key:
-            config["account_key"] = kwargs.get("account_key", self.account_key)
-        if "client_id" in kwargs or self.client_id:
-            config["client_id"] = kwargs.get("client_id", self.client_id)
-        if "client_secret" in kwargs or self.client_secret:
-            config["client_secret"] = kwargs.get("client_secret", self.client_secret)
-        if "tenant_id" in kwargs or self.tenant_id:
-            config["tenant_id"] = kwargs.get("tenant_id", self.tenant_id)
-
-        if anonymous:
-            config[self._KEY_SKIP_SIGNATURE] = True
-
-        client_options = {"timeout": "99999s", "allow_http": "true"}
-
-        if config:
-            kwargs["config"] = config
-        kwargs["client_options"] = client_options
-
-        return kwargs
-
-
 @dataclass(init=True, repr=True, eq=True, frozen=True, kw_only=True)
 class CommonInit:
     """
@@ -265,6 +32,7 @@ class CommonInit:
     org: str | None = None
     project: str | None = None
     domain: str | None = None
+    batch_size: int = 1000
 
 
 @dataclass(init=True, kw_only=True, repr=True, eq=True, frozen=True)
@@ -303,11 +71,10 @@ async def _initialize_client(
     """
     from flyte.remote._client.controlplane import ClientSet
 
-    if endpoint is not None:
+    if endpoint:
         return await ClientSet.for_endpoint(
             endpoint,
             insecure=insecure,
-            api_key=api_key,
             insecure_skip_verify=insecure_skip_verify,
             auth_type=auth_type,
             headless=headless,
@@ -320,10 +87,29 @@ async def _initialize_client(
             rpc_retries=rpc_retries,
             http_proxy_url=http_proxy_url,
         )
-    raise NotImplementedError("Currently only endpoints are supported.")
+    elif api_key:
+        return await ClientSet.for_api_key(
+            api_key,
+            insecure=insecure,
+            insecure_skip_verify=insecure_skip_verify,
+            auth_type=auth_type,
+            headless=headless,
+            ca_cert_file_path=ca_cert_file_path,
+            command=command,
+            proxy_command=proxy_command,
+            client_id=client_id,
+            client_credentials_secret=client_credentials_secret,
+            client_config=client_config,
+            rpc_retries=rpc_retries,
+            http_proxy_url=http_proxy_url,
+        )
 
+    raise InitializationError(
+        "MissingEndpointOrApiKeyError", "user", "Either endpoint or api_key must be provided to initialize the client."
+    )
 
-@syncer.wrap
+
+@syncify
 async def init(
     org: str | None = None,
     project: str | None = None,
@@ -345,6 +131,7 @@ async def init(
     rpc_retries: int = 3,
     http_proxy_url: str | None = None,
     storage: Storage | None = None,
+    batch_size: int = 1000,
 ) -> None:
     """
     Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
@@ -353,14 +140,13 @@ async def init(
     :param project: Optional project name (not used in this implementation)
     :param domain: Optional domain name (not used in this implementation)
     :param root_dir: Optional root directory from which to determine how to load files, and find paths to files.
+      This is useful for determining the root directory for the current project, and for locating files like config etc.
+      also use to determine all the code that needs to be copied to the remote location.
       defaults to the editable install directory if the cwd is in a Python editable install, else just the cwd.
     :param log_level: Optional logging level for the logger, default is set using the default initialization policies
     :param api_key: Optional API key for authentication
     :param endpoint: Optional API endpoint URL
     :param headless: Optional Whether to run in headless mode
-    :param mode: Optional execution model (local, remote). Default is local. When local is used,
-           the execution will be done locally. When remote is used, the execution will be sent to a remote server,
-           In the remote case, the endpoint or api_key must be set.
     :param insecure_skip_verify: Whether to skip SSL certificate verification
     :param auth_client_config: Optional client configuration for authentication
     :param auth_type: The authentication type to use (Pkce, ClientSecret, ExternalCommand, DeviceFlow)
@@ -378,10 +164,12 @@ async def init(
     :param insecure: insecure flag for the client
     :param storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
     :param org: Optional organization override for the client. Should be set by auth instead.
+    :param batch_size: Optional batch size for operations that use listings, defaults to 1000, so limit larger than
+      batch_size will be split into multiple requests.
 
     :return: None
     """
-    from flyte._utils import get_cwd_editable_install
+    from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
 
     interactive_mode = ipython_check()
 
@@ -391,6 +179,8 @@ async def init(
 
     global _init_config  # noqa: PLW0603
 
+    endpoint = sanitize_endpoint(endpoint)
+
     with _init_lock:
         client = None
         if endpoint or api_key:
@@ -418,10 +208,67 @@ async def init(
             domain=domain,
             client=client,
             storage=storage,
-            org=org,
+            org=org or org_from_endpoint(endpoint),
+            batch_size=batch_size,
         )
 
 
+@syncify
+async def init_from_config(
+    path_or_config: str | Config | None = None, root_dir: Path | None = None, log_level: int | None = None
+) -> None:
+    """
+    Initialize the Flyte system using a configuration file or Config object. This method should be called before any
+    other Flyte remote API methods are called. Thread-safe implementation.
+
+    :param path_or_config: Path to the configuration file or Config object
+    :param root_dir: Optional root directory from which to determine how to load files, and find paths to
+        files like config etc. For example if one uses the copy-style=="all", it is essential to determine the
+        root directory for the current project. If not provided, it defaults to the editable install directory or
+        if not available, the current working directory.
+    :param log_level: Optional logging level for the framework logger,
+        default is set using the default initialization policies
+    :return: None
+    """
+    import flyte.config as config
+
+    cfg: config.Config
+    if path_or_config is None or isinstance(path_or_config, str):
+        # If a string is passed, treat it as a path to the config file
+        if path_or_config:
+            if not Path(path_or_config).exists():
+                raise InitializationError(
+                    "ConfigFileNotFoundError",
+                    "user",
+                    f"Configuration file '{path_or_config}' does not exist., current working directory is {Path.cwd()}",
+                )
+        if root_dir and path_or_config:
+            cfg = config.auto(str(root_dir / path_or_config))
+        else:
+            cfg = config.auto(path_or_config)
+    else:
+        # If a Config object is passed, use it directly
+        cfg = path_or_config
+
+    logger.debug(f"Flyte config initialized as {cfg}")
+    await init.aio(
+        org=cfg.task.org,
+        project=cfg.task.project,
+        domain=cfg.task.domain,
+        endpoint=cfg.platform.endpoint,
+        insecure=cfg.platform.insecure,
+        insecure_skip_verify=cfg.platform.insecure_skip_verify,
+        ca_cert_file_path=cfg.platform.ca_cert_file_path,
+        auth_type=cfg.platform.auth_mode,
+        command=cfg.platform.command,
+        proxy_command=cfg.platform.proxy_command,
+        client_id=cfg.platform.client_id,
+        client_credentials_secret=cfg.platform.client_credentials_secret,
+        root_dir=root_dir,
+        log_level=log_level,
+    )
+
+
 def _get_init_config() -> Optional[_InitConfig]:
     """
     Get the current initialization configuration. Thread-safe implementation.
@@ -449,7 +296,7 @@ def get_common_config() -> CommonInit:
     return cfg
 
 
-def get_storage() -> Storage:
+def get_storage() -> Storage | None:
     """
     Get the current storage configuration. Thread-safe implementation.
 
@@ -463,9 +310,6 @@ def get_storage() -> Storage:
             "Configuration has not been initialized. Call flyte.init() with a valid endpoint or",
             " api-key before using this function.",
         )
-    if cfg.storage is None:
-        # return default local storage
-        return typing.cast(Storage, cfg.replace(storage=Storage()).storage)
     return cfg.storage
 
 
@@ -495,19 +339,33 @@ def is_initialized() -> bool:
     return _get_init_config() is not None
 
 
-def initialize_in_cluster(storage: Storage | None = None) -> None:
+def initialize_in_cluster() -> None:
     """
     Initialize the system for in-cluster execution. This is a placeholder function and does not perform any actions.
 
     :return: None
     """
-    init(storage=storage)
+    init()
 
 
 # Define a generic type variable for the decorated function
 T = TypeVar("T", bound=Callable)
 
 
+def ensure_client():
+    """
+    Ensure that the client is initialized. If not, raise an InitializationError.
+    This function is used to check if the client is initialized before executing any Flyte remote API methods.
+    """
+    if _get_init_config() is None or _get_init_config().client is None:
+        raise InitializationError(
+            "ClientNotInitializedError",
+            "user",
+            "Client has not been initialized. Call flyte.init() with a valid endpoint"
+            " or api-key before using this function.",
+        )
+
+
 def requires_client(func: T) -> T:
     """
     Decorator that checks if the client has been initialized before executing the function.
@@ -518,7 +376,7 @@ def requires_client(func: T) -> T:
     """
 
     @functools.wraps(func)
-    def wrapper(*args, **kwargs) -> T:
+    async def wrapper(*args, **kwargs) -> T:
         init_config = _get_init_config()
         if init_config is None or init_config.client is None:
             raise InitializationError(

flyte/_interface.py

@@ -62,12 +62,12 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
         # Options 1 and 2
         bases = return_annotation.__bases__  # type: ignore
         if len(bases) == 1 and bases[0] is tuple and hasattr(return_annotation, "_fields"):
-            logger.debug(f"Task returns named tuple {return_annotation}")
+            # Task returns named tuple
             return dict(get_type_hints(cast(Type, return_annotation), include_extras=True))
 
     if hasattr(return_annotation, "__origin__") and return_annotation.__origin__ is tuple:  # type: ignore
         # Handle option 3
-        logger.debug(f"Task returns unnamed typing.Tuple {return_annotation}")
+        # Task returns unnamed typing.Tuple
         if len(return_annotation.__args__) == 1:  # type: ignore
             raise TypeError("Tuples should be used to indicate multiple return values, found only one return variable.")
         ra = get_args(return_annotation)
@@ -80,5 +80,5 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
 
     else:
         # Handle all other single return types
-        logger.debug(f"Task returns unnamed native tuple {return_annotation}")
+        # Task returns unnamed native tuple
         return {default_output_name(): cast(Type, return_annotation)}

flyte/_internal/controllers/__init__.py

@@ -1,8 +1,9 @@
+import concurrent.futures
 import threading
-from typing import Any, Literal, Optional, Protocol, Tuple, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol, Tuple, TypeVar
 
-from flyte._datastructures import ActionID, NativeInterface
 from flyte._task import TaskTemplate
+from flyte.models import ActionID, NativeInterface
 
 from ._trace import TraceInfo
 
@@ -10,6 +11,9 @@ __all__ = ["Controller", "ControllerType", "TraceInfo", "create_controller", "ge
 
 from ..._protos.workflow import task_definition_pb2
 
+if TYPE_CHECKING:
+    import concurrent.futures
+
 ControllerType = Literal["local", "remote"]
 
 R = TypeVar("R")
@@ -28,6 +32,15 @@ class Controller(Protocol):
         """
         ...
 
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        """
+        This should call the async submit method above, but return a concurrent Future object that can be
+        used in a blocking wait or wrapped in an async future. This is called when
+          a) a synchronous task is kicked off locally,
+          b) a running task (of either kind) kicks off a downstream synchronous task.
+        """
+        ...
+
     async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
         """
         Submit a task reference to the controller asynchronously and wait for the result. This is async and will block
@@ -47,12 +60,12 @@ class Controller(Protocol):
     async def watch_for_errors(self): ...
 
     async def get_action_outputs(
-        self, _interface: NativeInterface, _func_name: str, *args, **kwargs
+        self, _interface: NativeInterface, _func: Callable, *args, **kwargs
     ) -> Tuple[TraceInfo, bool]:
         """
         This method returns the outputs of the action, if it is available.
         :param _interface: NativeInterface
-        :param _func_name: Function name
+        :param _func: Function name
         :param args: Arguments
         :param kwargs: Keyword arguments
         :return: TraceInfo object and a boolean indicating if the action was found.
@@ -81,13 +94,13 @@ class _ControllerState:
     lock = threading.Lock()
 
 
-async def get_controller() -> Controller:
+def get_controller() -> Controller:
     """
     Get the controller instance. Raise an error if it has not been created.
     """
     if _ControllerState.controller is not None:
         return _ControllerState.controller
-    raise RuntimeError("Controller is not initialized. Please call get_or_create_controller() first.")
+    raise RuntimeError("Controller is not initialized. Please call create_controller() first.")
 
 
 def create_controller(