Fix get_task_inputs in flytekitpluugins-interactive (#2398)

=Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · May 9, 2024 · 28f1e84 · 28f1e84
1 parent 4dd4d22
commit 28f1e84
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/utils.py b/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/utils.py
@@ -53,6 +53,7 @@ def get_task_inputs(task_module_name, task_name, context_working_dir):
     local_inputs_file = os.path.join(context_working_dir, "inputs.pb")
     input_proto = utils.load_proto_from_file(_literals_pb2.LiteralMap, local_inputs_file)
     idl_input_literals = _literal_models.LiteralMap.from_flyte_idl(input_proto)
+
     task_module = load_module_from_path(task_module_name, os.path.join(context_working_dir, f"{task_module_name}.py"))
     task_def = getattr(task_module, task_name)
     native_inputs = TypeEngine.literal_map_to_kwargs(

diff --git a/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/vscode_lib/decorator.py b/plugins/flytekit-flyteinteractive/flytekitplugins/flyteinteractive/vscode_lib/decorator.py
@@ -261,7 +261,7 @@ def prepare_interactive_python(task_function):
 
 if __name__ == "__main__":
     inputs = get_task_inputs(
-        task_module_name="{task_module_name}",
+        task_module_name="{task_module_name.split('.')[-1]}",
         task_name="{task_name}",
         context_working_dir="{context_working_dir}",
     )