pytorch-lightning
63 строки · 2.4 Кб
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
15import os
16from typing import Optional
17
18from lightning_cloud.openapi import V1Membership
19
20import lightning.app
21from lightning.app.core.constants import LIGHTNING_CLOUD_PROJECT_ID
22from lightning.app.utilities.enum import AppStage
23from lightning.app.utilities.network import LightningClient
24
25
26def _get_project(client: LightningClient, project_id: Optional[str] = None, verbose: bool = True) -> V1Membership:
27"""Get a project membership for the user from the backend."""
28if project_id is None:
29project_id = LIGHTNING_CLOUD_PROJECT_ID
30
31if project_id is not None:
32project = client.projects_service_get_project(project_id)
33if not project:
34raise ValueError(
35"Environment variable `LIGHTNING_CLOUD_PROJECT_ID` is set but could not find an associated project."
36)
37return V1Membership(
38name=project.name,
39display_name=project.display_name,
40description=project.description,
41created_at=project.created_at,
42project_id=project.id,
43owner_id=project.owner_id,
44owner_type=project.owner_type,
45quotas=project.quotas,
46updated_at=project.updated_at,
47)
48
49projects = client.projects_service_list_memberships()
50if len(projects.memberships) == 0:
51raise ValueError("No valid projects found. Please reach out to lightning.ai team to create a project")
52if len(projects.memberships) > 1 and verbose:
53print(f"Defaulting to the project: {projects.memberships[0].name}")
54return projects.memberships[0]
55
56
57def _sigterm_flow_handler(*_, app: "lightning.app.LightningApp"):
58app.stage = AppStage.STOPPING
59
60
61def is_running_in_cloud() -> bool:
62"""Returns True if the Lightning App is running in the cloud."""
63return bool(int(os.environ.get("LAI_RUNNING_IN_CLOUD", "0"))) or "LIGHTNING_APP_STATE_URL" in os.environ
64