Add --raw-output-data-prefix to the pyflyte-execute command

flytekit/bin/entrypoint.py

@@ -53,7 +53,7 @@ def _map_job_index_to_child_index(local_input_dir, datadir, index):
 
 
 @_scopes.system_entry_point
-def _execute_task(task_module, task_name, inputs, output_prefix, test):
+def _execute_task(task_module, task_name, inputs, output_prefix, raw_output_data_prefix, test):
     with _TemporaryConfiguration(_internal_config.CONFIGURATION_PATH.get()):
         with _utils.AutoDeletingTempDir("input_dir") as input_dir:
             # Load user code
@@ -83,7 +83,10 @@ def _execute_task(task_module, task_name, inputs, output_prefix, test):
                 _data_proxy.Data.get_data(inputs, local_inputs_file)
                 input_proto = _utils.load_proto_from_file(_literals_pb2.LiteralMap, local_inputs_file)
                 _engine_loader.get_engine().get_task(task_def).execute(
-                    _literal_models.LiteralMap.from_flyte_idl(input_proto), context={"output_prefix": output_prefix},
+                    _literal_models.LiteralMap.from_flyte_idl(input_proto), context={
+                        "output_prefix": output_prefix,
+                        "raw_output_data_prefix": raw_output_data_prefix,
+                    },
                 )
 
 
@@ -97,10 +100,11 @@ def _pass_through():
 @_click.option("--task-name", required=True)
 @_click.option("--inputs", required=True)
 @_click.option("--output-prefix", required=True)
+@_click.option("--raw-data-output-prefix", required=False)
 @_click.option("--test", is_flag=True)
-def execute_task_cmd(task_module, task_name, inputs, output_prefix, test):
+def execute_task_cmd(task_module, task_name, inputs, output_prefix, raw_output_data_prefix, test):
     _click.echo(_utils.get_version_message())
-    _execute_task(task_module, task_name, inputs, output_prefix, test)
+    _execute_task(task_module, task_name, inputs, output_prefix, raw_output_data_prefix, test)
 
 
 if __name__ == "__main__":

flytekit/common/tasks/sdk_runnable.py

@@ -447,6 +447,8 @@ def _get_container_definition(
                 "{{.input}}",
                 "--output-prefix",
                 "{{.outputPrefix}}",
+                "--raw-data-output-prefix",
+                "{{.outputPrefix}}"
             ],
             resources=_task_models.Resources(limits=limits, requests=requests),
             env=environment,

flytekit/engines/common.py

@@ -396,12 +396,13 @@ def fetch_workflow(self, workflow_id):
 
 
 class EngineContext(object):
-    def __init__(self, execution_date, tmp_dir, stats, execution_id, logging):
+    def __init__(self, execution_date, tmp_dir, stats, execution_id, logging, raw_output_data_prefix=None):
         self._stats = stats
         self._execution_date = execution_date
         self._working_directory = tmp_dir
         self._execution_id = execution_id
         self._logging = logging
+        self._raw_output_data_prefix = raw_output_data_prefix
 
     @property
     def stats(self):
@@ -437,3 +438,7 @@ def execution_id(self):
         :rtype: flytekit.models.core.identifier.WorkflowExecutionIdentifier
         """
         return self._execution_id
+
+    @property
+    def raw_output_data_prefix(self) -> str:
+        return self._raw_output_data_prefix

flytekit/engines/flyte/engine.py

@@ -270,11 +270,11 @@ def execute(self, inputs, context=None):
         :param dict[Text, Text] context:
         :rtype: dict[Text, flytekit.models.common.FlyteIdlEntity]
         """
-
         with _common_utils.AutoDeletingTempDir("engine_dir") as temp_dir:
             with _common_utils.AutoDeletingTempDir("task_dir") as task_dir:
                 with _data_proxy.LocalWorkingDirectoryContext(task_dir):
-                    with _data_proxy.RemoteDataContext():
+                    raw_output_data_prefix = context.get("raw_output_data_prefix", None)
+                    with _data_proxy.RemoteDataContext(raw_output_data_prefix_override=raw_output_data_prefix):
                         output_file_dict = dict()
 
                         # This sets the logging level for user code and is the only place an sdk setting gets
@@ -311,6 +311,7 @@ def execute(self, inputs, context=None):
                                     ),
                                     logging=_logging,
                                     tmp_dir=task_dir,
+                                    raw_output_data_prefix=context['raw_output_data_prefix'] if "raw_output_data_prefix" in context else None,
                                 ),
                                 inputs,
                             )

flytekit/interfaces/data/data_proxy.py

@@ -64,22 +64,23 @@ def __init__(self, sandbox):
 class RemoteDataContext(_OutputDataContext):
 
     _CLOUD_PROVIDER_TO_PROXIES = {
-        _constants.CloudProvider.AWS: _s3proxy.AwsS3Proxy(),
-        _constants.CloudProvider.GCP: _gcs_proxy.GCSProxy(),
+        _constants.CloudProvider.AWS: _s3proxy.AwsS3Proxy,
+        _constants.CloudProvider.GCP: _gcs_proxy.GCSProxy,
     }
 
-    def __init__(self, cloud_provider=None):
+    def __init__(self, cloud_provider=None, raw_output_data_prefix_override=None):
         """
         :param Optional[Text] cloud_provider: From flytekit.common.constants.CloudProvider enum
         """
         cloud_provider = cloud_provider or _platform_config.CLOUD_PROVIDER.get()
-        proxy = type(self)._CLOUD_PROVIDER_TO_PROXIES.get(cloud_provider, None)
-        if proxy is None:
+        proxy_class = type(self)._CLOUD_PROVIDER_TO_PROXIES.get(cloud_provider, None)
+        if proxy_class is None:
             raise _user_exception.FlyteAssertion(
                 "Configured cloud provider is not supported for data I/O.  Received: {}, expected one of: {}".format(
                     cloud_provider, list(type(self)._CLOUD_PROVIDER_TO_PROXIES.keys())
                 )
             )
+        proxy = proxy_class(raw_output_data_prefix_override)
         super(RemoteDataContext, self).__init__(proxy)
 
 

flytekit/interfaces/data/gcs/gcs_proxy.py

@@ -28,6 +28,19 @@ def _amend_path(path):
 class GCSProxy(_common_data.DataProxy):
     _GS_UTIL_CLI = "gsutil"
 
+    def __init__(self, raw_output_data_prefix_override: str = None):
+        """
+        :param raw_output_data_prefix_override: Instead of relying on the AWS or GCS configuration (see
+            S3_SHARD_FORMATTER for AWS and GCS_PREFIX for GCP) setting when computing the shard
+            path (_get_shard_path), use this prefix instead as a base. This code assumes that the
+            path passed in is correct. That is, an S3 path won't be passed in when running on GCP.
+        """
+        self._raw_output_data_prefix_override = raw_output_data_prefix_override
+
+    @property
+    def raw_output_data_prefix_override(self) -> str:
+        return self._raw_output_data_prefix_override
+
     @staticmethod
     def _check_binary():
         """
@@ -119,12 +132,14 @@ def upload_directory(self, local_path, remote_path):
         )
         return _update_cmd_config_and_execute(cmd)
 
-    def get_random_path(self):
+    def get_random_path(self) -> str:
         """
-        :rtype: Text
+        If this object was created with a raw output data prefix, usually set by Propeller/Plugins at execution time
+        and piped all the way here, it will be used instead of referencing the GCS_PREFIX configuration.
         """
         key = _uuid.UUID(int=_flyte_random.random.getrandbits(128)).hex
-        return _os.path.join(_gcp_config.GCS_PREFIX.get(), key)
+        prefix = self.raw_output_data_prefix_override or _gcp_config.GCS_PREFIX.get()
+        return _os.path.join(prefix, key)
 
     def get_random_directory(self):
         """

flytekit/interfaces/data/s3/s3proxy.py

@@ -41,6 +41,19 @@ class AwsS3Proxy(_common_data.DataProxy):
     _AWS_CLI = "aws"
     _SHARD_CHARACTERS = [_text_type(x) for x in _six_moves.range(10)] + list(_string.ascii_lowercase)
 
+    def __init__(self, raw_output_data_prefix_override: str = None):
+        """
+        :param raw_output_data_prefix_override: Instead of relying on the AWS or GCS configuration (see
+            S3_SHARD_FORMATTER for AWS and GCS_PREFIX for GCP) setting when computing the shard
+            path (_get_shard_path), use this prefix instead as a base. This code assumes that the
+            path passed in is correct. That is, an S3 path won't be passed in when running on GCP.
+        """
+        self._raw_output_data_prefix_override = raw_output_data_prefix_override
+
+    @property
+    def raw_output_data_prefix_override(self) -> str:
+        return self._raw_output_data_prefix_override
+
     @staticmethod
     def _check_binary():
         """
@@ -179,10 +192,14 @@ def get_random_directory(self):
         """
         return self.get_random_path() + "/"
 
-    def _get_shard_path(self):
+    def _get_shard_path(self) -> str:
         """
-        :rtype: Text
+        If this object was created with a raw output data prefix, usually set by Propeller/Plugins at execution time
+        and piped all the way here, it will be used instead of referencing the S3 shard configuration.
         """
+        if self.raw_output_data_prefix_override:
+            return self.raw_output_data_prefix_override
+
         shard = ""
         for _ in _six_moves.range(_aws_config.S3_SHARD_STRING_LENGTH.get()):
             shard += _flyte_random.random.choice(self._SHARD_CHARACTERS)