add remote file download

Signed-off-by: ChungYujoyce <[email protected]>

flytekit/clis/sdk_in_container/run.py

@@ -24,6 +24,7 @@
     CTX_MODULE,
     CTX_PROJECT,
     CTX_PROJECT_ROOT,
+    CTX_FILENAME,
 )
 from flytekit.clis.sdk_in_container.helpers import (
     FLYTE_REMOTE_INSTANCE_KEY,
@@ -708,18 +709,22 @@ def _run(*args, **kwargs):
 
         console_url = remote.generate_console_url(execution)
         click.secho(f"Go to {console_url} to see execution in the console.")
-        if remote_entity:
-            file = str(remote_entity.id.name).split('.')[0]
-            file += ".py"
-            if os.path.exists(file):
-                os.remove(file)
 
         if run_level_params.get("dump_snippet"):
             dump_flyte_remote_snippet(execution, project, domain)
 
+        return ctx
+
     return _run
 
 
+def result_command(ctx: click.Context):
+    if ctx is not None:
+        filename = ctx.obj[RUN_LEVEL_PARAMS_KEY].get('filename')
+        if os.path.exists(filename):
+            os.remove(filename)
+
+
 class WorkflowCommand(click.RichGroup):
     """
     click multicommand at the python file layer, subcommands should be all the workflows in the file.
@@ -729,17 +734,20 @@ def __init__(self, filename: str, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         ctx = context_manager.FlyteContextManager.current_context()
+        self.file = self.download(ctx, filename)
+        self._filename = pathlib.Path(self.file).resolve()
+
+    def download(self, ctx, filename):
         if ctx.file_access.is_remote(filename):
             local_path = '.'
             ctx.file_access.get(filename, local_path)
             filename = filename.rsplit('/', 1)[1]
-
-        self._filename = pathlib.Path(filename).resolve()
-
+        return filename
+
     def list_commands(self, ctx):
         entities = get_entities_in_file(self._filename)
         return entities.all()
-
+    
     def get_command(self, ctx, exe_entity):
         """
         This command uses the filename with which this command was created, and the string name of the entity passed
@@ -767,7 +775,7 @@ def get_command(self, ctx, exe_entity):
 
         ctx.obj[RUN_LEVEL_PARAMS_KEY][CTX_PROJECT_ROOT] = project_root
         ctx.obj[RUN_LEVEL_PARAMS_KEY][CTX_MODULE] = module
-
+        ctx.obj[RUN_LEVEL_PARAMS_KEY][CTX_FILENAME] = self.file
         entity = load_naive_entity(module, exe_entity, project_root)
 
         # If this is a remote execution, which we should know at this point, then create the remote object
@@ -810,8 +818,11 @@ def list_commands(self, ctx):
     def get_command(self, ctx, filename):
         if ctx.obj:
             ctx.obj[RUN_LEVEL_PARAMS_KEY] = ctx.params
-        return WorkflowCommand(filename, name=filename, help="Run a [workflow|task] in a file using script mode")
-
+        return WorkflowCommand(filename, 
+                               name=filename, 
+                               result_callback=result_command,
+                               help="Run a [workflow|task] in a file using script mode")
+
 
 _run_help = """
 This command can execute either a workflow or a task from the command line, for fully self-contained scripts.

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -85,9 +85,6 @@ def test_remote_files():
         ["run", "--remote", REMOTE_WORKFLOW_FILE, "my_wf", "--a", "1", "--b", "Hello"],
         catch_exceptions=False,
     )
-    file = "basic_workflow.py"
-    if os.path.exists(file):
-        os.remove(file)
     assert result.exit_code == 0