lightning-sdk 2025.9.16__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 (78) hide show
  1. lightning_sdk/__init__.py +4 -3
  2. lightning_sdk/api/cloud_account_api.py +12 -1
  3. lightning_sdk/api/job_api.py +12 -11
  4. lightning_sdk/api/mmt_api.py +1 -1
  5. lightning_sdk/api/studio_api.py +1 -1
  6. lightning_sdk/api/teamspace_api.py +18 -0
  7. lightning_sdk/api/user_api.py +8 -2
  8. lightning_sdk/cli/entrypoint.py +3 -1
  9. lightning_sdk/cli/groups.py +8 -1
  10. lightning_sdk/cli/legacy/entrypoint.py +1 -1
  11. lightning_sdk/cli/studio/create.py +19 -5
  12. lightning_sdk/cli/studio/delete.py +9 -5
  13. lightning_sdk/cli/studio/list.py +5 -1
  14. lightning_sdk/cli/studio/ssh.py +9 -3
  15. lightning_sdk/cli/studio/start.py +26 -3
  16. lightning_sdk/cli/studio/stop.py +7 -3
  17. lightning_sdk/cli/studio/switch.py +21 -5
  18. lightning_sdk/cli/utils/owner_selection.py +110 -0
  19. lightning_sdk/cli/utils/studio_selection.py +22 -15
  20. lightning_sdk/cli/utils/teamspace_selection.py +63 -62
  21. lightning_sdk/cli/vm/__init__.py +20 -0
  22. lightning_sdk/cli/vm/create.py +33 -0
  23. lightning_sdk/cli/vm/delete.py +25 -0
  24. lightning_sdk/cli/vm/list.py +30 -0
  25. lightning_sdk/cli/vm/ssh.py +31 -0
  26. lightning_sdk/cli/vm/start.py +60 -0
  27. lightning_sdk/cli/vm/stop.py +25 -0
  28. lightning_sdk/cli/vm/switch.py +38 -0
  29. lightning_sdk/lightning_cloud/openapi/__init__.py +20 -1
  30. lightning_sdk/lightning_cloud/openapi/api/assistants_service_api.py +2 -95
  31. lightning_sdk/lightning_cloud/openapi/api/billing_service_api.py +24 -8
  32. lightning_sdk/lightning_cloud/openapi/api/cluster_service_api.py +420 -0
  33. lightning_sdk/lightning_cloud/openapi/api/jobs_service_api.py +121 -0
  34. lightning_sdk/lightning_cloud/openapi/api/storage_service_api.py +655 -0
  35. lightning_sdk/lightning_cloud/openapi/models/__init__.py +20 -1
  36. lightning_sdk/lightning_cloud/openapi/models/create_machine_request_represents_the_request_to_create_a_machine.py +435 -0
  37. lightning_sdk/lightning_cloud/openapi/models/job_id_reportroutingtelemetry_body.py +123 -0
  38. lightning_sdk/lightning_cloud/openapi/models/project_id_storagetransfers_body.py +149 -0
  39. lightning_sdk/lightning_cloud/openapi/models/user_id_affiliatelinks_body.py +107 -3
  40. lightning_sdk/lightning_cloud/openapi/models/v1_abort_storage_transfer_response.py +97 -0
  41. lightning_sdk/lightning_cloud/openapi/models/v1_assistant_session_daily_aggregated.py +27 -1
  42. lightning_sdk/lightning_cloud/openapi/models/v1_cloud_provider.py +2 -0
  43. lightning_sdk/lightning_cloud/openapi/models/v1_cluster_accelerator.py +27 -1
  44. lightning_sdk/lightning_cloud/openapi/models/v1_cluster_spec.py +53 -1
  45. lightning_sdk/lightning_cloud/openapi/models/v1_cluster_type.py +1 -0
  46. lightning_sdk/lightning_cloud/openapi/models/v1_create_machine_response.py +123 -0
  47. lightning_sdk/lightning_cloud/openapi/models/v1_create_project_request.py +27 -1
  48. lightning_sdk/lightning_cloud/openapi/models/v1_delete_machine_response.py +97 -0
  49. lightning_sdk/lightning_cloud/openapi/models/v1_external_cluster_spec.py +27 -1
  50. lightning_sdk/lightning_cloud/openapi/models/v1_get_machine_response.py +123 -0
  51. lightning_sdk/lightning_cloud/openapi/models/v1_get_user_response.py +27 -1
  52. lightning_sdk/lightning_cloud/openapi/models/v1_kubernetes_direct_v1.py +27 -1
  53. lightning_sdk/lightning_cloud/openapi/models/v1_lightning_elastic_cluster_v1.py +97 -0
  54. lightning_sdk/lightning_cloud/openapi/models/{v1_get_model_total_usage_metrics_response.py → v1_list_machines_response.py} +37 -37
  55. lightning_sdk/lightning_cloud/openapi/models/v1_list_storage_transfers_response.py +123 -0
  56. lightning_sdk/lightning_cloud/openapi/models/v1_machine.py +539 -0
  57. lightning_sdk/lightning_cloud/openapi/models/v1_machine_direct_v1.py +123 -0
  58. lightning_sdk/lightning_cloud/openapi/models/v1_pause_storage_transfer_response.py +97 -0
  59. lightning_sdk/lightning_cloud/openapi/models/v1_purchase_annual_upsell_request.py +123 -0
  60. lightning_sdk/lightning_cloud/openapi/models/v1_report_deployment_routing_telemetry_response.py +97 -0
  61. lightning_sdk/lightning_cloud/openapi/models/v1_resume_storage_transfer_response.py +97 -0
  62. lightning_sdk/lightning_cloud/openapi/models/v1_routing_telemetry.py +79 -1
  63. lightning_sdk/lightning_cloud/openapi/models/v1_rule_resource.py +1 -0
  64. lightning_sdk/lightning_cloud/openapi/models/v1_slack_notifier.py +149 -0
  65. lightning_sdk/lightning_cloud/openapi/models/v1_slack_notifier_type.py +105 -0
  66. lightning_sdk/lightning_cloud/openapi/models/v1_storage_transfer.py +435 -0
  67. lightning_sdk/lightning_cloud/openapi/models/v1_storage_transfer_status.py +108 -0
  68. lightning_sdk/lightning_cloud/openapi/models/v1_update_user_request.py +27 -1
  69. lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py +105 -79
  70. lightning_sdk/machine.py +16 -1
  71. lightning_sdk/studio.py +55 -11
  72. lightning_sdk/teamspace.py +65 -2
  73. {lightning_sdk-2025.9.16.dist-info → lightning_sdk-2025.9.29.dist-info}/METADATA +1 -1
  74. {lightning_sdk-2025.9.16.dist-info → lightning_sdk-2025.9.29.dist-info}/RECORD +78 -50
  75. {lightning_sdk-2025.9.16.dist-info → lightning_sdk-2025.9.29.dist-info}/LICENSE +0 -0
  76. {lightning_sdk-2025.9.16.dist-info → lightning_sdk-2025.9.29.dist-info}/WHEEL +0 -0
  77. {lightning_sdk-2025.9.16.dist-info → lightning_sdk-2025.9.29.dist-info}/entry_points.txt +0 -0
  78. {lightning_sdk-2025.9.16.dist-info → lightning_sdk-2025.9.29.dist-info}/top_level.txt +0 -0
lightning_sdk/__init__.py CHANGED
@@ -10,14 +10,16 @@ from lightning_sdk.organization import Organization
10
10
  from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin, SlurmJobsPlugin
11
11
  from lightning_sdk.status import Status
12
12
  from lightning_sdk.studio import Studio
13
- from lightning_sdk.teamspace import FolderLocation, Teamspace
13
+ from lightning_sdk.teamspace import ConnectionType, FolderLocation, Teamspace
14
14
  from lightning_sdk.user import User
15
15
 
16
16
  __all__ = [
17
17
  "AIHub",
18
18
  "Agent",
19
19
  "CloudProvider",
20
+ "ConnectionType",
20
21
  "Deployment",
22
+ "FolderLocation",
21
23
  "Job",
22
24
  "JobsPlugin",
23
25
  "Machine",
@@ -29,10 +31,9 @@ __all__ = [
29
31
  "Status",
30
32
  "Studio",
31
33
  "Teamspace",
32
- "FolderLocation",
33
34
  "User",
34
35
  ]
35
36
 
36
- __version__ = "2025.09.16"
37
+ __version__ = "2025.09.29"
37
38
  _check_version_and_prompt_upgrade(__version__)
38
39
  _set_tqdm_envvars_noninteractive()
@@ -13,6 +13,7 @@ from lightning_sdk.lightning_cloud.rest_client import LightningClient
13
13
 
14
14
  if TYPE_CHECKING:
15
15
  from lightning_sdk.machine import CloudProvider
16
+ from lightning_sdk.teamspace import ConnectionType
16
17
 
17
18
 
18
19
  class CloudAccountApi:
@@ -133,7 +134,7 @@ class CloudAccountApi:
133
134
 
134
135
  def get_cloud_account_provider_mapping(self, teamspace_id: str) -> Dict["CloudProvider", V1ExternalCluster]:
135
136
  """Gets the cloud account <-> provider mapping."""
136
- res = self.list_global_cloud_accounts(teamspace_id=teamspace_id)
137
+ res = self.list_cloud_accounts(teamspace_id=teamspace_id)
137
138
  cloud_accounts = {cloud_account.id: cloud_account for cloud_account in res}
138
139
  providers = {cloud_account.id: self._get_cloud_account_provider(cloud_account) for cloud_account in res}
139
140
 
@@ -214,3 +215,13 @@ class CloudAccountApi:
214
215
  return default_cloud_account
215
216
 
216
217
  return None
218
+
219
+ @staticmethod
220
+ def get_cloud_provider_for_connection_type(connection_type: "ConnectionType") -> "CloudProvider":
221
+ from lightning_sdk.machine import CloudProvider
222
+ from lightning_sdk.teamspace import ConnectionType
223
+
224
+ if connection_type == ConnectionType.EFS:
225
+ return CloudProvider.AWS
226
+
227
+ raise ValueError(f"ConnectionType {ConnectionType} currently not supported!")
@@ -104,22 +104,23 @@ class JobApiV1:
104
104
  org_id=org_id,
105
105
  )
106
106
 
107
- identifier = None
107
+ identifiers = []
108
108
 
109
109
  if user_requested_compute_config and user_requested_compute_config.name:
110
- identifier = user_requested_compute_config.name
110
+ identifiers.append(user_requested_compute_config.name)
111
111
  else:
112
- identifier = spec.compute_config.instance_type
112
+ identifiers.append(spec.compute_config.instance_type)
113
113
 
114
114
  for accelerator in accelerators:
115
- if identifier in (
116
- accelerator.slug,
117
- accelerator.slug_multi_cloud,
118
- accelerator.instance_id,
119
- ):
120
- return Machine.from_str(accelerator.slug_multi_cloud)
115
+ for ident in identifiers:
116
+ if ident in (
117
+ accelerator.slug,
118
+ accelerator.slug_multi_cloud,
119
+ accelerator.instance_id,
120
+ ):
121
+ return Machine._from_accelerator(accelerator)
121
122
 
122
- return Machine.from_str(identifier)
123
+ return Machine.from_str(identifiers[0])
123
124
 
124
125
  def _get_machines_for_cloud_account(
125
126
  self, teamspace_id: str, cloud_account_id: str, org_id: str
@@ -431,7 +432,7 @@ class JobApiV2:
431
432
  if (spec.instance_name and spec.instance_name in possible_identifiers) or (
432
433
  spec.instance_type and spec.instance_type in possible_identifiers
433
434
  ):
434
- return Machine.from_str(accelerator.slug_multi_cloud)
435
+ return Machine._from_accelerator(accelerator)
435
436
 
436
437
  return Machine.from_str(spec.instance_name or spec.instance_type)
437
438
 
@@ -265,7 +265,7 @@ class MMTApiV2:
265
265
  if (spec.instance_name and spec.instance_name in possible_identifiers) or (
266
266
  spec.instance_type and spec.instance_type in possible_identifiers
267
267
  ):
268
- return Machine.from_str(accelerator.slug_multi_cloud)
268
+ return Machine._from_accelerator(accelerator)
269
269
 
270
270
  return Machine.from_str(spec.instance_name or spec.instance_type)
271
271
 
@@ -426,7 +426,7 @@ class StudioApi:
426
426
  accelerator.slug_multi_cloud,
427
427
  accelerator.instance_id,
428
428
  ):
429
- return Machine.from_str(accelerator.slug_multi_cloud)
429
+ return Machine._from_accelerator(accelerator)
430
430
 
431
431
  return Machine.from_str(response.compute_config.name)
432
432
 
@@ -28,6 +28,7 @@ from lightning_sdk.lightning_cloud.openapi import (
28
28
  V1Assistant,
29
29
  V1CloudSpace,
30
30
  V1ClusterAccelerator,
31
+ V1EfsConfig,
31
32
  V1Endpoint,
32
33
  V1ExternalCluster,
33
34
  V1GCSFolderDataConnection,
@@ -515,3 +516,20 @@ class TeamspaceApi:
515
516
  create_request.gcs_folder = V1GCSFolderDataConnection()
516
517
 
517
518
  self._client.data_connection_service_create_data_connection(create_request, teamspace_id)
519
+
520
+ def new_connection(
521
+ self, teamspace_id: str, name: str, source: str, cluster: V1ExternalCluster, writable: bool, region: str
522
+ ) -> None:
523
+ create_request = Create(
524
+ name=name,
525
+ create_resources=False,
526
+ force=True,
527
+ writable=writable,
528
+ cluster_id=cluster.id,
529
+ access_cluster_ids=[cluster.id],
530
+ )
531
+
532
+ # TODO: Add support for other connection types
533
+ create_request.efs = V1EfsConfig(file_system_id=source, region=region)
534
+
535
+ self._client.data_connection_service_create_data_connection(create_request, teamspace_id)
@@ -1,5 +1,5 @@
1
1
  import re
2
- from typing import Dict, List, Union
2
+ from typing import Dict, List, Optional, Union
3
3
 
4
4
  from lightning_sdk.lightning_cloud.login import Auth
5
5
  from lightning_sdk.lightning_cloud.openapi import (
@@ -61,8 +61,14 @@ class UserApi:
61
61
  def _get_all_teamspace_memberships(
62
62
  self,
63
63
  user_id: str, # todo: this is unused, but still required
64
+ org_id: Optional[str] = None,
64
65
  ) -> List[V1Membership]:
65
- return self._client.projects_service_list_memberships(filter_by_user_id=True).memberships
66
+ kwargs: Dict[str, Union[bool, str]] = {"filter_by_user_id": True}
67
+
68
+ if org_id is not None:
69
+ kwargs["organization_id"] = org_id
70
+
71
+ return self._client.projects_service_list_memberships(**kwargs).memberships
66
72
 
67
73
  def _get_authed_user_name(self) -> str:
68
74
  """Gets the currently logged-in user."""
@@ -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")
@@ -0,0 +1,110 @@
1
+ import os
2
+ from contextlib import suppress
3
+ from typing import Dict, List, Optional, TypedDict
4
+
5
+ import click
6
+ from simple_term_menu import TerminalMenu
7
+
8
+ from lightning_sdk.cli.legacy.exceptions import StudioCliError
9
+ from lightning_sdk.organization import Organization
10
+ from lightning_sdk.owner import Owner
11
+ from lightning_sdk.user import User
12
+ from lightning_sdk.utils.resolve import ApiException, _get_authed_user, _resolve_org, _resolve_user
13
+
14
+
15
+ class _OwnerMenuType(TypedDict):
16
+ name: str
17
+ is_org: bool
18
+
19
+
20
+ class OwnerMenu:
21
+ """This class is used to select a teamspace owner (org/user) from a list of possible owners.
22
+
23
+ It can be used to select an owner from a list of possible owners, or to resolve an owner from a name.
24
+ """
25
+
26
+ def _get_owner_from_interactive_menu(self, possible_owners: Dict[str, _OwnerMenuType]) -> _OwnerMenuType:
27
+ owner_ids = sorted(possible_owners.keys())
28
+ terminal_menu = self._prepare_terminal_menu_owners([possible_owners[k] for k in owner_ids])
29
+ terminal_menu.show()
30
+
31
+ selected_id = owner_ids[terminal_menu.chosen_menu_index]
32
+ return possible_owners[selected_id]
33
+
34
+ def _get_owner_from_name(self, owner: str, possible_owners: Dict[str, _OwnerMenuType]) -> _OwnerMenuType:
35
+ for _, ts in possible_owners.items():
36
+ if ts["name"]:
37
+ return ts
38
+
39
+ click.echo(f"Could not find Owner {owner}, please select it from the list:")
40
+ return self._get_owner_from_interactive_menu(possible_owners)
41
+
42
+ @staticmethod
43
+ def _prepare_terminal_menu_owners(
44
+ possible_owners: List[_OwnerMenuType], title: Optional[str] = None
45
+ ) -> TerminalMenu:
46
+ if title is None:
47
+ title = "Please select a Teamspace-Owner out of the following:"
48
+
49
+ return TerminalMenu(
50
+ [f"{to['name']} ({'Organization' if to['is_org'] else 'User'})" for to in possible_owners],
51
+ title=title,
52
+ clear_menu_on_exit=True,
53
+ )
54
+
55
+ @staticmethod
56
+ def _get_possible_owners(user: User) -> Dict[str, _OwnerMenuType]:
57
+ user_api = user._user_api
58
+
59
+ orgs = user_api._get_organizations_for_authed_user()
60
+ owners: Dict[str, _OwnerMenuType] = {user.id: {"name": user.name, "is_org": False}}
61
+
62
+ for org in orgs:
63
+ owners[org.id] = {"name": org.name, "is_org": True}
64
+
65
+ return owners
66
+
67
+ def __call__(self, owner: Optional[str] = None) -> Owner:
68
+ try:
69
+ # try to resolve the teamspace from the name, environment or config
70
+ resolved_owner = None
71
+ with suppress(ApiException, ValueError, RuntimeError):
72
+ resolved_owner = _resolve_org(owner)
73
+
74
+ if resolved_owner is not None:
75
+ return resolved_owner
76
+
77
+ with suppress(ApiException, ValueError, RuntimeError):
78
+ resolved_owner = _resolve_user(owner)
79
+
80
+ if resolved_owner is not None:
81
+ return resolved_owner
82
+
83
+ if os.environ.get("LIGHTNING_NON_INTERACTIVE", "0") == "1":
84
+ raise ValueError(
85
+ "Owner selection is not supported in non-interactive mode. Please provide an owner name."
86
+ )
87
+
88
+ # if the owner is not resolved, try to get the owner from the interactive menu
89
+ # this could mean that either no owner was provided or the provided owner is not valid
90
+ user = _get_authed_user()
91
+
92
+ possible_owners = self._get_possible_owners(user)
93
+ if owner is None:
94
+ owner_dict = self._get_owner_from_interactive_menu(possible_owners=possible_owners)
95
+
96
+ else:
97
+ owner_dict = self._get_owner_from_name(owner=owner, possible_owners=possible_owners)
98
+
99
+ if owner_dict.get("is_org", False):
100
+ return Organization(owner_dict.get("name", None))
101
+
102
+ return User(owner_dict.get("name", None))
103
+ except KeyboardInterrupt:
104
+ raise KeyboardInterrupt from None
105
+
106
+ except Exception as e:
107
+ raise StudioCliError(
108
+ f"Could not find the given Teamspace-Owner {owner}. "
109
+ "Please contact Lightning AI directly to resolve this issue."
110
+ ) from e