Skip to content

Commit 67a47d4

Browse files
authored
[App] Add automatic conversion to structures (#15961)
1 parent 772d121 commit 67a47d4

File tree

3 files changed

+38
-2
lines changed

3 files changed

+38
-2
lines changed

src/lightning_app/CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
2828

2929
- Added a `configure_layout` method to the `LightningWork` which can be used to control how the work is handled in the layout of a parent flow ([#15926](https://github.com/Lightning-AI/lightning/pull/15926))
3030

31+
- Added automatic conversion of list and dict of works and flows to structures ([#15961](https://github.com/Lightning-AI/lightning/pull/15961))
32+
3133

3234
### Changed
3335

src/lightning_app/core/flow.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,14 @@ def __setattr__(self, name: str, value: Any) -> None:
142142
if name in self._works and value != getattr(self, name):
143143
raise AttributeError(f"Cannot set attributes as the work can't be changed once defined: {name}")
144144

145+
if isinstance(value, (list, dict)) and value:
146+
_type = (LightningFlow, LightningWork, List, Dict)
147+
if isinstance(value, list) and all(isinstance(va, _type) for va in value):
148+
value = List(*value)
149+
150+
if isinstance(value, dict) and all(isinstance(va, _type) for va in value.values()):
151+
value = Dict(**value)
152+
145153
if isinstance(value, LightningFlow):
146154
self._flows.add(name)
147155
_set_child_name(self, value, name)
@@ -163,10 +171,10 @@ def __setattr__(self, name: str, value: Any) -> None:
163171
value._register_cloud_compute()
164172

165173
elif isinstance(value, (Dict, List)):
166-
value._backend = self._backend
167174
self._structures.add(name)
168175
_set_child_name(self, value, name)
169-
if self._backend:
176+
if getattr(self, "_backend", None) is not None:
177+
value._backend = self._backend
170178
for flow in value.flows:
171179
LightningFlow._attach_backend(flow, self._backend)
172180
for work in value.works:

tests/tests_app/structures/test_structures.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -518,3 +518,29 @@ def __init__(self):
518518
LightningApp(flow) # wrap in app to init all component names
519519
assert flow.list_structure[0].name == "root.list_structure.0"
520520
assert flow.dict_structure["dict_child"].name == "root.dict_structure.dict_child"
521+
522+
523+
class FlowWiStructures(LightningFlow):
524+
def __init__(self):
525+
super().__init__()
526+
527+
self.ws = [EmptyFlow(), EmptyFlow()]
528+
529+
self.ws1 = {"a": EmptyFlow(), "b": EmptyFlow()}
530+
531+
self.ws2 = {
532+
"a": EmptyFlow(),
533+
"b": EmptyFlow(),
534+
"c": List(EmptyFlow(), EmptyFlow()),
535+
"d": Dict(**{"a": EmptyFlow()}),
536+
}
537+
538+
def run(self):
539+
pass
540+
541+
542+
def test_flow_without_structures():
543+
544+
flow = FlowWiStructures()
545+
assert isinstance(flow.ws, List)
546+
assert isinstance(flow.ws1, Dict)

0 commit comments

Comments
 (0)