From cf16ea6ecd6f5537cb7f2a14cb7c9a50be58248e Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Thu, 31 Mar 2022 12:07:58 -0400 Subject: [PATCH 01/27] checkpoint Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 73 ++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 70990647f1..56e2a82eac 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -237,3 +237,76 @@ def execute(self, **kwargs) -> typing.Any: def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any: return self._config_task_instance.post_execute(user_params, rval) + + +class PortableShellTask(ShellTask): + + def __init__( + self, + name: str, + script_file: str, + debug: bool = False, + task_config: T = None, + inputs: typing.Optional[typing.Dict[str, typing.Type]] = None, + output_locs: typing.Optional[typing.List[OutputLocation]] = None, + **kwargs, + ): + self._script = """ + #!/bin/bash + + set -uexo pipefail + + {inputs.export_env} + + bash {inputs.script_file} {inputs.args} + """ + + def execute(self, **kwargs) -> typing.Any: + """ + Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem + """ + logger.info(f"Running shell script as type {self.task_type}") + if self.script_file: + with open(self.script_file) as f: + self._script = f.read() + + outputs: typing.Dict[str, str] = {} + if self._output_locs: + for v in self._output_locs: + outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=kwargs) + + if os.name == "nt": + self._script = self._script.lstrip().rstrip().replace("\n", "&&") + + gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) + if self._debug: + print("\n==============================================\n") + print(gen_script) + print("\n==============================================\n") + + try: + subprocess.check_call(gen_script, shell=True) + except subprocess.CalledProcessError as e: + files = os.listdir(".") + fstr = "\n-".join(files) + logger.error( + f"Failed to Execute Script, return-code {e.returncode} \n" + f"StdErr: {e.stderr}\n" + f"StdOut: {e.stdout}\n" + f" Current directory contents: .\n-{fstr}" + ) + raise + + final_outputs = [] + for v in self._output_locs: + if issubclass(v.var_type, FlyteFile): + final_outputs.append(FlyteFile(outputs[v.var])) + if issubclass(v.var_type, FlyteDirectory): + final_outputs.append(FlyteDirectory(outputs[v.var])) + if len(final_outputs) == 1: + return final_outputs[0] + if len(final_outputs) > 1: + return tuple(final_outputs) + return None + + From eff0a7dc0e4dbd59540257e876c1256d44325bf1 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Thu, 31 Mar 2022 13:05:01 -0400 Subject: [PATCH 02/27] Experimental implementation works well. Instead of messing with the class, we utilize the interpolation and context to dynamically generate a wrapper around the desired script. In the wrapper, we cd to the ctx.working_directory, export the env variables (handled by an additional method to convert dict to str), and then pass the arguments to the script Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 95 ++++++++++------------------------ 1 file changed, 28 insertions(+), 67 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 56e2a82eac..014e7a2576 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -184,6 +184,12 @@ def script(self) -> typing.Optional[str]: def script_file(self) -> typing.Optional[os.PathLike]: return self._script_file + def make_export_string_from_env_dict(self, d): + items = [] + for k, v in d.items(): + items.append(f"export {k}={v}") + return "\n".join(items) + def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters: return self._config_task_instance.pre_execute(user_params) @@ -204,6 +210,9 @@ def execute(self, **kwargs) -> typing.Any: if os.name == "nt": self._script = self._script.lstrip().rstrip().replace("\n", "&&") + if "env" in kwargs: + kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"]) + gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) if self._debug: print("\n==============================================\n") @@ -239,74 +248,26 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty return self._config_task_instance.post_execute(user_params, rval) -class PortableShellTask(ShellTask): - - def __init__( - self, - name: str, - script_file: str, - debug: bool = False, - task_config: T = None, - inputs: typing.Optional[typing.Dict[str, typing.Type]] = None, - output_locs: typing.Optional[typing.List[OutputLocation]] = None, - **kwargs, - ): - self._script = """ - #!/bin/bash - - set -uexo pipefail - - {inputs.export_env} - - bash {inputs.script_file} {inputs.args} - """ +portable_shell_task = ShellTask( + name="portable_shell_task_instance", + debug=True, + inputs=kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), + output_locs=[ + OutputLocation( + var="k", + var_type=FlyteFile, + location="{ctx.working_directory}/test_output.txt", + ) + ], + script=""" + #!/bin/bash - def execute(self, **kwargs) -> typing.Any: - """ - Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem - """ - logger.info(f"Running shell script as type {self.task_type}") - if self.script_file: - with open(self.script_file) as f: - self._script = f.read() - - outputs: typing.Dict[str, str] = {} - if self._output_locs: - for v in self._output_locs: - outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=kwargs) - - if os.name == "nt": - self._script = self._script.lstrip().rstrip().replace("\n", "&&") - - gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) - if self._debug: - print("\n==============================================\n") - print(gen_script) - print("\n==============================================\n") - - try: - subprocess.check_call(gen_script, shell=True) - except subprocess.CalledProcessError as e: - files = os.listdir(".") - fstr = "\n-".join(files) - logger.error( - f"Failed to Execute Script, return-code {e.returncode} \n" - f"StdErr: {e.stderr}\n" - f"StdOut: {e.stdout}\n" - f" Current directory contents: .\n-{fstr}" - ) - raise + set -uexo pipefail - final_outputs = [] - for v in self._output_locs: - if issubclass(v.var_type, FlyteFile): - final_outputs.append(FlyteFile(outputs[v.var])) - if issubclass(v.var_type, FlyteDirectory): - final_outputs.append(FlyteDirectory(outputs[v.var])) - if len(final_outputs) == 1: - return final_outputs[0] - if len(final_outputs) > 1: - return tuple(final_outputs) - return None + cd {ctx.working_directory} + {inputs.export_env} + bash {inputs.script_file} {inputs.script_args} + """ +) \ No newline at end of file From afa9c5f482faec37ef8228071dc65c22079767ff Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Thu, 31 Mar 2022 13:08:26 -0400 Subject: [PATCH 03/27] fix spacing Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 014e7a2576..d424004315 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -212,7 +212,7 @@ def execute(self, **kwargs) -> typing.Any: if "env" in kwargs: kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"]) - + breakpoint() gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) if self._debug: print("\n==============================================\n") @@ -251,7 +251,7 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty portable_shell_task = ShellTask( name="portable_shell_task_instance", debug=True, - inputs=kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), + inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), output_locs=[ OutputLocation( var="k", @@ -260,14 +260,14 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty ) ], script=""" - #!/bin/bash +#!/bin/bash - set -uexo pipefail +set -uexo pipefail - cd {ctx.working_directory} +cd {ctx.working_directory} - {inputs.export_env} +{inputs.export_env} - bash {inputs.script_file} {inputs.script_args} +bash {inputs.script_file} {inputs.script_args} """ ) \ No newline at end of file From 369eca0ed05ba1323a17aaed6edcfe55eef3c338 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Thu, 31 Mar 2022 13:10:51 -0400 Subject: [PATCH 04/27] remove breakpoint Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index d424004315..79d3927374 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -212,7 +212,7 @@ def execute(self, **kwargs) -> typing.Any: if "env" in kwargs: kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"]) - breakpoint() + gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) if self._debug: print("\n==============================================\n") From a29367e6b1c746c23f6cd4bcde3a0c7c57ee2ea9 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Thu, 31 Mar 2022 13:23:10 -0400 Subject: [PATCH 05/27] Output ctx.working_directory as single output Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 79d3927374..a98d2b2b1a 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -255,8 +255,8 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty output_locs=[ OutputLocation( var="k", - var_type=FlyteFile, - location="{ctx.working_directory}/test_output.txt", + var_type=FlyteDirectory, + location="{ctx.working_directory}", ) ], script=""" From a64774b5f35305f0a5f9d6bab4d282a46843c631 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Thu, 31 Mar 2022 13:39:41 -0400 Subject: [PATCH 06/27] more doc strings Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index a98d2b2b1a..a18b44f882 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -184,7 +184,15 @@ def script(self) -> typing.Optional[str]: def script_file(self) -> typing.Optional[os.PathLike]: return self._script_file - def make_export_string_from_env_dict(self, d): + def make_export_string_from_env_dict(self, d) -> str: + """ + Utility function to convert a dictionary of desired environment variable key: value pairs into a string of + ` + export k1=v1 + export k2=v2 + ... + ` + """ items = [] for k, v in d.items(): items.append(f"export {k}={v}") @@ -210,7 +218,7 @@ def execute(self, **kwargs) -> typing.Any: if os.name == "nt": self._script = self._script.lstrip().rstrip().replace("\n", "&&") - if "env" in kwargs: + if "env" in kwargs and isinstance(kwargs["env"], dict): kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"]) gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) From 78292811334c241245d7d803fcde2d476467745e Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Thu, 31 Mar 2022 13:42:57 -0400 Subject: [PATCH 07/27] more comments Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index a18b44f882..452c9af10c 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -255,7 +255,9 @@ def execute(self, **kwargs) -> typing.Any: def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any: return self._config_task_instance.post_execute(user_params, rval) - +# The portable_shell_task is an instance of ShellTask which wraps a 'pure' shell script +# This utility function allows for the specification of env variables, arguments, and the actual script within the +# workflow definition rather than at `ShellTask` instantiation portable_shell_task = ShellTask( name="portable_shell_task_instance", debug=True, From 15453aa93202e54039d5c30336ea70a5536be6c4 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Fri, 1 Apr 2022 11:01:17 -0400 Subject: [PATCH 08/27] Added comments Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 452c9af10c..7edb5dd6c7 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -219,6 +219,8 @@ def execute(self, **kwargs) -> typing.Any: self._script = self._script.lstrip().rstrip().replace("\n", "&&") if "env" in kwargs and isinstance(kwargs["env"], dict): + # This supports the portable_shell_task by adding an additional key:value pair to kwargs/input + # This will cause collisions if a user tries to use `env` AND `export_env` in their inputs kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"]) gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) From 0abcfdc3565759e99e369c2c9a5de709a0e81ba5 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Fri, 1 Apr 2022 15:40:17 -0400 Subject: [PATCH 09/27] Add tests and test files from other branch Signed-off-by: Mike Zhong --- .../flytekit/unit/extras/tasks/test_shell.py | 36 ++++++++++++++++++- .../extras/tasks/testdata/script_args_env.sh | 22 ++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index ad1d75e272..8ded46048b 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -9,7 +9,7 @@ import flytekit from flytekit import kwtypes -from flytekit.extras.tasks.shell import OutputLocation, ShellTask +from flytekit.extras.tasks.shell import OutputLocation, ShellTask, portable_shell_task from flytekit.types.directory import FlyteDirectory from flytekit.types.file import CSVFile, FlyteFile @@ -20,6 +20,7 @@ script_sh = os.path.join(testdata, "script.exe") else: script_sh = os.path.join(testdata, "script.sh") + script_sh_2 = os.path.join(testdata, "script_args_env.sh") def test_shell_task_no_io(): @@ -249,3 +250,36 @@ def test_shell_script(): assert t.script_file == script_sh t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) + + +def test_portable_shell_task_with_args(capfd): + portable_shell_task( + script_file=script_sh_2, + script_args="first_arg second_arg", + env={} + ) + cap = capfd.readouterr() + assert "first_arg" in cap.out + assert "second_arg" in cap.out + + +def test_shell_task_with_env(capfd): + portable_shell_task( + script_file=script_sh_2, + env={"A": "AAAA", "B": "BBBB"}, + script_args="" + ) + cap = capfd.readouterr() + assert "AAAA" in cap.out + assert "BBBB" in cap.out + + +def test_shell_task_properly_restores_env_after_execution(): + env_as_dict = os.environ.copy() + portable_shell_task( + script_file=script_sh_2, + env={"A": "AAAA", "B": "BBBB"}, + script_args="" + ) + env_as_dict_after = os.environ.copy() + assert env_as_dict == env_as_dict_after \ No newline at end of file diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh new file mode 100644 index 0000000000..b63d4c0a29 --- /dev/null +++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +set -exo pipefail + +echo "A: $A" +echo "B: $B" + +if [ ! -z $1 ]; then + echo $1 +else + echo "Unset first positional argument" +fi + +if [ ! -z $2 ]; then + echo $2 +else + echo "Unset second positional argument" +fi + +SOME_VAR="This var is set" + +echo "Reading SOME_VAR: ${SOME_VAR}" \ No newline at end of file From e0efff65e093bc9e6c5ac35f44cadd7127adc5b3 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Wed, 6 Apr 2022 15:28:33 -0400 Subject: [PATCH 10/27] minor fix, have function return the instance rather than create. It seems to get registered when flyte packages your project Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 7edb5dd6c7..653fbdb25e 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -257,10 +257,12 @@ def execute(self, **kwargs) -> typing.Any: def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any: return self._config_task_instance.post_execute(user_params, rval) + # The portable_shell_task is an instance of ShellTask which wraps a 'pure' shell script # This utility function allows for the specification of env variables, arguments, and the actual script within the # workflow definition rather than at `ShellTask` instantiation -portable_shell_task = ShellTask( +def get_portable_shell_task() -> ShellTask: + return ShellTask( name="portable_shell_task_instance", debug=True, inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), From 4abd44707a34c15845513e59f256c71f8beb272e Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Mon, 11 Apr 2022 13:00:44 -0400 Subject: [PATCH 11/27] fix tests Signed-off-by: Mike Zhong --- tests/flytekit/unit/extras/tasks/test_shell.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 8ded46048b..1f443bf47b 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -9,7 +9,7 @@ import flytekit from flytekit import kwtypes -from flytekit.extras.tasks.shell import OutputLocation, ShellTask, portable_shell_task +from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_portable_shell_task from flytekit.types.directory import FlyteDirectory from flytekit.types.file import CSVFile, FlyteFile @@ -253,7 +253,8 @@ def test_shell_script(): def test_portable_shell_task_with_args(capfd): - portable_shell_task( + pst = get_portable_shell_task() + pst( script_file=script_sh_2, script_args="first_arg second_arg", env={} @@ -264,7 +265,8 @@ def test_portable_shell_task_with_args(capfd): def test_shell_task_with_env(capfd): - portable_shell_task( + pst = get_portable_shell_task() + pst( script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="" @@ -276,7 +278,8 @@ def test_shell_task_with_env(capfd): def test_shell_task_properly_restores_env_after_execution(): env_as_dict = os.environ.copy() - portable_shell_task( + pst = get_portable_shell_task() + pst( script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="" From 6e2964b5075b2e1684daf6f222a5f950a1d3ffae Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Tue, 12 Apr 2022 09:12:32 -0400 Subject: [PATCH 12/27] remove set flags not supported by sh Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 653fbdb25e..1701ff5ae2 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -276,7 +276,7 @@ def get_portable_shell_task() -> ShellTask: script=""" #!/bin/bash -set -uexo pipefail +set -uex cd {ctx.working_directory} From d16b54ed20d7701051ea8da0b16bb196b76d6d88 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Tue, 12 Apr 2022 18:57:01 -0400 Subject: [PATCH 13/27] fix spellcheck lint errors Signed-off-by: Mike Zhong --- .../unit/extras/tasks/testdata/script_args_env.sh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh index b63d4c0a29..dd77131e22 100644 --- a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh +++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh @@ -5,14 +5,14 @@ set -exo pipefail echo "A: $A" echo "B: $B" -if [ ! -z $1 ]; then - echo $1 +if [ ! -n "$1" ]; then + echo "$1" else echo "Unset first positional argument" fi -if [ ! -z $2 ]; then - echo $2 +if [ ! -n "$2" ]; then + echo "$2" else echo "Unset second positional argument" fi From 9a989de47d051a27d4a9aa358af5b7528a6b02fd Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Fri, 15 Apr 2022 12:42:49 -0400 Subject: [PATCH 14/27] don't have a windows equivalent test script so bypassing those tests for now Signed-off-by: Mike Zhong --- tests/flytekit/unit/extras/tasks/test_shell.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 1f443bf47b..2eafe90468 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -18,6 +18,7 @@ test_csv = os.path.join(testdata, "test.csv") if os.name == "nt": script_sh = os.path.join(testdata, "script.exe") + script_sh_2 = None else: script_sh = os.path.join(testdata, "script.sh") script_sh_2 = os.path.join(testdata, "script_args_env.sh") @@ -253,6 +254,8 @@ def test_shell_script(): def test_portable_shell_task_with_args(capfd): + if script_sh_2 is None: + return pst = get_portable_shell_task() pst( script_file=script_sh_2, @@ -265,6 +268,8 @@ def test_portable_shell_task_with_args(capfd): def test_shell_task_with_env(capfd): + if script_sh_2 is None: + return pst = get_portable_shell_task() pst( script_file=script_sh_2, @@ -277,6 +282,8 @@ def test_shell_task_with_env(capfd): def test_shell_task_properly_restores_env_after_execution(): + if script_sh_2 is None: + return env_as_dict = os.environ.copy() pst = get_portable_shell_task() pst( From e217806820b9ad64307b197b3babb1c214ccbd10 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Mon, 18 Apr 2022 09:20:39 -0400 Subject: [PATCH 15/27] Add typing to make_export_string_from_env_dict params Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 1701ff5ae2..d22dc06691 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -184,7 +184,7 @@ def script(self) -> typing.Optional[str]: def script_file(self) -> typing.Optional[os.PathLike]: return self._script_file - def make_export_string_from_env_dict(self, d) -> str: + def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str: """ Utility function to convert a dictionary of desired environment variable key: value pairs into a string of ` From 2bce64e95fe18e7951e497787efde9b83faf1129 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Mon, 18 Apr 2022 09:52:49 -0400 Subject: [PATCH 16/27] fixup doc string Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index d22dc06691..ec49ffeb32 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -187,11 +187,11 @@ def script_file(self) -> typing.Optional[os.PathLike]: def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str: """ Utility function to convert a dictionary of desired environment variable key: value pairs into a string of - ` + ``` export k1=v1 export k2=v2 ... - ` + ``` """ items = [] for k, v in d.items(): From 3d26075dfaf023711e36a3960d0778d583d95eed Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Wed, 27 Apr 2022 15:43:35 -0400 Subject: [PATCH 17/27] Refactored the new behavior out into a new class. Did not change implementation details at all Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 144 +++++++++++++++--- .../flytekit/unit/extras/tasks/test_shell.py | 8 +- 2 files changed, 123 insertions(+), 29 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index ec49ffeb32..71b33a5bc8 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -184,6 +184,106 @@ def script(self) -> typing.Optional[str]: def script_file(self) -> typing.Optional[os.PathLike]: return self._script_file + def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters: + return self._config_task_instance.pre_execute(user_params) + + def execute(self, **kwargs) -> typing.Any: + """ + Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem + """ + logger.info(f"Running shell script as type {self.task_type}") + if self.script_file: + with open(self.script_file) as f: + self._script = f.read() + + outputs: typing.Dict[str, str] = {} + if self._output_locs: + for v in self._output_locs: + outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=kwargs) + + if os.name == "nt": + self._script = self._script.lstrip().rstrip().replace("\n", "&&") + + gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) + if self._debug: + print("\n==============================================\n") + print(gen_script) + print("\n==============================================\n") + + try: + subprocess.check_call(gen_script, shell=True) + except subprocess.CalledProcessError as e: + files = os.listdir(".") + fstr = "\n-".join(files) + logger.error( + f"Failed to Execute Script, return-code {e.returncode} \n" + f"StdErr: {e.stderr}\n" + f"StdOut: {e.stdout}\n" + f" Current directory contents: .\n-{fstr}" + ) + raise + + final_outputs = [] + for v in self._output_locs: + if issubclass(v.var_type, FlyteFile): + final_outputs.append(FlyteFile(outputs[v.var])) + if issubclass(v.var_type, FlyteDirectory): + final_outputs.append(FlyteDirectory(outputs[v.var])) + if len(final_outputs) == 1: + return final_outputs[0] + if len(final_outputs) > 1: + return tuple(final_outputs) + return None + + def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any: + return self._config_task_instance.post_execute(user_params, rval) + + +class _RawShellTask(ShellTask): + """ """ + + def __init__( + self, + name: str, + debug: bool = False, + script: typing.Optional[str] = None, + script_file: typing.Optional[str] = None, + task_config: T = None, + inputs: typing.Optional[typing.Dict[str, typing.Type]] = None, + output_locs: typing.Optional[typing.List[OutputLocation]] = None, + **kwargs, + ): + """ + The `_RawShellTask` is a minimal extension of the existing `ShellTask`. It's purpose is to support wrapping a + "raw" or "pure" shell script which needs to be executed with some environment variables set, and some arguments, + which may not be known until execution time. + + This class is not meant to be instantiated into tasks by users, but used with the factory function + `get_raw_shell_task()`. An instance of this class will be returned with either user-specified or default + template. The template itself will export the desired environment variables, and subsequently execute the + desired "raw" script with the specified arguments. + + .. note:: + This means that within your workflow, you can dynamically control the env variables, arguments, and even the + actual script you want to run. + + .. note:: + The downside is that a dynamic workflow will be required. The "raw" script passed in at execution time must + be at the specified location. + + These args are forwarded directly to the parent `ShellTask` constructor as behavior does not diverge + """ + super().__init__( + name=name, + debug=debug, + script=script, + script_file=script_file, + task_config=task_config, + inputs=inputs, + output_locs=output_locs, + **kwargs, + ) + def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str: """ Utility function to convert a dictionary of desired environment variable key: value pairs into a string of @@ -198,9 +298,6 @@ def make_export_string_from_env_dict(self, d: typing.Dict[str, str]) -> str: items.append(f"export {k}={v}") return "\n".join(items) - def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters: - return self._config_task_instance.pre_execute(user_params) - def execute(self, **kwargs) -> typing.Any: """ Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem @@ -219,8 +316,6 @@ def execute(self, **kwargs) -> typing.Any: self._script = self._script.lstrip().rstrip().replace("\n", "&&") if "env" in kwargs and isinstance(kwargs["env"], dict): - # This supports the portable_shell_task by adding an additional key:value pair to kwargs/input - # This will cause collisions if a user tries to use `env` AND `export_env` in their inputs kwargs["export_env"] = self.make_export_string_from_env_dict(kwargs["env"]) gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) @@ -254,26 +349,25 @@ def execute(self, **kwargs) -> typing.Any: return tuple(final_outputs) return None - def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any: - return self._config_task_instance.post_execute(user_params, rval) - -# The portable_shell_task is an instance of ShellTask which wraps a 'pure' shell script +# The raw_shell_task is an instance of _RawShellTask and wraps a 'pure' shell script # This utility function allows for the specification of env variables, arguments, and the actual script within the -# workflow definition rather than at `ShellTask` instantiation -def get_portable_shell_task() -> ShellTask: - return ShellTask( - name="portable_shell_task_instance", - debug=True, - inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), - output_locs=[ - OutputLocation( - var="k", - var_type=FlyteDirectory, - location="{ctx.working_directory}", - ) - ], - script=""" +# workflow definition rather than at `_RawShellTask` instantiation +def get_raw_shell_task(name: str = None) -> _RawShellTask: + _name = name if name else "raw_shell_task_instance" + + return _RawShellTask( + name=_name, + debug=True, + inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), + output_locs=[ + OutputLocation( + var="out", + var_type=FlyteDirectory, + location="{ctx.working_directory}", + ) + ], + script=""" #!/bin/bash set -uex @@ -283,5 +377,5 @@ def get_portable_shell_task() -> ShellTask: {inputs.export_env} bash {inputs.script_file} {inputs.script_args} - """ -) \ No newline at end of file +""" + ) \ No newline at end of file diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 2eafe90468..197296eaea 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -9,7 +9,7 @@ import flytekit from flytekit import kwtypes -from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_portable_shell_task +from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task from flytekit.types.directory import FlyteDirectory from flytekit.types.file import CSVFile, FlyteFile @@ -256,7 +256,7 @@ def test_shell_script(): def test_portable_shell_task_with_args(capfd): if script_sh_2 is None: return - pst = get_portable_shell_task() + pst = get_raw_shell_task() pst( script_file=script_sh_2, script_args="first_arg second_arg", @@ -270,7 +270,7 @@ def test_portable_shell_task_with_args(capfd): def test_shell_task_with_env(capfd): if script_sh_2 is None: return - pst = get_portable_shell_task() + pst = get_raw_shell_task() pst( script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, @@ -285,7 +285,7 @@ def test_shell_task_properly_restores_env_after_execution(): if script_sh_2 is None: return env_as_dict = os.environ.copy() - pst = get_portable_shell_task() + pst = get_raw_shell_task() pst( script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, From ce4603fe59cf4a7ff2b2d3e1d4f86470357898b3 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Mon, 2 May 2022 09:47:06 -0400 Subject: [PATCH 18/27] Address linter issues and up test cov Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 2 +- .../flytekit/unit/extras/tasks/test_shell.py | 45 ++++++++++++++++--- .../extras/tasks/testdata/script_args_env.sh | 4 +- 3 files changed, 43 insertions(+), 8 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 71b33a5bc8..1e249d0cb1 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -378,4 +378,4 @@ def get_raw_shell_task(name: str = None) -> _RawShellTask: bash {inputs.script_file} {inputs.script_args} """ - ) \ No newline at end of file + ) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 197296eaea..cef54e22c9 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -1,6 +1,7 @@ import datetime import os import tempfile +import typing from dataclasses import dataclass from subprocess import CalledProcessError @@ -9,7 +10,7 @@ import flytekit from flytekit import kwtypes -from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task +from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task, _RawShellTask from flytekit.types.directory import FlyteDirectory from flytekit.types.file import CSVFile, FlyteFile @@ -253,7 +254,7 @@ def test_shell_script(): t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) -def test_portable_shell_task_with_args(capfd): +def test_raw_shell_task_with_args(capfd): if script_sh_2 is None: return pst = get_raw_shell_task() @@ -267,7 +268,7 @@ def test_portable_shell_task_with_args(capfd): assert "second_arg" in cap.out -def test_shell_task_with_env(capfd): +def test_raw_shell_task_with_env(capfd): if script_sh_2 is None: return pst = get_raw_shell_task() @@ -281,7 +282,7 @@ def test_shell_task_with_env(capfd): assert "BBBB" in cap.out -def test_shell_task_properly_restores_env_after_execution(): +def test_raw_shell_task_properly_restores_env_after_execution(): if script_sh_2 is None: return env_as_dict = os.environ.copy() @@ -292,4 +293,38 @@ def test_shell_task_properly_restores_env_after_execution(): script_args="" ) env_as_dict_after = os.environ.copy() - assert env_as_dict == env_as_dict_after \ No newline at end of file + assert env_as_dict == env_as_dict_after + + +def test_raw_shell_task_instantiation(capfd): + pst = _RawShellTask( + name="test", + debug=True, + inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), + output_locs=[ + OutputLocation( + var="out", + var_type=FlyteDirectory, + location="{ctx.working_directory}", + ) + ], + script=""" +#!/bin/bash + +set -uex + +cd {ctx.working_directory} + +{inputs.export_env} + +bash {inputs.script_file} {inputs.script_args} +""" + ) + pst( + script_file=script_sh_2, + script_args="first_arg second_arg", + env={} + ) + cap = capfd.readouterr() + assert "first_arg" in cap.out + assert "second_arg" in cap.out diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh index dd77131e22..ceb4bdf975 100644 --- a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh +++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh @@ -5,13 +5,13 @@ set -exo pipefail echo "A: $A" echo "B: $B" -if [ ! -n "$1" ]; then +if [ ! -z "$1" ]; then echo "$1" else echo "Unset first positional argument" fi -if [ ! -n "$2" ]; then +if [ ! -z "$2" ]; then echo "$2" else echo "Unset second positional argument" From ead64955b19f23101621d5f4a0a43ccbc703472c Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Mon, 2 May 2022 11:02:42 -0400 Subject: [PATCH 19/27] Skip test on windows, no equivalent script Signed-off-by: Mike Zhong --- tests/flytekit/unit/extras/tasks/test_shell.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index cef54e22c9..d6ab80eb94 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -297,6 +297,8 @@ def test_raw_shell_task_properly_restores_env_after_execution(): def test_raw_shell_task_instantiation(capfd): + if script_sh_2 is None: + return pst = _RawShellTask( name="test", debug=True, From 2ebeecd59d381d6270431c928c3d5cd9b140c569 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Tue, 3 May 2022 10:41:04 -0400 Subject: [PATCH 20/27] Run black and isort, address SC2236 Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 2 +- .../flytekit/unit/extras/tasks/test_shell.py | 28 ++++--------------- .../extras/tasks/testdata/script_args_env.sh | 6 ++-- 3 files changed, 10 insertions(+), 26 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 1e249d0cb1..e3de1756fb 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -377,5 +377,5 @@ def get_raw_shell_task(name: str = None) -> _RawShellTask: {inputs.export_env} bash {inputs.script_file} {inputs.script_args} -""" +""", ) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index d6ab80eb94..c369787761 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -10,7 +10,7 @@ import flytekit from flytekit import kwtypes -from flytekit.extras.tasks.shell import OutputLocation, ShellTask, get_raw_shell_task, _RawShellTask +from flytekit.extras.tasks.shell import OutputLocation, ShellTask, _RawShellTask, get_raw_shell_task from flytekit.types.directory import FlyteDirectory from flytekit.types.file import CSVFile, FlyteFile @@ -258,11 +258,7 @@ def test_raw_shell_task_with_args(capfd): if script_sh_2 is None: return pst = get_raw_shell_task() - pst( - script_file=script_sh_2, - script_args="first_arg second_arg", - env={} - ) + pst(script_file=script_sh_2, script_args="first_arg second_arg", env={}) cap = capfd.readouterr() assert "first_arg" in cap.out assert "second_arg" in cap.out @@ -272,11 +268,7 @@ def test_raw_shell_task_with_env(capfd): if script_sh_2 is None: return pst = get_raw_shell_task() - pst( - script_file=script_sh_2, - env={"A": "AAAA", "B": "BBBB"}, - script_args="" - ) + pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="") cap = capfd.readouterr() assert "AAAA" in cap.out assert "BBBB" in cap.out @@ -287,11 +279,7 @@ def test_raw_shell_task_properly_restores_env_after_execution(): return env_as_dict = os.environ.copy() pst = get_raw_shell_task() - pst( - script_file=script_sh_2, - env={"A": "AAAA", "B": "BBBB"}, - script_args="" - ) + pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="") env_as_dict_after = os.environ.copy() assert env_as_dict == env_as_dict_after @@ -320,13 +308,9 @@ def test_raw_shell_task_instantiation(capfd): {inputs.export_env} bash {inputs.script_file} {inputs.script_args} -""" - ) - pst( - script_file=script_sh_2, - script_args="first_arg second_arg", - env={} +""", ) + pst(script_file=script_sh_2, script_args="first_arg second_arg", env={}) cap = capfd.readouterr() assert "first_arg" in cap.out assert "second_arg" in cap.out diff --git a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh index ceb4bdf975..0927512a05 100644 --- a/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh +++ b/tests/flytekit/unit/extras/tasks/testdata/script_args_env.sh @@ -5,13 +5,13 @@ set -exo pipefail echo "A: $A" echo "B: $B" -if [ ! -z "$1" ]; then +if [ -n "$1" ]; then echo "$1" else echo "Unset first positional argument" fi -if [ ! -z "$2" ]; then +if [ -n "$2" ]; then echo "$2" else echo "Unset second positional argument" @@ -19,4 +19,4 @@ fi SOME_VAR="This var is set" -echo "Reading SOME_VAR: ${SOME_VAR}" \ No newline at end of file +echo "Reading SOME_VAR: ${SOME_VAR}" From e4c121dc69437fa091bd4ff541792641519831f5 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Tue, 3 May 2022 10:45:00 -0400 Subject: [PATCH 21/27] Refactored class name to remove _, make utility function require name so multiple uses in a workflow don't break Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 12 ++++++------ tests/flytekit/unit/extras/tasks/test_shell.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index e3de1756fb..8cd2e39422 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -239,7 +239,7 @@ def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> ty return self._config_task_instance.post_execute(user_params, rval) -class _RawShellTask(ShellTask): +class RawShellTask(ShellTask): """ """ def __init__( @@ -254,7 +254,7 @@ def __init__( **kwargs, ): """ - The `_RawShellTask` is a minimal extension of the existing `ShellTask`. It's purpose is to support wrapping a + The `RawShellTask` is a minimal extension of the existing `ShellTask`. It's purpose is to support wrapping a "raw" or "pure" shell script which needs to be executed with some environment variables set, and some arguments, which may not be known until execution time. @@ -350,13 +350,13 @@ def execute(self, **kwargs) -> typing.Any: return None -# The raw_shell_task is an instance of _RawShellTask and wraps a 'pure' shell script +# The raw_shell_task is an instance of RawShellTask and wraps a 'pure' shell script # This utility function allows for the specification of env variables, arguments, and the actual script within the -# workflow definition rather than at `_RawShellTask` instantiation -def get_raw_shell_task(name: str = None) -> _RawShellTask: +# workflow definition rather than at `RawShellTask` instantiation +def get_raw_shell_task(name) -> RawShellTask: _name = name if name else "raw_shell_task_instance" - return _RawShellTask( + return RawShellTask( name=_name, debug=True, inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index c369787761..45f4801a80 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -10,7 +10,7 @@ import flytekit from flytekit import kwtypes -from flytekit.extras.tasks.shell import OutputLocation, ShellTask, _RawShellTask, get_raw_shell_task +from flytekit.extras.tasks.shell import OutputLocation, ShellTask, RawShellTask, get_raw_shell_task from flytekit.types.directory import FlyteDirectory from flytekit.types.file import CSVFile, FlyteFile @@ -287,7 +287,7 @@ def test_raw_shell_task_properly_restores_env_after_execution(): def test_raw_shell_task_instantiation(capfd): if script_sh_2 is None: return - pst = _RawShellTask( + pst = RawShellTask( name="test", debug=True, inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), From 215f48c51f6f39e900dd0a08d583de4e2b62fe58 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Tue, 3 May 2022 11:15:54 -0400 Subject: [PATCH 22/27] fix utility function Signed-off-by: Mike Zhong --- flytekit/extras/tasks/shell.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 8cd2e39422..812c0a3749 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -353,11 +353,10 @@ def execute(self, **kwargs) -> typing.Any: # The raw_shell_task is an instance of RawShellTask and wraps a 'pure' shell script # This utility function allows for the specification of env variables, arguments, and the actual script within the # workflow definition rather than at `RawShellTask` instantiation -def get_raw_shell_task(name) -> RawShellTask: - _name = name if name else "raw_shell_task_instance" +def get_raw_shell_task(name: str) -> RawShellTask: return RawShellTask( - name=_name, + name=name, debug=True, inputs=flytekit.kwtypes(env=typing.Dict[str, str], script_args=str, script_file=str), output_locs=[ From 5b756919dbad6f974d4cb0c4f16bcc978d90572a Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Tue, 3 May 2022 11:16:26 -0400 Subject: [PATCH 23/27] fix utility function call in tests Signed-off-by: Mike Zhong --- tests/flytekit/unit/extras/tasks/test_shell.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 45f4801a80..11e9fa3d6f 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -257,7 +257,7 @@ def test_shell_script(): def test_raw_shell_task_with_args(capfd): if script_sh_2 is None: return - pst = get_raw_shell_task() + pst = get_raw_shell_task(name="test") pst(script_file=script_sh_2, script_args="first_arg second_arg", env={}) cap = capfd.readouterr() assert "first_arg" in cap.out @@ -267,7 +267,7 @@ def test_raw_shell_task_with_args(capfd): def test_raw_shell_task_with_env(capfd): if script_sh_2 is None: return - pst = get_raw_shell_task() + pst = get_raw_shell_task(name="test") pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="") cap = capfd.readouterr() assert "AAAA" in cap.out @@ -278,7 +278,7 @@ def test_raw_shell_task_properly_restores_env_after_execution(): if script_sh_2 is None: return env_as_dict = os.environ.copy() - pst = get_raw_shell_task() + pst = get_raw_shell_task(name="test") pst(script_file=script_sh_2, env={"A": "AAAA", "B": "BBBB"}, script_args="") env_as_dict_after = os.environ.copy() assert env_as_dict == env_as_dict_after From 53e88f60dbd36d6b257c3ad997655eedde9b2838 Mon Sep 17 00:00:00 2001 From: Mike Zhong Date: Wed, 4 May 2022 09:40:34 -0400 Subject: [PATCH 24/27] Fix isort on test_shell.py Signed-off-by: Mike Zhong --- tests/flytekit/unit/extras/tasks/test_shell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 11e9fa3d6f..7081d8d65f 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -10,7 +10,7 @@ import flytekit from flytekit import kwtypes -from flytekit.extras.tasks.shell import OutputLocation, ShellTask, RawShellTask, get_raw_shell_task +from flytekit.extras.tasks.shell import OutputLocation, RawShellTask, ShellTask, get_raw_shell_task from flytekit.types.directory import FlyteDirectory from flytekit.types.file import CSVFile, FlyteFile From 6d5e67af7be4352d1d50bc9baa19d4dd231f5afd Mon Sep 17 00:00:00 2001 From: Calvin Leather Date: Tue, 9 Aug 2022 21:29:24 -0400 Subject: [PATCH 25/27] adding logging settings to papermill plugin Signed-off-by: Calvin Leather --- .../flytekitplugins/papermill/task.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py index 304932a828..fcceb7d70b 100644 --- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py +++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py @@ -1,5 +1,7 @@ import json +import logging import os +import sys import typing from typing import Any @@ -112,6 +114,7 @@ def __init__( name: str, notebook_path: str, render_deck: bool = False, + stream_logs: bool = False, task_config: T = None, inputs: typing.Optional[typing.Dict[str, typing.Type]] = None, outputs: typing.Optional[typing.Dict[str, typing.Type]] = None, @@ -132,6 +135,13 @@ def __init__( self._notebook_path = os.path.abspath(notebook_path) self._render_deck = render_deck + self._stream_logs = stream_logs + + # Send the papermill logger to stdout so that it appears in pod logs. Note that papermill doesn't allow + # injecting a logger, so we cannot redirect logs to the flyte child loggers (e.g., the userspace logger) + # and we instead send to stdout directly + if self._stream_logs: + logging.getLogger("papermill").addHandler(logging.StreamHandler(sys.stdout)) if not os.path.exists(self._notebook_path): raise ValueError(f"Illegal notebook path passed in {self._notebook_path}") @@ -207,7 +217,7 @@ def execute(self, **kwargs) -> Any: """ logger.info(f"Hijacking the call for task-type {self.task_type}, to call notebook.") # Execute Notebook via Papermill. - pm.execute_notebook(self._notebook_path, self.output_notebook_path, parameters=kwargs) # type: ignore + pm.execute_notebook(self._notebook_path, self.output_notebook_path, parameters=kwargs, log_output=self._stream_logs) # type: ignore outputs = self.extract_outputs(self.output_notebook_path) self.render_nb_html(self.output_notebook_path, self.rendered_output_path) From 8b85e13c18e7860d30027abfb54db2f2e21f2778 Mon Sep 17 00:00:00 2001 From: Calvin Leather Date: Wed, 10 Aug 2022 06:29:27 -0400 Subject: [PATCH 26/27] set level to info logging --- plugins/flytekit-papermill/flytekitplugins/papermill/task.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py index fcceb7d70b..00ec942fcd 100644 --- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py +++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py @@ -141,7 +141,9 @@ def __init__( # injecting a logger, so we cannot redirect logs to the flyte child loggers (e.g., the userspace logger) # and we instead send to stdout directly if self._stream_logs: - logging.getLogger("papermill").addHandler(logging.StreamHandler(sys.stdout)) + papermill_logger = logging.getLogger("papermill") + papermill_logger.addHandler(logging.StreamHandler(sys.stdout)) + papermill_logger.setLevel(logging.INFO) if not os.path.exists(self._notebook_path): raise ValueError(f"Illegal notebook path passed in {self._notebook_path}") From c854024ff9f49e2da39732111550606531b2a7c3 Mon Sep 17 00:00:00 2001 From: Calvin Leather Date: Wed, 10 Aug 2022 06:40:00 -0400 Subject: [PATCH 27/27] improved comments/docs --- .../flytekitplugins/papermill/task.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py index 00ec942fcd..9e9535dec2 100644 --- a/plugins/flytekit-papermill/flytekitplugins/papermill/task.py +++ b/plugins/flytekit-papermill/flytekitplugins/papermill/task.py @@ -86,6 +86,13 @@ class NotebookTask(PythonInstanceTask[T]): Users can access these notebooks after execution of the task locally or from remote servers. + .. note: + + By default, print statements in your notebook won't be transmitted to the pod logs/stdout. If you would + like to have logs forwarded as the notebook executes, pass the stream_logs argument. Note that notebook + logs can be quite verbose, so ensure you are prepared for any downstream log ingestion costs + (e.g., cloudwatch) + .. todo: Implicit extraction of SparkConfiguration from the notebook is not supported. @@ -139,10 +146,11 @@ def __init__( # Send the papermill logger to stdout so that it appears in pod logs. Note that papermill doesn't allow # injecting a logger, so we cannot redirect logs to the flyte child loggers (e.g., the userspace logger) - # and we instead send to stdout directly + # and inherit their settings, but we instead must send logs to stdout directly if self._stream_logs: papermill_logger = logging.getLogger("papermill") papermill_logger.addHandler(logging.StreamHandler(sys.stdout)) + # Papermill leaves the default level of DEBUG. We increase it here. papermill_logger.setLevel(logging.INFO) if not os.path.exists(self._notebook_path):