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

flyte/_initialize.py

@@ -1,14 +1,11 @@
 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
 
@@ -20,242 +17,11 @@ 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:
     """
@@ -304,11 +70,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,
@@ -321,7 +86,26 @@ 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
@@ -396,9 +180,9 @@ async def init(
 
     with _init_lock:
         if config is None:
-            from flyte.config import Config
+            import flyte.config as _f_cfg
 
-            config = Config.auto()
+            config = _f_cfg.Config()
         platform_cfg = config.platform
         task_cfg = config.task
         client = None
@@ -458,7 +242,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.
 
@@ -472,9 +256,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
 
 
@@ -504,13 +285,13 @@ 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

flyte/_internal/controllers/__init__.py

@@ -1,8 +1,8 @@
 import threading
-from typing import Any, Literal, Optional, Protocol, Tuple, TypeVar
+from typing import 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
 
@@ -47,12 +47,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 +81,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(

flyte/_internal/controllers/_local_controller.py

@@ -1,14 +1,14 @@
-from typing import Any, Tuple, TypeVar
+from typing import Any, Callable, Tuple, TypeVar
 
 import flyte.errors
 from flyte._context import internal_ctx
-from flyte._datastructures import ActionID, NativeInterface, RawDataPath
 from flyte._internal.controllers import TraceInfo
 from flyte._internal.runtime import convert
 from flyte._internal.runtime.entrypoints import direct_dispatch
 from flyte._logging import log, logger
 from flyte._protos.workflow import task_definition_pb2
 from flyte._task import TaskTemplate
+from flyte.models import ActionID, NativeInterface, RawDataPath
 
 R = TypeVar("R")
 
@@ -28,7 +28,11 @@ class LocalController:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
         inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _task.name, inputs)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, _task.name, serialized_inputs, 0
+        )
         sub_action_raw_data_path = RawDataPath(path=sub_action_output_path)
 
         out, err = await direct_dispatch(
@@ -64,7 +68,7 @@ class LocalController:
         pass
 
     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.
@@ -79,8 +83,13 @@ class LocalController:
         if _interface.inputs:
             converted_inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
             assert converted_inputs
+
+        serialized_inputs = converted_inputs.proto_inputs.SerializeToString(deterministic=True)
         action_id, action_output_path = convert.generate_sub_action_id_and_output_path(
-            tctx, _func_name, converted_inputs
+            tctx,
+            _func.__name__,
+            serialized_inputs,
+            0,
         )
         assert action_output_path
         return (

flyte/_internal/controllers/_trace.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any, Optional
 
-from flyte._datastructures import ActionID, NativeInterface
+from flyte.models import ActionID, NativeInterface
 
 
 @dataclass

flyte/_internal/controllers/remote/__init__.py

@@ -10,13 +10,13 @@ __all__ = ["RemoteController", "create_remote_controller"]
 def create_remote_controller(
     *,
     api_key: str | None = None,
-    auth_type: AuthType = "Pkce",
-    endpoint: str,
-    client_config: ClientConfig | None = None,
-    headless: bool = False,
+    endpoint: str | None = None,
     insecure: bool = False,
     insecure_skip_verify: bool = False,
     ca_cert_file_path: str | None = None,
+    client_config: ClientConfig | None = None,
+    auth_type: AuthType = "Pkce",
+    headless: bool = False,
     command: List[str] | None = None,
     proxy_command: List[str] | None = None,
     client_id: str | None = None,
@@ -27,13 +27,33 @@ def create_remote_controller(
     """
     Create a new instance of the remote controller.
     """
+    assert endpoint or api_key, "Either endpoint or api_key must be provided when initializing remote controller"
     from ._client import ControllerClient
     from ._controller import RemoteController
 
+    if endpoint:
+        client_coro = ControllerClient.for_endpoint(
+            endpoint,
+            insecure=insecure,
+            insecure_skip_verify=insecure_skip_verify,
+            ca_cert_file_path=ca_cert_file_path,
+            client_id=client_id,
+            client_credentials_secret=client_credentials_secret,
+            auth_type=auth_type,
+        )
+    elif api_key:
+        client_coro = ControllerClient.for_api_key(
+            api_key,
+            insecure=insecure,
+            insecure_skip_verify=insecure_skip_verify,
+            ca_cert_file_path=ca_cert_file_path,
+            client_id=client_id,
+            client_credentials_secret=client_credentials_secret,
+            auth_type=auth_type,
+        )
+
     controller = RemoteController(
-        client_coro=ControllerClient.for_endpoint(
-            endpoint=endpoint, insecure=insecure, insecure_skip_verify=insecure_skip_verify
-        ),
+        client_coro=client_coro,
         workers=10,
         max_system_retries=5,
     )

flyte/_internal/controllers/remote/_action.py

@@ -4,8 +4,8 @@ from dataclasses import dataclass
 
 from flyteidl.core import execution_pb2
 
-from flyte._datastructures import GroupData
 from flyte._protos.workflow import run_definition_pb2, state_service_pb2, task_definition_pb2
+from flyte.models import GroupData
 
 
 @dataclass

flyte/_internal/controllers/remote/_client.py

@@ -20,7 +20,11 @@ class ControllerClient:
 
     @classmethod
     async def for_endpoint(cls, endpoint: str, insecure: bool = False, **kwargs) -> ControllerClient:
-        return cls(await create_channel(endpoint, insecure=insecure, **kwargs))
+        return cls(await create_channel(endpoint, None, insecure=insecure, **kwargs))
+
+    @classmethod
+    async def for_api_key(cls, api_key: str, insecure: bool = False, **kwargs) -> ControllerClient:
+        return cls(await create_channel(None, api_key, insecure=insecure, **kwargs))
 
     @property
     def state_service(self) -> StateService: