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 os
16
import socket
17
from typing import Optional
18

19
from lightning_cloud.openapi import AppinstancesIdBody, Externalv1LightningappInstance, V1NetworkConfig
20

21
from lightning.app.utilities.network import LightningClient, find_free_network_port
22

23

24
def is_port_in_use(port: int) -> bool:
25
    """Checks if the given port is in use."""
26
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
27
        return s.connect_ex(("localhost", port)) == 0
28

29

30
def _find_lit_app_port(default_port: int) -> int:
31
    """Make a request to the cloud controlplane to find a disabled port of the flow, enable it and return it."""
32
    app_id = os.getenv("LIGHTNING_CLOUD_APP_ID", None)
33
    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID", None)
34
    enable_multiple_works_in_default_container = bool(int(os.getenv("ENABLE_MULTIPLE_WORKS_IN_DEFAULT_CONTAINER", "0")))
35

36
    if not app_id or not project_id or not enable_multiple_works_in_default_container:
37
        app_port = default_port
38

39
        # If the default port is not available, picks any other available one
40
        if is_port_in_use(default_port):
41
            app_port = find_free_network_port()
42

43
        return app_port
44

45
    client = LightningClient()
46
    list_apps_resp = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project_id)
47
    lit_app: Optional[Externalv1LightningappInstance] = None
48

49
    for lapp in list_apps_resp.lightningapps:
50
        if lapp.id == app_id:
51
            lit_app = lapp
52

53
    if not lit_app:
54
        raise RuntimeError(
55
            "App was not found. Please open an issue at https://github.com/lightning-AI/lightning/issues."
56
        )
57

58
    found_nc = None
59

60
    for nc in lit_app.spec.network_config:
61
        if not nc.enable:
62
            found_nc = nc
63
            nc.enable = True
64
            break
65

66
    client.lightningapp_instance_service_update_lightningapp_instance(
67
        project_id=project_id,
68
        id=lit_app.id,
69
        body=AppinstancesIdBody(name=lit_app.name, spec=lit_app.spec),
70
    )
71

72
    if not found_nc:
73
        raise RuntimeError(
74
            "No available port was found. Please open an issue at https://github.com/lightning-AI/lightning/issues."
75
        )
76

77
    # Note: This is required for the framework to know we need to use the CloudMultiProcessRuntime.
78
    os.environ["APP_SERVER_HOST"] = f"https://{found_nc.host}"
79

80
    return found_nc.port
81

82

83
def enable_port() -> V1NetworkConfig:
84
    """Make a request to the cloud controlplane to open a port of the flow."""
85
    app_id = os.getenv("LIGHTNING_CLOUD_APP_ID", None)
86
    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID", None)
87

88
    if not app_id or not project_id:
89
        raise Exception("The app_id and project_id should be defined.")
90

91
    client = LightningClient()
92
    list_apps_resp = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project_id)
93
    lit_app: Optional[Externalv1LightningappInstance] = None
94

95
    for lapp in list_apps_resp.lightningapps:
96
        if lapp.id == app_id:
97
            lit_app = lapp
98

99
    if not lit_app:
100
        raise RuntimeError(
101
            "App was not found. Please open an issue at https://github.com/lightning-AI/lightning/issues."
102
        )
103

104
    found_nc = None
105

106
    for nc in lit_app.spec.network_config:
107
        if not nc.enable:
108
            found_nc = nc
109
            nc.enable = True
110
            break
111

112
    client.lightningapp_instance_service_update_lightningapp_instance(
113
        project_id=project_id,
114
        id=lit_app.id,
115
        body=AppinstancesIdBody(name=lit_app.name, spec=lit_app.spec),
116
    )
117

118
    if not found_nc:
119
        raise RuntimeError(
120
            "No available port was found. Please open an issue at https://github.com/lightning-AI/lightning/issues."
121
        )
122

123
    return found_nc
124

125

126
def disable_port(port: int, ignore_disabled: bool = True) -> None:
127
    """Make a request to the cloud controlplane to close a port of the flow."""
128
    app_id = os.getenv("LIGHTNING_CLOUD_APP_ID", None)
129
    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID", None)
130

131
    if not app_id or not project_id:
132
        raise Exception("The app_id and project_id should be defined.")
133

134
    client = LightningClient()
135
    list_apps_resp = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project_id)
136
    lit_app: Optional[Externalv1LightningappInstance] = None
137

138
    for lapp in list_apps_resp.lightningapps:
139
        if lapp.id == app_id:
140
            lit_app = lapp
141

142
    if not lit_app:
143
        raise RuntimeError(
144
            "App was not found. Please open an issue at https://github.com/lightning-AI/lightning/issues."
145
        )
146

147
    found_nc = None
148

149
    for nc in lit_app.spec.network_config:
150
        if nc.port == port:
151
            if not nc.enable and not ignore_disabled:
152
                raise RuntimeError(f"The port {port} was already disabled.")
153

154
            nc.enable = False
155
            found_nc = nc
156
            break
157

158
    client.lightningapp_instance_service_update_lightningapp_instance(
159
        project_id=project_id,
160
        id=lit_app.id,
161
        body=AppinstancesIdBody(name=lit_app.name, spec=lit_app.spec),
162
    )
163

164
    if not found_nc:
165
        ports = [nc.port for nc in lit_app.spec.network_config]
166
        raise ValueError(f"The provided port doesn't exists. Available ports are {ports}.")
167

168
    assert found_nc
169