Swap out inspect file location

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:

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))

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]
 
 

tests/flytekit/unit/core/functools/decorator_source.py

@@ -0,0 +1,23 @@
+"""Script used for testing local execution of functool.wraps-wrapped tasks for stacked decorators"""
+
+from functools import wraps
+from typing import List
+
+
+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

tests/flytekit/unit/core/functools/decorator_usage.py

@@ -0,0 +1,11 @@
+from flytekit import task
+from flytekit.core.tracker import extract_task_module
+
+from .decorator_source import task_setup
+
+
+@task
+@task_setup
+def get_data(x: int) -> int:
+    print("running get_data")
+    return x + 1

tests/flytekit/unit/core/functools/test_decorator_location.py

@@ -0,0 +1,11 @@
+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")
+    get_data_task = getattr(m, "get_data")
+    assert "decorator_source" not in get_data_task.name
+    assert "decorator_usage" in get_data_task.name