flyte 2.0.0b23__py3-none-any.whl → 2.0.0b25__py3-none-any.whl

flyte/io/_dataframe/basic_dfs.py

@@ -3,7 +3,7 @@ import typing
 from pathlib import Path
 from typing import TypeVar
 
-from flyteidl.core import literals_pb2, types_pb2
+from flyteidl2.core import literals_pb2, types_pb2
 from fsspec.core import split_protocol, strip_protocol
 
 import flyte.storage as storage

flyte/io/_dataframe/dataframe.py

@@ -8,7 +8,7 @@ from abc import ABC, abstractmethod
 from dataclasses import is_dataclass
 from typing import Any, ClassVar, Coroutine, Dict, Generic, List, Optional, Type, Union
 
-from flyteidl.core import literals_pb2, types_pb2
+from flyteidl2.core import literals_pb2, types_pb2
 from fsspec.utils import get_protocol
 from mashumaro.types import SerializableType
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_serializer, model_validator
@@ -647,17 +647,21 @@ class DataFrameTransformerEngine(TypeTransformer[DataFrame]):
                 f"Already registered a handler for {(h.python_type, protocol, h.supported_format)}"
             )
         lowest_level[h.supported_format] = h
-        logger.debug(f"Registered {h} as handler for {h.python_type}, protocol {protocol}, fmt {h.supported_format}")
+        logger.debug(
+            f"Registered {h.__class__.__name__} as handler for {h.python_type.__class__.__name__},"
+            f" protocol {protocol}, fmt {h.supported_format}"
+        )
 
         if (default_format_for_type or default_for_type) and h.supported_format != GENERIC_FORMAT:
             if h.python_type in cls.DEFAULT_FORMATS and not override:
                 if cls.DEFAULT_FORMATS[h.python_type] != h.supported_format:
                     logger.info(
-                        f"Not using handler {h} with format {h.supported_format}"
-                        f" as default for {h.python_type}, {cls.DEFAULT_FORMATS[h.python_type]} already specified."
+                        f"Not using handler {h.__class__.__name__} with format {h.supported_format}"
+                        f" as default for {h.python_type.__class__.__name__},"
+                        f" {cls.DEFAULT_FORMATS[h.python_type]} already specified."
                     )
             else:
-                logger.debug(f"Use {type(h).__name__} as default handler for {h.python_type}.")
+                logger.debug(f"Use {type(h).__name__} as default handler for {h.python_type.__class__.__name__}.")
                 cls.DEFAULT_FORMATS[h.python_type] = h.supported_format
         if default_storage_for_type or default_for_type:
             if h.protocol in cls.DEFAULT_PROTOCOLS and not override:
@@ -685,7 +689,7 @@ class DataFrameTransformerEngine(TypeTransformer[DataFrame]):
         expected: types_pb2.LiteralType,
     ) -> literals_pb2.Literal:
         # Make a copy in case we need to hand off to encoders, since we can't be sure of mutations.
-        python_type, *attrs = extract_cols_and_format(python_type)
+        python_type, *_attrs = extract_cols_and_format(python_type)
         sdt = types_pb2.StructuredDatasetType(format=self.DEFAULT_FORMATS.get(python_type, GENERIC_FORMAT))
 
         if issubclass(python_type, DataFrame) and not isinstance(python_val, DataFrame):
@@ -876,7 +880,7 @@ class DataFrameTransformerEngine(TypeTransformer[DataFrame]):
             raise TypeTransformerFailedError("Attribute access unsupported.")
 
         # Detect annotations and extract out all the relevant information that the user might supply
-        expected_python_type, column_dict, storage_fmt, pa_schema = extract_cols_and_format(expected_python_type)
+        expected_python_type, column_dict, _storage_fmt, _pa_schema = extract_cols_and_format(expected_python_type)
 
         # Start handling for DataFrame scalars, first look at the columns
         incoming_columns = lv.scalar.structured_dataset.metadata.structured_dataset_type.columns
@@ -986,7 +990,7 @@ class DataFrameTransformerEngine(TypeTransformer[DataFrame]):
         return converted_cols
 
     def _get_dataset_type(self, t: typing.Union[Type[DataFrame], typing.Any]) -> types_pb2.StructuredDatasetType:
-        original_python_type, column_map, storage_format, pa_schema = extract_cols_and_format(t)  # type: ignore
+        _original_python_type, column_map, storage_format, pa_schema = extract_cols_and_format(t)  # type: ignore
 
         # Get the column information
         converted_cols: typing.List[types_pb2.StructuredDatasetType.DatasetColumn] = (

flyte/io/_dir.py

@@ -4,12 +4,14 @@ import os
 from pathlib import Path
 from typing import AsyncIterator, Dict, Generic, Iterator, List, Optional, Type, TypeVar, Union
 
-from flyteidl.core import literals_pb2, types_pb2
+from flyteidl2.core import literals_pb2, types_pb2
 from fsspec.asyn import AsyncFileSystem
+from fsspec.utils import get_protocol
 from mashumaro.types import SerializableType
 from pydantic import BaseModel, model_validator
 
 import flyte.storage as storage
+from flyte._context import internal_ctx
 from flyte.io._file import File
 from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
 
@@ -462,7 +464,15 @@ class Dir(BaseModel, Generic[T], SerializableType):
         Returns:
             The absolute path to the downloaded directory
         """
-        local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
+        # If no local_path specified, create a unique path + append source directory name
+        if local_path is None:
+            unique_path = storage.get_random_local_path()
+            source_dirname = Path(self.path).name  # will need to be updated for windows
+            local_dest = str(Path(unique_path) / source_dirname)
+        else:
+            # If local_path is specified, use it directly (contents go into it)
+            local_dest = str(local_path)
+
         if not storage.is_remote(self.path):
             if not local_path or local_path == self.path:
                 # Skip copying
@@ -477,6 +487,7 @@ class Dir(BaseModel, Generic[T], SerializableType):
                     await loop.run_in_executor(None, lambda: shutil.copytree(self.path, local_dest, dirs_exist_ok=True))
 
                 await copy_tree()
+                return local_dest
         return await storage.get(self.path, local_dest, recursive=True)
 
     def download_sync(self, local_path: Optional[Union[str, Path]] = None) -> str:
@@ -510,7 +521,15 @@ class Dir(BaseModel, Generic[T], SerializableType):
         Returns:
             The absolute path to the downloaded directory
         """
-        local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
+        # If no local_path specified, create a unique path + append source directory name
+        if local_path is None:
+            unique_path = storage.get_random_local_path()
+            source_dirname = Path(self.path).name
+            local_dest = str(Path(unique_path) / source_dirname)
+        else:
+            # If local_path is specified, use it directly (contents go into it)
+            local_dest = str(local_path)
+
         if not storage.is_remote(self.path):
             if not local_path or local_path == self.path:
                 # Skip copying
@@ -520,6 +539,7 @@ class Dir(BaseModel, Generic[T], SerializableType):
                 import shutil
 
                 shutil.copytree(self.path, local_dest, dirs_exist_ok=True)
+                return local_dest
 
         fs = storage.get_underlying_filesystem(path=self.path)
         fs.get(self.path, local_dest, recursive=True)
@@ -529,7 +549,7 @@ class Dir(BaseModel, Generic[T], SerializableType):
     async def from_local(
         cls,
         local_path: Union[str, Path],
-        remote_path: Optional[str] = None,
+        remote_destination: Optional[str] = None,
         dir_cache_key: Optional[str] = None,
     ) -> Dir[T]:
         """
@@ -571,7 +591,8 @@ class Dir(BaseModel, Generic[T], SerializableType):
         ```
         Args:
             local_path: Path to the local directory
-            remote_path: Optional remote path to store the directory. If None, a path will be automatically generated.
+            remote_destination: Optional remote path to store the directory. If None, a path will be automatically
+              generated.
             dir_cache_key: Optional precomputed hash value to use for cache key computation when this Dir is used
                           as an input to discoverable tasks. If not specified, the cache key will be based on
                           directory attributes.
@@ -581,15 +602,23 @@ class Dir(BaseModel, Generic[T], SerializableType):
         """
         local_path_str = str(local_path)
         dirname = os.path.basename(os.path.normpath(local_path_str))
+        resolved_remote_path = remote_destination or internal_ctx().raw_data.get_random_remote_path(dirname)
+        protocol = get_protocol(resolved_remote_path)
+
+        # Shortcut for local, don't copy and just return
+        if "file" in protocol and remote_destination is None:
+            output_path = str(Path(local_path).absolute())
+            return cls(path=output_path, name=dirname, hash=dir_cache_key)
 
-        output_path = await storage.put(from_path=local_path_str, to_path=remote_path, recursive=True)
+        # todo: in the future, mirror File and set the file to_path here
+        output_path = await storage.put(from_path=local_path_str, to_path=remote_destination, recursive=True)
         return cls(path=output_path, name=dirname, hash=dir_cache_key)
 
     @classmethod
     def from_local_sync(
         cls,
         local_path: Union[str, Path],
-        remote_path: Optional[str] = None,
+        remote_destination: Optional[str] = None,
         dir_cache_key: Optional[str] = None,
     ) -> Dir[T]:
         """
@@ -632,7 +661,8 @@ class Dir(BaseModel, Generic[T], SerializableType):
 
         Args:
             local_path: Path to the local directory
-            remote_path: Optional remote path to store the directory. If None, a path will be automatically generated.
+            remote_destination: Optional remote path to store the directory. If None, a path will be automatically
+              generated.
             dir_cache_key: Optional precomputed hash value to use for cache key computation when this Dir is used
                           as an input to discoverable tasks. If not specified, the cache key will be based on
                           directory attributes.
@@ -643,14 +673,17 @@ class Dir(BaseModel, Generic[T], SerializableType):
         local_path_str = str(local_path)
         dirname = os.path.basename(os.path.normpath(local_path_str))
 
-        if not remote_path:
-            from flyte._context import internal_ctx
+        resolved_remote_path = remote_destination or internal_ctx().raw_data.get_random_remote_path(dirname)
+        protocol = get_protocol(resolved_remote_path)
+
+        # Shortcut for local, don't copy and just return
+        if "file" in protocol and remote_destination is None:
+            output_path = str(Path(local_path).absolute())
+            return cls(path=output_path, name=dirname, hash=dir_cache_key)
 
-            ctx = internal_ctx()
-            remote_path = ctx.raw_data.get_random_remote_path(dirname)
-        fs = storage.get_underlying_filesystem(path=remote_path)
-        fs.put(local_path_str, remote_path, recursive=True)
-        return cls(path=remote_path, name=dirname, hash=dir_cache_key)
+        fs = storage.get_underlying_filesystem(path=resolved_remote_path)
+        fs.put(local_path_str, resolved_remote_path, recursive=True)
+        return cls(path=resolved_remote_path, name=dirname, hash=dir_cache_key)
 
     @classmethod
     def from_existing_remote(cls, remote_path: str, dir_cache_key: Optional[str] = None) -> Dir[T]:

flyte/io/_file.py

@@ -20,7 +20,7 @@ from typing import (
 )
 
 import aiofiles
-from flyteidl.core import literals_pb2, types_pb2
+from flyteidl2.core import literals_pb2, types_pb2
 from fsspec.utils import get_protocol
 from mashumaro.types import SerializableType
 from pydantic import BaseModel, Field, model_validator
@@ -33,6 +33,9 @@ from flyte._initialize import requires_initialization
 from flyte.io._hashing_io import AsyncHashingReader, HashingWriter, HashMethod, PrecomputedValue
 from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
 
+if typing.TYPE_CHECKING:
+    from obstore import AsyncReadableFile, AsyncWritableFile
+
 if typing.TYPE_CHECKING:
     from obstore import AsyncReadableFile, AsyncWritableFile
 
@@ -524,23 +527,40 @@ class File(BaseModel, Generic[T], SerializableType):
             The absolute path to the downloaded file
         """
         if local_path is None:
-            local_path = storage.get_random_local_path(file_path_or_file_name=local_path)
+            local_path = storage.get_random_local_path(file_path_or_file_name=self.path)
         else:
+            # Preserve trailing separator if present (Path.absolute() strips it)
+            local_path_str = str(local_path)
+            has_trailing_sep = local_path_str.endswith(os.sep)
             local_path = str(Path(local_path).absolute())
+            if has_trailing_sep:
+                local_path = local_path + os.sep
 
         fs = storage.get_underlying_filesystem(path=self.path)
 
         # If it's already a local file, just copy it
         if "file" in fs.protocol:
+            # Apply directory logic for local-to-local copies
+            local_path_for_copy = local_path
+            if isinstance(local_path, str):
+                local_path_obj = Path(local_path)
+                # Check if it's a directory or ends with separator
+                if local_path.endswith(os.sep) or (local_path_obj.exists() and local_path_obj.is_dir()):
+                    remote_filename = Path(self.path).name
+                    local_path_for_copy = str(local_path_obj / remote_filename)
+
+            # Ensure parent directory exists
+            Path(local_path_for_copy).parent.mkdir(parents=True, exist_ok=True)
+
             # Use aiofiles for async copy
             async with aiofiles.open(self.path, "rb") as src:
-                async with aiofiles.open(local_path, "wb") as dst:
+                async with aiofiles.open(local_path_for_copy, "wb") as dst:
                     await dst.write(await src.read())
-            return str(local_path)
+            return str(local_path_for_copy)
 
         # Otherwise download from remote using async functionality
-        await storage.get(self.path, str(local_path))
-        return str(local_path)
+        result_path = await storage.get(self.path, str(local_path))
+        return result_path
 
     def download_sync(self, local_path: Optional[Union[str, Path]] = None) -> str:
         """
@@ -576,19 +596,36 @@ class File(BaseModel, Generic[T], SerializableType):
             The absolute path to the downloaded file
         """
         if local_path is None:
-            local_path = storage.get_random_local_path(file_path_or_file_name=local_path)
+            local_path = storage.get_random_local_path(file_path_or_file_name=self.path)
         else:
+            # Preserve trailing separator if present (Path.absolute() strips it)
+            local_path_str = str(local_path)
+            has_trailing_sep = local_path_str.endswith(os.sep)
             local_path = str(Path(local_path).absolute())
+            if has_trailing_sep:
+                local_path = local_path + os.sep
 
         fs = storage.get_underlying_filesystem(path=self.path)
 
         # If it's already a local file, just copy it
         if "file" in fs.protocol:
+            # Apply directory logic for local-to-local copies
+            local_path_for_copy = local_path
+            if isinstance(local_path, str):
+                local_path_obj = Path(local_path)
+                # Check if it's a directory or ends with separator
+                if local_path.endswith(os.sep) or (local_path_obj.exists() and local_path_obj.is_dir()):
+                    remote_filename = Path(self.path).name
+                    local_path_for_copy = str(local_path_obj / remote_filename)
+
+            # Ensure parent directory exists
+            Path(local_path_for_copy).parent.mkdir(parents=True, exist_ok=True)
+
             # Use standard file operations for sync copy
             import shutil
 
-            shutil.copy2(self.path, local_path)
-            return str(local_path)
+            shutil.copy2(self.path, local_path_for_copy)
+            return str(local_path_for_copy)
 
         # Otherwise download from remote using sync functionality
         # Use the sync version of storage operations
@@ -756,9 +793,9 @@ class File(BaseModel, Generic[T], SerializableType):
         if not os.path.exists(local_path):
             raise ValueError(f"File not found: {local_path}")
 
-        remote_path = remote_destination or internal_ctx().raw_data.get_random_remote_path()
-        protocol = get_protocol(remote_path)
         filename = Path(local_path).name
+        remote_path = remote_destination or internal_ctx().raw_data.get_random_remote_path(filename)
+        protocol = get_protocol(remote_path)
 
         # If remote_destination was not set by the user, and the configured raw data path is also local,
         # then let's optimize by not uploading.

flyte/models.py

@@ -14,7 +14,7 @@ from flyte._interface import extract_return_annotation, literal_to_enum
 from flyte._logging import logger
 
 if TYPE_CHECKING:
-    from flyteidl.core import literals_pb2
+    from flyteidl2.core import literals_pb2
 
     from flyte._internal.imagebuild.image_builder import ImageCache
     from flyte.report import Report
@@ -144,7 +144,7 @@ class RawDataPath:
 
     def get_random_remote_path(self, file_name: Optional[str] = None) -> str:
         """
-        Returns a random path for uploading a file/directory to.
+        Returns a random path for uploading a file/directory to. This file/folder will not be created, it's just a path.
 
         :param file_name: If given, will be joined after a randomly generated portion.
         :return:
@@ -160,13 +160,14 @@ class RawDataPath:
 
         protocol = get_protocol(file_prefix)
         if "file" in protocol:
-            local_path = pathlib.Path(file_prefix) / random_string
+            parent_folder = pathlib.Path(file_prefix)
+            parent_folder.mkdir(exist_ok=True, parents=True)
             if file_name:
-                # Only if file name is given do we create the parent, because it may be needed as a folder otherwise
-                local_path = local_path / file_name
-                if not local_path.exists():
-                    local_path.parent.mkdir(exist_ok=True, parents=True)
-                    local_path.touch()
+                random_folder = parent_folder / random_string
+                random_folder.mkdir()
+                local_path = random_folder / file_name
+            else:
+                local_path = parent_folder / random_string
             return str(local_path.absolute())
 
         fs = fsspec.filesystem(protocol)
@@ -194,6 +195,8 @@ class TaskContext:
     :param action: The action ID of the current execution. This is always set, within a run.
     :param version: The version of the executed task. This is set when the task is executed by an action and will be
       set on all sub-actions.
+    :param custom_context: Context metadata for the action. If an action receives context, it'll automatically pass it
+      to any actions it spawns. Context will not be used for cache key computation.
     """
 
     action: ActionID
@@ -210,6 +213,7 @@ class TaskContext:
     data: Dict[str, Any] = field(default_factory=dict)
     mode: Literal["local", "remote", "hybrid"] = "remote"
     interactive_mode: bool = False
+    custom_context: Dict[str, str] = field(default_factory=dict)
 
     def replace(self, **kwargs) -> TaskContext:
         if "data" in kwargs:

flyte/remote/_action.py

@@ -20,14 +20,15 @@ from typing import (
 import grpc
 import rich.pretty
 import rich.repr
+from flyteidl2.common import identifier_pb2, list_pb2
+from flyteidl2.task import common_pb2
+from flyteidl2.workflow import run_definition_pb2, run_service_pb2
+from flyteidl2.workflow.run_service_pb2 import WatchActionDetailsResponse
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 
 from flyte import types
-from flyte._initialize import ensure_client, get_client, get_common_config
-from flyte._protos.common import identifier_pb2, list_pb2
-from flyte._protos.workflow import run_definition_pb2, run_service_pb2
-from flyte._protos.workflow.run_service_pb2 import WatchActionDetailsResponse
+from flyte._initialize import ensure_client, get_client, get_init_config
 from flyte.remote._common import ToJSONMixin
 from flyte.remote._logs import Logs
 from flyte.syncify import syncify
@@ -67,14 +68,14 @@ def _action_rich_repr(action: run_definition_pb2.Action) -> rich.repr.Result:
     """
     Rich representation of the action.
     """
-    yield "run", action.id.run.name
+    yield "name", action.id.run.name
     if action.metadata.HasField("task"):
-        yield "task", action.metadata.task.id.name
+        yield "task name", action.metadata.task.id.name
         yield "type", action.metadata.task.task_type
     elif action.metadata.HasField("trace"):
         yield "trace", action.metadata.trace.name
         yield "type", "trace"
-    yield "name", action.id.name
+    yield "action name", action.id.name
     yield from _action_time_phase(action)
     yield "group", action.metadata.group
     yield "parent", action.metadata.parent
@@ -98,9 +99,10 @@ def _action_details_rich_repr(
     """
     yield "name", action.id.run.name
     yield from _action_time_phase(action)
-    yield "task", action.resolved_task_spec.task_template.id.name
-    yield "task_type", action.resolved_task_spec.task_template.type
-    yield "task_version", action.resolved_task_spec.task_template.id.version
+    if action.HasField("task"):
+        yield "task", action.task.task_template.id.name
+        yield "task_type", action.task.task_template.type
+        yield "task_version", action.task.task_template.id.version
     yield "attempts", action.attempts
     yield "error", (f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA")
     yield "phase", run_definition_pb2.Phase.Name(action.status.phase)
@@ -152,7 +154,7 @@ class Action(ToJSONMixin):
             key=sort_by[0],
             direction=(list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING),
         )
-        cfg = get_common_config()
+        cfg = get_init_config()
         while True:
             req = list_pb2.ListRequest(
                 limit=100,
@@ -193,7 +195,7 @@ class Action(ToJSONMixin):
         :param name: The name of the action.
         """
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         details: ActionDetails = await ActionDetails.get_details.aio(
             identifier_pb2.ActionIdentifier(
                 run=identifier_pb2.RunIdentifier(
@@ -458,7 +460,7 @@ class ActionDetails(ToJSONMixin):
         ensure_client()
         if not uri:
             assert name is not None and run_name is not None, "Either uri or name and run_name must be provided"
-        cfg = get_common_config()
+        cfg = get_init_config()
         return await cls.get_details.aio(
             identifier_pb2.ActionIdentifier(
                 run=identifier_pb2.RunIdentifier(
@@ -706,7 +708,7 @@ class ActionInputs(UserDict, ToJSONMixin):
     remote Union API.
     """
 
-    pb2: run_definition_pb2.Inputs
+    pb2: common_pb2.Inputs
     data: Dict[str, Any]
 
     def __repr__(self):
@@ -723,14 +725,14 @@ class ActionOutputs(tuple, ToJSONMixin):
     remote Union API.
     """
 
-    def __new__(cls, pb2: run_definition_pb2.Outputs, data: Tuple[Any, ...]):
+    def __new__(cls, pb2: common_pb2.Outputs, data: Tuple[Any, ...]):
         # Create the tuple part
         obj = super().__new__(cls, data)
         # Store extra data (you can't do this here directly since it's immutable)
         obj.pb2 = pb2
         return obj
 
-    def __init__(self, pb2: run_definition_pb2.Outputs, data: Tuple[Any, ...]):
+    def __init__(self, pb2: common_pb2.Outputs, data: Tuple[Any, ...]):
         # Normally you'd set instance attributes here,
         # but we've already set `pb2` in `__new__`
         self.pb2 = pb2

flyte/remote/_client/_protocols.py

@@ -2,12 +2,13 @@ from typing import AsyncIterator, Protocol
 
 from flyteidl.admin import project_attributes_pb2, project_pb2, version_pb2
 from flyteidl.service import dataproxy_pb2, identity_pb2
+from flyteidl2.secret import payload_pb2
+from flyteidl2.task import task_service_pb2
+from flyteidl2.trigger import trigger_service_pb2
+from flyteidl2.workflow import run_logs_service_pb2, run_service_pb2
 from grpc.aio import UnaryStreamCall
 from grpc.aio._typing import RequestType
 
-from flyte._protos.secret import payload_pb2
-from flyte._protos.workflow import run_logs_service_pb2, run_service_pb2, task_service_pb2, trigger_service_pb2
-
 
 class MetadataServiceProtocol(Protocol):
     async def GetVersion(self, request: version_pb2.GetVersionRequest) -> version_pb2.GetVersionResponse: ...

flyte/remote/_client/auth/_channel.py

@@ -112,7 +112,7 @@ async def create_channel(
     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)
+        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

flyte/remote/_client/controlplane.py

@@ -16,14 +16,10 @@ if "GRPC_VERBOSITY" not in os.environ:
 
 import grpc
 from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc, identity_pb2_grpc
-
-from flyte._protos.secret import secret_pb2_grpc
-from flyte._protos.workflow import (
-    run_logs_service_pb2_grpc,
-    run_service_pb2_grpc,
-    task_service_pb2_grpc,
-    trigger_service_pb2_grpc,
-)
+from flyteidl2.secret import secret_pb2_grpc
+from flyteidl2.task import task_service_pb2_grpc
+from flyteidl2.trigger import trigger_service_pb2_grpc
+from flyteidl2.workflow import run_logs_service_pb2_grpc, run_service_pb2_grpc
 
 from ._protocols import (
     DataProxyService,

flyte/remote/_data.py

@@ -15,7 +15,7 @@ import httpx
 from flyteidl.service import dataproxy_pb2
 from google.protobuf import duration_pb2
 
-from flyte._initialize import CommonInit, ensure_client, get_client, get_common_config
+from flyte._initialize import CommonInit, ensure_client, get_client, get_init_config, require_project_and_domain
 from flyte.errors import InitializationError, RuntimeSystemError
 from flyte.syncify import syncify
 
@@ -54,6 +54,7 @@ def hash_file(file_path: typing.Union[os.PathLike, str]) -> Tuple[bytes, str, in
     return h.digest(), h.hexdigest(), size
 
 
+@require_project_and_domain
 async def _upload_single_file(
     cfg: CommonInit, fp: Path, verify: bool = True, basedir: str | None = None
 ) -> Tuple[str, str]:
@@ -124,7 +125,7 @@ async def upload_file(fp: Path, verify: bool = True) -> Tuple[str, str]:
     """
     # This is a placeholder implementation. Replace with actual upload logic.
     ensure_client()
-    cfg = get_common_config()
+    cfg = get_init_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)
@@ -140,7 +141,7 @@ async def upload_dir(dir_path: Path, verify: bool = True) -> str:
     """
     # This is a placeholder implementation. Replace with actual upload logic.
     ensure_client()
-    cfg = get_common_config()
+    cfg = get_init_config()
     if not dir_path.is_dir():
         raise ValueError(f"{dir_path} is not a directory, upload arg must be a directory.")
 

flyte/remote/_logs.py

@@ -4,6 +4,9 @@ from dataclasses import dataclass
 from typing import AsyncGenerator, AsyncIterator
 
 import grpc
+from flyteidl2.common import identifier_pb2
+from flyteidl2.logs.dataplane import payload_pb2
+from flyteidl2.workflow import run_logs_service_pb2
 from rich.console import Console
 from rich.live import Live
 from rich.panel import Panel
@@ -11,9 +14,6 @@ from rich.text import Text
 
 from flyte._initialize import ensure_client, get_client
 from flyte._logging import logger
-from flyte._protos.common import identifier_pb2
-from flyte._protos.logs.dataplane import payload_pb2
-from flyte._protos.workflow import run_logs_service_pb2
 from flyte._tools import ipython_check, ipywidgets_check
 from flyte.errors import LogsNotYetAvailableError
 from flyte.syncify import syncify

flyte/remote/_run.py

@@ -5,11 +5,11 @@ from typing import AsyncGenerator, AsyncIterator, Literal, Tuple
 
 import grpc
 import rich.repr
+from flyteidl2.common import identifier_pb2, list_pb2
+from flyteidl2.workflow import run_definition_pb2, run_service_pb2
 
-from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._initialize import ensure_client, get_client, get_init_config
 from flyte._logging import logger
-from flyte._protos.common import identifier_pb2, list_pb2
-from flyte._protos.workflow import run_definition_pb2, run_service_pb2
 from flyte.syncify import syncify
 
 from . import Action, ActionDetails, ActionInputs, ActionOutputs
@@ -97,7 +97,7 @@ class Run(ToJSONMixin):
                 ),
             )
 
-        cfg = get_common_config()
+        cfg = get_init_config()
         i = 0
         while True:
             req = list_pb2.ListRequest(
@@ -316,7 +316,7 @@ class RunDetails(ToJSONMixin):
         :param name: The name of the run.
         """
         ensure_client()
-        cfg = get_common_config()
+        cfg = get_init_config()
         return await RunDetails.get_details.aio(
             run_id=identifier_pb2.RunIdentifier(
                 org=cfg.org,