From f378c38b6f4a9537006d10f84733d8071f8ce508 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Thu, 5 May 2022 16:03:29 -0700 Subject: [PATCH 1/9] testing change Signed-off-by: Yee Hing Tong --- flytekit/clis/sdk_in_container/pyflyte.py | 5 ++++- flytekit/core/tracker.py | 6 +++--- flytekit/tools/serialize_helpers.py | 5 +++++ flytekit/tools/translator.py | 2 ++ 4 files changed, 14 insertions(+), 4 deletions(-) diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index f27128d108..9ce86e9f2a 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -42,7 +42,10 @@ def main(ctx, pkgs=None, config=None): Entrypoint for all the user commands. """ ctx.obj = dict() - + # import importlib + # importlib.import_module("core.flyte_basics.hello_world") + # + # raise Exception("jfkdlsa") # Handle package management - get from config if not specified on the command line pkgs = pkgs or [] if config: diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py index 9db8233a05..0fad8335c2 100644 --- a/flytekit/core/tracker.py +++ b/flytekit/core/tracker.py @@ -199,7 +199,7 @@ def _resolve_abs_module_name(self, path: str, package_root: str) -> str: if "__init__.py" not in os.listdir(dirname): return basename - # Now recurse down such that we can extract the absolute module path + # Now recurse down such that we can extract the absolute module path mod_name = self._resolve_abs_module_name(dirname, package_root) final_mod_name = f"{mod_name}.{basename}" if mod_name else basename self._module_cache[path] = final_mod_name @@ -243,8 +243,8 @@ def extract_task_module(f: Union[Callable, TrackedInstance]) -> Tuple[str, str, package_root = ( FeatureFlags.FLYTE_PYTHON_PACKAGE_ROOT if FeatureFlags.FLYTE_PYTHON_PACKAGE_ROOT != "auto" else None ) - new_mod_name = _mod_sanitizer.get_absolute_module_name(inspect.getabsfile(f), package_root) + new_mod_name = _mod_sanitizer.get_absolute_module_name(inspect.getabsfile(mod), package_root) # We only replace the mod_name if it is more specific, else we already have a fully resolved path if len(new_mod_name) > len(mod_name): mod_name = new_mod_name - return f"{mod_name}.{name}", mod_name, name, os.path.abspath(inspect.getfile(f)) + return f"{mod_name}.{name}", mod_name, name, os.path.abspath(inspect.getfile(mod)) diff --git a/flytekit/tools/serialize_helpers.py b/flytekit/tools/serialize_helpers.py index 7c1969afa7..8c99a4e757 100644 --- a/flytekit/tools/serialize_helpers.py +++ b/flytekit/tools/serialize_helpers.py @@ -96,6 +96,11 @@ def get_registrable_entities( f"Multiple definitions of the following tasks were found: {duplicate_task_names}" ) + # from flytekit.models.task import TaskSpec + # for xx in entities_to_be_serialized: + # if isinstance(xx, TaskSpec) and "use_setup" in xx.template.id.name: + # import ipdb; ipdb.set_trace() + return [v.to_flyte_idl() for v in entities_to_be_serialized] diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 8f3deba22b..c2006b5f0e 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -572,6 +572,8 @@ def get_serializable( elif isinstance(entity, PythonTask): cp_entity = get_serializable_task(entity_mapping, settings, entity) + if "use_setup" in entity.name: + import ipdb; ipdb.set_trace() elif isinstance(entity, WorkflowBase): cp_entity = get_serializable_workflow(entity_mapping, settings, entity, options) From 821a0ecad1b086ef5ed0ee11d9d1748b8a01ca23 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Thu, 5 May 2022 16:04:07 -0700 Subject: [PATCH 2/9] remove debug Signed-off-by: Yee Hing Tong --- flytekit/tools/translator.py | 2 - .../flytekit/unit/core/functools/__init__.py | 0 .../unit/core/functools/decorator_source.py | 49 +++++++++++++++++++ .../core/functools/test_decorator_location.py | 23 +++++++++ 4 files changed, 72 insertions(+), 2 deletions(-) create mode 100644 tests/flytekit/unit/core/functools/__init__.py create mode 100644 tests/flytekit/unit/core/functools/decorator_source.py create mode 100644 tests/flytekit/unit/core/functools/test_decorator_location.py diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index c2006b5f0e..8f3deba22b 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -572,8 +572,6 @@ def get_serializable( elif isinstance(entity, PythonTask): cp_entity = get_serializable_task(entity_mapping, settings, entity) - if "use_setup" in entity.name: - import ipdb; ipdb.set_trace() elif isinstance(entity, WorkflowBase): cp_entity = get_serializable_workflow(entity_mapping, settings, entity, options) diff --git a/tests/flytekit/unit/core/functools/__init__.py b/tests/flytekit/unit/core/functools/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit/unit/core/functools/decorator_source.py b/tests/flytekit/unit/core/functools/decorator_source.py new file mode 100644 index 0000000000..f66c30374b --- /dev/null +++ b/tests/flytekit/unit/core/functools/decorator_source.py @@ -0,0 +1,49 @@ +"""Script used for testing local execution of functool.wraps-wrapped tasks for stacked decorators""" + +import os +from functools import wraps +from typing import List + + +def task_decorator_1(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_1") + return fn(*args, **kwargs) + + return wrapper + + +def task_decorator_2(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + print("running task_decorator_2") + return fn(*args, **kwargs) + + return wrapper + + +def task_setup( + function: callable = None, *, + integration_requests: List = None +) -> None: + integration_requests = integration_requests or [] + + @wraps(function) + def wrapper(*args, **kwargs): + # Preprocessing of task + print("preprocessing") + + # Execute function + output = function(*args, **kwargs) + + # Postprocessing of output + print("postprocessing") + + return output + + return ( + functools.partial( + task_setup, integration_requests=integration_requests) + if function is None else wrapper + ) diff --git a/tests/flytekit/unit/core/functools/test_decorator_location.py b/tests/flytekit/unit/core/functools/test_decorator_location.py new file mode 100644 index 0000000000..c4c0ca63e9 --- /dev/null +++ b/tests/flytekit/unit/core/functools/test_decorator_location.py @@ -0,0 +1,23 @@ +from flytekit import task +from .decorator_source import task_decorator_1, task_decorator_2, task_setup +from flytekit.core.tracker import extract_task_module + +@task +@task_decorator_1 +@task_decorator_2 +def my_task(x: int) -> int: + print("running my_task") + return x + 1 + + +@task +@task_setup +def get_data(x: int) -> int: + print("running get_data") + return x + 1 + + +def test_fjdskla(): + res = extract_task_module(get_data) + print(res) + print(get_data.name) From 1b750843b5c3098505cf7c49aada0acb2b65ea52 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 6 May 2022 13:22:59 -0700 Subject: [PATCH 3/9] unit test Signed-off-by: Yee Hing Tong --- .../unit/core/functools/decorator_source.py | 30 ++----------------- .../unit/core/functools/decorator_usage.py | 10 +++++++ .../core/functools/test_decorator_location.py | 26 +++++----------- 3 files changed, 19 insertions(+), 47 deletions(-) create mode 100644 tests/flytekit/unit/core/functools/decorator_usage.py diff --git a/tests/flytekit/unit/core/functools/decorator_source.py b/tests/flytekit/unit/core/functools/decorator_source.py index f66c30374b..9c92364649 100644 --- a/tests/flytekit/unit/core/functools/decorator_source.py +++ b/tests/flytekit/unit/core/functools/decorator_source.py @@ -1,32 +1,10 @@ """Script used for testing local execution of functool.wraps-wrapped tasks for stacked decorators""" -import os from functools import wraps from typing import List -def task_decorator_1(fn): - @wraps(fn) - def wrapper(*args, **kwargs): - print("running task_decorator_1") - return fn(*args, **kwargs) - - return wrapper - - -def task_decorator_2(fn): - @wraps(fn) - def wrapper(*args, **kwargs): - print("running task_decorator_2") - return fn(*args, **kwargs) - - return wrapper - - -def task_setup( - function: callable = None, *, - integration_requests: List = None -) -> None: +def task_setup(function: callable = None, *, integration_requests: List = None) -> None: integration_requests = integration_requests or [] @wraps(function) @@ -42,8 +20,4 @@ def wrapper(*args, **kwargs): return output - return ( - functools.partial( - task_setup, integration_requests=integration_requests) - if function is None else wrapper - ) + return functools.partial(task_setup, integration_requests=integration_requests) if function is None else wrapper diff --git a/tests/flytekit/unit/core/functools/decorator_usage.py b/tests/flytekit/unit/core/functools/decorator_usage.py new file mode 100644 index 0000000000..8b6e8a9ecd --- /dev/null +++ b/tests/flytekit/unit/core/functools/decorator_usage.py @@ -0,0 +1,10 @@ +from flytekit import task +from .decorator_source import task_setup +from flytekit.core.tracker import extract_task_module + + +@task +@task_setup +def get_data(x: int) -> int: + print("running get_data") + return x + 1 diff --git a/tests/flytekit/unit/core/functools/test_decorator_location.py b/tests/flytekit/unit/core/functools/test_decorator_location.py index c4c0ca63e9..34b41ad1f8 100644 --- a/tests/flytekit/unit/core/functools/test_decorator_location.py +++ b/tests/flytekit/unit/core/functools/test_decorator_location.py @@ -1,23 +1,11 @@ +import importlib + from flytekit import task -from .decorator_source import task_decorator_1, task_decorator_2, task_setup from flytekit.core.tracker import extract_task_module -@task -@task_decorator_1 -@task_decorator_2 -def my_task(x: int) -> int: - print("running my_task") - return x + 1 - - -@task -@task_setup -def get_data(x: int) -> int: - print("running get_data") - return x + 1 - -def test_fjdskla(): - res = extract_task_module(get_data) - print(res) - print(get_data.name) +def test_dont_use_wrapper_location(): + m = importlib.import_module("tests.flytekit.unit.core.functools.decorator_usage") + get_data_task = getattr(m, "get_data") + assert "decorator_source" not in get_data_task.name + assert "decorator_usage" in get_data_task.name From a9f090854cd1b2556bc16118711b702cded9a538 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 6 May 2022 14:32:43 -0700 Subject: [PATCH 4/9] nit Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/core/functools/decorator_usage.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/functools/decorator_usage.py b/tests/flytekit/unit/core/functools/decorator_usage.py index 8b6e8a9ecd..c9866e1004 100644 --- a/tests/flytekit/unit/core/functools/decorator_usage.py +++ b/tests/flytekit/unit/core/functools/decorator_usage.py @@ -1,7 +1,8 @@ from flytekit import task -from .decorator_source import task_setup from flytekit.core.tracker import extract_task_module +from .decorator_source import task_setup + @task @task_setup From 2f22e5bce84acb3df7fb639dc28d7e3aecc8abe3 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 6 May 2022 15:17:57 -0700 Subject: [PATCH 5/9] nit Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/core/functools/decorator_usage.py | 1 - tests/flytekit/unit/core/functools/test_decorator_location.py | 3 --- 2 files changed, 4 deletions(-) diff --git a/tests/flytekit/unit/core/functools/decorator_usage.py b/tests/flytekit/unit/core/functools/decorator_usage.py index c9866e1004..09006f55ad 100644 --- a/tests/flytekit/unit/core/functools/decorator_usage.py +++ b/tests/flytekit/unit/core/functools/decorator_usage.py @@ -1,5 +1,4 @@ from flytekit import task -from flytekit.core.tracker import extract_task_module from .decorator_source import task_setup diff --git a/tests/flytekit/unit/core/functools/test_decorator_location.py b/tests/flytekit/unit/core/functools/test_decorator_location.py index 34b41ad1f8..0439485f9d 100644 --- a/tests/flytekit/unit/core/functools/test_decorator_location.py +++ b/tests/flytekit/unit/core/functools/test_decorator_location.py @@ -1,8 +1,5 @@ import importlib -from flytekit import task -from flytekit.core.tracker import extract_task_module - def test_dont_use_wrapper_location(): m = importlib.import_module("tests.flytekit.unit.core.functools.decorator_usage") From 45345d929a82af7f7e9de8e2d191af38b283a170 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 6 May 2022 15:36:15 -0700 Subject: [PATCH 6/9] nit Signed-off-by: Yee Hing Tong --- flytekit/clis/sdk_in_container/pyflyte.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index 9ce86e9f2a..b386866e30 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -42,11 +42,7 @@ def main(ctx, pkgs=None, config=None): Entrypoint for all the user commands. """ ctx.obj = dict() - # import importlib - # importlib.import_module("core.flyte_basics.hello_world") - # - # raise Exception("jfkdlsa") - # Handle package management - get from config if not specified on the command line + pkgs = pkgs or [] if config: ctx.obj[CTX_CONFIG_FILE] = config From 61b636b04b1d7cae3e45e14ed24d7a1454a19cae Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 6 May 2022 15:36:37 -0700 Subject: [PATCH 7/9] nit Signed-off-by: Yee Hing Tong --- flytekit/clis/sdk_in_container/pyflyte.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index b386866e30..f27128d108 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -43,6 +43,7 @@ def main(ctx, pkgs=None, config=None): """ ctx.obj = dict() + # Handle package management - get from config if not specified on the command line pkgs = pkgs or [] if config: ctx.obj[CTX_CONFIG_FILE] = config From ac44b59990cdc0818458165c7d8d7721d4ba80ba Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 6 May 2022 15:36:59 -0700 Subject: [PATCH 8/9] nit Signed-off-by: Yee Hing Tong --- flytekit/tools/serialize_helpers.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/flytekit/tools/serialize_helpers.py b/flytekit/tools/serialize_helpers.py index 8c99a4e757..7c1969afa7 100644 --- a/flytekit/tools/serialize_helpers.py +++ b/flytekit/tools/serialize_helpers.py @@ -96,11 +96,6 @@ def get_registrable_entities( f"Multiple definitions of the following tasks were found: {duplicate_task_names}" ) - # from flytekit.models.task import TaskSpec - # for xx in entities_to_be_serialized: - # if isinstance(xx, TaskSpec) and "use_setup" in xx.template.id.name: - # import ipdb; ipdb.set_trace() - return [v.to_flyte_idl() for v in entities_to_be_serialized] From 1b0b1e00c567fcadb16d1331d6b57eea5594cf2e Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 6 May 2022 15:39:57 -0700 Subject: [PATCH 9/9] add one test Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/core/functools/decorator_usage.py | 1 - .../unit/core/functools/test_decorator_location.py | 9 +++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/functools/decorator_usage.py b/tests/flytekit/unit/core/functools/decorator_usage.py index 09006f55ad..bdd4fca2d9 100644 --- a/tests/flytekit/unit/core/functools/decorator_usage.py +++ b/tests/flytekit/unit/core/functools/decorator_usage.py @@ -6,5 +6,4 @@ @task @task_setup def get_data(x: int) -> int: - print("running get_data") return x + 1 diff --git a/tests/flytekit/unit/core/functools/test_decorator_location.py b/tests/flytekit/unit/core/functools/test_decorator_location.py index 0439485f9d..d896d09592 100644 --- a/tests/flytekit/unit/core/functools/test_decorator_location.py +++ b/tests/flytekit/unit/core/functools/test_decorator_location.py @@ -1,8 +1,17 @@ import importlib +from flytekit.core.tracker import extract_task_module + def test_dont_use_wrapper_location(): m = importlib.import_module("tests.flytekit.unit.core.functools.decorator_usage") get_data_task = getattr(m, "get_data") assert "decorator_source" not in get_data_task.name assert "decorator_usage" in get_data_task.name + + a, b, c, _ = extract_task_module(get_data_task) + assert (a, b, c) == ( + "tests.flytekit.unit.core.functools.decorator_usage.get_data", + "tests.flytekit.unit.core.functools.decorator_usage", + "get_data", + )