Flyte flyte:// file system and improve remote file handling

dev-requirements.in

@@ -32,4 +32,5 @@ types-protobuf
 types-croniter
 types-mock
 autoflake
+types-requests
 prometheus-client

flytekit/clients/friendly.py

@@ -978,7 +978,13 @@ def list_matchable_attributes(self, resource_type):
         )
 
     def get_upload_signed_url(
-        self, project: str, domain: str, content_md5: bytes, filename: str = None, expires_in: datetime.timedelta = None
+        self,
+        project: str,
+        domain: str,
+        content_md5: typing.Optional[bytes] = None,
+        filename: typing.Optional[str] = None,
+        expires_in: typing.Optional[datetime.timedelta] = None,
+        filename_root: typing.Optional[str] = None,
     ) -> _data_proxy_pb2.CreateUploadLocationResponse:
         """
         Get a signed url to be used during fast registration.
@@ -990,6 +996,8 @@ def get_upload_signed_url(
         :param str filename: [Optional] If provided this specifies a desired suffix for the generated location
         :param datetime.timedelta expires_in: [Optional] If provided this defines a requested expiration duration for
             the generated url
+        :param filename_root: If provided will be used as the root of the filename.  If not, Admin will use a hash
+          This option is useful when uploading a series of files that you want to be grouped together.
         :rtype: flyteidl.service.dataproxy_pb2.CreateUploadLocationResponse
         """
         expires_in_pb = None
@@ -1003,12 +1011,13 @@ def get_upload_signed_url(
                 content_md5=content_md5,
                 filename=filename,
                 expires_in=expires_in_pb,
+                filename_root=filename_root,
             )
         )
 
     def get_download_signed_url(
         self, native_url: str, expires_in: datetime.timedelta = None
-    ) -> _data_proxy_pb2.CreateDownloadLocationRequest:
+    ) -> _data_proxy_pb2.CreateDownloadLocationResponse:
         expires_in_pb = None
         if expires_in:
             expires_in_pb = Duration()

flytekit/clis/flyte_cli/main.py

@@ -1822,7 +1822,7 @@ def fast_register_files(
     version = version if version else digest
     full_remote_path = _get_additional_distribution_loc(additional_distribution_dir, version)
     ctx = FlyteContextManager.current_context()
-    ctx.file_access.put_data(compressed_source, full_remote_path)
+    full_remote_path = ctx.file_access.put_data(compressed_source, full_remote_path)
     _click.secho(f"Uploaded compressed code archive {compressed_source} to {full_remote_path}", fg="green")
 
     def fast_register_task(entity: _GeneratedProtocolMessageType) -> _GeneratedProtocolMessageType:

flytekit/clis/sdk_in_container/fetch.py

@@ -0,0 +1,21 @@
+import rich
+import rich_click as click
+
+from flytekit.clis.sdk_in_container.helpers import get_and_save_remote_with_click_context
+from flytekit.remote import FlyteRemote
+
+
+@click.command("fetch")
+@click.argument("flyte_data_uri", type=str, required=True, metavar="FLYTE-DATA-URI (of the form flyte://...)")
+@click.pass_context
+def fetch(ctx: click.Context, flyte_data_uri: str):
+    """
+    Retrieve Inputs/Outputs for a Flyte Execution or any of the inner node executions from the remote server.
+
+    The URI can be retrieved from the Flyte Console, or by invoking the get_data API.
+    """
+
+    remote: FlyteRemote = get_and_save_remote_with_click_context(ctx, project="flytesnacks", domain="development")
+    click.secho(f"Fetching data from {flyte_data_uri}...", dim=True)
+    data = remote.get(flyte_data_uri)
+    rich.print(data.literals)

flytekit/clis/sdk_in_container/helpers.py

@@ -12,7 +12,9 @@
 FLYTE_REMOTE_INSTANCE_KEY = "flyte_remote"
 
 
-def get_remote(cfg_file_path: typing.Optional[str], project: str, domain: str) -> FlyteRemote:
+def get_remote(
+    cfg_file_path: typing.Optional[str], project: str, domain: str, data_upload_location: Optional[str] = None
+) -> FlyteRemote:
     cfg_file = get_config_file(cfg_file_path)
     if cfg_file is None:
         cfg_obj = Config.for_sandbox()
@@ -22,11 +24,17 @@ def get_remote(cfg_file_path: typing.Optional[str], project: str, domain: str) -
         cli_logger.info(
             f"Creating remote with config {cfg_obj}" + (f" with file {cfg_file_path}" if cfg_file_path else "")
         )
-    return FlyteRemote(cfg_obj, default_project=project, default_domain=domain)
+    return FlyteRemote(
+        cfg_obj, default_project=project, default_domain=domain, data_upload_location=data_upload_location
+    )
 
 
 def get_and_save_remote_with_click_context(
-    ctx: click.Context, project: str, domain: str, save: bool = True
+    ctx: click.Context,
+    project: str,
+    domain: str,
+    save: bool = True,
+    data_upload_location: Optional[str] = None,
 ) -> FlyteRemote:
     """
     NB: This function will by default mutate the click Context.obj dictionary, adding a remote key with value
@@ -36,12 +44,13 @@ def get_and_save_remote_with_click_context(
     :param project: default project for the remote instance
     :param domain: default domain
     :param save: If false, will not mutate the context.obj dict
+    :param data_upload_location: if specified, will set the data upload location for the remote instance
     :return: FlyteRemote instance
     """
     if ctx.obj.get(FLYTE_REMOTE_INSTANCE_KEY) is not None:
         return ctx.obj[FLYTE_REMOTE_INSTANCE_KEY]
     cfg_file_location = ctx.obj.get(CTX_CONFIG_FILE)
-    r = get_remote(cfg_file_location, project, domain)
+    r = get_remote(cfg_file_location, project, domain, data_upload_location)
     if save:
         ctx.obj[FLYTE_REMOTE_INSTANCE_KEY] = r
     return r

flytekit/clis/sdk_in_container/launchplan.py

@@ -44,7 +44,7 @@ def launchplan(
     launchplan: str,
     launchplan_version: str,
 ):
-    remote = get_and_save_remote_with_click_context(ctx, project, domain)
+    remote = get_and_save_remote_with_click_context(ctx, project, domain, data_upload_location="flyte://data")
     with Progress() as progress:
         t1 = progress.add_task(f"[cyan] {'Activating' if activate else 'Deactivating'}...", total=1)
         try:

flytekit/clis/sdk_in_container/pyflyte.py

@@ -7,6 +7,7 @@
 from flytekit.clis.sdk_in_container.backfill import backfill
 from flytekit.clis.sdk_in_container.build import build
 from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE, CTX_PACKAGES, CTX_VERBOSE
+from flytekit.clis.sdk_in_container.fetch import fetch
 from flytekit.clis.sdk_in_container.init import init
 from flytekit.clis.sdk_in_container.launchplan import launchplan
 from flytekit.clis.sdk_in_container.local_cache import local_cache
@@ -17,6 +18,7 @@
 from flytekit.clis.sdk_in_container.serialize import serialize
 from flytekit.clis.sdk_in_container.serve import serve
 from flytekit.clis.sdk_in_container.utils import ErrorHandlingCommand, validate_package
+from flytekit.clis.version import info
 from flytekit.configuration.file import FLYTECTL_CONFIG_ENV_VAR, FLYTECTL_CONFIG_ENV_VAR_OVERRIDE
 from flytekit.configuration.internal import LocalSDK
 from flytekit.loggers import cli_logger
@@ -80,6 +82,8 @@ def main(ctx, pkgs: typing.List[str], config: str, verbose: bool):
 main.add_command(build)
 main.add_command(metrics)
 main.add_command(launchplan)
+main.add_command(fetch)
+main.add_command(info)
 main.epilog
 
 if __name__ == "__main__":

flytekit/clis/sdk_in_container/register.py

@@ -157,7 +157,7 @@ def register(
     )
 
     # Create and save FlyteRemote,
-    remote = get_and_save_remote_with_click_context(ctx, project, domain)
+    remote = get_and_save_remote_with_click_context(ctx, project, domain, data_upload_location="flyte://data")
     click.secho(f"Registering against {remote.config.platform.endpoint}")
     try:
         repo.register(

flytekit/clis/sdk_in_container/run.py

@@ -1,5 +1,4 @@
 import asyncio
-import functools
 import importlib
 import inspect
 import json
@@ -34,7 +33,7 @@
 from flytekit.models.common import RawOutputDataConfig
 from flytekit.models.interface import Parameter, Variable
 from flytekit.models.types import SimpleType
-from flytekit.remote import FlyteLaunchPlan, FlyteRemote, FlyteTask, FlyteWorkflow
+from flytekit.remote import FlyteLaunchPlan, FlyteRemote, FlyteTask, FlyteWorkflow, remote_fs
 from flytekit.remote.executions import FlyteWorkflowExecution
 from flytekit.tools import module_loader
 from flytekit.tools.script_mode import _find_project_root
@@ -257,7 +256,10 @@ class RunLevelParams(PyFlyteParams):
 
     def remote_instance(self) -> FlyteRemote:
         if self._remote is None:
-            self._remote = get_remote(self.config_file, self.project, self.domain)
+            data_upload_location = None
+            if self.is_remote:
+                data_upload_location = remote_fs.REMOTE_PLACEHOLDER
+            self._remote = get_remote(self.config_file, self.project, self.domain, data_upload_location)
         return self._remote
 
     @property
@@ -350,7 +352,6 @@ def to_click_option(
     literal_var: Variable,
     python_type: typing.Type,
     default_val: typing.Any,
-    get_upload_url_fn: typing.Callable,
     required: bool,
 ) -> click.Option:
     """
@@ -363,9 +364,7 @@ def to_click_option(
         flyte_ctx,
         literal_type=literal_var.type,
         python_type=python_type,
-        get_upload_url_fn=get_upload_url_fn,
         is_remote=run_level_params.is_remote,
-        remote_instance_accessor=run_level_params.remote_instance,
     )
 
     if literal_converter.is_bool() and not default_val:
@@ -487,26 +486,27 @@ def _run(*args, **kwargs):
             image_config = run_level_params.image_config
             image_config = patch_image_config(config_file, image_config)
 
-            remote_entity = remote.register_script(
-                entity,
-                project=run_level_params.project,
-                domain=run_level_params.domain,
-                image_config=image_config,
-                destination_dir=run_level_params.destination_dir,
-                source_path=run_level_params.computed_params.project_root,
-                module_name=run_level_params.computed_params.module,
-                copy_all=run_level_params.copy_all,
-            )
+            with context_manager.FlyteContextManager.with_context(remote.context.new_builder()):
+                remote_entity = remote.register_script(
+                    entity,
+                    project=run_level_params.project,
+                    domain=run_level_params.domain,
+                    image_config=image_config,
+                    destination_dir=run_level_params.destination_dir,
+                    source_path=run_level_params.computed_params.project_root,
+                    module_name=run_level_params.computed_params.module,
+                    copy_all=run_level_params.copy_all,
+                )
 
-            run_remote(
-                remote,
-                remote_entity,
-                run_level_params.project,
-                run_level_params.domain,
-                inputs,
-                run_level_params,
-                type_hints=entity.python_interface.inputs,
-            )
+                run_remote(
+                    remote,
+                    remote_entity,
+                    run_level_params.project,
+                    run_level_params.domain,
+                    inputs,
+                    run_level_params,
+                    type_hints=entity.python_interface.inputs,
+                )
         finally:
             if run_level_params.computed_params.temp_file_name:
                 os.remove(run_level_params.computed_params.temp_file_name)
@@ -542,22 +542,14 @@ def _get_params(
         defaults: typing.Dict[str, Parameter],
     ) -> typing.List["click.Parameter"]:
         params = []
-        run_level_params: RunLevelParams = ctx.obj
-        r = run_level_params.remote_instance()
-
-        get_upload_url_fn = functools.partial(
-            r.client.get_upload_signed_url, project=run_level_params.project, domain=run_level_params.domain
-        )
         flyte_ctx = context_manager.FlyteContextManager.current_context()
         for name, var in inputs.items():
             if fixed and name in fixed:
                 continue
             required = True
             if defaults and name in defaults:
                 required = False
-            params.append(
-                to_click_option(ctx, flyte_ctx, name, var, native_inputs[name], None, get_upload_url_fn, required)
-            )
+            params.append(to_click_option(ctx, flyte_ctx, name, var, native_inputs[name], None, required))
         return params
 
     def get_params(self, ctx: click.Context) -> typing.List["click.Parameter"]:
@@ -675,23 +667,15 @@ def _create_command(
 
         # If this is a remote execution, which we should know at this point, then create the remote object
         r = run_level_params.remote_instance()
-        get_upload_url_fn = functools.partial(
-            r.client.get_upload_signed_url, project=run_level_params.project, domain=run_level_params.domain
-        )
-
-        flyte_ctx = context_manager.FlyteContextManager.current_context()
+        flyte_ctx = r.context
 
         # Add options for each of the workflow inputs
         params = []
         for input_name, input_type_val in loaded_entity.python_interface.inputs_with_defaults.items():
             literal_var = loaded_entity.interface.inputs.get(input_name)
             python_type, default_val = input_type_val
             required = type(None) not in get_args(python_type) and default_val is None
-            params.append(
-                to_click_option(
-                    ctx, flyte_ctx, input_name, literal_var, python_type, default_val, get_upload_url_fn, required
-                )
-            )
+            params.append(to_click_option(ctx, flyte_ctx, input_name, literal_var, python_type, default_val, required))
 
         entity_type = "Workflow" if is_workflow else "Task"
         h = f"{click.style(entity_type, bold=True)} ({run_level_params.computed_params.module}.{entity_name})"

flytekit/clis/version.py

@@ -0,0 +1,27 @@
+import rich
+import rich_click as click
+from rich.panel import Panel
+
+from flytekit.clis.sdk_in_container.helpers import get_and_save_remote_with_click_context
+from flytekit.remote import FlyteRemote
+
+Content = """
+This CLI is meant to be used within a virtual environment that has Flytekit installed. Ideally it is used to iterate on your Flyte workflows and tasks.
+
+Flytekit Version: [cyan]{version}[reset]
+Flyte Backend Endpoint: [cyan]{endpoint}
+"""
+
+
+@click.command("info")
+@click.pass_context
+def info(ctx: click.Context):
+    """
+    Print out information about the current Flyte Python CLI environment - like the version of Flytekit, backend endpoint
+    currently configured, etc.
+    """
+    import flytekit
+
+    remote: FlyteRemote = get_and_save_remote_with_click_context(ctx, project="flytesnacks", domain="development")
+    c = Content.format(version=flytekit.__version__, endpoint=remote.client.url)
+    rich.print(Panel(c, title="Flytekit CLI Info", border_style="purple", padding=(1, 1, 1, 1)))

flytekit/configuration/__init__.py

@@ -614,7 +614,7 @@ def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> DataConfig:
 @dataclass(init=True, repr=True, eq=True, frozen=True)
 class Config(object):
     """
-    This is the parent configuration object and holds all the underlying configuration object types. An instance of
+    This the parent configuration object and holds all the underlying configuration object types. An instance of
     this object holds all the config necessary to
 
     1. Interactive session with Flyte backend

flytekit/core/array_node_map_task.py

@@ -9,7 +9,6 @@
 from flytekit.configuration import SerializationSettings
 from flytekit.core import tracker
 from flytekit.core.base_task import PythonTask, TaskResolverMixin
-from flytekit.core.constants import SdkTaskType
 from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager
 from flytekit.core.interface import transform_interface_to_list_interface
 from flytekit.core.python_function_task import PythonFunctionTask, PythonInstanceTask
@@ -94,7 +93,7 @@ def __init__(
         super().__init__(
             name=self.name,
             interface=collection_interface,
-            task_type=SdkTaskType.PYTHON_TASK,
+            task_type=self._run_task.task_type,
             task_config=None,
             task_type_version=1,
             **kwargs,