pytorch-lightning

1
# Copyright The Lightning AI team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
import inspect
16
import warnings
17
from typing import Dict, List, Union
18

19
import lightning.app
20
from lightning.app.frontend.frontend import Frontend
21
from lightning.app.utilities.app_helpers import _MagicMockJsonSerializable, is_overridden
22
from lightning.app.utilities.cloud import is_running_in_cloud
23

24

25
def _add_comment_to_literal_code(method, contains, comment):
26
    """Inspects a method's code and adds a message to it.
27

28
    This is a nice to have, so if it fails for some reason, it shouldn't affect the program.
29

30
    """
31
    try:
32
        lines = inspect.getsource(method)
33
        lines = lines.split("\n")
34
        idx_list = [i for i, x in enumerate(lines) if contains in x]
35
        for i in idx_list:
36
            line = lines[i]
37
            line += comment
38
            lines[i] = line
39

40
        return "\n".join(lines)
41

42
    except Exception:
43
        return ""
44

45

46
def _collect_layout(app: "lightning.app.LightningApp", flow: "lightning.app.LightningFlow") -> Union[Dict, List[Dict]]:
47
    """Process the layout returned by the ``configure_layout()`` method in each flow."""
48
    layout = flow.configure_layout()
49

50
    if isinstance(layout, Frontend):
51
        frontend = layout
52
        frontend.flow = flow
53
        app.frontends.setdefault(flow.name, frontend)
54

55
        # When running locally, the target will get overwritten by the dispatcher when launching the frontend servers
56
        # When running in the cloud, the frontend code will construct the URL based on the flow name
57
        return flow._layout
58
    if isinstance(layout, _MagicMockJsonSerializable):
59
        # The import was mocked, we set a dummy `Frontend` so that `is_headless` knows there is a UI
60
        app.frontends.setdefault(flow.name, "mock")
61
        return flow._layout
62
    if isinstance(layout, dict):
63
        layout = _collect_content_layout([layout], app, flow)
64
    elif isinstance(layout, (list, tuple)) and all(isinstance(item, dict) for item in layout):
65
        layout = _collect_content_layout(layout, app, flow)
66
    else:
67
        lines = _add_comment_to_literal_code(flow.configure_layout, contains="return", comment="  <------- this guy")
68
        raise TypeError(
69
            f"""
70
        The return value of configure_layout() in `{flow.__class__.__name__}`  is an unsupported layout format:
71
        \n{lines}
72

73
        Return either an object of type {Frontend} (e.g., StreamlitFrontend, StaticWebFrontend):
74
            def configure_layout(self):
75
                return la.frontend.Frontend(...)
76

77
        OR a single dict:
78
            def configure_layout(self):
79
                tab1 = {{'name': 'tab name', 'content': self.a_component}}
80
                return tab1
81

82
        OR a list of dicts:
83
            def configure_layout(self):
84
                tab1 = {{'name': 'tab name 1', 'content': self.component_a}}
85
                tab2 = {{'name': 'tab name 2', 'content': self.component_b}}
86
                return [tab1, tab2]
87

88
        (see the docs for `LightningFlow.configure_layout`).
89
        """
90
        )
91

92
    return layout
93

94

95
def _collect_content_layout(
96
    layout: List[Dict], app: "lightning.app.LightningApp", flow: "lightning.app.LightningFlow"
97
) -> Union[List[Dict], Dict]:
98
    """Process the layout returned by the ``configure_layout()`` method if the returned format represents an
99
    aggregation of child layouts."""
100
    for entry in layout:
101
        if "content" not in entry:
102
            raise ValueError(
103
                f"A dictionary returned by `{flow.__class__.__name__}.configure_layout()` is missing a key 'content'."
104
                f" For the value, choose either a reference to a child flow or a URla."
105
            )
106
        if isinstance(entry["content"], str):  # assume this is a URL
107
            url = entry["content"]
108
            if url.startswith("/"):
109
                # The URL isn't fully defined yet. Looks something like ``self.work.url + /something``.
110
                entry["target"] = ""
111
            else:
112
                entry["target"] = url
113
            if url.startswith("http://") and is_running_in_cloud():
114
                warnings.warn(
115
                    f"You configured an http link {url[:32]}... but it won't be accessible in the cloud."
116
                    f" Consider replacing 'http' with 'https' in the link above."
117
                )
118

119
        elif isinstance(entry["content"], lightning.app.LightningFlow):
120
            entry["content"] = entry["content"].name
121

122
        elif isinstance(entry["content"], lightning.app.LightningWork):
123
            work = entry["content"]
124
            work_layout = _collect_work_layout(work)
125

126
            if work_layout is None:
127
                entry["content"] = ""
128
            elif isinstance(work_layout, str):
129
                entry["content"] = work_layout
130
                entry["target"] = work_layout
131
            elif isinstance(work_layout, (Frontend, _MagicMockJsonSerializable)):
132
                if len(layout) > 1:
133
                    lines = _add_comment_to_literal_code(
134
                        flow.configure_layout, contains="return", comment="  <------- this guy"
135
                    )
136
                    m = f"""
137
                    The return value of configure_layout() in `{flow.__class__.__name__}`  is an
138
                    unsupported format:
139
                    \n{lines}
140

141
                    The tab containing a `{work.__class__.__name__}` must be the only tab in the
142
                    layout of this flow.
143

144
                    (see the docs for `LightningWork.configure_layout`).
145
                    """
146
                    raise TypeError(m)
147

148
                if isinstance(work_layout, Frontend):
149
                    # If the work returned a frontend, treat it as belonging to the flow.
150
                    # NOTE: This could evolve in the future to run the Frontend directly in the work machine.
151
                    frontend = work_layout
152
                    frontend.flow = flow
153
                elif isinstance(work_layout, _MagicMockJsonSerializable):
154
                    # The import was mocked, we set a dummy `Frontend` so that `is_headless` knows there is a UI.
155
                    frontend = "mock"
156

157
                app.frontends.setdefault(flow.name, frontend)
158
                return flow._layout
159

160
        elif isinstance(entry["content"], _MagicMockJsonSerializable):
161
            # The import was mocked, we just record dummy content so that `is_headless` knows there is a UI
162
            entry["content"] = "mock"
163
            entry["target"] = "mock"
164
        else:
165
            m = f"""
166
            A dictionary returned by `{flow.__class__.__name__}.configure_layout()` contains an unsupported entry.
167

168
            {{'content': {repr(entry['content'])}}}
169

170
            Set the `content` key to a child flow or a URL, for example:
171

172
            class {flow.__class__.__name__}(LightningFlow):
173
                def configure_layout(self):
174
                    return {{'content': childFlow OR childWork OR 'http://some/url'}}
175
            """
176
            raise ValueError(m)
177
    return layout
178

179

180
def _collect_work_layout(work: "lightning.app.LightningWork") -> Union[None, str, Frontend, _MagicMockJsonSerializable]:
181
    """Check if ``configure_layout`` is overridden on the given work and return the work layout (either a string, a
182
    ``Frontend`` object, or an instance of a mocked import).
183

184
    Args:
185
        work: The work to collect the layout for.
186

187
    Raises:
188
        TypeError: If the value returned by ``configure_layout`` is not of a supported format.
189

190
    """
191
    work_layout = work.configure_layout() if is_overridden("configure_layout", work) else work.url
192

193
    if work_layout is None:
194
        return None
195
    if isinstance(work_layout, str):
196
        url = work_layout
197
        # The URL isn't fully defined yet. Looks something like ``self.work.url + /something``.
198
        if url and not url.startswith("/"):
199
            return url
200
        return ""
201
    if isinstance(work_layout, (Frontend, _MagicMockJsonSerializable)):
202
        return work_layout
203
    raise TypeError(
204
        f"""
205
    The value returned by `{work.__class__.__name__}.configure_layout()` is of an unsupported type.
206

207
    {repr(work_layout)}
208

209
    Return a `Frontend` or a URL string, for example:
210

211
    class {work.__class__.__name__}(LightningWork):
212
        def configure_layout(self):
213
            return MyFrontend() OR 'http://some/url'
214
    """
215
    )
216