lightning-sdk 2025.9.23__py3-none-any.whl → 2025.9.29__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.
Files changed (69) hide show
  1. lightning_sdk/__init__.py +1 -1
  2. lightning_sdk/cli/entrypoint.py +3 -1
  3. lightning_sdk/cli/groups.py +8 -1
  4. lightning_sdk/cli/legacy/entrypoint.py +1 -1
  5. lightning_sdk/cli/studio/create.py +19 -5
  6. lightning_sdk/cli/studio/delete.py +9 -5
  7. lightning_sdk/cli/studio/list.py +5 -1
  8. lightning_sdk/cli/studio/ssh.py +9 -3
  9. lightning_sdk/cli/studio/start.py +26 -3
  10. lightning_sdk/cli/studio/stop.py +7 -3
  11. lightning_sdk/cli/studio/switch.py +21 -5
  12. lightning_sdk/cli/utils/studio_selection.py +22 -15
  13. lightning_sdk/cli/vm/__init__.py +20 -0
  14. lightning_sdk/cli/vm/create.py +33 -0
  15. lightning_sdk/cli/vm/delete.py +25 -0
  16. lightning_sdk/cli/vm/list.py +30 -0
  17. lightning_sdk/cli/vm/ssh.py +31 -0
  18. lightning_sdk/cli/vm/start.py +60 -0
  19. lightning_sdk/cli/vm/stop.py +25 -0
  20. lightning_sdk/cli/vm/switch.py +38 -0
  21. lightning_sdk/lightning_cloud/openapi/__init__.py +20 -1
  22. lightning_sdk/lightning_cloud/openapi/api/assistants_service_api.py +2 -95
  23. lightning_sdk/lightning_cloud/openapi/api/billing_service_api.py +24 -8
  24. lightning_sdk/lightning_cloud/openapi/api/cluster_service_api.py +420 -0
  25. lightning_sdk/lightning_cloud/openapi/api/jobs_service_api.py +121 -0
  26. lightning_sdk/lightning_cloud/openapi/api/storage_service_api.py +655 -0
  27. lightning_sdk/lightning_cloud/openapi/models/__init__.py +20 -1
  28. lightning_sdk/lightning_cloud/openapi/models/create_machine_request_represents_the_request_to_create_a_machine.py +435 -0
  29. lightning_sdk/lightning_cloud/openapi/models/job_id_reportroutingtelemetry_body.py +123 -0
  30. lightning_sdk/lightning_cloud/openapi/models/project_id_storagetransfers_body.py +149 -0
  31. lightning_sdk/lightning_cloud/openapi/models/user_id_affiliatelinks_body.py +107 -3
  32. lightning_sdk/lightning_cloud/openapi/models/v1_abort_storage_transfer_response.py +97 -0
  33. lightning_sdk/lightning_cloud/openapi/models/v1_assistant_session_daily_aggregated.py +27 -1
  34. lightning_sdk/lightning_cloud/openapi/models/v1_cloud_provider.py +2 -0
  35. lightning_sdk/lightning_cloud/openapi/models/v1_cluster_accelerator.py +27 -1
  36. lightning_sdk/lightning_cloud/openapi/models/v1_cluster_spec.py +53 -1
  37. lightning_sdk/lightning_cloud/openapi/models/v1_cluster_type.py +1 -0
  38. lightning_sdk/lightning_cloud/openapi/models/v1_create_machine_response.py +123 -0
  39. lightning_sdk/lightning_cloud/openapi/models/v1_create_project_request.py +27 -1
  40. lightning_sdk/lightning_cloud/openapi/models/v1_delete_machine_response.py +97 -0
  41. lightning_sdk/lightning_cloud/openapi/models/v1_external_cluster_spec.py +27 -1
  42. lightning_sdk/lightning_cloud/openapi/models/v1_get_machine_response.py +123 -0
  43. lightning_sdk/lightning_cloud/openapi/models/v1_get_user_response.py +27 -1
  44. lightning_sdk/lightning_cloud/openapi/models/v1_kubernetes_direct_v1.py +27 -1
  45. lightning_sdk/lightning_cloud/openapi/models/v1_lightning_elastic_cluster_v1.py +97 -0
  46. lightning_sdk/lightning_cloud/openapi/models/{v1_get_model_total_usage_metrics_response.py → v1_list_machines_response.py} +37 -37
  47. lightning_sdk/lightning_cloud/openapi/models/v1_list_storage_transfers_response.py +123 -0
  48. lightning_sdk/lightning_cloud/openapi/models/v1_machine.py +539 -0
  49. lightning_sdk/lightning_cloud/openapi/models/v1_machine_direct_v1.py +123 -0
  50. lightning_sdk/lightning_cloud/openapi/models/v1_pause_storage_transfer_response.py +97 -0
  51. lightning_sdk/lightning_cloud/openapi/models/v1_purchase_annual_upsell_request.py +123 -0
  52. lightning_sdk/lightning_cloud/openapi/models/v1_report_deployment_routing_telemetry_response.py +97 -0
  53. lightning_sdk/lightning_cloud/openapi/models/v1_resume_storage_transfer_response.py +97 -0
  54. lightning_sdk/lightning_cloud/openapi/models/v1_routing_telemetry.py +79 -1
  55. lightning_sdk/lightning_cloud/openapi/models/v1_rule_resource.py +1 -0
  56. lightning_sdk/lightning_cloud/openapi/models/v1_slack_notifier.py +149 -0
  57. lightning_sdk/lightning_cloud/openapi/models/v1_slack_notifier_type.py +105 -0
  58. lightning_sdk/lightning_cloud/openapi/models/v1_storage_transfer.py +435 -0
  59. lightning_sdk/lightning_cloud/openapi/models/v1_storage_transfer_status.py +108 -0
  60. lightning_sdk/lightning_cloud/openapi/models/v1_update_user_request.py +27 -1
  61. lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py +105 -79
  62. lightning_sdk/studio.py +55 -11
  63. lightning_sdk/teamspace.py +11 -2
  64. {lightning_sdk-2025.9.23.dist-info → lightning_sdk-2025.9.29.dist-info}/METADATA +1 -1
  65. {lightning_sdk-2025.9.23.dist-info → lightning_sdk-2025.9.29.dist-info}/RECORD +69 -42
  66. {lightning_sdk-2025.9.23.dist-info → lightning_sdk-2025.9.29.dist-info}/LICENSE +0 -0
  67. {lightning_sdk-2025.9.23.dist-info → lightning_sdk-2025.9.29.dist-info}/WHEEL +0 -0
  68. {lightning_sdk-2025.9.23.dist-info → lightning_sdk-2025.9.29.dist-info}/entry_points.txt +0 -0
  69. {lightning_sdk-2025.9.23.dist-info → lightning_sdk-2025.9.29.dist-info}/top_level.txt +0 -0
lightning_sdk/__init__.py CHANGED
@@ -34,6 +34,6 @@ __all__ = [
34
34
  "User",
35
35
  ]
36
36
 
37
- __version__ = "2025.09.23"
37
+ __version__ = "2025.09.29"
38
38
  _check_version_and_prompt_upgrade(__version__)
39
39
  _set_tqdm_envvars_noninteractive()
@@ -21,6 +21,7 @@ from lightning_sdk.cli.groups import (
21
21
  # job,
22
22
  # mmt,
23
23
  studio,
24
+ vm,
24
25
  )
25
26
  from lightning_sdk.cli.utils import CustomHelpFormatter, rich_to_str
26
27
  from lightning_sdk.constants import _LIGHTNING_DEBUG
@@ -40,7 +41,7 @@ def _notify_exception(exception_type: Type[BaseException], value: BaseException,
40
41
  if _LIGHTNING_DEBUG:
41
42
  tb_text = "".join(traceback.format_exception(exception_type, value, tb))
42
43
  renderables.append(Text("\n\nFull traceback:\n", style="bold yellow"))
43
- renderables.append(Syntax(tb_text, "python", theme="monokai", line_numbers=False, word_wrap=True))
44
+ renderables.append(Syntax(tb_text, "python", theme="monokai light", line_numbers=False, word_wrap=True))
44
45
  else:
45
46
  renderables.append(Text("\n\n🐞 To view the full traceback, set: LIGHTNING_DEBUG=1"))
46
47
 
@@ -83,6 +84,7 @@ main_cli.add_command(config)
83
84
  # main_cli.add_command(job)
84
85
  # main_cli.add_command(mmt)
85
86
  main_cli.add_command(studio)
87
+ main_cli.add_command(vm)
86
88
  if os.environ.get("LIGHTNING_EXPERIMENTAL_CLI_ONLY", "0") != "1":
87
89
  #### LEGACY COMMANDS ####
88
90
  # these commands are currently supported for backwards compatibility, but will potentially be removed in the future.
@@ -6,6 +6,7 @@ from lightning_sdk.cli.config import register_commands as register_config_comman
6
6
  from lightning_sdk.cli.job import register_commands as register_job_commands
7
7
  from lightning_sdk.cli.mmt import register_commands as register_mmt_commands
8
8
  from lightning_sdk.cli.studio import register_commands as register_studio_commands
9
+ from lightning_sdk.cli.vm import register_commands as register_vm_commands
9
10
 
10
11
 
11
12
  @click.group(name="studio")
@@ -25,7 +26,12 @@ def mmt() -> None:
25
26
 
26
27
  @click.group(name="config")
27
28
  def config() -> None:
28
- """Manage Lightning SDK and CLIconfiguration."""
29
+ """Manage Lightning SDK and CLI configuration."""
30
+
31
+
32
+ @click.group(name="vm")
33
+ def vm() -> None:
34
+ """Manage Lightning AI VMs."""
29
35
 
30
36
 
31
37
  # Register config commands with the main config group
@@ -33,3 +39,4 @@ register_job_commands(job)
33
39
  register_mmt_commands(mmt)
34
40
  register_studio_commands(studio)
35
41
  register_config_commands(config)
42
+ register_vm_commands(vm)
@@ -48,7 +48,7 @@ def _notify_exception(exception_type: Type[BaseException], value: BaseException,
48
48
  if _LIGHTNING_DEBUG:
49
49
  tb_text = "".join(traceback.format_exception(exception_type, value, tb))
50
50
  renderables.append(Text("\n\nFull traceback:\n", style="bold yellow"))
51
- renderables.append(Syntax(tb_text, "python", theme="monokai", line_numbers=False, word_wrap=True))
51
+ renderables.append(Syntax(tb_text, "python", theme="monokai light", line_numbers=False, word_wrap=True))
52
52
  else:
53
53
  renderables.append(Text("\n\n🐞 To view the full traceback, set: LIGHTNING_DEBUG=1"))
54
54
 
@@ -9,7 +9,7 @@ from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
9
9
  from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
10
10
  from lightning_sdk.lightning_cloud.openapi.rest import ApiException
11
11
  from lightning_sdk.machine import CloudProvider
12
- from lightning_sdk.studio import Studio
12
+ from lightning_sdk.studio import VM, Studio
13
13
 
14
14
 
15
15
  @click.command("create")
@@ -36,6 +36,16 @@ def create_studio(
36
36
  Example:
37
37
  lightning studio create
38
38
  """
39
+ create_impl(name=name, teamspace=teamspace, cloud_provider=cloud_provider, cloud_account=cloud_account, vm=False)
40
+
41
+
42
+ def create_impl(
43
+ name: Optional[str],
44
+ teamspace: Optional[str],
45
+ cloud_provider: Optional[str],
46
+ cloud_account: Optional[str],
47
+ vm: bool,
48
+ ) -> None:
39
49
  menu = TeamspacesMenu()
40
50
 
41
51
  resolved_teamspace = menu(teamspace)
@@ -44,8 +54,12 @@ def create_studio(
44
54
  if cloud_provider is not None:
45
55
  cloud_provider = CloudProvider(cloud_provider)
46
56
 
57
+ create_cls = VM if vm else Studio
58
+ cls_name = create_cls.__qualname__
59
+
47
60
  try:
48
- studio = Studio(
61
+ create_cls = VM if vm else Studio
62
+ studio = create_cls(
49
63
  name=name,
50
64
  teamspace=resolved_teamspace,
51
65
  create_ok=True,
@@ -54,7 +68,7 @@ def create_studio(
54
68
  )
55
69
  except (RuntimeError, ValueError, ApiException):
56
70
  if name:
57
- raise ValueError(f"Could not create Studio: '{name}'. Does the Studio exist?") from None
58
- raise ValueError(f"Could not create Studio: '{name}'. Please provide a Studio name") from None
71
+ raise ValueError(f"Could not create {cls_name}: '{name}'. Does the {cls_name} exist?") from None
72
+ raise ValueError(f"Could not create {cls_name}: '{name}'. Please provide a {cls_name} name") from None
59
73
 
60
- click.echo(f"Studio {studio_name_link(studio)} created successfully")
74
+ click.echo(f"{cls_name} {studio_name_link(studio)} created successfully")
@@ -12,7 +12,7 @@ from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
12
12
  @click.option(
13
13
  "--name",
14
14
  help=(
15
- "The name of the studio to start. "
15
+ "The name of the studio to delete. "
16
16
  "If not provided, will try to infer from environment, "
17
17
  "use the default value from the config or prompt for interactive selection."
18
18
  ),
@@ -25,21 +25,25 @@ def delete_studio(name: Optional[str] = None, teamspace: Optional[str] = None) -
25
25
  lightning studio delete --name my-studio
26
26
 
27
27
  """
28
+ return delete_impl(name=name, teamspace=teamspace, vm=False)
29
+
30
+
31
+ def delete_impl(name: Optional[str], teamspace: Optional[str], vm: bool) -> None:
28
32
  menu = TeamspacesMenu()
29
33
  resolved_teamspace = menu(teamspace=teamspace)
30
34
 
31
- menu = StudiosMenu(resolved_teamspace)
35
+ menu = StudiosMenu(resolved_teamspace, vm=vm)
32
36
  studio = menu(studio=name)
33
37
 
34
38
  studio_name = f"{studio.teamspace.owner.name}/{studio.teamspace.name}/{studio.name}"
35
39
  confirmed = click.confirm(
36
- f"Are you sure you want to delete studio '{studio_name}'?",
40
+ f"Are you sure you want to delete {studio._cls_name} '{studio_name}'?",
37
41
  abort=True,
38
42
  )
39
43
  if not confirmed:
40
- click.echo("Studio deletion cancelled")
44
+ click.echo(f"{studio._cls_name} deletion cancelled")
41
45
  return
42
46
 
43
47
  studio.delete()
44
48
 
45
- click.echo(f"Studio '{studio.name}' deleted successfully")
49
+ click.echo(f"{studio._cls_name} '{studio.name}' deleted successfully")
@@ -35,13 +35,17 @@ def list_studios(teamspace: Optional[str] = None, all: bool = False, sort_by: Op
35
35
  lightning studio list --teamspace owner/teamspace
36
36
 
37
37
  """
38
+ return list_impl(teamspace=teamspace, all=all, sort_by=sort_by, vm=False)
39
+
40
+
41
+ def list_impl(teamspace: Optional[str], all: bool, sort_by: Optional[str], vm: bool) -> None: # noqa: A002
38
42
  menu = TeamspacesMenu()
39
43
  teamspace_resolved = menu(teamspace=teamspace)
40
44
  save_teamspace_to_config(teamspace_resolved, overwrite=False)
41
45
 
42
46
  user = _get_authed_user()
43
47
 
44
- studios = teamspace_resolved.studios
48
+ studios = teamspace_resolved.vms if vm else teamspace_resolved.studios
45
49
 
46
50
  table = Table(
47
51
  pad_edge=True,
@@ -20,7 +20,7 @@ from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
20
20
  @click.option(
21
21
  "--name",
22
22
  help=(
23
- "The name of the studio to start. "
23
+ "The name of the studio to ssh into. "
24
24
  "If not provided, will try to infer from environment, "
25
25
  "use the default value from the config or prompt for interactive selection."
26
26
  ),
@@ -39,6 +39,10 @@ def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, opti
39
39
  Example:
40
40
  lightning studio ssh --name my-studio
41
41
  """
42
+ return ssh_impl(name=name, teamspace=teamspace, option=option, vm=False)
43
+
44
+
45
+ def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[List[str]], vm: bool) -> None:
42
46
  auth = Auth()
43
47
  auth.authenticate()
44
48
  ssh_private_key_path = _download_ssh_keys(auth.api_key, force_download=False)
@@ -46,8 +50,10 @@ def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, opti
46
50
  menu = TeamspacesMenu()
47
51
  resolved_teamspace = menu(teamspace=teamspace)
48
52
 
49
- menu = StudiosMenu(resolved_teamspace)
50
- studio = menu(studio=name)
53
+ menu = StudiosMenu(resolved_teamspace, vm=vm)
54
+ studio = menu(
55
+ studio=name,
56
+ )
51
57
  save_studio_to_config(studio)
52
58
 
53
59
  ssh_options = " -o " + " -o ".join(option) if option else ""
@@ -9,7 +9,7 @@ from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
9
9
  from lightning_sdk.cli.utils.studio_selection import StudiosMenu
10
10
  from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
11
11
  from lightning_sdk.machine import CloudProvider, Machine
12
- from lightning_sdk.studio import Studio
12
+ from lightning_sdk.studio import VM, Studio
13
13
 
14
14
 
15
15
  @click.command("start")
@@ -57,6 +57,28 @@ def start_studio(
57
57
  lightning studio start --name my-studio
58
58
 
59
59
  """
60
+ return start_impl(
61
+ name=name,
62
+ teamspace=teamspace,
63
+ create=create,
64
+ machine=machine,
65
+ interruptible=interruptible,
66
+ cloud_provider=cloud_provider,
67
+ cloud_account=cloud_account,
68
+ vm=False,
69
+ )
70
+
71
+
72
+ def start_impl(
73
+ name: Optional[str],
74
+ teamspace: Optional[str],
75
+ create: bool,
76
+ machine: str,
77
+ interruptible: bool,
78
+ cloud_provider: Optional[str],
79
+ cloud_account: Optional[str],
80
+ vm: bool,
81
+ ) -> None:
60
82
  menu = TeamspacesMenu()
61
83
  resolved_teamspace = menu(teamspace=teamspace)
62
84
 
@@ -64,10 +86,11 @@ def start_studio(
64
86
  cloud_provider = CloudProvider(cloud_provider)
65
87
 
66
88
  if not create:
67
- menu = StudiosMenu(resolved_teamspace)
89
+ menu = StudiosMenu(resolved_teamspace, vm=vm)
68
90
  studio = menu(studio=name)
69
91
  else:
70
- studio = Studio(
92
+ create_cls = VM if vm else Studio
93
+ studio = create_cls(
71
94
  name=name,
72
95
  teamspace=resolved_teamspace,
73
96
  create_ok=create,
@@ -14,7 +14,7 @@ from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
14
14
  @click.option(
15
15
  "--name",
16
16
  help=(
17
- "The name of the studio to start. "
17
+ "The name of the studio to stop. "
18
18
  "If not provided, will try to infer from environment, "
19
19
  "use the default value from the config or prompt for interactive selection."
20
20
  ),
@@ -27,15 +27,19 @@ def stop_studio(name: Optional[str] = None, teamspace: Optional[str] = None) ->
27
27
  lightning studio stop --name my-studio
28
28
 
29
29
  """
30
+ return stop_impl(name=name, teamspace=teamspace, vm=False)
31
+
32
+
33
+ def stop_impl(name: Optional[str], teamspace: Optional[str], vm: bool) -> None:
30
34
  # missing studio_name and teamspace are handled by the studio class
31
35
  menu = TeamspacesMenu()
32
36
  resolved_teamspace = menu(teamspace=teamspace)
33
37
 
34
- menu = StudiosMenu(resolved_teamspace)
38
+ menu = StudiosMenu(resolved_teamspace, vm=vm)
35
39
  studio = menu(studio=name)
36
40
 
37
41
  studio.stop()
38
42
 
39
43
  save_studio_to_config(studio)
40
44
 
41
- click.echo(f"Studio {studio_name_link(studio)} stopped successfully")
45
+ click.echo(f"{studio._cls_name} {studio_name_link(studio)} stopped successfully")
@@ -9,14 +9,13 @@ from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
9
9
  from lightning_sdk.cli.utils.studio_selection import StudiosMenu
10
10
  from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
11
11
  from lightning_sdk.machine import Machine
12
- from lightning_sdk.studio import Studio
13
12
 
14
13
 
15
14
  @click.command("switch")
16
15
  @click.option(
17
16
  "--name",
18
17
  help=(
19
- "The name of the studio to start. "
18
+ "The name of the studio to switch to a different machine. "
20
19
  "If not provided, will try to infer from environment, "
21
20
  "use the default value from the config or prompt for interactive selection."
22
21
  ),
@@ -35,16 +34,33 @@ def switch_studio(
35
34
  interruptible: bool = False,
36
35
  ) -> None:
37
36
  """Switch a Studio to a different machine type."""
37
+ return switch_impl(
38
+ name=name,
39
+ teamspace=teamspace,
40
+ machine=machine,
41
+ interruptible=interruptible,
42
+ vm=False,
43
+ )
44
+
45
+
46
+ def switch_impl(
47
+ name: Optional[str],
48
+ teamspace: Optional[str],
49
+ machine: Optional[str],
50
+ interruptible: bool,
51
+ vm: bool,
52
+ ) -> None:
38
53
  menu = TeamspacesMenu()
39
54
  resolved_teamspace = menu(teamspace=teamspace)
40
55
 
41
- menu = StudiosMenu(resolved_teamspace)
56
+ menu = StudiosMenu(resolved_teamspace, vm=vm)
42
57
  studio = menu(studio=name)
43
58
 
44
59
  resolved_machine = Machine.from_str(machine)
45
- Studio.show_progress = True
60
+
61
+ studio.__class__.show_progress = True
46
62
  studio.switch_machine(resolved_machine, interruptible=interruptible)
47
63
 
48
64
  save_studio_to_config(studio)
49
65
 
50
- click.echo(f"Studio {studio_name_link(studio)} switched to machine '{resolved_machine}' successfully")
66
+ click.echo(f"{studio._cls_name} {studio_name_link(studio)} switched to machine '{resolved_machine}' successfully")
@@ -1,12 +1,12 @@
1
1
  import os
2
2
  from contextlib import suppress
3
- from typing import Dict, List, Optional
3
+ from typing import Dict, List, Optional, Union
4
4
 
5
5
  import click
6
6
  from simple_term_menu import TerminalMenu
7
7
 
8
8
  from lightning_sdk.cli.legacy.exceptions import StudioCliError
9
- from lightning_sdk.studio import Studio
9
+ from lightning_sdk.studio import VM, Studio
10
10
  from lightning_sdk.teamspace import Teamspace
11
11
  from lightning_sdk.utils.resolve import _get_authed_user
12
12
 
@@ -17,54 +17,58 @@ class StudiosMenu:
17
17
  It can be used to select a studio from a list of possible studios, or to resolve a studio from a name.
18
18
  """
19
19
 
20
- def __init__(self, teamspace: Teamspace) -> None:
20
+ def __init__(self, teamspace: Teamspace, vm: bool = False) -> None:
21
21
  """Initialize the StudiosMenu with a teamspace.
22
22
 
23
23
  Args:
24
24
  teamspace: The teamspace to list studios from
25
25
  """
26
26
  self.teamspace = teamspace
27
+ self.vm = vm
27
28
 
28
- def _get_studio_from_interactive_menu(self, possible_studios: Dict[str, Studio]) -> Studio:
29
+ def _get_studio_from_interactive_menu(self, possible_studios: Dict[str, Union[Studio, VM]]) -> Union[Studio, VM]:
29
30
  studio_names = sorted(possible_studios.keys())
30
- terminal_menu = self._prepare_terminal_menu_studios(studio_names)
31
+ terminal_menu = self._prepare_terminal_menu_studios(studio_names, vm=self.vm)
31
32
  terminal_menu.show()
32
33
 
33
34
  selected_name = studio_names[terminal_menu.chosen_menu_index]
34
35
  return possible_studios[selected_name]
35
36
 
36
- def _get_studio_from_name(self, studio: str, possible_studios: Dict[str, Studio]) -> Studio:
37
+ def _get_studio_from_name(self, studio: str, possible_studios: Dict[str, Union[Studio, VM]]) -> Union[Studio, VM]:
37
38
  if studio in possible_studios:
38
39
  return possible_studios[studio]
39
40
 
40
- click.echo(f"Could not find Studio {studio}, please select it from the list:")
41
+ click.echo(f"Could not find {'VM' if self.vm else 'Studio'} {studio}, please select it from the list:")
41
42
  return self._get_studio_from_interactive_menu(possible_studios)
42
43
 
43
44
  @staticmethod
44
- def _prepare_terminal_menu_studios(studio_names: List[str], title: Optional[str] = None) -> TerminalMenu:
45
+ def _prepare_terminal_menu_studios(
46
+ studio_names: List[str], title: Optional[str] = None, vm: bool = False
47
+ ) -> TerminalMenu:
45
48
  if title is None:
46
- title = "Please select a Studio out of the following:"
49
+ title = f"Please select a {'VM' if vm else 'Studio'} out of the following:"
47
50
 
48
51
  return TerminalMenu(studio_names, title=title, clear_menu_on_exit=True)
49
52
 
50
- def _get_possible_studios(self) -> Dict[str, Studio]:
53
+ def _get_possible_studios(self) -> Dict[str, Union[Studio, VM]]:
51
54
  """Get all available studios in the teamspace."""
52
55
  studios = {}
53
56
 
54
57
  user = _get_authed_user()
58
+ studios = self.teamspace.vms if self.vm else self.teamspace.studios
55
59
  for studio in self.teamspace.studios:
56
60
  if studio._studio.user_id == user.id:
57
61
  studios[studio.name] = studio
58
62
  return studios
59
63
 
60
- def __call__(self, studio: Optional[str] = None) -> Studio:
64
+ def __call__(self, studio: Optional[str] = None) -> Union[Studio, VM]:
61
65
  """Select a studio from the teamspace.
62
66
 
63
67
  Args:
64
68
  studio: Optional studio name to select. If not provided, will show interactive menu.
65
69
 
66
70
  Returns:
67
- Selected Studio object
71
+ Selected Studio/VM object
68
72
 
69
73
  Raises:
70
74
  StudioCliError: If studio selection fails
@@ -73,15 +77,18 @@ class StudiosMenu:
73
77
  # try to resolve the studio from the name, environment or config
74
78
  resolved_studio = None
75
79
 
80
+ selected_cls = VM if self.vm else Studio
81
+
76
82
  with suppress(Exception):
77
- resolved_studio = Studio(name=studio, teamspace=self.teamspace, create_ok=False)
83
+ resolved_studio = selected_cls(name=studio, teamspace=self.teamspace, create_ok=False)
78
84
 
79
85
  if resolved_studio is not None:
80
86
  return resolved_studio
81
87
 
82
88
  if os.environ.get("LIGHTNING_NON_INTERACTIVE", "0") == "1" and studio is None:
83
89
  raise ValueError(
84
- "Studio selection is not supported in non-interactive mode. Please provide a studio name."
90
+ f"{'VM' if self.vm else 'Studio'} selection is not supported in non-interactive mode. "
91
+ "Please provide a studio name."
85
92
  )
86
93
 
87
94
  click.echo(f"Listing studios in teamspace {self.teamspace.owner.name}/{self.teamspace.name}...")
@@ -101,6 +108,6 @@ class StudiosMenu:
101
108
 
102
109
  except Exception as e:
103
110
  raise StudioCliError(
104
- "Could not resolve a Studio. "
111
+ f"Could not resolve a {'VM' if self.vm else 'Studio'}. "
105
112
  "Please pass it as an argument or contact Lightning AI directly to resolve this issue."
106
113
  ) from e
@@ -0,0 +1,20 @@
1
+ import click
2
+
3
+
4
+ def register_commands(group: click.Group) -> None:
5
+ """Register studio commands with the given group."""
6
+ from lightning_sdk.cli.vm.create import create_vm
7
+ from lightning_sdk.cli.vm.delete import delete_vm
8
+ from lightning_sdk.cli.vm.list import list_vms
9
+ from lightning_sdk.cli.vm.ssh import ssh_vm
10
+ from lightning_sdk.cli.vm.start import start_vm
11
+ from lightning_sdk.cli.vm.stop import stop_vm
12
+ from lightning_sdk.cli.vm.switch import switch_vm
13
+
14
+ group.add_command(create_vm)
15
+ group.add_command(delete_vm)
16
+ group.add_command(list_vms)
17
+ group.add_command(ssh_vm)
18
+ group.add_command(start_vm)
19
+ group.add_command(stop_vm)
20
+ group.add_command(switch_vm)
@@ -0,0 +1,33 @@
1
+ from typing import Optional
2
+
3
+ import click
4
+
5
+ from lightning_sdk.cli.studio.create import create_impl
6
+ from lightning_sdk.machine import CloudProvider
7
+
8
+
9
+ @click.command("create")
10
+ @click.option("--name", help="The name of the VM to create. If not provided, a random name will be generated.")
11
+ @click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
12
+ @click.option(
13
+ "--cloud-provider",
14
+ help="The cloud provider to start the VM on. Defaults to teamspace default.",
15
+ type=click.Choice(m.name for m in list(CloudProvider)),
16
+ )
17
+ @click.option(
18
+ "--cloud-account",
19
+ help="The cloud account to create the VM on. Defaults to teamspace default.",
20
+ type=click.STRING,
21
+ )
22
+ def create_vm(
23
+ name: Optional[str] = None,
24
+ teamspace: Optional[str] = None,
25
+ cloud_provider: Optional[str] = None,
26
+ cloud_account: Optional[str] = None,
27
+ ) -> None:
28
+ """Create a new VM.
29
+
30
+ Example:
31
+ lightning vm create
32
+ """
33
+ create_impl(name=name, teamspace=teamspace, cloud_provider=cloud_provider, cloud_account=cloud_account, vm=True)
@@ -0,0 +1,25 @@
1
+ from typing import Optional
2
+
3
+ import click
4
+
5
+ from lightning_sdk.cli.studio.delete import delete_impl
6
+
7
+
8
+ @click.command("delete")
9
+ @click.option(
10
+ "--name",
11
+ help=(
12
+ "The name of the VM to delete. "
13
+ "If not provided, will try to infer from environment, "
14
+ "use the default value from the config or prompt for interactive selection."
15
+ ),
16
+ )
17
+ @click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
18
+ def delete_vm(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
19
+ """Delete a VM.
20
+
21
+ Example:
22
+ lightning vm delete --name my-vm
23
+
24
+ """
25
+ return delete_impl(name=name, teamspace=teamspace, vm=True)
@@ -0,0 +1,30 @@
1
+ from typing import Optional
2
+
3
+ import click
4
+
5
+ from lightning_sdk.cli.studio.list import list_impl
6
+
7
+
8
+ @click.command("list")
9
+ @click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
10
+ @click.option(
11
+ "--all",
12
+ is_flag=True,
13
+ flag_value=True,
14
+ default=False,
15
+ help="List all VMs, not just the ones belonging to the authed user",
16
+ )
17
+ @click.option(
18
+ "--sort-by",
19
+ default=None,
20
+ type=click.Choice(["name", "teamspace", "status", "machine", "cloud-account"], case_sensitive=False),
21
+ help="the attribute to sort the VMs by.",
22
+ )
23
+ def list_vms(teamspace: Optional[str] = None, all: bool = False, sort_by: Optional[str] = None) -> None: # noqa: A002
24
+ """List VMs in a teamspace.
25
+
26
+ Example:
27
+ lightning vm list --teamspace owner/teamspace
28
+
29
+ """
30
+ return list_impl(teamspace=teamspace, all=all, sort_by=sort_by, vm=True)
@@ -0,0 +1,31 @@
1
+ from typing import List, Optional
2
+
3
+ import click
4
+
5
+ from lightning_sdk.cli.studio.ssh import ssh_impl
6
+
7
+
8
+ @click.command("ssh")
9
+ @click.option(
10
+ "--name",
11
+ help=(
12
+ "The name of the VM to ssh into. "
13
+ "If not provided, will try to infer from environment, "
14
+ "use the default value from the config or prompt for interactive selection."
15
+ ),
16
+ )
17
+ @click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)", type=click.STRING)
18
+ @click.option(
19
+ "--option",
20
+ "-o",
21
+ help="Additional options to pass to the SSH command. Can be specified multiple times.",
22
+ multiple=True,
23
+ type=click.STRING,
24
+ )
25
+ def ssh_vm(name: Optional[str] = None, teamspace: Optional[str] = None, option: Optional[List[str]] = None) -> None:
26
+ """SSH into a VM.
27
+
28
+ Example:
29
+ lightning vm ssh --name my-vm
30
+ """
31
+ return ssh_impl(name=name, teamspace=teamspace, option=option, vm=False)
@@ -0,0 +1,60 @@
1
+ from typing import Optional
2
+
3
+ import click
4
+
5
+ from lightning_sdk.cli.studio.start import start_impl
6
+ from lightning_sdk.machine import CloudProvider, Machine
7
+
8
+
9
+ @click.command("start")
10
+ @click.option(
11
+ "--name",
12
+ help=(
13
+ "The name of the VM to start. "
14
+ "If not provided, will try to infer from environment, "
15
+ "use the default value from the config or prompt for interactive selection."
16
+ ),
17
+ )
18
+ @click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
19
+ @click.option("--create", is_flag=True, help="Create the VM if it doesn't exist")
20
+ @click.option(
21
+ "--machine",
22
+ help="The machine type to start the VM on. Defaults to CPU-4",
23
+ type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli),
24
+ )
25
+ @click.option("--interruptible", is_flag=True, help="Start the VM on an interruptible instance.")
26
+ @click.option(
27
+ "--cloud-provider",
28
+ help=("The cloud provider to start the VM on. Defaults to teamspace default. Only used if --create is specified."),
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 VM on. Defaults to teamspace default. Only used if --create is specified.",
34
+ type=click.STRING,
35
+ )
36
+ def start_vm(
37
+ name: Optional[str] = None,
38
+ teamspace: Optional[str] = None,
39
+ create: bool = False,
40
+ machine: str = "CPU",
41
+ interruptible: bool = False,
42
+ cloud_provider: Optional[str] = None,
43
+ cloud_account: Optional[str] = None,
44
+ ) -> None:
45
+ """Start a VM.
46
+
47
+ Example:
48
+ lightning vm start --name my-vm
49
+
50
+ """
51
+ return start_impl(
52
+ name=name,
53
+ teamspace=teamspace,
54
+ create=create,
55
+ machine=machine,
56
+ interruptible=interruptible,
57
+ cloud_provider=cloud_provider,
58
+ cloud_account=cloud_account,
59
+ vm=True,
60
+ )