Merge pull request #530 from latchbio/ayush/fix-lps

fix launchplans

Author: ayushkamat

CHANGELOG.md

@@ -16,6 +16,12 @@ Types of changes
 
 # Latch SDK Changelog
 
+## 2.58.2 - 2024-04-03
+
+### Fixed
+
+* Fix issue where `LaunchPlan`s would not be interpreted correctly in custom workflow modules
+
 ## 2.58.1 - 2024-04-03
 
 ### Fixed

pyproject.toml

@@ -12,7 +12,7 @@ include = ["src/**/*.py", "src/latch_cli/services/init/*"]
 
 [project]
 name = "latch"
-version = "2.58.1"
+version = "2.58.2"
 description = "The Latch SDK"
 authors = [{ name = "Kenny Workman", email = "[email protected]" }]
 maintainers = [

src/latch/resources/launch_plan.py

@@ -1,7 +1,7 @@
-import inspect
-from typing import Any, Callable, Dict
+from typing import Any
 
-from flytekit import LaunchPlan as _LaunchPlan
+from flytekit.core.launch_plan import LaunchPlan as _LaunchPlan
+from flytekit.core.workflow import PythonFunctionWorkflow
 
 
 class LaunchPlan:
@@ -26,18 +26,19 @@ class LaunchPlan:
         )
     """
 
-    def __init__(self, workflow: Callable, name: str, default_params: Dict[str, Any]):
-        # This constructor is invoked twice when task code is executed.
-        #   1. When the pyflyte-execute entrypoint is invoked to start task.
-        #      `mod.__name__`  of caller is `wf`
-        #   2. When the PythonAutoContainer loads our module to call our task.
-        #      `mod.__name__`  of caller is `wf.__init__`
-        # LaunchPlans are stored in a global array and redundant additions will
-        # throw an exception, so we want to create on the first constructor
-        # call only.
-
-        frame = inspect.stack()[1]
-        mod = inspect.getmodule(frame[0])
-        if mod.__name__ == "wf":
-            str_repr = f"wf.__init__.{workflow.__name__}.{name}"
-            _LaunchPlan.create(str_repr, workflow, default_params)
+    def __init__(
+        self,
+        workflow: PythonFunctionWorkflow,
+        name: str,
+        default_params: dict[str, Any],
+    ):
+        try:  # noqa: SIM105
+            _LaunchPlan.create(
+                f"{workflow.__module__}.{workflow.__name__}.{name}",
+                workflow,
+                default_params,
+            )
+
+        # if the launchplan already exists, the `create` method throws an AssertionError
+        except AssertionError:
+            pass