lightning-sdk 2025.8.14.post0__py3-none-any.whl → 2025.8.18.post0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- lightning_sdk/__init__.py +1 -1
- lightning_sdk/api/cloud_account_api.py +5 -0
- lightning_sdk/cli/__init__.py +1 -0
- lightning_sdk/cli/config/__init__.py +14 -0
- lightning_sdk/cli/config/get.py +57 -0
- lightning_sdk/cli/config/set.py +104 -0
- lightning_sdk/cli/config/show.py +9 -0
- lightning_sdk/cli/entrypoint.py +60 -41
- lightning_sdk/cli/groups.py +35 -0
- lightning_sdk/cli/job/__init__.py +7 -0
- lightning_sdk/cli/{configure.py → legacy/configure.py} +2 -2
- lightning_sdk/cli/{connect.py → legacy/connect.py} +2 -2
- lightning_sdk/cli/{create.py → legacy/create.py} +1 -1
- lightning_sdk/cli/{delete.py → legacy/delete.py} +3 -3
- lightning_sdk/cli/legacy/deploy/__init__.py +0 -0
- lightning_sdk/cli/{deploy → legacy/deploy}/_auth.py +1 -1
- lightning_sdk/cli/{deploy → legacy/deploy}/devbox.py +8 -2
- lightning_sdk/cli/{deploy → legacy/deploy}/serve.py +3 -3
- lightning_sdk/cli/{download.py → legacy/download.py} +3 -3
- lightning_sdk/cli/legacy/entrypoint.py +110 -0
- lightning_sdk/cli/{generate.py → legacy/generate.py} +1 -1
- lightning_sdk/cli/{inspection.py → legacy/inspection.py} +1 -1
- lightning_sdk/cli/{job_and_mmt_action.py → legacy/job_and_mmt_action.py} +3 -3
- lightning_sdk/cli/{jobs_menu.py → legacy/jobs_menu.py} +1 -1
- lightning_sdk/cli/{list.py → legacy/list.py} +2 -2
- lightning_sdk/cli/{mmts_menu.py → legacy/mmts_menu.py} +1 -1
- lightning_sdk/cli/{open.py → legacy/open.py} +2 -2
- lightning_sdk/cli/{stop.py → legacy/stop.py} +1 -1
- lightning_sdk/cli/{teamspace_menu.py → legacy/teamspace_menu.py} +1 -1
- lightning_sdk/cli/{upload.py → legacy/upload.py} +3 -3
- lightning_sdk/cli/mmt/__init__.py +7 -0
- lightning_sdk/cli/studio/__init__.py +22 -0
- lightning_sdk/cli/studio/create.py +68 -0
- lightning_sdk/cli/studio/delete.py +44 -0
- lightning_sdk/cli/studio/list.py +67 -0
- lightning_sdk/cli/studio/ssh.py +112 -0
- lightning_sdk/cli/studio/start.py +81 -0
- lightning_sdk/cli/studio/stop.py +44 -0
- lightning_sdk/cli/studio/switch.py +52 -0
- lightning_sdk/cli/utils/__init__.py +7 -0
- lightning_sdk/cli/utils/cloud_account_map.py +10 -0
- lightning_sdk/cli/utils/resolve.py +28 -0
- lightning_sdk/cli/utils/richt_print.py +11 -0
- lightning_sdk/lightning_cloud/openapi/api/billing_service_api.py +5 -1
- lightning_sdk/lightning_cloud/openapi/api/k8_s_cluster_service_api.py +117 -0
- lightning_sdk/lightning_cloud/openapi/models/v1_managed_model.py +29 -3
- lightning_sdk/lightning_cloud/openapi/models/v1_notification_type.py +1 -0
- lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py +1 -27
- lightning_sdk/lightning_cloud/utils/data_connection.py +51 -1
- lightning_sdk/studio.py +17 -5
- lightning_sdk/teamspace.py +14 -0
- lightning_sdk/utils/config.py +158 -0
- lightning_sdk/utils/resolve.py +49 -3
- {lightning_sdk-2025.8.14.post0.dist-info → lightning_sdk-2025.8.18.post0.dist-info}/METADATA +2 -1
- {lightning_sdk-2025.8.14.post0.dist-info → lightning_sdk-2025.8.18.post0.dist-info}/RECORD +69 -47
- /lightning_sdk/cli/{deploy → legacy}/__init__.py +0 -0
- /lightning_sdk/cli/{ai_hub.py → legacy/ai_hub.py} +0 -0
- /lightning_sdk/cli/{clusters_menu.py → legacy/clusters_menu.py} +0 -0
- /lightning_sdk/cli/{docker_cli.py → legacy/docker_cli.py} +0 -0
- /lightning_sdk/cli/{exceptions.py → legacy/exceptions.py} +0 -0
- /lightning_sdk/cli/{run.py → legacy/run.py} +0 -0
- /lightning_sdk/cli/{start.py → legacy/start.py} +0 -0
- /lightning_sdk/cli/{studios_menu.py → legacy/studios_menu.py} +0 -0
- /lightning_sdk/cli/{switch.py → legacy/switch.py} +0 -0
- /lightning_sdk/cli/{coloring.py → utils/coloring.py} +0 -0
- {lightning_sdk-2025.8.14.post0.dist-info → lightning_sdk-2025.8.18.post0.dist-info}/LICENSE +0 -0
- {lightning_sdk-2025.8.14.post0.dist-info → lightning_sdk-2025.8.18.post0.dist-info}/WHEEL +0 -0
- {lightning_sdk-2025.8.14.post0.dist-info → lightning_sdk-2025.8.18.post0.dist-info}/entry_points.txt +0 -0
- {lightning_sdk-2025.8.14.post0.dist-info → lightning_sdk-2025.8.18.post0.dist-info}/top_level.txt +0 -0
|
@@ -7,8 +7,8 @@ from rich.table import Table
|
|
|
7
7
|
from typing_extensions import Literal
|
|
8
8
|
|
|
9
9
|
from lightning_sdk import Job, Machine, Studio, Teamspace
|
|
10
|
-
from lightning_sdk.cli.clusters_menu import _ClustersMenu
|
|
11
|
-
from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
|
|
10
|
+
from lightning_sdk.cli.legacy.clusters_menu import _ClustersMenu
|
|
11
|
+
from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
|
|
12
12
|
from lightning_sdk.lightning_cloud.openapi import V1MultiMachineJob
|
|
13
13
|
from lightning_sdk.lit_container import LitContainer
|
|
14
14
|
from lightning_sdk.utils.resolve import _get_authed_user
|
|
@@ -3,7 +3,7 @@ from typing import Dict, List, Optional
|
|
|
3
3
|
from rich.console import Console
|
|
4
4
|
from simple_term_menu import TerminalMenu
|
|
5
5
|
|
|
6
|
-
from lightning_sdk.cli.exceptions import StudioCliError
|
|
6
|
+
from lightning_sdk.cli.legacy.exceptions import StudioCliError
|
|
7
7
|
from lightning_sdk.mmt import MMT
|
|
8
8
|
from lightning_sdk.teamspace import Teamspace
|
|
9
9
|
|
|
@@ -6,8 +6,8 @@ from typing import Optional
|
|
|
6
6
|
import click
|
|
7
7
|
from rich.console import Console
|
|
8
8
|
|
|
9
|
-
from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
|
|
10
|
-
from lightning_sdk.cli.upload import _upload_folder
|
|
9
|
+
from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
|
|
10
|
+
from lightning_sdk.cli.legacy.upload import _upload_folder
|
|
11
11
|
from lightning_sdk.studio import Studio
|
|
12
12
|
from lightning_sdk.teamspace import Teamspace
|
|
13
13
|
from lightning_sdk.utils.resolve import _get_studio_url
|
|
@@ -3,7 +3,7 @@ from typing import Optional
|
|
|
3
3
|
import click
|
|
4
4
|
from rich.console import Console
|
|
5
5
|
|
|
6
|
-
from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
|
|
6
|
+
from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction
|
|
7
7
|
from lightning_sdk.lightning_cloud.openapi.rest import ApiException
|
|
8
8
|
from lightning_sdk.studio import Studio
|
|
9
9
|
|
|
@@ -4,7 +4,7 @@ from rich.console import Console
|
|
|
4
4
|
from simple_term_menu import TerminalMenu
|
|
5
5
|
|
|
6
6
|
from lightning_sdk.api import OrgApi
|
|
7
|
-
from lightning_sdk.cli.exceptions import StudioCliError
|
|
7
|
+
from lightning_sdk.cli.legacy.exceptions import StudioCliError
|
|
8
8
|
from lightning_sdk.teamspace import Teamspace
|
|
9
9
|
from lightning_sdk.user import User
|
|
10
10
|
from lightning_sdk.utils.resolve import _get_authed_user
|
|
@@ -14,9 +14,9 @@ from tqdm import tqdm
|
|
|
14
14
|
|
|
15
15
|
from lightning_sdk.api.lit_container_api import DockerNotRunningError, LCRAuthFailedError, LitContainerApi
|
|
16
16
|
from lightning_sdk.api.utils import _get_cloud_url
|
|
17
|
-
from lightning_sdk.cli.exceptions import StudioCliError
|
|
18
|
-
from lightning_sdk.cli.studios_menu import _StudiosMenu
|
|
19
|
-
from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
|
|
17
|
+
from lightning_sdk.cli.legacy.exceptions import StudioCliError
|
|
18
|
+
from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
|
|
19
|
+
from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
|
|
20
20
|
from lightning_sdk.constants import _LIGHTNING_DEBUG
|
|
21
21
|
from lightning_sdk.models import upload_model as _upload_model
|
|
22
22
|
from lightning_sdk.studio import Studio
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
"""Studio CLI commands."""
|
|
2
|
+
|
|
3
|
+
import click
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def register_commands(group: click.Group) -> None:
|
|
7
|
+
"""Register studio commands with the given group."""
|
|
8
|
+
from lightning_sdk.cli.studio.create import create_studio
|
|
9
|
+
from lightning_sdk.cli.studio.delete import delete_studio
|
|
10
|
+
from lightning_sdk.cli.studio.list import list_studios
|
|
11
|
+
from lightning_sdk.cli.studio.ssh import ssh_studio
|
|
12
|
+
from lightning_sdk.cli.studio.start import start_studio
|
|
13
|
+
from lightning_sdk.cli.studio.stop import stop_studio
|
|
14
|
+
from lightning_sdk.cli.studio.switch import switch_studio
|
|
15
|
+
|
|
16
|
+
group.add_command(delete_studio)
|
|
17
|
+
group.add_command(create_studio)
|
|
18
|
+
group.add_command(list_studios)
|
|
19
|
+
group.add_command(ssh_studio)
|
|
20
|
+
group.add_command(start_studio)
|
|
21
|
+
group.add_command(stop_studio)
|
|
22
|
+
group.add_command(switch_studio)
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""Studio start command."""
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import click
|
|
6
|
+
|
|
7
|
+
from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
|
|
8
|
+
from lightning_sdk.lightning_cloud.openapi.rest import ApiException
|
|
9
|
+
from lightning_sdk.machine import CloudProvider
|
|
10
|
+
from lightning_sdk.studio import Studio
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@click.command("create")
|
|
14
|
+
@click.argument("studio_name", required=False)
|
|
15
|
+
@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
|
|
16
|
+
@click.option(
|
|
17
|
+
"--cloud-provider",
|
|
18
|
+
help="The cloud provider to start the studio on. Defaults to teamspace default.",
|
|
19
|
+
type=click.Choice(m.name for m in list(CloudProvider)),
|
|
20
|
+
)
|
|
21
|
+
@click.option(
|
|
22
|
+
"--cloud-account",
|
|
23
|
+
help="The cloud account to create the studio on. Defaults to teamspace default.",
|
|
24
|
+
type=click.STRING,
|
|
25
|
+
)
|
|
26
|
+
def create_studio(
|
|
27
|
+
studio_name: Optional[str] = None,
|
|
28
|
+
teamspace: Optional[str] = None,
|
|
29
|
+
cloud_provider: Optional[str] = None,
|
|
30
|
+
cloud_account: Optional[str] = None,
|
|
31
|
+
) -> None:
|
|
32
|
+
"""Create a new Studio.
|
|
33
|
+
|
|
34
|
+
Example:
|
|
35
|
+
lightning studio create [STUDIO_NAME]
|
|
36
|
+
|
|
37
|
+
STUDIO_NAME: the name of the studio to create.
|
|
38
|
+
|
|
39
|
+
If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
|
|
40
|
+
"""
|
|
41
|
+
if teamspace is not None:
|
|
42
|
+
resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
|
|
43
|
+
if resolved_teamspace is None:
|
|
44
|
+
raise ValueError(
|
|
45
|
+
f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
|
|
46
|
+
"Does the teamspace exist?"
|
|
47
|
+
)
|
|
48
|
+
else:
|
|
49
|
+
resolved_teamspace = None
|
|
50
|
+
|
|
51
|
+
if cloud_provider is not None:
|
|
52
|
+
cloud_provider = CloudProvider(cloud_provider)
|
|
53
|
+
|
|
54
|
+
try:
|
|
55
|
+
studio = Studio(
|
|
56
|
+
studio_name,
|
|
57
|
+
teamspace=resolved_teamspace,
|
|
58
|
+
create_ok=True,
|
|
59
|
+
cloud_provider=cloud_provider,
|
|
60
|
+
cloud_account=cloud_account,
|
|
61
|
+
)
|
|
62
|
+
except (RuntimeError, ValueError, ApiException) as e:
|
|
63
|
+
print(e)
|
|
64
|
+
if studio_name:
|
|
65
|
+
raise ValueError(f"Could not create Studio: '{studio_name}'. Does the Studio exist?") from None
|
|
66
|
+
raise ValueError(f"Could not create Studio: '{studio_name}'. Please provide a Studio name") from None
|
|
67
|
+
|
|
68
|
+
click.echo(f"Studio '{studio.name}' created successfully")
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
"""Studio delete command."""
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import click
|
|
6
|
+
|
|
7
|
+
from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
|
|
8
|
+
from lightning_sdk.studio import Studio
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@click.command("delete")
|
|
12
|
+
@click.argument("studio_name", required=False)
|
|
13
|
+
@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
|
|
14
|
+
def delete_studio(studio_name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
|
|
15
|
+
"""Delete a Studio.
|
|
16
|
+
|
|
17
|
+
Example:
|
|
18
|
+
lightning studio delete [STUDIO_NAME]
|
|
19
|
+
|
|
20
|
+
STUDIO_NAME: the name of the studio to delete.
|
|
21
|
+
|
|
22
|
+
If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
|
|
23
|
+
"""
|
|
24
|
+
# missing studio_name and teamspace are handled by the studio class
|
|
25
|
+
if teamspace is not None:
|
|
26
|
+
resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
|
|
27
|
+
if resolved_teamspace is None:
|
|
28
|
+
raise ValueError(
|
|
29
|
+
f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
|
|
30
|
+
"Does the teamspace exist?"
|
|
31
|
+
)
|
|
32
|
+
else:
|
|
33
|
+
resolved_teamspace = None
|
|
34
|
+
|
|
35
|
+
try:
|
|
36
|
+
studio = Studio(studio_name, teamspace=resolved_teamspace, create_ok=False)
|
|
37
|
+
studio.delete()
|
|
38
|
+
except Exception:
|
|
39
|
+
# TODO: make this a generic CLI error
|
|
40
|
+
if studio_name:
|
|
41
|
+
raise ValueError(f"Could not delete Studio: '{studio_name}'. Does the Studio exist?") from None
|
|
42
|
+
raise ValueError("No studio name provided. Use 'lightning studio delete <name>' to delete a studio.") from None
|
|
43
|
+
|
|
44
|
+
click.echo(f"Studio '{studio.name}' deleted successfully")
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
"""Studio list command."""
|
|
2
|
+
|
|
3
|
+
from typing import Callable, Optional
|
|
4
|
+
|
|
5
|
+
import click
|
|
6
|
+
from rich.table import Table
|
|
7
|
+
|
|
8
|
+
from lightning_sdk.cli.utils.cloud_account_map import cloud_account_to_display_name
|
|
9
|
+
from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
|
|
10
|
+
from lightning_sdk.cli.utils.richt_print import rich_to_str
|
|
11
|
+
from lightning_sdk.studio import Studio
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@click.command("list")
|
|
15
|
+
@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
|
|
16
|
+
@click.option(
|
|
17
|
+
"--sort-by",
|
|
18
|
+
default=None,
|
|
19
|
+
type=click.Choice(["name", "teamspace", "status", "machine", "cloud-account"], case_sensitive=False),
|
|
20
|
+
help="the attribute to sort the studios by.",
|
|
21
|
+
)
|
|
22
|
+
def list_studios(teamspace: Optional[str] = None, sort_by: Optional[str] = None) -> None:
|
|
23
|
+
"""List Studios in a teamspace.
|
|
24
|
+
|
|
25
|
+
Example:
|
|
26
|
+
lightning studio list --teamspace owner/teamspace
|
|
27
|
+
|
|
28
|
+
"""
|
|
29
|
+
teamspace_resolved = resolve_teamspace_owner_name_format(teamspace)
|
|
30
|
+
|
|
31
|
+
if teamspace_resolved is None:
|
|
32
|
+
# TODO: make this a generic CLI error
|
|
33
|
+
raise ValueError(f"Could not resolve teamspace: {teamspace}")
|
|
34
|
+
|
|
35
|
+
studios = teamspace_resolved.studios
|
|
36
|
+
|
|
37
|
+
table = Table(
|
|
38
|
+
pad_edge=True,
|
|
39
|
+
)
|
|
40
|
+
table.add_column("Name")
|
|
41
|
+
table.add_column("Teamspace")
|
|
42
|
+
table.add_column("Status")
|
|
43
|
+
table.add_column("Machine")
|
|
44
|
+
table.add_column("Cloud account")
|
|
45
|
+
|
|
46
|
+
for studio in sorted(studios, key=_sort_studios_key(sort_by)):
|
|
47
|
+
table.add_row(
|
|
48
|
+
studio.name,
|
|
49
|
+
f"{studio.teamspace.owner.name}/{studio.teamspace.name}",
|
|
50
|
+
str(studio.status),
|
|
51
|
+
str(studio.machine) if studio.machine is not None else None, # when None the cell is empty
|
|
52
|
+
str(cloud_account_to_display_name(studio.cloud_account, studio.teamspace.id)),
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
click.echo(rich_to_str(table), color=True)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def _sort_studios_key(sort_by: str) -> Callable[[Studio], str]:
|
|
59
|
+
"""Return a key function to sort studios by a given attribute."""
|
|
60
|
+
sort_key_map = {
|
|
61
|
+
"name": lambda s: str(s.name or ""),
|
|
62
|
+
"teamspace": lambda s: str(s.teamspace.name or ""),
|
|
63
|
+
"status": lambda s: str(s.status or ""),
|
|
64
|
+
"machine": lambda s: str(s.machine or ""),
|
|
65
|
+
"cloud-account": lambda s: str(cloud_account_to_display_name(s.cloud_account or "", s.teamspace.id)),
|
|
66
|
+
}
|
|
67
|
+
return sort_key_map.get(sort_by, lambda s: s.name)
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
"""Studio SSH command."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import platform
|
|
5
|
+
import subprocess
|
|
6
|
+
import uuid
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from typing import List, Optional
|
|
9
|
+
|
|
10
|
+
import click
|
|
11
|
+
|
|
12
|
+
from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
|
|
13
|
+
from lightning_sdk.lightning_cloud.login import Auth
|
|
14
|
+
from lightning_sdk.studio import Studio
|
|
15
|
+
from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@click.command("ssh")
|
|
19
|
+
@click.argument("studio_name", required=False)
|
|
20
|
+
@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)", type=click.STRING)
|
|
21
|
+
@click.option(
|
|
22
|
+
"--option",
|
|
23
|
+
"-o",
|
|
24
|
+
help="Additional options to pass to the SSH command. Can be specified multiple times.",
|
|
25
|
+
multiple=True,
|
|
26
|
+
type=click.STRING,
|
|
27
|
+
)
|
|
28
|
+
def ssh_studio(
|
|
29
|
+
studio_name: Optional[str] = None, teamspace: Optional[str] = None, option: Optional[List[str]] = None
|
|
30
|
+
) -> None:
|
|
31
|
+
"""SSH into a Studio.
|
|
32
|
+
|
|
33
|
+
Example:
|
|
34
|
+
lightning studio ssh [STUDIO_NAME]
|
|
35
|
+
|
|
36
|
+
STUDIO_NAME: the name of the studio to SSH into.
|
|
37
|
+
|
|
38
|
+
If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
|
|
39
|
+
"""
|
|
40
|
+
auth = Auth()
|
|
41
|
+
auth.authenticate()
|
|
42
|
+
ssh_private_key_path = _download_ssh_keys(auth.api_key, force_download=False)
|
|
43
|
+
|
|
44
|
+
if teamspace is not None:
|
|
45
|
+
resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
|
|
46
|
+
if resolved_teamspace is None:
|
|
47
|
+
raise ValueError(
|
|
48
|
+
f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
|
|
49
|
+
"Does the teamspace exist?"
|
|
50
|
+
)
|
|
51
|
+
else:
|
|
52
|
+
resolved_teamspace = None
|
|
53
|
+
|
|
54
|
+
studio = Studio(studio_name, teamspace=resolved_teamspace)
|
|
55
|
+
|
|
56
|
+
ssh_options = " -o " + " -o ".join(option) if option else ""
|
|
57
|
+
ssh_command = f"ssh -i {ssh_private_key_path}{ssh_options} s_{studio._studio.id}@ssh.lightning.ai"
|
|
58
|
+
|
|
59
|
+
try:
|
|
60
|
+
subprocess.run(ssh_command.split())
|
|
61
|
+
except Exception:
|
|
62
|
+
# redownload the keys to be sure they are up to date
|
|
63
|
+
_download_ssh_keys(auth.api_key, force_download=True)
|
|
64
|
+
try:
|
|
65
|
+
subprocess.run(ssh_command.split())
|
|
66
|
+
except Exception:
|
|
67
|
+
# TODO: make this a generic CLI error
|
|
68
|
+
raise RuntimeError("Failed to establish SSH connection") from None
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def _download_ssh_keys(
|
|
72
|
+
api_key: str,
|
|
73
|
+
force_download: bool = False,
|
|
74
|
+
ssh_key_name: str = "lightning_rsa",
|
|
75
|
+
) -> None:
|
|
76
|
+
"""Download the SSH key for a User."""
|
|
77
|
+
ssh_private_key_path = os.path.join(os.path.expanduser(os.path.dirname(_DEFAULT_CONFIG_FILE_PATH)), ssh_key_name)
|
|
78
|
+
|
|
79
|
+
os.makedirs(os.path.dirname(ssh_private_key_path), exist_ok=True)
|
|
80
|
+
|
|
81
|
+
if not os.path.isfile(ssh_private_key_path) or force_download:
|
|
82
|
+
key_id = str(uuid.uuid4())
|
|
83
|
+
_download_file(
|
|
84
|
+
f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
|
|
85
|
+
ssh_private_key_path,
|
|
86
|
+
overwrite=True,
|
|
87
|
+
chmod=0o600,
|
|
88
|
+
)
|
|
89
|
+
_download_file(
|
|
90
|
+
f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}",
|
|
91
|
+
ssh_private_key_path + ".pub",
|
|
92
|
+
overwrite=True,
|
|
93
|
+
)
|
|
94
|
+
|
|
95
|
+
return ssh_private_key_path
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
|
|
99
|
+
"""Download a file from a URL."""
|
|
100
|
+
import requests
|
|
101
|
+
|
|
102
|
+
if os.path.isfile(local_path) and not overwrite:
|
|
103
|
+
raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.")
|
|
104
|
+
|
|
105
|
+
response = requests.get(url, stream=True)
|
|
106
|
+
response.raise_for_status()
|
|
107
|
+
|
|
108
|
+
with open(local_path, "wb") as file:
|
|
109
|
+
for chunk in response.iter_content(chunk_size=8192):
|
|
110
|
+
file.write(chunk)
|
|
111
|
+
if chmod is not None:
|
|
112
|
+
os.chmod(local_path, 0o600)
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
"""Studio start command."""
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import click
|
|
6
|
+
|
|
7
|
+
from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
|
|
8
|
+
from lightning_sdk.lightning_cloud.openapi.rest import ApiException
|
|
9
|
+
from lightning_sdk.machine import CloudProvider, Machine
|
|
10
|
+
from lightning_sdk.studio import Studio
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@click.command("start")
|
|
14
|
+
@click.argument("studio_name", required=False)
|
|
15
|
+
@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
|
|
16
|
+
@click.option("--create", is_flag=True, help="Create the studio if it doesn't exist")
|
|
17
|
+
@click.option(
|
|
18
|
+
"--machine",
|
|
19
|
+
help="The machine type to start the studio on. Defaults to CPU-4",
|
|
20
|
+
type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli),
|
|
21
|
+
)
|
|
22
|
+
@click.option("--interruptible", is_flag=True, help="Start the studio on an interruptible instance.")
|
|
23
|
+
@click.option(
|
|
24
|
+
"--cloud-provider",
|
|
25
|
+
help=(
|
|
26
|
+
"The cloud provider to start the studio on. Defaults to teamspace default. "
|
|
27
|
+
"Only used if --create is specified."
|
|
28
|
+
),
|
|
29
|
+
type=click.Choice(m.name for m in list(CloudProvider)),
|
|
30
|
+
)
|
|
31
|
+
@click.option(
|
|
32
|
+
"--cloud-account",
|
|
33
|
+
help="The cloud account to start the studio on. Defaults to teamspace default. Only used if --create is specified.",
|
|
34
|
+
type=click.STRING,
|
|
35
|
+
)
|
|
36
|
+
def start_studio(
|
|
37
|
+
studio_name: Optional[str] = None,
|
|
38
|
+
teamspace: Optional[str] = None,
|
|
39
|
+
create: bool = False,
|
|
40
|
+
machine: Optional[str] = None,
|
|
41
|
+
interruptible: bool = False,
|
|
42
|
+
cloud_provider: Optional[str] = None,
|
|
43
|
+
cloud_account: Optional[str] = None,
|
|
44
|
+
) -> None:
|
|
45
|
+
"""Start a Studio.
|
|
46
|
+
|
|
47
|
+
Example:
|
|
48
|
+
lightning studio start [STUDIO_NAME]
|
|
49
|
+
|
|
50
|
+
STUDIO_NAME: the name of the studio to start.
|
|
51
|
+
|
|
52
|
+
If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
|
|
53
|
+
"""
|
|
54
|
+
if teamspace is not None:
|
|
55
|
+
resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
|
|
56
|
+
if resolved_teamspace is None:
|
|
57
|
+
raise ValueError(
|
|
58
|
+
f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
|
|
59
|
+
"Does the teamspace exist?"
|
|
60
|
+
)
|
|
61
|
+
else:
|
|
62
|
+
resolved_teamspace = None
|
|
63
|
+
|
|
64
|
+
if cloud_provider is not None:
|
|
65
|
+
cloud_provider = CloudProvider(cloud_provider)
|
|
66
|
+
|
|
67
|
+
try:
|
|
68
|
+
studio = Studio(
|
|
69
|
+
studio_name,
|
|
70
|
+
teamspace=resolved_teamspace,
|
|
71
|
+
create_ok=create,
|
|
72
|
+
cloud_provider=cloud_provider,
|
|
73
|
+
cloud_account=cloud_account,
|
|
74
|
+
)
|
|
75
|
+
except (RuntimeError, ValueError, ApiException):
|
|
76
|
+
if studio_name:
|
|
77
|
+
raise ValueError(f"Could not start Studio: '{studio_name}'. Does the Studio exist?") from None
|
|
78
|
+
raise ValueError(f"Could not start Studio: '{studio_name}'. Please provide a Studio name") from None
|
|
79
|
+
|
|
80
|
+
studio.start(machine, interruptible=interruptible)
|
|
81
|
+
click.echo(f"Studio '{studio.name}' started successfully")
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
"""Studio stop command."""
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import click
|
|
6
|
+
|
|
7
|
+
from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
|
|
8
|
+
from lightning_sdk.studio import Studio
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@click.command("stop")
|
|
12
|
+
@click.argument("studio_name", required=False)
|
|
13
|
+
@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
|
|
14
|
+
def stop_studio(studio_name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
|
|
15
|
+
"""Stop a Studio.
|
|
16
|
+
|
|
17
|
+
Example:
|
|
18
|
+
lightning studio stop [STUDIO_NAME]
|
|
19
|
+
|
|
20
|
+
STUDIO_NAME: the name of the studio to stop.
|
|
21
|
+
|
|
22
|
+
If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
|
|
23
|
+
"""
|
|
24
|
+
# missing studio_name and teamspace are handled by the studio class
|
|
25
|
+
if teamspace is not None:
|
|
26
|
+
resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
|
|
27
|
+
if resolved_teamspace is None:
|
|
28
|
+
raise ValueError(
|
|
29
|
+
f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
|
|
30
|
+
"Does the teamspace exist?"
|
|
31
|
+
)
|
|
32
|
+
else:
|
|
33
|
+
resolved_teamspace = None
|
|
34
|
+
|
|
35
|
+
try:
|
|
36
|
+
studio = Studio(studio_name, teamspace=resolved_teamspace)
|
|
37
|
+
studio.stop()
|
|
38
|
+
except Exception:
|
|
39
|
+
# TODO: make this a generic CLI error
|
|
40
|
+
if studio_name:
|
|
41
|
+
raise ValueError(f"Could not stop studio: '{studio_name}'. Does the studio exist?") from None
|
|
42
|
+
raise ValueError("No studio name provided. Use 'lightning studio stop <name>' to stop a studio.") from None
|
|
43
|
+
|
|
44
|
+
click.echo(f"Studio '{studio.name}' stopped successfully")
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
"""Studio switch command."""
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import click
|
|
6
|
+
|
|
7
|
+
from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
|
|
8
|
+
from lightning_sdk.lightning_cloud.openapi.rest import ApiException
|
|
9
|
+
from lightning_sdk.machine import Machine
|
|
10
|
+
from lightning_sdk.studio import Studio
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@click.command("switch")
|
|
14
|
+
@click.argument("studio_name", required=False)
|
|
15
|
+
@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
|
|
16
|
+
@click.option(
|
|
17
|
+
"--machine",
|
|
18
|
+
help="The machine type to switch the studio to.",
|
|
19
|
+
type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine)),
|
|
20
|
+
)
|
|
21
|
+
@click.option("--interruptible", is_flag=True, help="Switch the studio to an interruptible instance.")
|
|
22
|
+
def switch_studio(
|
|
23
|
+
studio_name: Optional[str] = None,
|
|
24
|
+
teamspace: Optional[str] = None,
|
|
25
|
+
machine: Optional[str] = None,
|
|
26
|
+
interruptible: bool = False,
|
|
27
|
+
) -> None:
|
|
28
|
+
"""Switch a Studio to a different machine type."""
|
|
29
|
+
if teamspace is not None:
|
|
30
|
+
resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
|
|
31
|
+
if resolved_teamspace is None:
|
|
32
|
+
raise ValueError(
|
|
33
|
+
f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
|
|
34
|
+
"Does the teamspace exist?"
|
|
35
|
+
)
|
|
36
|
+
else:
|
|
37
|
+
resolved_teamspace = None
|
|
38
|
+
|
|
39
|
+
try:
|
|
40
|
+
studio = Studio(
|
|
41
|
+
studio_name,
|
|
42
|
+
teamspace=resolved_teamspace,
|
|
43
|
+
)
|
|
44
|
+
except (RuntimeError, ValueError, ApiException):
|
|
45
|
+
if studio_name:
|
|
46
|
+
raise ValueError(f"Could not switch Studio: '{studio_name}'. Does the Studio exist?") from None
|
|
47
|
+
raise ValueError(f"Could not switch Studio: '{studio_name}'. Please provide a Studio name") from None
|
|
48
|
+
|
|
49
|
+
resolved_machine = Machine.from_str(machine)
|
|
50
|
+
studio.switch_machine(resolved_machine, interruptible=interruptible)
|
|
51
|
+
|
|
52
|
+
click.echo(f"Studio '{studio.name}' switched to machine '{resolved_machine}' successfully")
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
from lightning_sdk.api.cloud_account_api import CloudAccountApi
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def cloud_account_to_display_name(cloud_account: str, teamspace_id: str) -> str:
|
|
5
|
+
api = CloudAccountApi()
|
|
6
|
+
cloud_accounts = api.list_global_cloud_accounts(teamspace_id=teamspace_id)
|
|
7
|
+
for global_cloud_account in cloud_accounts:
|
|
8
|
+
if global_cloud_account.id == cloud_account:
|
|
9
|
+
return "Lightning AI"
|
|
10
|
+
return cloud_account
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
from typing import Optional
|
|
2
|
+
|
|
3
|
+
from lightning_sdk.teamspace import Teamspace
|
|
4
|
+
from lightning_sdk.utils.resolve import _resolve_teamspace
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def resolve_teamspace_owner_name_format(teamspace_name: str) -> Optional[Teamspace]:
|
|
8
|
+
teamspace_resolved = None
|
|
9
|
+
if teamspace_name is None:
|
|
10
|
+
return _resolve_teamspace(None, None, None)
|
|
11
|
+
|
|
12
|
+
splits = teamspace_name.split("/")
|
|
13
|
+
if len(splits) == 1:
|
|
14
|
+
try:
|
|
15
|
+
teamspace_resolved = _resolve_teamspace(teamspace_name, None, None)
|
|
16
|
+
except Exception:
|
|
17
|
+
teamspace_resolved = None
|
|
18
|
+
|
|
19
|
+
elif len(splits) == 2:
|
|
20
|
+
try:
|
|
21
|
+
try:
|
|
22
|
+
teamspace_resolved = _resolve_teamspace(splits[1], splits[0], None)
|
|
23
|
+
except Exception:
|
|
24
|
+
teamspace_resolved = _resolve_teamspace(splits[1], None, splits[0])
|
|
25
|
+
except Exception:
|
|
26
|
+
teamspace_resolved = None
|
|
27
|
+
|
|
28
|
+
return teamspace_resolved
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from typing import Any
|
|
3
|
+
|
|
4
|
+
from rich.console import Console
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def rich_to_str(*renderables: Any) -> str:
|
|
8
|
+
with open(os.devnull, "w") as f:
|
|
9
|
+
console = Console(file=f, record=True)
|
|
10
|
+
console.print(*renderables)
|
|
11
|
+
return console.export_text(styles=True)
|
|
@@ -663,6 +663,7 @@ class BillingServiceApi(object):
|
|
|
663
663
|
|
|
664
664
|
:param async_req bool
|
|
665
665
|
:param str org_id:
|
|
666
|
+
:param str project_id:
|
|
666
667
|
:return: V1BillingSubscription
|
|
667
668
|
If the method is called asynchronously,
|
|
668
669
|
returns the request thread.
|
|
@@ -684,12 +685,13 @@ class BillingServiceApi(object):
|
|
|
684
685
|
|
|
685
686
|
:param async_req bool
|
|
686
687
|
:param str org_id:
|
|
688
|
+
:param str project_id:
|
|
687
689
|
:return: V1BillingSubscription
|
|
688
690
|
If the method is called asynchronously,
|
|
689
691
|
returns the request thread.
|
|
690
692
|
"""
|
|
691
693
|
|
|
692
|
-
all_params = ['org_id'] # noqa: E501
|
|
694
|
+
all_params = ['org_id', 'project_id'] # noqa: E501
|
|
693
695
|
all_params.append('async_req')
|
|
694
696
|
all_params.append('_return_http_data_only')
|
|
695
697
|
all_params.append('_preload_content')
|
|
@@ -712,6 +714,8 @@ class BillingServiceApi(object):
|
|
|
712
714
|
query_params = []
|
|
713
715
|
if 'org_id' in params:
|
|
714
716
|
query_params.append(('orgId', params['org_id'])) # noqa: E501
|
|
717
|
+
if 'project_id' in params:
|
|
718
|
+
query_params.append(('projectId', params['project_id'])) # noqa: E501
|
|
715
719
|
|
|
716
720
|
header_params = {}
|
|
717
721
|
|