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

flyte/_initialize.py

@@ -1,261 +1,28 @@
 from __future__ import annotations
 
-import datetime
 import functools
-import os
+import sys
 import threading
 import typing
-from dataclasses import dataclass, replace
-from datetime import timedelta
+from dataclasses import dataclass, field, replace
 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 ._tools import ipython_check
+from ._logging import LogFormat, initialize_logger, logger
 
 if TYPE_CHECKING:
+    from flyte._internal.imagebuild import ImageBuildEngine
     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:
     """
@@ -266,12 +33,17 @@ class CommonInit:
     org: str | None = None
     project: str | None = None
     domain: str | None = None
+    batch_size: int = 1000
+    source_config_path: Optional[Path] = None  # Only used for documentation
+    sync_local_sys_paths: bool = True
 
 
 @dataclass(init=True, kw_only=True, repr=True, eq=True, frozen=True)
 class _InitConfig(CommonInit):
     client: Optional[ClientSet] = None
     storage: Optional[Storage] = None
+    image_builder: "ImageBuildEngine.ImageBuilderType" = "local"
+    images: typing.Dict[str, str] = field(default_factory=dict)
 
     def replace(self, **kwargs) -> _InitConfig:
         return replace(self, **kwargs)
@@ -304,11 +76,19 @@ async def _initialize_client(
     """
     from flyte.remote._client.controlplane import ClientSet
 
-    if endpoint is not None:
+    # https://grpc.io/docs/guides/keepalive/#keepalive-configuration-specification
+    channel_options = [
+        ("grpc.keepalive_permit_without_calls", 1),
+        ("grpc.keepalive_time_ms", 30000),  # Send keepalive ping every 30 seconds
+        ("grpc.keepalive_timeout_ms", 10000),  # Wait 10 seconds for keepalive response
+        ("grpc.http2.max_pings_without_data", 0),  # Allow unlimited pings without data
+        ("grpc.http2.min_ping_interval_without_data_ms", 30000),  # Min 30s between pings
+    ]
+
+    if endpoint and api_key is None:
         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,17 +100,46 @@ async def _initialize_client(
             client_config=client_config,
             rpc_retries=rpc_retries,
             http_proxy_url=http_proxy_url,
+            grpc_options=channel_options,
         )
-    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,
+            grpc_options=channel_options,
+        )
+
+    raise InitializationError(
+        "MissingEndpointOrApiKeyError", "user", "Either endpoint or api_key must be provided to initialize the client."
+    )
 
 
-@syncer.wrap
+def _initialize_logger(
+    log_level: int | None = None, log_format: LogFormat | None = None, reset_root_logger: bool = False
+) -> None:
+    initialize_logger(log_level=log_level, log_format=log_format, enable_rich=True, reset_root_logger=reset_root_logger)
+
+
+@syncify
 async def init(
     org: str | None = None,
     project: str | None = None,
     domain: str | None = None,
     root_dir: Path | None = None,
     log_level: int | None = None,
+    log_format: LogFormat | None = None,
+    reset_root_logger: bool = False,
     endpoint: str | None = None,
     headless: bool = False,
     insecure: bool = False,
@@ -346,7 +155,12 @@ async def init(
     rpc_retries: int = 3,
     http_proxy_url: str | None = None,
     storage: Storage | None = None,
-    config: Config | None = None,
+    batch_size: int = 1000,
+    image_builder: ImageBuildEngine.ImageBuilderType = "local",
+    images: typing.Dict[str, str] | None = None,
+    source_config_path: Optional[Path] = None,
+    sync_local_sys_paths: bool = True,
+    load_plugin_type_transformers: bool = True,
 ) -> None:
     """
     Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
@@ -355,14 +169,15 @@ 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 log_format: Optional logging format for the logger, default is "console"
+    :param reset_root_logger: By default, we clear out root logger handlers and set up our own.
     :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)
@@ -376,61 +191,338 @@ async def init(
     :param ca_cert_file_path: [optional] str Root Cert to be loaded and used to verify admin
     :param http_proxy_url: [optional] HTTP Proxy to be used for OAuth requests
     :param rpc_retries: [optional] int Number of times to retry the platform calls
-    :param audience: oauth2 audience for the token request. This is used to validate the token
     :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 config: Optional config to override the init parameters
-
+    :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.
+    :param image_builder: Optional image builder configuration, if not provided, the default image builder will be used.
+    :param images: Optional dict of images that can be used by referencing the image name.
+    :param source_config_path: Optional path to the source configuration file (This is only used for documentation)
+    :param sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
+      into the remote container (default: True).
+    :param load_plugin_type_transformers: If enabled (default True), load the type transformer plugins registered under
+      the "flyte.plugins.types" entry point group.
     :return: None
     """
-    from flyte._utils import get_cwd_editable_install
+    from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
+    from flyte.types import _load_custom_type_transformers
 
-    interactive_mode = ipython_check()
-
-    initialize_logger(enable_rich=interactive_mode)
-    if log_level:
-        initialize_logger(log_level=log_level, enable_rich=interactive_mode)
+    _initialize_logger(log_level=log_level, log_format=log_format, reset_root_logger=reset_root_logger)
+    if load_plugin_type_transformers:
+        _load_custom_type_transformers()
 
     global _init_config  # noqa: PLW0603
 
-    with _init_lock:
-        if config is None:
-            from flyte.config import Config
+    endpoint = sanitize_endpoint(endpoint)
 
-            config = Config.auto()
-        platform_cfg = config.platform
-        task_cfg = config.task
+    with _init_lock:
         client = None
-        if endpoint or platform_cfg.endpoint or api_key:
+        if endpoint or api_key:
             client = await _initialize_client(
                 api_key=api_key,
-                auth_type=auth_type or platform_cfg.auth_mode,
-                endpoint=endpoint or platform_cfg.endpoint,
+                auth_type=auth_type,
+                endpoint=endpoint,
                 headless=headless,
-                insecure=insecure or platform_cfg.insecure,
-                insecure_skip_verify=insecure_skip_verify or platform_cfg.insecure_skip_verify,
-                ca_cert_file_path=ca_cert_file_path or platform_cfg.ca_cert_file_path,
-                command=command or platform_cfg.command,
-                proxy_command=proxy_command or platform_cfg.proxy_command,
-                client_id=client_id or platform_cfg.client_id,
-                client_credentials_secret=client_credentials_secret or platform_cfg.client_credentials_secret,
+                insecure=insecure,
+                insecure_skip_verify=insecure_skip_verify,
+                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=auth_client_config,
-                rpc_retries=rpc_retries or platform_cfg.rpc_retries,
-                http_proxy_url=http_proxy_url or platform_cfg.http_proxy_url,
+                rpc_retries=rpc_retries,
+                http_proxy_url=http_proxy_url,
             )
 
-        root_dir = root_dir or get_cwd_editable_install() or Path.cwd()
+        if not root_dir:
+            editable_root = get_cwd_editable_install()
+            if editable_root:
+                logger.info(f"Using editable install as root directory: {editable_root}")
+                root_dir = editable_root
+            else:
+                logger.info("No editable install found, using current working directory as root directory.")
+                root_dir = Path.cwd()
+        # We will inject the root_dir into the sys,path for module resolution
+        sys.path.append(str(root_dir))
+
         _init_config = _InitConfig(
             root_dir=root_dir,
-            project=project or task_cfg.project,
-            domain=domain or task_cfg.domain,
+            project=project,
+            domain=domain,
             client=client,
             storage=storage,
-            org=org or task_cfg.org,
+            org=org or org_from_endpoint(endpoint),
+            batch_size=batch_size,
+            image_builder=image_builder,
+            images=images or {},
+            source_config_path=source_config_path,
+            sync_local_sys_paths=sync_local_sys_paths,
         )
 
 
+@syncify
+async def init_from_config(
+    path_or_config: str | Path | Config | None = None,
+    root_dir: Path | None = None,
+    log_level: int | None = None,
+    log_format: LogFormat = "console",
+    project: str | None = None,
+    domain: str | None = None,
+    storage: Storage | None = None,
+    batch_size: int = 1000,
+    image_builder: ImageBuildEngine.ImageBuilderType | None = None,
+    images: tuple[str, ...] | None = None,
+    sync_local_sys_paths: bool = True,
+) -> 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 project: Project name, this will override any project names in the configuration file
+    :param domain: Domain name, this will override any domain names in the configuration file
+    :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
+    :param log_format: Optional logging format for the logger, default is "console"
+    :param storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
+    :param images: List of image strings in format "imagename=imageuri" or just "imageuri".
+    :param sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
+     into the remote container (default: True).
+    :param batch_size: Optional batch size for operations that use listings, defaults to 1000
+    :param image_builder: Optional image builder configuration, if provided,
+        will override any defaults set in the configuration.
+    :return: None
+    """
+    from rich.highlighter import ReprHighlighter
+
+    import flyte.config as config
+    from flyte.cli._common import parse_images
+
+    cfg: config.Config
+    cfg_path: Optional[Path] = None
+    if path_or_config is None:
+        # If no path is provided, use the default config file
+        cfg = config.auto()
+    elif isinstance(path_or_config, (str, Path)):
+        if root_dir:
+            cfg_path = root_dir.expanduser() / path_or_config
+        else:
+            cfg_path = Path(path_or_config).expanduser()
+        if not Path(cfg_path).exists():
+            raise InitializationError(
+                "ConfigFileNotFoundError",
+                "user",
+                f"Configuration file '{cfg_path}' does not exist., current working directory is {Path.cwd()}",
+            )
+        cfg = config.auto(cfg_path)
+    else:
+        cfg = path_or_config
+
+    logger.info(f"Flyte config initialized as {cfg}", extra={"highlighter": ReprHighlighter()})
+
+    # parse image, this will overwrite the image_refs set in the config file
+    parse_images(cfg, images)
+
+    await init.aio(
+        org=cfg.task.org,
+        project=project or cfg.task.project,
+        domain=domain or 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,
+        log_format=log_format,
+        image_builder=image_builder or cfg.image.builder,
+        batch_size=batch_size,
+        images=cfg.image.image_refs,
+        storage=storage,
+        source_config_path=cfg_path,
+        sync_local_sys_paths=sync_local_sys_paths,
+    )
+
+
+@syncify
+async def init_from_api_key(
+    api_key: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+    root_dir: Path | None = None,
+    log_level: int | None = None,
+    log_format: LogFormat | None = None,
+    storage: Storage | None = None,
+    batch_size: int = 1000,
+    image_builder: ImageBuildEngine.ImageBuilderType = "local",
+    images: typing.Dict[str, str] | None = None,
+    sync_local_sys_paths: bool = True,
+) -> None:
+    """
+    Initialize the Flyte system using an API key for authentication. This is a convenience
+    method for API key-based authentication. Thread-safe implementation.
+
+    The API key should be an encoded API key that contains the endpoint, client ID, client secret,
+    and organization information. You can obtain this encoded API key from your Flyte administrator
+    or cloud provider.
+
+    :param api_key: Optional encoded API key for authentication. If None, reads from FLYTE_API_KEY
+        environment variable. The API key is a base64-encoded string containing endpoint, client_id,
+        client_secret, and org information.
+    :param project: Optional project name
+    :param domain: Optional domain name
+    :param root_dir: Optional root directory from which to determine how to load files, and find paths to files.
+      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
+    :param log_format: Optional logging format for the logger, default is "console"
+    :param storage: Optional blob store (S3, GCS, Azure) configuration
+    :param batch_size: Optional batch size for operations that use listings, defaults to 1000
+    :param image_builder: Optional image builder configuration
+    :param images: Optional dict of images that can be used by referencing the image name
+    :param sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
+      into the remote container (default: True)
+    :return: None
+    """
+    import os
+
+    from flyte._utils import sanitize_endpoint
+    from flyte.remote._client.auth._auth_utils import decode_api_key
+
+    # If api_key is not provided, read from environment variable
+    if api_key is None:
+        api_key = os.getenv("FLYTE_API_KEY")
+        if api_key is None:
+            raise InitializationError(
+                "MissingApiKeyError",
+                "user",
+                "API key must be provided either as a parameter or via the FLYTE_API_KEY environment variable.",
+            )
+
+    # Decode the API key to extract endpoint, client_id, client_secret, and org
+    endpoint, client_id, client_secret, org = decode_api_key(api_key)
+
+    # Sanitize the endpoint
+    endpoint = sanitize_endpoint(endpoint)  # type: ignore[assignment]
+
+    await init.aio(
+        org=None if org == "None" else org,
+        project=project,
+        domain=domain,
+        endpoint=endpoint,
+        api_key=api_key,
+        client_id=client_id,
+        client_credentials_secret=client_secret,
+        auth_type="ClientSecret",  # API keys use client credentials flow
+        root_dir=root_dir,
+        log_level=log_level,
+        log_format=log_format,
+        insecure=False,
+        insecure_skip_verify=False,
+        storage=storage,
+        batch_size=batch_size,
+        image_builder=image_builder,
+        images=images,
+        sync_local_sys_paths=sync_local_sys_paths,
+    )
+
+
+@syncify
+async def init_in_cluster(
+    org: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+    api_key: str | None = None,
+    endpoint: str | None = None,
+    insecure: bool = False,
+) -> dict[str, typing.Any]:
+    import os
+
+    from flyte._utils import str2bool
+
+    PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
+    DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
+    ORG_NAME = "_U_ORG_NAME"
+    ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+    INSECURE_SKIP_VERIFY_OVERRIDE = "_U_INSECURE_SKIP_VERIFY"
+    INSECURE_OVERRIDE = "_U_INSECURE"
+    _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+    EAGER_API_KEY = "EAGER_API_KEY"
+
+    org = org or os.getenv(ORG_NAME)
+    project = project or os.getenv(PROJECT_NAME)
+    domain = domain or os.getenv(DOMAIN_NAME)
+    api_key = api_key or os.getenv(_UNION_EAGER_API_KEY_ENV_VAR) or os.getenv(EAGER_API_KEY)
+
+    remote_kwargs: dict[str, typing.Any] = {"insecure": insecure}
+    if api_key:
+        logger.info("Using api key from environment")
+        remote_kwargs["api_key"] = api_key
+    else:
+        ep = endpoint or os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+        remote_kwargs["endpoint"] = ep
+        if not insecure:
+            if "localhost" in ep or "docker" in ep:
+                remote_kwargs["insecure"] = True
+        if str2bool(os.getenv(INSECURE_OVERRIDE, "")):
+            remote_kwargs["insecure"] = True
+        logger.debug(f"Using controller endpoint: {ep} with kwargs: {remote_kwargs}")
+
+    # Check for insecure_skip_verify override (e.g. for self-signed certs)
+    insecure_skip_verify_str = os.getenv(INSECURE_SKIP_VERIFY_OVERRIDE, "")
+    if str2bool(insecure_skip_verify_str):
+        remote_kwargs["insecure_skip_verify"] = True
+        logger.info("SSL certificate verification disabled (insecure_skip_verify=True)")
+
+    await init.aio(
+        org=org, project=project, domain=domain, root_dir=Path.cwd(), image_builder="remote", **remote_kwargs
+    )
+    return remote_kwargs
+
+
+@syncify
+async def init_passthrough(
+    endpoint: str | None = None,
+    org: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+    insecure: bool = False,
+) -> dict[str, typing.Any]:
+    """
+    Initialize the Flyte system with passthrough authentication.
+
+    This authentication mode allows you to pass custom authentication metadata
+    using the `flyte.remote.auth_metadata()` context manager.
+
+    :param org: Optional organization name
+    :param project: Optional project name
+    :param domain: Optional domain name
+    :param endpoint: Optional API endpoint URL
+    :param insecure: Whether to use an insecure channel
+    :return: Dictionary of remote kwargs used for initialization
+    """
+    await init.aio(
+        org=org,
+        project=project,
+        domain=domain,
+        root_dir=Path.cwd(),
+        image_builder="remote",
+        endpoint=endpoint,
+        insecure=insecure,
+        auth_type="Passthrough",
+    )
+    return {"endpoint": endpoint, "insecure": insecure}
+
+
 def _get_init_config() -> Optional[_InitConfig]:
     """
     Get the current initialization configuration. Thread-safe implementation.
@@ -441,7 +533,7 @@ def _get_init_config() -> Optional[_InitConfig]:
         return _init_config
 
 
-def get_common_config() -> CommonInit:
+def get_init_config() -> _InitConfig:
     """
     Get the current initialization configuration. Thread-safe implementation.
 
@@ -450,15 +542,15 @@ def get_common_config() -> CommonInit:
     cfg = _get_init_config()
     if cfg is None:
         raise InitializationError(
-            "StorageNotInitializedError",
+            "ClientNotInitializedError",
             "user",
-            "Configuration has not been initialized. Call flyte.init() with a valid endpoint or",
-            " api-key before using this function.",
+            "Configuration has not been initialized. Call flyte.init() with a valid endpoint/api-key before",
+            " using this function or Call flyte.init_from_config() with a valid path to the config file",
         )
     return cfg
 
 
-def get_storage() -> Storage:
+def get_storage() -> Storage | None:
     """
     Get the current storage configuration. Thread-safe implementation.
 
@@ -469,12 +561,10 @@ def get_storage() -> Storage:
         raise InitializationError(
             "StorageNotInitializedError",
             "user",
-            "Configuration has not been initialized. Call flyte.init() with a valid endpoint or",
-            " api-key before using this function.",
+            "Configuration has not been initialized. Call flyte.init() with a valid"
+            " storage configuration before using this function or Call flyte.init_from_config()" 
+            " with a valid path to the config file",
         )
-    if cfg.storage is None:
-        # return default local storage
-        return typing.cast(Storage, cfg.replace(storage=Storage()).storage)
     return cfg.storage
 
 
@@ -489,8 +579,8 @@ def get_client() -> ClientSet:
         raise InitializationError(
             "ClientNotInitializedError",
             "user",
-            "Client has not been initialized. Call flyte.init() with a valid endpoint or"
-            " api-key before using this function.",
+            "Client has not been initialized. Call flyte.init() with a valid endpoint/api-key "
+            "before using this function or Call flyte.init_from_config() with a valid path to the config file",
         )
     return cfg.client
 
@@ -504,41 +594,31 @@ 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 requires_client(func: T) -> T:
+def ensure_client():
     """
-    Decorator that checks if the client has been initialized before executing the function.
-    Raises InitializationError if the client is not initialized.
-
-    :param func: Function to decorate
-    :return: Decorated function that checks for initialization
+    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.
     """
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs) -> T:
-        init_config = _get_init_config()
-        if init_config is None or init_config.client is None:
-            raise InitializationError(
-                "ClientNotInitializedError",
-                "user",
-                f"Function '{func.__name__}' requires client to be initialized. "
-                f"Call flyte.init() with a valid endpoint or api-key before using this function.",
-            )
-        return func(*args, **kwargs)
-
-    return typing.cast(T, wrapper)
+    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/api-key before using"
+            " this function or Call flyte.init_from_config() with a valid path to the config file",
+        )
 
 
 def requires_storage(func: T) -> T:
@@ -557,7 +637,8 @@ def requires_storage(func: T) -> T:
                 "StorageNotInitializedError",
                 "user",
                 f"Function '{func.__name__}' requires storage to be initialized. "
-                f"Call flyte.init() with a valid storage configuration before using this function.",
+                "Call flyte.init() with a valid storage configuration before using this function."
+                "or Call flyte.init_from_config() with a valid path to the config file",
             )
         return func(*args, **kwargs)
 
@@ -583,7 +664,8 @@ def requires_upload_location(func: T) -> T:
                 "No upload path configured",
                 "user",
                 f"Function '{func.__name__}' requires client to be initialized. "
-                f"Call flyte.init() with storage configuration before using this function.",
+                "Call flyte.init() with storage configuration before using this function."
+                "or Call flyte.init_from_config() with a valid path to the config file."
             )
         return func(*args, **kwargs)
 
@@ -605,13 +687,42 @@ def requires_initialization(func: T) -> T:
             raise InitializationError(
                 "NotInitConfiguredError",
                 "user",
-                f"Function '{func.__name__}' requires initialization. Call flyte.init() before using this function.",
+                f"Function '{func.__name__}' requires initialization. Call flyte.init() before using this function"
+                " or Call flyte.init_from_config() with a valid path to the config file."
             )
         return func(*args, **kwargs)
 
     return typing.cast(T, wrapper)
 
 
+def require_project_and_domain(func):
+    """
+    Decorator that ensures the current Flyte configuration defines
+    both 'project' and 'domain'. Raises a clear error if not found.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        cfg = get_init_config()
+        if cfg.project is None:
+            raise ValueError(
+                "Project must be provided to initialize the client. "
+                "Please set 'project' in the 'task' section of your config file, "
+                "or pass it directly to flyte.init(project='your-project-name')."
+            )
+
+        if cfg.domain is None:
+            raise ValueError(
+                "Domain must be provided to initialize the client. "
+                "Please set 'domain' in the 'task' section of your config file, "
+                "or pass it directly to flyte.init(domain='your-domain-name')."
+            )
+
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 async def _init_for_testing(
     project: str | None = None,
     domain: str | None = None,
@@ -641,3 +752,32 @@ def replace_client(client):
 
     with _init_lock:
         _init_config = _init_config.replace(client=client)
+
+
+def current_domain() -> str:
+    """
+    Returns the current domain from Runtime environment (on the cluster) or from the initialized configuration.
+    This is safe to be used during `deploy`, `run` and within `task` code.
+
+    NOTE: This will not work if you deploy a task to a domain and then run it in another domain.
+
+    Raises InitializationError if the configuration is not initialized or domain is not set.
+    :return: The current domain
+    """
+    from ._context import ctx
+
+    tctx = ctx()
+    if tctx is not None:
+        domain = tctx.action.domain
+        if domain is not None:
+            return domain
+
+    cfg = _get_init_config()
+    if cfg is None or cfg.domain is None:
+        raise InitializationError(
+            "DomainNotInitializedError",
+            "user",
+            "Domain has not been initialized. Call flyte.init() with a valid domain before using this function"
+            " or Call flyte.init_from_config() with a valid path to the config file"
+        )
+    return cfg.domain