def sanitize(obj: Union[Path, Drive]) -> Union[Path, Dict]:
if isinstance(obj, Path):
# create a copy of the Path and erase the consumer
# the LightningWork on the receiving end of the caller queue will become the new consumer
# this is necessary to make the Path deepdiff-hashable
path_copy = Path(obj)
path_copy._sanitize()
path_copy._consumer = None
return path_copy
return obj.to_dict()
def _transfer_path_attributes(self) -> None:
"""Transfer all Path attributes in the Work if they have an origin and exist."""
for name in self.work._paths:
path = getattr(self.work, name)
if isinstance(path, str):
path = Path(path)
path._attach_work(self.work)
if path.origin_name and path.origin_name != self.work.name and path.exists_remote(
):
path.get(overwrite=True)
def test_flow_state_change_with_path():
"""Test that type changes to a Path attribute are properly reflected within the state."""
class Flow(LightningFlow):
def __init__(self):
super().__init__()
self.none_to_path = None
self.path_to_none = Path()
self.path_to_path = Path()
def run(self):
self.none_to_path = "lit://none/to/path"
self.path_to_none = None
self.path_to_path = "lit://path/to/path"
self._exit()
flow = Flow()
MultiProcessRuntime(LightningApp(flow)).dispatch()
assert flow.none_to_path == Path("lit://none/to/path")
assert flow.path_to_none is None
assert flow.path_to_path == Path("lit://path/to/path")
assert "path_to_none" not in flow._paths
assert "path_to_none" in flow._state
assert flow._paths["none_to_path"] == Path("lit://none/to/path").to_dict()
assert flow._paths["path_to_path"] == Path("lit://path/to/path").to_dict()
assert flow.state["vars"]["none_to_path"] == Path("lit://none/to/path")
assert flow.state["vars"]["path_to_none"] is None
assert flow.state["vars"]["path_to_path"] == Path("lit://path/to/path")
def test_convert_paths_after_init():
"""Test that we can convert paths after the Flow/Work initialization, i.e., when the LightningApp is fully
instantiated."""
# TODO: Add a test case for the Lightning List and Dict containers
class Flow1(EmptyFlow):
def __init__(self):
super().__init__()
self.path1 = Path("a")
self.path2 = Path("b")
flow1 = Flow1()
assert flow1._paths == {}
_convert_paths_after_init(flow1)
assert flow1._paths == {
"path1": Path("a").to_dict(),
"path2": Path("b").to_dict()
}
class Work1(EmptyWork):
def __init__(self):
super().__init__()
self.path3 = Path("c")
class Flow2(EmptyFlow):
def __init__(self):
super().__init__()
self.work1 = Work1()
self.path4 = Path("d")
flow2 = Flow2()
assert flow2._paths == {}
assert flow2.work1._paths == {}
_convert_paths_after_init(flow2)
assert flow2._paths == {"path4": Path("d").to_dict()}
assert set(flow2.work1._paths.keys()) == {"path3"}
assert flow2.work1._paths["path3"]["origin_name"] == "root.work1"
assert flow2.work1._paths["path3"]["consumer_name"] == "root.work1"
def __init__(self):
super().__init__()
self.none_to_path = None
self.path_to_none = Path()
self.path_to_path = Path()
def __init__(self):
super().__init__()
self.no_path = "a/b/c"
self.path = Path("lit://x/y/z")
self.lit_path = "lit://x/y/z"
def __init__(self, raise_exception, enable_exception=True):
super().__init__(raise_exception=raise_exception)
self.enable_exception = enable_exception
self.dummy_path = Path("test")
def __init__(self):
super().__init__()
self.work1 = Work1()
self.path4 = Path("d")
def __init__(self):
super().__init__()
self.path3 = Path("c")
def __init__(self):
super().__init__()
self.path1 = Path("a")
self.path2 = Path("b")
def __getattr__(self, item):
if item in self.__dict__.get("_paths", {}):
return Path.from_dict(self._paths[item])
return self.__getattribute__(item)
def __setattr__(self, name, value):
from lightning_app.structures import Dict, List
if (
not _is_init_context(self)
and name not in self._state
and name not in self._paths
and (
not isinstance(value, (LightningWork, LightningFlow))
or (isinstance(value, (LightningWork, LightningFlow)) and not _is_run_context(self))
)
and name not in self._works.union(self._flows)
and self._is_state_attribute(name)
):
raise AttributeError(f"Cannot set attributes that were not defined in __init__: {name}")
if isinstance(value, str) and value.startswith("lit://"):
value = Path(value)
if self._is_state_attribute(name):
if hasattr(self, name):
if name in self._flows and value != getattr(self, name):
raise AttributeError(f"Cannot set attributes as the flow can't be changed once defined: {name}")
if name in self._works and value != getattr(self, name):
raise AttributeError(f"Cannot set attributes as the work can't be changed once defined: {name}")
if isinstance(value, LightningFlow):
self._flows.add(name)
_set_child_name(self, value, name)
if name in self._state:
self._state.remove(name)
# Attach the backend to the flow and its children work.
if self._backend:
LightningFlow._attach_backend(value, self._backend)
elif isinstance(value, LightningWork):
self._works.add(name)
_set_child_name(self, value, name)
if name in self._state:
self._state.remove(name)
if self._backend:
self._backend._wrap_run_method(_LightningAppRef().get_current(), value)
elif isinstance(value, (Dict, List)):
value._backend = self._backend
self._structures.add(name)
_set_child_name(self, value, name)
if self._backend:
for flow in value.flows:
LightningFlow._attach_backend(flow, self._backend)
for work in value.works:
self._backend._wrap_run_method(_LightningAppRef().get_current(), work)
elif isinstance(value, Path):
# In the init context, the full name of the Flow and Work is not known, i.e., we can't serialize
# the path without losing the information of origin and consumer. Hence, we delay the serialization
# of the path object until the app is instantiated.
if not _is_init_context(self):
self._paths[name] = value.to_dict()
self._state.add(name)
elif isinstance(value, Drive):
value = deepcopy(value)
value.component_name = self.name
self._state.add(name)
elif _is_json_serializable(value):
self._state.add(name)
if not isinstance(value, Path) and hasattr(self, "_paths") and name in self._paths:
# The attribute changed type from Path to another
self._paths.pop(name)
else:
raise AttributeError(
f"Only JSON-serializable attributes are currently supported"
f" (str, int, float, bool, tuple, list, dict etc.) to be part of {self} state. "
f"Found the attribute {name} with {value} instead. \n"
"HINT: Private attributes defined as follows `self._x = y` won't be shared between components "
"and therefore don't need to be JSON-serializable."
)
super().__setattr__(name, value)
def sanitize_path(path: Path) -> Path:
path_copy = Path(path)
path_copy._sanitize()
return path_copy
