[App] Multiprocessing-safe work pickling

src/lightning_app/utilities/safe_pickle.py

@@ -0,0 +1,70 @@
+import pickle
+import sys
+import types
+from copy import deepcopy
+from pathlib import Path
+
+from lightning_app.core.work import LightningWork
+from lightning_app.utilities.app_helpers import _LightningAppRef
+
+
+def get_picklable_work(work: LightningWork):
+    """Pickling a LightningWork instance fails if it doesn from the work process
+    itself. This function is safe to call from the work process within both MultiprocessRuntime
+    and Cloud.
+    Note: This function modifies the module information of the work object. Specifically, it injects
+    the relative module path into the __module__ attribute of the work object. If the object is not
+    importable from the CWD, then the pickle load will fail.
+
+    Example:
+        for a directory structure like below and the work class is defined in the app.py where
+        the app.py is the entrypoint for the app, it will inject `foo.bar.app` into the
+        __module__ attribute
+
+        └── foo
+            ├── __init__.py
+            └── bar
+                └── app.py
+    """
+
+    # pickling the user work class - pickling `self` will cause issue because the
+    # work is running under a process, in local
+    for w in _LightningAppRef.get_current().works:
+        if work.name == w.name:
+            # copying the work object to avoid modifying the original work object
+            copied_work = deepcopy(w)
+            break
+    else:
+        raise ValueError(f"Work with name {work.name} not found in the app references")
+    # if work is defined in the __main__ or __mp__main__ (the entrypoint file for `lightning run app` command),
+    # pickling/unpickling will fail, hence we need patch the module information
+    if "_main__" in copied_work.__class__.__module__:
+        work_class_module = sys.modules[copied_work.__class__.__module__]
+        relative_path = Path(work_class_module.__file__).relative_to(Path.cwd())
+        expected_module_name = relative_path.as_posix().replace(".py", "").replace("/", ".")
+        # TODO @sherin: also check if the module is importable from the CWD
+        fake_module = types.ModuleType(expected_module_name)
+        fake_module.__dict__.update(work_class_module.__dict__)
+        fake_module.__dict__["__name__"] = expected_module_name
+        sys.modules[expected_module_name] = fake_module
+        for k, v in fake_module.__dict__.items():
+            if not k.startswith("__") and hasattr(v, "__module__"):
+                if "_main__" in v.__module__:
+                    v.__module__ = expected_module_name
+
+    # removing reference to backend; backend is not picklable because of openapi client reference in it
+    copied_work._backend = None
+    return copied_work
+
+
+def dump(work: LightningWork, f):
+    picklable_work = get_picklable_work(work)
+    pickle.dump(picklable_work, f)
+
+
+def load(f):
+    # inject current working directory to sys.path
+    sys.path.insert(1, str(Path.cwd()))
+    work = pickle.load(f)
+    sys.path.pop(1)
+    return work

tests/tests_app/utilities/test_safe_pickle.py

@@ -0,0 +1,10 @@
+import subprocess
+from pathlib import Path
+
+
+def test_safe_pickle_app():
+    test_dir = Path(__file__).parent / "testdata"
+    proc = subprocess.Popen(
+        ["lightning", "run", "app", "safe_pickle_app.py", "--open-ui", "false"], stdout=subprocess.PIPE, cwd=test_dir
+    )
+    assert "Exiting the pickling app successfully!!" in proc.stdout.read().decode("UTF-8")

tests/tests_app/utilities/testdata/safe_pickle_app.py

@@ -0,0 +1,63 @@
+"""
+This app tests three things
+1. Can a work pickle `self`
+2. Can the pickled work be unpickled in another work
+3. Can the pickled work be unpickled from a script
+"""
+
+import subprocess
+from pathlib import Path
+
+from lightning_app import LightningApp, LightningFlow, LightningWork
+from lightning_app.utilities import safe_pickle
+
+
+class SelfPicklingWork(LightningWork):
+    def run(self):
+        with open("work.pkl", "wb") as f:
+            safe_pickle.dump(self, f)
+
+    def get_test_string(self):
+        return f"Hello from {self.__class__.__name__}!"
+
+
+class WorkThatLoadsPickledWork(LightningWork):
+    def run(self):
+        with open("work.pkl", "rb") as f:
+            work = safe_pickle.load(f)
+        assert work.get_test_string() == "Hello from SelfPicklingWork!"
+
+
+script_load_pickled_work = """
+import pickle
+work = pickle.load(open("work.pkl", "rb"))
+print(work.get_test_string())
+"""
+
+
+class RootFlow(LightningFlow):
+    def __init__(self):
+        super().__init__()
+        self.self_pickling_work = SelfPicklingWork()
+        self.work_that_loads_pickled_work = WorkThatLoadsPickledWork()
+
+    def run(self):
+        self.self_pickling_work.run()
+        self.work_that_loads_pickled_work.run()
+
+        with open("script_that_loads_pickled_work.py", "w") as f:
+            f.write(script_load_pickled_work)
+
+        # read the output from subprocess
+        proc = subprocess.Popen(["python", "script_that_loads_pickled_work.py"], stdout=subprocess.PIPE)
+        assert "Hello from SelfPicklingWork" in proc.stdout.read().decode("UTF-8")
+
+        # deleting the script
+        Path("script_that_loads_pickled_work.py").unlink()
+        # deleting the pkl file
+        Path("work.pkl").unlink()
+
+        self._exit("Exiting the pickling app successfully!!")
+
+
+app = LightningApp(RootFlow())