Added alt prefix head to FlyteFile.new_remote

flytekit/core/data_persistence.py

@@ -454,6 +454,27 @@ def join(
             f = fs.unstrip_protocol(f)
         return f
 
+    def get_new_path(
+        self,
+        fs: typing.Optional[fsspec.AbstractFileSystem] = None,
+        alt: typing.Optional[str] = None,
+        stem: typing.Optional[str] = None,
+        unstrip: bool = False,
+    ) -> str:
+        fs = fs or self.raw_output_fs
+        pref = self.raw_output_prefix
+        s_pref = pref.split(fs.sep)
+        if alt:
+            s_pref[2] = alt
+        if stem:
+            s_pref.append(stem)
+        else:
+            s_pref.append(self.get_random_string())
+        p = fs.sep.join(s_pref)
+        if unstrip:
+            p = fs.unstrip_protocol(p)
+        return p
+
     def get_random_local_path(self, file_path_or_file_name: typing.Optional[str] = None) -> str:
         """
         Use file_path_or_file_name, when you want a random directory, but want to preserve the leaf file name

flytekit/types/file/file.py

@@ -179,13 +179,15 @@ def extension(cls) -> str:
         return ""
 
     @classmethod
-    def new_remote_file(cls, name: typing.Optional[str] = None) -> FlyteFile:
+    def new_remote_file(cls, name: typing.Optional[str] = None, alt: typing.Optional[str] = None) -> FlyteFile:
         """
         Create a new FlyteFile object with a remote path.
+
+        :param name: If you want to specify a different name for the file, you can specify it here.
+        :param alt: If you want to specify a different prefix head than the default one, you can specify it here.
         """
         ctx = FlyteContextManager.current_context()
-        r = name or ctx.file_access.get_random_string()
-        remote_path = ctx.file_access.join(ctx.file_access.raw_output_prefix, r)
+        remote_path = ctx.file_access.get_new_path(alt=alt, stem=name)
         return cls(path=remote_path)
 
     @classmethod

tests/flytekit/unit/core/test_data_persistence.py

@@ -136,6 +136,18 @@ def test_write_known_location():
         assert f.read() == arbitrary_text.encode("utf-8")
 
 
+def test_get_new_path():
+    """
+    Test that a new path given alternate bucket and name is generated correctly
+    """
+    random_dir = tempfile.mkdtemp()
+    raw = os.path.join(random_dir, "raw")
+    fs = FileAccessProvider(local_sandbox_dir=random_dir, raw_output_prefix=raw)
+    np = fs.get_new_path(alt="foo", stem="bar.txt")
+    assert "foo" in np
+    assert np.endswith("bar.txt")
+
+
 def test_initialise_azure_file_provider_with_account_key():
     with mock.patch.dict(
         os.environ,

tests/flytekit/unit/types/directory/test_types.py

@@ -22,6 +22,9 @@ def test_new_remote_dir():
     fd = FlyteDirectory.new_remote()
     assert FlyteContext.current_context().file_access.raw_output_prefix in fd.path
 
+# def test_new_remote_dir_alt():
+#     fd = FlyteDirectory.new_remote(alt_bucket="my-alt-bucket")
+#     assert "s3://my-alt-bucket" in fd.path
 
 @mock.patch("flytekit.types.directory.types.os.name", "nt")
 def test_sep_nt():

tests/flytekit/unit/types/file/test_types.py

@@ -0,0 +1,6 @@
+from flytekit.types.file import FlyteFile
+from flytekit import FlyteContextManager
+
+def test_new_remote_alt():
+    ff = FlyteFile.new_remote_file(alt="my-alt-prefix", name="my-file.txt")
+    assert "my-alt-prefix" in ff.path