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

flyte/extras/_container.py

@@ -5,24 +5,19 @@ 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"
 
-
-def _extract_command_key(cmd: str, **kwargs) -> Any:
+def _extract_command_key(cmd: str, **kwargs) -> List[Any] | None:
     """
     Extract the key from the command using regex.
     """
     import re
 
-    input_regex = r"^\{\{\s*\.inputs\.(.*?)\s*\}\}$"
-    match = re.match(input_regex, cmd)
-    if match:
-        return match.group(1)
-    return None
+    input_regex = r"\{\{\.inputs\.([a-zA-Z0-9_]+)\}\}"
+    return re.findall(input_regex, cmd)
 
 
 def _extract_path_command_key(cmd: str, input_data_dir: Optional[str]) -> Optional[str]:
@@ -70,7 +65,7 @@ class ContainerTask(TaskTemplate):
         input_data_dir: str | pathlib.Path = "/var/inputs",
         output_data_dir: str | pathlib.Path = "/var/outputs",
         metadata_format: MetadataFormat = "JSON",
-        local_logs: bool = False,
+        local_logs: bool = True,
         **kwargs,
     ):
         super().__init__(
@@ -106,34 +101,33 @@ class ContainerTask(TaskTemplate):
         For FlyteFile and FlyteDirectory commands, e.g., "/var/inputs/inputs", we extract the key from strings that
          begin with the specified `input_data_dir`.
         """
-        # from flytekit.types.directory import FlyteDirectory
-        # from flytekit.types.file import FlyteFile
+        from flyte.io import Dir, File
 
         volume_binding: Dict[str, Dict[str, str]] = {}
         path_k = _extract_path_command_key(cmd, str(self._input_data_dir))
-        k = path_k if path_k else _extract_command_key(cmd)
-
-        if k:
-            input_val = kwargs.get(k)
-            # TODO: Add support file and directory transformer first
-            # if type(input_val) in [FlyteFile, FlyteDirectory]:
-            #     if not path_k:
-            #         raise AssertionError(
-            #             "FlyteFile and FlyteDirectory commands should not use the template syntax like this:
-            #             {{.inputs.infile}}\n"
-            #             "Please use a path-like syntax, such as: /var/inputs/infile.\n"
-            #             "This requirement is due to how Flyte Propeller processes template syntax inputs."
-            #         )
-            #     local_flyte_file_or_dir_path = str(input_val)
-            #     remote_flyte_file_or_dir_path = os.path.join(self._input_data_dir, k)  # type: ignore
-            #     volume_binding[local_flyte_file_or_dir_path] = {
-            #         "bind": remote_flyte_file_or_dir_path,
-            #         "mode": "rw",
-            #     }
-            #     command = remote_flyte_file_or_dir_path
-            command = str(input_val)
-        else:
-            command = cmd
+        keys = path_k if path_k else _extract_command_key(cmd)
+
+        if keys:
+            for k in keys:
+                input_val = kwargs.get(k)
+                # TODO: Add support file and directory transformer first
+                if type(input_val) in [File, Dir]:
+                    if not path_k:
+                        raise AssertionError(
+                            "File and Directory commands should not use the template syntax "
+                            "like this: {{.inputs.infile}}\n"
+                            "Please use a path-like syntax, such as: /var/inputs/infile.\n"
+                            "This requirement is due to how Flyte Propeller processes template syntax inputs."
+                        )
+                    local_flyte_file_or_dir_path = str(input_val)
+                    remote_flyte_file_or_dir_path = os.path.join(self._input_data_dir, k)  # type: ignore
+                    volume_binding[local_flyte_file_or_dir_path] = {
+                        "bind": remote_flyte_file_or_dir_path,
+                        "mode": "rw",
+                    }
+                    command = remote_flyte_file_or_dir_path
+                else:
+                    command = cmd
 
         return command, volume_binding
 
@@ -215,7 +209,7 @@ class ContainerTask(TaskTemplate):
                 output_dict[k] = self._convert_output_val_to_correct_type(output_val, output_type)
         return output_dict
 
-    def execute(self, **kwargs) -> Any:
+    async def execute(self, **kwargs) -> Any:
         try:
             import docker
         except ImportError:
@@ -235,6 +229,7 @@ class ContainerTask(TaskTemplate):
             raise AssertionError(f"Only Image objects are supported, not strings. Got {self._image} instead.")
         uri = self._image.uri
         self._pull_image_if_not_exists(client, uri)
+        print(f"Command: {commands!r}")
 
         container = client.containers.run(uri, command=commands, remove=True, volumes=volume_bindings, detach=True)
 
@@ -258,16 +253,11 @@ 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"),
         )
 
     def container_args(self, sctx: SerializationContext) -> List[str]:
         return self._cmd + (self._args if self._args else [])
-
-    def config(self, sctx: SerializationContext) -> Dict[str, str]:
-        if self.pod_template is None:
-            return {}
-        return {_PRIMARY_CONTAINER_NAME_FIELD: self.primary_container_name}

flyte/io/__init__.py

@@ -3,9 +3,25 @@
 
 This package contains additional data types beyond the primitive data types in python to abstract data flow
 of large datasets in Union.
+
 """
 
-__all__ = ["Dir", "File"]
+__all__ = [
+    "Dir",
+    "File",
+    "StructuredDataset",
+    "StructuredDatasetDecoder",
+    "StructuredDatasetEncoder",
+    "StructuredDatasetTransformerEngine",
+    "lazy_import_structured_dataset_handler",
+]
 
 from ._dir import Dir
 from ._file import File
+from ._structured_dataset import (
+    StructuredDataset,
+    StructuredDatasetDecoder,
+    StructuredDatasetEncoder,
+    StructuredDatasetTransformerEngine,
+    lazy_import_structured_dataset_handler,
+)

flyte/io/_dir.py

@@ -244,7 +244,7 @@ class Dir(BaseModel, Generic[T], SerializableType):
 
                 shutil.copytree(self.path, local_dest, dirs_exist_ok=True)
 
-        # Figure this out when we figure out the final synchronicity story
+        # Figure this out when we figure out the final sync story
         raise NotImplementedError("Sync download is not implemented for remote paths")
 
     @classmethod
@@ -287,7 +287,7 @@ class Dir(BaseModel, Generic[T], SerializableType):
             remote_dir = Dir[DataFrame].from_local_sync('/tmp/data_dir/', 's3://bucket/data/')
             ```
         """
-        # Implement this after we figure out the final synchronicity story
+        # Implement this after we figure out the final sync story
         raise NotImplementedError("Sync upload is not implemented for remote paths")
 
     async def exists(self) -> bool:

flyte/io/_file.py

@@ -22,7 +22,6 @@ from fsspec.asyn import AsyncFileSystem
 from fsspec.utils import get_protocol
 from mashumaro.types import SerializableType
 from pydantic import BaseModel, model_validator
-from synchronicity import Synchronizer
 
 import flyte.storage as storage
 from flyte._context import internal_ctx
@@ -33,8 +32,6 @@ from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
 # Type variable for the file format
 T = TypeVar("T")
 
-synced = Synchronizer()
-
 
 class File(BaseModel, Generic[T], SerializableType):
     """
@@ -235,6 +232,8 @@ class File(BaseModel, Generic[T], SerializableType):
             # This code is broadly similar to what storage.get_stream does, but without actually reading from the stream
             file_handle = None
             try:
+                if "b" not in mode:
+                    raise ValueError("Mode must include 'b' for binary access, when using remote files.")
                 if isinstance(fs, AsyncFileSystem):
                     file_handle = await fs.open_async(self.path, mode)
                     yield file_handle
@@ -314,7 +313,7 @@ class File(BaseModel, Generic[T], SerializableType):
         with fs.open(self.path, **open_kwargs) as f:
             yield f
 
-    # @synced.wrap - enabling this did not work - synchronicity/pydantic issue
+    # TODO sync needs to be implemented
     async def download(self, local_path: Optional[Union[str, Path]] = None) -> str:
         """
         Asynchronously download the file to a local path.

flyte/io/{structured_dataset → _structured_dataset}/basic_dfs.py

@@ -9,7 +9,7 @@ from fsspec.core import split_protocol, strip_protocol
 import flyte.storage as storage
 from flyte._logging import logger
 from flyte._utils import lazy_module
-from flyte.io.structured_dataset.structured_dataset import (
+from flyte.io._structured_dataset.structured_dataset import (
     CSV,
     PARQUET,
     StructuredDataset,

flyte/io/{structured_dataset → _structured_dataset}/structured_dataset.py

@@ -168,7 +168,7 @@ class StructuredDataset(SerializableType, DataClassJSONMixin):
         return self._literal_sd
 
     def open(self, dataframe_type: Type[DF]):
-        from flyte.io.structured_dataset import lazy_import_structured_dataset_handler
+        from flyte.io._structured_dataset import lazy_import_structured_dataset_handler
 
         """
         Load the handler if needed. For the use case like:

flyte/{_datastructures.py → models.py}

@@ -5,7 +5,9 @@ import os
 import pathlib
 import tempfile
 from dataclasses import dataclass, field, replace
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Literal, Optional, Tuple, Type
+
+import rich.repr
 
 from flyte._docstring import Docstring
 from flyte._interface import extract_return_annotation
@@ -13,6 +15,8 @@ from flyte._logging import logger
 from flyte._utils.helpers import base36_encode
 
 if TYPE_CHECKING:
+    from flyteidl.core import literals_pb2
+
     from flyte._internal.imagebuild.image_builder import ImageCache
     from flyte.report import Report
 
@@ -27,6 +31,7 @@ def generate_random_name() -> str:
     return str(uuid4())  # Placeholder for actual random name generation logic
 
 
+@rich.repr.auto
 @dataclass(frozen=True, kw_only=True)
 class ActionID:
     """
@@ -56,17 +61,19 @@ 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)
         return self.new_sub_action(new_name)
 
 
+@rich.repr.auto
 @dataclass(frozen=True, kw_only=True)
 class RawDataPath:
     """
@@ -90,6 +97,7 @@ class RawDataPath:
                 # Create a temporary directory for data storage
                 p = tempfile.mkdtemp()
                 logger.debug(f"Creating temporary directory for data storage: {p}")
+                pathlib.Path(p).mkdir(parents=True, exist_ok=True)
                 return RawDataPath(path=p)
             case str():
                 return RawDataPath(path=local_folder)
@@ -132,11 +140,13 @@ class RawDataPath:
         return remote_path
 
 
+@rich.repr.auto
 @dataclass(frozen=True)
 class GroupData:
     name: str
 
 
+@rich.repr.auto
 @dataclass(frozen=True, kw_only=True)
 class TaskContext:
     """
@@ -159,6 +169,7 @@ class TaskContext:
     code_bundle: CodeBundle | None = None
     compiled_image_cache: ImageCache | None = None
     data: Dict[str, Any] = field(default_factory=dict)
+    mode: Literal["local", "remote", "hybrid"] = "remote"
 
     def replace(self, **kwargs) -> TaskContext:
         if "data" in kwargs:
@@ -176,6 +187,7 @@ class TaskContext:
         return self.data.get(key)
 
 
+@rich.repr.auto
 @dataclass(frozen=True, kw_only=True)
 class CodeBundle:
     """
@@ -210,6 +222,7 @@ class CodeBundle:
         return replace(self, downloaded_path=path)
 
 
+@rich.repr.auto
 @dataclass(frozen=True)
 class Checkpoints:
     """
@@ -220,6 +233,13 @@ class Checkpoints:
     checkpoint_path: str | None
 
 
+class _has_default:
+    """
+    A marker class to indicate that a specific input has a default value or not.
+    This is used to determine if the input is required or not.
+    """
+
+
 @dataclass(frozen=True)
 class NativeInterface:
     """
@@ -229,7 +249,14 @@ class NativeInterface:
 
     inputs: Dict[str, Tuple[Type, Any]]
     outputs: Dict[str, Type]
-    docstring: Optional[Docstring] = field(default=None)
+    docstring: Optional[Docstring] = None
+
+    # This field is used to indicate that the task has a default value for the input, but already in the
+    # remote form.
+    _remote_defaults: Optional[Dict[str, literals_pb2.Literal]] = field(default=None, repr=False)
+
+    has_default: ClassVar[Type[_has_default]] = _has_default  # This can be used to indicate if a specific input
+    # has a default value or not, in the case when the default value is not known. An example would be remote tasks.
 
     def has_outputs(self) -> bool:
         """
@@ -237,12 +264,31 @@ class NativeInterface:
         """
         return self.outputs is not None and len(self.outputs) > 0
 
+    def num_required_inputs(self) -> int:
+        """
+        Get the number of required inputs for the task. This is used to determine how many inputs are required for the
+        task execution.
+        """
+        return sum(1 for t in self.inputs.values() if t[1] is inspect.Parameter.empty)
+
     @classmethod
-    def from_types(cls, inputs: Dict[str, Type], outputs: Dict[str, Type]) -> NativeInterface:
+    def from_types(
+        cls,
+        inputs: Dict[str, Tuple[Type, Type[_has_default] | Type[inspect._empty]]],
+        outputs: Dict[str, Type],
+        default_inputs: Optional[Dict[str, literals_pb2.Literal]] = None,
+    ) -> NativeInterface:
         """
         Create a new NativeInterface from the given types. This is used to create a native interface for the task.
+        :param inputs: A dictionary of input names and their types and a value indicating if they have a default value.
+        :param outputs: A dictionary of output names and their types.
+        :param default_inputs: Optional dictionary of default inputs for remote tasks.
+        :return: A NativeInterface object with the given inputs and outputs.
         """
-        return cls(inputs={k: (v, inspect.Parameter.empty) for k, v in inputs.items()}, outputs=outputs)
+        for k, v in inputs.items():
+            if v[1] is cls.has_default and (default_inputs is None or k not in default_inputs):
+                raise ValueError(f"Input {k} has a default value but no default input provided for remote task.")
+        return cls(inputs=inputs, outputs=outputs, _remote_defaults=default_inputs)
 
     @classmethod
     def from_callable(cls, func: Callable) -> NativeInterface:
@@ -290,7 +336,10 @@ class NativeInterface:
                 tp = tpe[0] if isinstance(tpe[0], str) else tpe[0].__name__
                 i += f"{key}: {tp}"
                 if tpe[1] is not inspect.Parameter.empty:
-                    i += f" = {tpe[1]}"
+                    if tpe[1] is self.has_default:
+                        i += " = ..."
+                    else:
+                        i += f" = {tpe[1]}"
         i += ")"
         if self.outputs:
             initial = True

flyte/remote/__init__.py

@@ -11,6 +11,7 @@ __all__ = [
     "Run",
     "RunDetails",
     "Secret",
+    "SecretTypes",
     "Task",
     "create_channel",
     "upload_dir",
@@ -21,5 +22,5 @@ from ._client.auth import create_channel
 from ._data import upload_dir, upload_file
 from ._project import Project
 from ._run import Action, ActionDetails, ActionInputs, ActionOutputs, Run, RunDetails
-from ._secret import Secret
+from ._secret import Secret, SecretTypes
 from ._task import Task

flyte/remote/_client/_protocols.py

@@ -58,6 +58,8 @@ class TaskService(Protocol):
         self, request: task_service_pb2.GetTaskDetailsRequest
     ) -> task_service_pb2.GetTaskDetailsResponse: ...
 
+    async def ListTasks(self, request: task_service_pb2.ListTasksRequest) -> task_service_pb2.ListTasksResponse: ...
+
 
 class RunService(Protocol):
     async def CreateRun(self, request: run_service_pb2.CreateRunRequest) -> run_service_pb2.CreateRunResponse: ...

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

@@ -1,3 +1,4 @@
+import os
 import ssl
 import typing
 
@@ -15,6 +16,11 @@ from ._authenticators.factory import (
     get_async_proxy_authenticator,
 )
 
+# Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
+if "GRPC_VERBOSITY" not in os.environ:
+    os.environ["GRPC_VERBOSITY"] = "ERROR"
+    os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
+
 # Initialize gRPC AIO early enough so it can be used in the main thread
 init_grpc_aio()
 
@@ -44,7 +50,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 +74,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 +113,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 +140,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 +199,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/_console.py

@@ -7,7 +7,7 @@ def _get_http_domain(endpoint: str, insecure: bool) -> str:
     if parsed.scheme == "dns":
         domain = parsed.path.lstrip("/")
     else:
-        domain = parsed.netloc
+        domain = parsed.netloc or parsed.path
     # TODO: make console url configurable
     if domain.split(":")[0] == "localhost":
         domain = "localhost:8080"

flyte/remote/_data.py

@@ -15,8 +15,9 @@ import httpx
 from flyteidl.service import dataproxy_pb2
 from google.protobuf import duration_pb2
 
-from flyte._initialize import CommonInit, get_client, get_common_config, requires_client
-from flyte.errors import RuntimeSystemError
+from flyte._initialize import CommonInit, ensure_client, get_client, get_common_config
+from flyte._logging import make_hyperlink
+from flyte.errors import InitializationError, RuntimeSystemError
 
 _UPLOAD_EXPIRES_IN = timedelta(seconds=60)
 
@@ -83,11 +84,13 @@ async def _upload_single_file(
             raise RuntimeSystemError(
                 "PermissionDenied", f"Failed to get signed url for {fp}, please check your permissions."
             )
+        elif e.code() == grpc.StatusCode.UNAVAILABLE:
+            raise InitializationError("EndpointUnavailable", "user", "Service is unavailable.")
         else:
             raise RuntimeSystemError(e.code().value, f"Failed to get signed url for {fp}.")
     except Exception as e:
         raise RuntimeSystemError(type(e).__name__, f"Failed to get signed url for {fp}.") from e
-    logger.debug(f"Uploading to signed url {resp.signed_url} for {fp}")
+    logger.debug(f'Uploading to {make_hyperlink("signed url", resp.signed_url)} for {fp}')
     extra_headers = get_extra_headers_for_protocol(resp.native_url)
     extra_headers.update(resp.headers)
     encoded_md5 = b64encode(md5_bytes)
@@ -100,7 +103,8 @@ async def _upload_single_file(
             if put_resp.status_code != 200:
                 raise RuntimeSystemError(
                     "UploadFailed",
-                    f"Failed to upload {fp} to {resp.signed_url}, status code: {put_resp.status_code}",
+                    f"Failed to upload {fp} to {resp.signed_url}, status code: {put_resp.status_code}, "
+                    f"response: {put_resp.text}",
                 )
         # TODO in old code we did this
         #             if self._config.platform.insecure_skip_verify is True
@@ -109,7 +113,6 @@ async def _upload_single_file(
     return str_digest, resp.native_url
 
 
-@requires_client
 async def upload_file(fp: Path, verify: bool = True) -> Tuple[str, str]:
     """
     Uploads a file to a remote location and returns the remote URI.
@@ -119,13 +122,13 @@ async def upload_file(fp: Path, verify: bool = True) -> Tuple[str, str]:
     :return: A tuple containing the MD5 digest and the remote URI.
     """
     # This is a placeholder implementation. Replace with actual upload logic.
+    ensure_client()
     cfg = get_common_config()
     if not fp.is_file():
         raise ValueError(f"{fp} is not a single file, upload arg must be a single file.")
     return await _upload_single_file(cfg, fp, verify=verify)
 
 
-@requires_client
 async def upload_dir(dir_path: Path, verify: bool = True) -> str:
     """
     Uploads a directory to a remote location and returns the remote URI.
@@ -135,6 +138,7 @@ async def upload_dir(dir_path: Path, verify: bool = True) -> str:
     :return: The remote URI of the uploaded directory.
     """
     # This is a placeholder implementation. Replace with actual upload logic.
+    ensure_client()
     cfg = get_common_config()
     if not dir_path.is_dir():
         raise ValueError(f"{dir_path} is not a directory, upload arg must be a directory.")