From 46871eded57feed2f927bbba51a02b5c7a1a0024 Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 28 Nov 2022 00:34:07 +0530 Subject: [PATCH 01/16] safe work pickling --- src/lightning_app/utilities/safe_pickle.py | 70 +++++++++++++++++++ tests/tests_app/utilities/test_safe_pickle.py | 11 +++ .../utilities/testdata/safe_pickle_app.py | 63 +++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 src/lightning_app/utilities/safe_pickle.py create mode 100644 tests/tests_app/utilities/test_safe_pickle.py create mode 100644 tests/tests_app/utilities/testdata/safe_pickle_app.py diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py new file mode 100644 index 0000000000000..4768d99812172 --- /dev/null +++ b/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 diff --git a/tests/tests_app/utilities/test_safe_pickle.py b/tests/tests_app/utilities/test_safe_pickle.py new file mode 100644 index 0000000000000..ebe05d294979e --- /dev/null +++ b/tests/tests_app/utilities/test_safe_pickle.py @@ -0,0 +1,11 @@ +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") diff --git a/tests/tests_app/utilities/testdata/safe_pickle_app.py b/tests/tests_app/utilities/testdata/safe_pickle_app.py new file mode 100644 index 0000000000000..1213f5a74c2b9 --- /dev/null +++ b/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.utilities import safe_pickle +from lightning_app import LightningApp, LightningWork, LightningFlow + + +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()) From e40df251e905bc7f0f4fc6080d09f258e7a2c189 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 27 Nov 2022 19:07:41 +0000 Subject: [PATCH 02/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning_app/utilities/safe_pickle.py | 4 ++-- tests/tests_app/utilities/test_safe_pickle.py | 5 ++--- tests/tests_app/utilities/testdata/safe_pickle_app.py | 2 +- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 4768d99812172..26caf1854c5e8 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -9,7 +9,7 @@ def get_picklable_work(work: LightningWork): - """ Pickling a LightningWork instance fails if it doesn from the work process + """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 @@ -41,7 +41,7 @@ def get_picklable_work(work: LightningWork): 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('/', '.') + 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__) diff --git a/tests/tests_app/utilities/test_safe_pickle.py b/tests/tests_app/utilities/test_safe_pickle.py index ebe05d294979e..391f34d2c43f0 100644 --- a/tests/tests_app/utilities/test_safe_pickle.py +++ b/tests/tests_app/utilities/test_safe_pickle.py @@ -5,7 +5,6 @@ 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) + ["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") diff --git a/tests/tests_app/utilities/testdata/safe_pickle_app.py b/tests/tests_app/utilities/testdata/safe_pickle_app.py index 1213f5a74c2b9..f15344360d85f 100644 --- a/tests/tests_app/utilities/testdata/safe_pickle_app.py +++ b/tests/tests_app/utilities/testdata/safe_pickle_app.py @@ -8,8 +8,8 @@ import subprocess from pathlib import Path +from lightning_app import LightningApp, LightningFlow, LightningWork from lightning_app.utilities import safe_pickle -from lightning_app import LightningApp, LightningWork, LightningFlow class SelfPicklingWork(LightningWork): From fef035d023855c7b9080b2d2f1366993e5df0600 Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 28 Nov 2022 15:14:39 +0530 Subject: [PATCH 03/16] mypy --- src/lightning_app/utilities/safe_pickle.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 26caf1854c5e8..108489f95326c 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -1,6 +1,7 @@ import pickle import sys import types +import typing from copy import deepcopy from pathlib import Path @@ -8,7 +9,7 @@ from lightning_app.utilities.app_helpers import _LightningAppRef -def get_picklable_work(work: LightningWork): +def get_picklable_work(work: LightningWork) -> 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. @@ -29,7 +30,10 @@ def get_picklable_work(work: LightningWork): # 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: + app_ref = _LightningAppRef.get_current() + if app_ref is None: + raise RuntimeError("Cannot pickle LightningWork outside of a LightningApp") + for w in app_ref.works: if work.name == w.name: # copying the work object to avoid modifying the original work object copied_work = deepcopy(w) @@ -40,6 +44,12 @@ def get_picklable_work(work: LightningWork): # 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__] + work_class_file = work_class_module.__file__ + if not work_class_file: + raise ValueError( + f"Cannot pickle work class {copied_work.__class__.__name__} because we " + f"couldn't identify the module file" + ) 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 @@ -57,12 +67,12 @@ def get_picklable_work(work: LightningWork): return copied_work -def dump(work: LightningWork, f): +def dump(work: LightningWork, f: typing.BinaryIO) -> None: picklable_work = get_picklable_work(work) pickle.dump(picklable_work, f) -def load(f): +def load(f) -> typing.Any: # inject current working directory to sys.path sys.path.insert(1, str(Path.cwd())) work = pickle.load(f) From 9fe8bfd489597202e2312261b5be3655e7340e18 Mon Sep 17 00:00:00 2001 From: Sherin Thomas Date: Mon, 28 Nov 2022 15:16:17 +0530 Subject: [PATCH 04/16] Update tests/tests_app/utilities/test_safe_pickle.py Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> --- tests/tests_app/utilities/test_safe_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tests_app/utilities/test_safe_pickle.py b/tests/tests_app/utilities/test_safe_pickle.py index 391f34d2c43f0..8f8f086fcc877 100644 --- a/tests/tests_app/utilities/test_safe_pickle.py +++ b/tests/tests_app/utilities/test_safe_pickle.py @@ -7,4 +7,4 @@ def test_safe_pickle_app(): 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") + assert "Exiting the pickling app successfully!" in proc.stdout.read().decode("UTF-8") From b041d4ee301637cfacd4e19b593d2cc15c5222f4 Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 28 Nov 2022 15:19:33 +0530 Subject: [PATCH 05/16] changelog --- src/lightning_app/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/src/lightning_app/CHANGELOG.md b/src/lightning_app/CHANGELOG.md index 54219e0f31eed..e61c0572d5d5e 100644 --- a/src/lightning_app/CHANGELOG.md +++ b/src/lightning_app/CHANGELOG.md @@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Show a message when `BuildConfig(dockerfile="...")` is passed but a `Dockerfile` file is already present in the Work ([#15799](https://github.com/Lightning-AI/lightning/pull/15799)) - Added a CloudMultiProcessBackend which enables running a child App from within the Flow in the cloud ([#15800](https://github.com/Lightning-AI/lightning/pull/15800)) +- Utility for pickling work object safely even from a child process ([#15836](https://github.com/Lightning-AI/lightning/pull/15836)) ### Changed From 8732ef4b9f3573104a7404f036350cf6140bbec0 Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 28 Nov 2022 15:34:33 +0530 Subject: [PATCH 06/16] trying to fix ubuntu test --- tests/tests_app/utilities/test_safe_pickle.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/tests_app/utilities/test_safe_pickle.py b/tests/tests_app/utilities/test_safe_pickle.py index 8f8f086fcc877..1f6370cb8394f 100644 --- a/tests/tests_app/utilities/test_safe_pickle.py +++ b/tests/tests_app/utilities/test_safe_pickle.py @@ -7,4 +7,5 @@ def test_safe_pickle_app(): 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") + stdout, _ = proc.communicate() + assert "Exiting the pickling app successfully!!" in stdout.decode("UTF-8") From 4a7dd7d36fc9a094ae2a12bb16838aadf0cf3bf5 Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 28 Nov 2022 15:36:51 +0530 Subject: [PATCH 07/16] more typing fixes --- src/lightning_app/utilities/safe_pickle.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 108489f95326c..2289e830748bb 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -50,7 +50,7 @@ def get_picklable_work(work: LightningWork) -> LightningWork: f"Cannot pickle work class {copied_work.__class__.__name__} because we " f"couldn't identify the module file" ) - relative_path = Path(work_class_module.__file__).relative_to(Path.cwd()) + relative_path = Path(work_class_module.__file__).relative_to(Path.cwd()) # type: ignore 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) @@ -72,7 +72,7 @@ def dump(work: LightningWork, f: typing.BinaryIO) -> None: pickle.dump(picklable_work, f) -def load(f) -> typing.Any: +def load(f: typing.BinaryIO) -> typing.Any: # inject current working directory to sys.path sys.path.insert(1, str(Path.cwd())) work = pickle.load(f) From 5723995368c2c2860aecec21a3bb7274953e5855 Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 28 Nov 2022 15:44:33 +0530 Subject: [PATCH 08/16] trying to fix ubuntu test --- tests/tests_app/utilities/test_safe_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tests_app/utilities/test_safe_pickle.py b/tests/tests_app/utilities/test_safe_pickle.py index 1f6370cb8394f..473fe28ed22f7 100644 --- a/tests/tests_app/utilities/test_safe_pickle.py +++ b/tests/tests_app/utilities/test_safe_pickle.py @@ -8,4 +8,4 @@ def test_safe_pickle_app(): ["lightning", "run", "app", "safe_pickle_app.py", "--open-ui", "false"], stdout=subprocess.PIPE, cwd=test_dir ) stdout, _ = proc.communicate() - assert "Exiting the pickling app successfully!!" in stdout.decode("UTF-8") + assert "Exiting the pickling app successfully" in stdout.decode("UTF-8") From 22c1625b1a95a860aa43a4030d526692f709e823 Mon Sep 17 00:00:00 2001 From: Sherin Thomas Date: Mon, 28 Nov 2022 17:31:03 +0530 Subject: [PATCH 09/16] Update src/lightning_app/CHANGELOG.md Co-authored-by: Ethan Harris --- src/lightning_app/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/src/lightning_app/CHANGELOG.md b/src/lightning_app/CHANGELOG.md index e61c0572d5d5e..d88705571456a 100644 --- a/src/lightning_app/CHANGELOG.md +++ b/src/lightning_app/CHANGELOG.md @@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Show a message when `BuildConfig(dockerfile="...")` is passed but a `Dockerfile` file is already present in the Work ([#15799](https://github.com/Lightning-AI/lightning/pull/15799)) - Added a CloudMultiProcessBackend which enables running a child App from within the Flow in the cloud ([#15800](https://github.com/Lightning-AI/lightning/pull/15800)) + - Utility for pickling work object safely even from a child process ([#15836](https://github.com/Lightning-AI/lightning/pull/15836)) From 6a3fce3c2d53903369433b81c74c62d3d2a28fcd Mon Sep 17 00:00:00 2001 From: Sherin Thomas Date: Mon, 28 Nov 2022 17:35:05 +0530 Subject: [PATCH 10/16] Update src/lightning_app/utilities/safe_pickle.py Co-authored-by: Ethan Harris --- src/lightning_app/utilities/safe_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 2289e830748bb..e9b5aef04dc69 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -10,7 +10,7 @@ def get_picklable_work(work: LightningWork) -> LightningWork: - """Pickling a LightningWork instance fails if it doesn from the work process + """Pickling a LightningWork instance fails if done 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 From 1333e9bf41110fcab19fcd58a98c8fd40d8ba9d5 Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 5 Dec 2022 13:29:18 +0000 Subject: [PATCH 11/16] trim work args that has mp queue reference --- src/lightning_app/utilities/safe_pickle.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index e9b5aef04dc69..81145313d12a0 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -1,3 +1,4 @@ +import contextlib import pickle import sys import types @@ -9,6 +10,20 @@ from lightning_app.utilities.app_helpers import _LightningAppRef +NON_PICKLABLE_WORK_ARGS = ['_request_queue', '_response_queue', '_backend', '_setattr_replacement'] + + +@contextlib.contextmanager +def trimmed_work(work: LightningWork, to_trim: typing.List[str]) -> None: + holder = {} + for arg in to_trim: + holder[arg] = getattr(work, arg) + setattr(work, arg, None) + yield + for arg in to_trim: + setattr(work, arg, holder[arg]) + + def get_picklable_work(work: LightningWork) -> LightningWork: """Pickling a LightningWork instance fails if done from the work process itself. This function is safe to call from the work process within both MultiprocessRuntime @@ -36,10 +51,12 @@ def get_picklable_work(work: LightningWork) -> LightningWork: for w in app_ref.works: if work.name == w.name: # copying the work object to avoid modifying the original work object - copied_work = deepcopy(w) + with trimmed_work(w, to_trim=NON_PICKLABLE_WORK_ARGS): + 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__: From 1025aefee3a1162b00c56b4f15d8ab6386d89dfd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 5 Dec 2022 13:30:35 +0000 Subject: [PATCH 12/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning_app/utilities/safe_pickle.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 81145313d12a0..a80bd86addf61 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -9,8 +9,7 @@ from lightning_app.core.work import LightningWork from lightning_app.utilities.app_helpers import _LightningAppRef - -NON_PICKLABLE_WORK_ARGS = ['_request_queue', '_response_queue', '_backend', '_setattr_replacement'] +NON_PICKLABLE_WORK_ARGS = ["_request_queue", "_response_queue", "_backend", "_setattr_replacement"] @contextlib.contextmanager From 63f46ba2dc13c5f82bec247bd0960b41713f786f Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 5 Dec 2022 13:36:37 +0000 Subject: [PATCH 13/16] mypy --- src/lightning_app/utilities/safe_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 81145313d12a0..766d6aeaaa3cb 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -14,7 +14,7 @@ @contextlib.contextmanager -def trimmed_work(work: LightningWork, to_trim: typing.List[str]) -> None: +def trimmed_work(work: LightningWork, to_trim: typing.List[str]) -> typing.Iterator[None]: holder = {} for arg in to_trim: holder[arg] = getattr(work, arg) From 8c955a299ef5d0a9f88996072522dbbca4b93fee Mon Sep 17 00:00:00 2001 From: hhsecond Date: Mon, 5 Dec 2022 18:18:06 +0000 Subject: [PATCH 14/16] review --- src/lightning_app/utilities/safe_pickle.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 5a48c6a39a8b9..6e32a2d2572c6 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -9,11 +9,12 @@ from lightning_app.core.work import LightningWork from lightning_app.utilities.app_helpers import _LightningAppRef -NON_PICKLABLE_WORK_ARGS = ["_request_queue", "_response_queue", "_backend", "_setattr_replacement"] +NON_PICKLABLE_WORK_ATTRIBUTES = ["_request_queue", "_response_queue", "_backend", "_setattr_replacement"] @contextlib.contextmanager -def trimmed_work(work: LightningWork, to_trim: typing.List[str]) -> typing.Iterator[None]: +def _trimmed_work(work: LightningWork, to_trim: typing.List[str]) -> typing.Iterator[None]: + """ Context manager to trim the work object to remove attributes that are not picklable """ holder = {} for arg in to_trim: holder[arg] = getattr(work, arg) @@ -42,15 +43,16 @@ def get_picklable_work(work: LightningWork) -> LightningWork: └── app.py """ - # pickling the user work class - pickling `self` will cause issue because the - # work is running under a process, in local + # If the work object not taken from the app ref, there is a thread lock reference + # somewhere thats preventing it from being pickled. Investigate it later. We + # shouldn't be fetching the work object from the app ref. TODO @sherin app_ref = _LightningAppRef.get_current() if app_ref is None: raise RuntimeError("Cannot pickle LightningWork outside of a LightningApp") for w in app_ref.works: if work.name == w.name: # copying the work object to avoid modifying the original work object - with trimmed_work(w, to_trim=NON_PICKLABLE_WORK_ARGS): + with _trimmed_work(w, to_trim=NON_PICKLABLE_WORK_ATTRIBUTES): copied_work = deepcopy(w) break else: @@ -77,9 +79,6 @@ def get_picklable_work(work: LightningWork) -> LightningWork: 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 From f16735fe0a7c53506ffffca896fa8dd0b4333f68 Mon Sep 17 00:00:00 2001 From: Sherin Thomas Date: Mon, 5 Dec 2022 23:48:27 +0530 Subject: [PATCH 15/16] Update src/lightning_app/utilities/safe_pickle.py Co-authored-by: thomas chaton --- src/lightning_app/utilities/safe_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 6e32a2d2572c6..427048ef40527 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -51,7 +51,7 @@ def get_picklable_work(work: LightningWork) -> LightningWork: raise RuntimeError("Cannot pickle LightningWork outside of a LightningApp") for w in app_ref.works: if work.name == w.name: - # copying the work object to avoid modifying the original work object + # deep-copying the work object to avoid modifying the original work object with _trimmed_work(w, to_trim=NON_PICKLABLE_WORK_ATTRIBUTES): copied_work = deepcopy(w) break From 76b3ca67e9a6feb339a4b0ec93ca0efd472eea3e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 5 Dec 2022 18:19:29 +0000 Subject: [PATCH 16/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning_app/utilities/safe_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning_app/utilities/safe_pickle.py b/src/lightning_app/utilities/safe_pickle.py index 427048ef40527..8788ff22a3cb6 100644 --- a/src/lightning_app/utilities/safe_pickle.py +++ b/src/lightning_app/utilities/safe_pickle.py @@ -14,7 +14,7 @@ @contextlib.contextmanager def _trimmed_work(work: LightningWork, to_trim: typing.List[str]) -> typing.Iterator[None]: - """ Context manager to trim the work object to remove attributes that are not picklable """ + """Context manager to trim the work object to remove attributes that are not picklable.""" holder = {} for arg in to_trim: holder[arg] = getattr(work, arg)