lightning-sdk 2025.9.30__py3-none-any.whl → 2025.10.8__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 CHANGED
@@ -35,6 +35,6 @@ __all__ = [
35
35
  "VM",
36
36
  ]
37
37
 
38
- __version__ = "2025.09.30"
38
+ __version__ = "2025.10.08"
39
39
  _check_version_and_prompt_upgrade(__version__)
40
40
  _set_tqdm_envvars_noninteractive()
@@ -16,6 +16,8 @@ class BaseStudioInfo:
16
16
  name: str
17
17
  managed_id: str
18
18
  description: str
19
+ creator: str
20
+ enabled: bool
19
21
 
20
22
 
21
23
  class BaseStudio:
@@ -80,20 +82,29 @@ class BaseStudio:
80
82
  disabled=disabled,
81
83
  )
82
84
 
83
- def list(self, managed: bool = True) -> List[BaseStudioInfo]:
85
+ def list(self, managed: bool = True, include_disabled: bool = False) -> List[BaseStudioInfo]:
84
86
  """List all base studios in the organization.
85
87
 
88
+ Args:
89
+ managed: Whether to filter for managed base studios.
90
+ include_disabled: Whether to include disabled base studios in the results.
91
+
86
92
  Returns:
87
- List[V1CloudSpaceEnvironmentTemplate]: A list of base studio templates.
93
+ List[BaseStudioInfo]: A list of base studio templates.
88
94
  """
89
- result = []
90
- for template in self._base_studio_api.get_all_base_studios(self._org.id, managed).templates:
91
- result.append(
92
- BaseStudioInfo(
93
- id=template.id,
94
- name=template.name,
95
- managed_id=template.managed_id,
96
- description=template.description,
97
- ),
95
+ templates = self._base_studio_api.get_all_base_studios(self._org.id, managed).templates
96
+
97
+ return [
98
+ BaseStudioInfo(
99
+ id=template.id,
100
+ name=template.name,
101
+ managed_id=template.managed_id,
102
+ description=template.description,
103
+ creator="⚡ Lightning AI"
104
+ if template.managed_id
105
+ else UserApi()._get_user_by_id(template.user_id).username,
106
+ enabled=not template.disabled,
98
107
  )
99
- return result
108
+ for template in templates
109
+ if include_disabled or not template.disabled
110
+ ]
@@ -0,0 +1,10 @@
1
+ """Base Studio CLI commands."""
2
+
3
+ import click
4
+
5
+
6
+ def register_commands(group: click.Group) -> None:
7
+ """Register base studio commands with the given group."""
8
+ from lightning_sdk.cli.base_studio.list import list_base_studios
9
+
10
+ group.add_command(list_base_studios)
@@ -0,0 +1,45 @@
1
+ """Base Studio list command."""
2
+
3
+ import click
4
+ from rich.table import Table
5
+
6
+ from lightning_sdk.base_studio import BaseStudio
7
+ from lightning_sdk.cli.utils.richt_print import rich_to_str
8
+
9
+
10
+ @click.command("list")
11
+ @click.option("--include-disabled", help="Include disabled Base Studios in the list.", is_flag=True)
12
+ def list_base_studios(include_disabled: bool) -> None:
13
+ """List Base Studios in an org.
14
+
15
+ Example:
16
+ lightning base-studio list
17
+
18
+ """
19
+ return list_impl(include_disabled=include_disabled)
20
+
21
+
22
+ def list_impl(include_disabled: bool) -> None:
23
+ base_studio_cls = BaseStudio()
24
+ base_studios = base_studio_cls.list(include_disabled=include_disabled) + base_studio_cls.list(
25
+ managed=False, include_disabled=include_disabled
26
+ )
27
+
28
+ table = Table(
29
+ pad_edge=True,
30
+ )
31
+
32
+ table.add_column("Name")
33
+ table.add_column("Description")
34
+ table.add_column("Creator")
35
+ table.add_column("Enabled")
36
+
37
+ for base_studio in base_studios:
38
+ table.add_row(
39
+ base_studio.name.lower().replace(" ", "-"),
40
+ base_studio.description or "",
41
+ base_studio.creator,
42
+ "Yes" if base_studio.enabled else "No",
43
+ )
44
+
45
+ click.echo(rich_to_str(table), color=True)
@@ -17,6 +17,7 @@ from lightning_sdk.api.studio_api import _cloud_url
17
17
 
18
18
  # Import legacy groups directly from groups.py
19
19
  from lightning_sdk.cli.groups import (
20
+ base_studio,
20
21
  config,
21
22
  # job,
22
23
  # mmt,
@@ -85,6 +86,7 @@ main_cli.add_command(config)
85
86
  # main_cli.add_command(mmt)
86
87
  main_cli.add_command(studio)
87
88
  main_cli.add_command(vm)
89
+ main_cli.add_command(base_studio)
88
90
  if os.environ.get("LIGHTNING_EXPERIMENTAL_CLI_ONLY", "0") != "1":
89
91
  #### LEGACY COMMANDS ####
90
92
  # these commands are currently supported for backwards compatibility, but will potentially be removed in the future.
@@ -2,6 +2,7 @@
2
2
 
3
3
  import click
4
4
 
5
+ from lightning_sdk.cli.base_studio import register_commands as register_base_studio_commands
5
6
  from lightning_sdk.cli.config import register_commands as register_config_commands
6
7
  from lightning_sdk.cli.job import register_commands as register_job_commands
7
8
  from lightning_sdk.cli.mmt import register_commands as register_mmt_commands
@@ -34,9 +35,15 @@ def vm() -> None:
34
35
  """Manage Lightning AI VMs."""
35
36
 
36
37
 
38
+ @click.group(name="base-studio")
39
+ def base_studio() -> None:
40
+ """Manage Lightning AI Base Studios."""
41
+
42
+
37
43
  # Register config commands with the main config group
38
44
  register_job_commands(job)
39
45
  register_mmt_commands(mmt)
40
46
  register_studio_commands(studio)
41
47
  register_config_commands(config)
42
48
  register_vm_commands(vm)
49
+ register_base_studio_commands(base_studio)
@@ -2,20 +2,96 @@
2
2
 
3
3
  import subprocess
4
4
  import sys
5
- from typing import Optional
5
+ from typing import Dict, Optional, Set
6
6
 
7
7
  import click
8
8
 
9
+ from lightning_sdk.base_studio import BaseStudio
9
10
  from lightning_sdk.cli.utils.richt_print import studio_name_link
10
11
  from lightning_sdk.cli.utils.save_to_config import save_studio_to_config, save_teamspace_to_config
11
- from lightning_sdk.cli.utils.ssh_connection import download_ssh_keys
12
+ from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal
12
13
  from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
13
- from lightning_sdk.lightning_cloud.login import Auth
14
14
  from lightning_sdk.lightning_cloud.openapi.rest import ApiException
15
- from lightning_sdk.machine import CloudProvider
15
+ from lightning_sdk.machine import CloudProvider, Machine
16
16
  from lightning_sdk.studio import Studio
17
17
  from lightning_sdk.utils.names import random_unique_name
18
18
 
19
+ DEFAULT_MACHINE = "CPU"
20
+
21
+
22
+ def _split_gpus_spec(gpus: str) -> tuple[str, int]:
23
+ machine_name, machine_val = gpus.split(":", 1)
24
+ machine_name = machine_name.strip()
25
+ machine_val = machine_val.strip()
26
+
27
+ if not machine_val.isdigit() or int(machine_val) <= 0:
28
+ raise ValueError(f"Invalid GPU count '{machine_val}'. Must be a positive integer.")
29
+
30
+ machine_num = int(machine_val)
31
+ return machine_name, machine_num
32
+
33
+
34
+ def _construct_available_gpus(machine_options: Dict[str, str]) -> Set[str]:
35
+ # returns available gpus:count
36
+ available_gpus = set()
37
+ for v in machine_options.values():
38
+ if "_X_" in v:
39
+ gpu_type_num = v.replace("_X_", ":")
40
+ available_gpus.add(gpu_type_num)
41
+ else:
42
+ available_gpus.add(v)
43
+ return available_gpus
44
+
45
+
46
+ def _get_machine_from_gpus(gpus: str) -> Machine:
47
+ machine_name = gpus
48
+ machine_num = 1
49
+
50
+ if ":" in gpus:
51
+ machine_name, machine_num = _split_gpus_spec(gpus)
52
+
53
+ machine_options = {
54
+ m.name.lower(): m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli
55
+ }
56
+
57
+ if machine_num == 1:
58
+ # e.g. gpus=L4 or gpus=L4:1
59
+ gpu_key = machine_name.lower()
60
+ try:
61
+ return machine_options[gpu_key]
62
+ except KeyError:
63
+ available = ", ".join(_construct_available_gpus(machine_options))
64
+ raise ValueError(f"Invalid GPU type '{machine_name}'. Available options: {available}") from None
65
+
66
+ # Else: e.g. gpus=L4:4
67
+ gpu_key = f"{machine_name.lower()}_x_{machine_num}"
68
+ try:
69
+ return machine_options[gpu_key]
70
+ except KeyError:
71
+ available = ", ".join(_construct_available_gpus(machine_options))
72
+ raise ValueError(f"Invalid GPU configuration '{gpus}'. Available options: {available}") from None
73
+
74
+
75
+ def _get_base_studio_id(studio_type: Optional[str]) -> Optional[str]:
76
+ base_studios = BaseStudio()
77
+ base_studios = base_studios.list()
78
+ template_id = None
79
+
80
+ if base_studios and len(base_studios):
81
+ # if not specified by user, use the first existing template studio
82
+ template_id = base_studios[0].id
83
+ # else, try to match the provided studio_type to base studio name
84
+ if studio_type:
85
+ normalized_studio_type = studio_type.lower().replace(" ", "-")
86
+ match = next(
87
+ (s for s in base_studios if s.name.lower().replace(" ", "-") == normalized_studio_type),
88
+ None,
89
+ )
90
+ if match:
91
+ template_id = match.id
92
+
93
+ return template_id
94
+
19
95
 
20
96
  @click.command("connect")
21
97
  @click.argument("name", required=False)
@@ -30,13 +106,32 @@ from lightning_sdk.utils.names import random_unique_name
30
106
  help="The cloud account to create the studio on. Defaults to teamspace default.",
31
107
  type=click.STRING,
32
108
  )
33
- @click.option("--gpus", help="The number of GPUs to start the studio on. ", type=click.INT)
109
+ @click.option(
110
+ "--machine",
111
+ help="The machine type to start the studio on. Defaults to CPU-4",
112
+ type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli),
113
+ )
114
+ @click.option(
115
+ "--gpus",
116
+ help="The number and type of GPUs to start the studio on (format: TYPE:COUNT, e.g. L4:4)",
117
+ type=click.STRING,
118
+ )
119
+ @click.option(
120
+ "--studio-type",
121
+ help="The base studio template name to use for creating the studio. "
122
+ "Must be lowercase and hyphenated (use '-' instead of spaces). "
123
+ "Run 'lightning base-studio list' to see all available templates. "
124
+ "Defaults to the first available template.",
125
+ type=click.STRING,
126
+ )
34
127
  def connect_studio(
35
128
  name: Optional[str] = None,
36
129
  teamspace: Optional[str] = None,
37
130
  cloud_provider: Optional[str] = None,
38
131
  cloud_account: Optional[str] = None,
39
- gpus: Optional[int] = None,
132
+ machine: Optional[str] = None,
133
+ gpus: Optional[str] = None,
134
+ studio_type: Optional[str] = None,
40
135
  ) -> None:
41
136
  """Connect to a Studio.
42
137
 
@@ -53,6 +148,9 @@ def connect_studio(
53
148
 
54
149
  name = name or random_unique_name()
55
150
 
151
+ # check for available base studios
152
+ template_id = _get_base_studio_id(studio_type)
153
+
56
154
  try:
57
155
  studio = Studio(
58
156
  name=name,
@@ -60,35 +158,32 @@ def connect_studio(
60
158
  create_ok=True,
61
159
  cloud_provider=cloud_provider,
62
160
  cloud_account=cloud_account,
161
+ studio_type=template_id,
63
162
  )
64
163
  except (RuntimeError, ValueError, ApiException):
65
164
  raise ValueError(f"Could not create Studio: '{name}'") from None
66
165
 
67
166
  click.echo(f"Connecting to Studio '{studio_name_link(studio)}' ...")
68
167
 
69
- machine = "CPU"
70
168
  Studio.show_progress = True
71
- if gpus:
72
- # TODO: handle something like gpus=4:L4
73
- pass
169
+
170
+ if machine and gpus:
171
+ raise click.UsageError("Options --machine and --gpu are mutually exclusive. Provide only one.")
172
+ elif gpus:
173
+ machine = _get_machine_from_gpus(gpus.strip())
174
+ elif not machine:
175
+ machine = DEFAULT_MACHINE
74
176
 
75
177
  save_studio_to_config(studio)
76
- studio.start(machine, interruptible=True)
178
+ # by default, interruptible is False
179
+ studio.start(machine=machine, interruptible=False)
77
180
 
78
- ssh_private_key_path = _configure_ssh_internal()
181
+ ssh_private_key_path = configure_ssh_internal()
79
182
 
183
+ ssh_option = "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o LogLevel=ERROR"
80
184
  try:
81
- ssh_command = (
82
- f"ssh -i {ssh_private_key_path} -o UserKnownHostsFile=/dev/null s_{studio._studio.id}@ssh.lightning.ai"
83
- )
185
+ ssh_command = f"ssh -i {ssh_private_key_path} {ssh_option} s_{studio._studio.id}@ssh.lightning.ai"
84
186
  subprocess.run(ssh_command.split())
85
187
  except Exception as ex:
86
188
  print(f"Failed to establish SSH connection: {ex}")
87
189
  sys.exit(1)
88
-
89
-
90
- def _configure_ssh_internal() -> str:
91
- """Internal function to configure SSH without Click decorators."""
92
- auth = Auth()
93
- auth.authenticate()
94
- return download_ssh_keys(auth.api_key, force_download=False)
@@ -6,10 +6,9 @@ from typing import List, Optional
6
6
  import click
7
7
 
8
8
  from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
9
- from lightning_sdk.cli.utils.ssh_connection import download_ssh_keys
9
+ from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal
10
10
  from lightning_sdk.cli.utils.studio_selection import StudiosMenu
11
11
  from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
12
- from lightning_sdk.lightning_cloud.login import Auth
13
12
 
14
13
 
15
14
  @click.command("ssh")
@@ -39,9 +38,7 @@ def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, opti
39
38
 
40
39
 
41
40
  def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[List[str]], vm: bool) -> None:
42
- auth = Auth()
43
- auth.authenticate()
44
- ssh_private_key_path = download_ssh_keys(auth.api_key, force_download=False)
41
+ ssh_private_key_path = configure_ssh_internal()
45
42
 
46
43
  menu = TeamspacesMenu()
47
44
  resolved_teamspace = menu(teamspace=teamspace)
@@ -59,7 +56,7 @@ def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[Lis
59
56
  subprocess.run(ssh_command.split())
60
57
  except Exception:
61
58
  # redownload the keys to be sure they are up to date
62
- download_ssh_keys(auth.api_key, force_download=True)
59
+ ssh_private_key_path = configure_ssh_internal(force_download=True)
63
60
  try:
64
61
  subprocess.run(ssh_command.split())
65
62
  except Exception:
@@ -4,9 +4,17 @@ import uuid
4
4
  from pathlib import Path
5
5
  from typing import Optional
6
6
 
7
+ from lightning_sdk.lightning_cloud.login import Auth
7
8
  from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
8
9
 
9
10
 
11
+ def configure_ssh_internal(force_download: bool = False) -> str:
12
+ """Internal function to configure SSH without Click decorators."""
13
+ auth = Auth()
14
+ auth.authenticate()
15
+ return download_ssh_keys(auth.api_key, force_download=force_download)
16
+
17
+
10
18
  def download_ssh_keys(
11
19
  api_key: str | None,
12
20
  force_download: bool = False,
lightning_sdk/studio.py CHANGED
@@ -8,6 +8,7 @@ from tqdm.auto import tqdm
8
8
 
9
9
  from lightning_sdk.api.cloud_account_api import CloudAccountApi
10
10
  from lightning_sdk.api.studio_api import StudioApi
11
+ from lightning_sdk.base_studio import BaseStudio
11
12
  from lightning_sdk.constants import _LIGHTNING_DEBUG
12
13
  from lightning_sdk.lightning_cloud.openapi import V1ClusterType
13
14
  from lightning_sdk.machine import CloudProvider, Machine
@@ -51,6 +52,8 @@ class Studio:
51
52
  If not specified, falls backto the teamspace default cloud account.
52
53
  create_ok: whether the studio will be created if it does not yet exist. Defaults to True
53
54
  provider: the provider of the machine, the studio should be created on.
55
+ studio_type: Type of studio to create. Only effective during initial creation;
56
+ ignored for existing studios.
54
57
 
55
58
  Note:
56
59
  Since a teamspace can either be owned by an org or by a user directly,
@@ -78,6 +81,7 @@ class Studio:
78
81
  source: Optional[str] = None,
79
82
  disable_secrets: bool = False,
80
83
  provider: Optional[Union[CloudProvider, str]] = None, # deprecated in favor of cloud_provider
84
+ studio_type: Optional[str] = None, # for base studio templates
81
85
  ) -> None:
82
86
  self._studio_api = StudioApi()
83
87
  self._cloud_account_api = CloudAccountApi()
@@ -113,6 +117,25 @@ class Studio:
113
117
  default_cloud_account=self._teamspace.default_cloud_account,
114
118
  )
115
119
 
120
+ self._studio_type = None
121
+ if studio_type:
122
+ self._base_studio = BaseStudio()
123
+ self._available_base_studios = self._base_studio.list()
124
+ for bst in self._available_base_studios:
125
+ if (
126
+ bst.id == studio_type
127
+ or bst.name == studio_type
128
+ or bst.name.lower().replace(" ", "-") == studio_type
129
+ ):
130
+ self._studio_type = bst.id
131
+
132
+ if not self._studio_type:
133
+ raise ValueError(
134
+ f"Could not find studio type with ID or name '{studio_type}'. "
135
+ f"Available studio types: "
136
+ f"{[bst.name.lower().replace(' ', '-') for bst in self._available_base_studios]}"
137
+ )
138
+
116
139
  # Resolve studio name if not provided: explicit → env (LIGHTNING_CLOUD_SPACE_ID) → config defaults
117
140
  if name is None and not getattr(self._skip_init, "value", False):
118
141
  studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
@@ -150,6 +173,7 @@ class Studio:
150
173
  cloud_account=_cloud_account,
151
174
  source=source,
152
175
  disable_secrets=self._disable_secrets,
176
+ cloud_space_environment_template_id=self._studio_type,
153
177
  )
154
178
  else:
155
179
  raise e
@@ -264,10 +288,11 @@ class Studio:
264
288
  else:
265
289
  interruptible = self.teamspace.start_studios_on_interruptible
266
290
 
291
+ new_machine = machine
292
+ if not isinstance(machine, Machine):
293
+ new_machine = Machine.from_str(machine)
294
+
267
295
  if status == Status.Running:
268
- new_machine = machine
269
- if not isinstance(machine, Machine):
270
- new_machine = Machine.from_str(machine)
271
296
  if new_machine != self.machine:
272
297
  raise RuntimeError(
273
298
  f"Requested to start {self._cls_name} on {new_machine}, "
@@ -289,7 +314,11 @@ class Studio:
289
314
  with StudioProgressTracker("start", show_progress=True) as progress:
290
315
  # Start the studio without blocking
291
316
  self._studio_api.start_studio_async(
292
- self._studio.id, self._teamspace.id, machine, interruptible=interruptible, max_runtime=max_runtime
317
+ self._studio.id,
318
+ self._teamspace.id,
319
+ new_machine,
320
+ interruptible=interruptible,
321
+ max_runtime=max_runtime,
293
322
  )
294
323
 
295
324
  # Track progress through completion
@@ -299,7 +328,7 @@ class Studio:
299
328
  else:
300
329
  # Use the blocking version if no progress is needed
301
330
  self._studio_api.start_studio(
302
- self._studio.id, self._teamspace.id, machine, interruptible=interruptible, max_runtime=max_runtime
331
+ self._studio.id, self._teamspace.id, new_machine, interruptible=interruptible, max_runtime=max_runtime
303
332
  )
304
333
 
305
334
  self._setup()
@@ -17,30 +17,31 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_get_cloud_space_instance_st
17
17
  class StartupPhase(Enum):
18
18
  """Studio startup phase messages."""
19
19
 
20
- ALLOCATING_MACHINE = "Allocating machine from cloud provider..."
21
- STUDIO_STARTING = "Studio is starting up..."
22
- SETTING_UP_ENVIRONMENT = "Setting up Studio environment..."
23
- RESTORING_STATE = "Restoring Studio state..."
24
- FINALIZING_SETUP = "Finalizing Studio setup..."
25
- COMPLETED = "Studio started successfully"
20
+ STARTING_STUDIO = "Starting Studio..."
21
+ GETTING_MACHINE = "Getting a machine..."
26
22
 
23
+ SWITCHING_STUDIO = "Switching Studio..."
27
24
 
28
- def get_switching_progress_message(percentage: int, is_base_studio: bool, is_new_cloud_space: bool) -> str:
25
+ SETTING_UP_MACHINE = "Setting up machine..."
26
+ RESTORING_STUDIO = "Restoring Studio..."
27
+ PREPARING_STUDIO = "Preparing Studio..."
28
+ RESTORING_BASE_STUDIO = "Restoring Base Studio..."
29
+ SETTING_UP_BASE_STUDIO = "Setting up Base Studio..."
30
+ DONE = "Done"
31
+
32
+
33
+ def get_switching_progress_message(percentage: int, is_base_studio: bool) -> str:
29
34
  """Get progress message for switching studios."""
30
35
  percentage = max(0, min(100, round(percentage)))
31
36
 
32
37
  if percentage > 98:
33
- message = "Done"
38
+ message = StartupPhase.DONE.value
34
39
  elif percentage > 80:
35
- if is_new_cloud_space:
36
- message = "Setting up Base Studio..." if is_base_studio else "Preparing Studio..."
37
- else:
38
- message = "Restoring Studio..."
40
+ message = StartupPhase.RESTORING_BASE_STUDIO.value if is_base_studio else StartupPhase.RESTORING_STUDIO.value
39
41
  elif percentage > 60:
40
- message = "Setting up machine from cloud provider"
42
+ message = StartupPhase.SETTING_UP_MACHINE.value
41
43
  else:
42
- message = "Allocating machine from cloud provider"
43
-
44
+ message = StartupPhase.SWITCHING_STUDIO.value
44
45
  return f"({percentage}%) {message}"
45
46
 
46
47
 
@@ -110,15 +111,13 @@ class StudioProgressTracker:
110
111
  if self.progress:
111
112
  self.progress.stop()
112
113
 
113
- def update_progress(
114
- self, percentage: int, message: str = "", is_base_studio: bool = False, is_new_cloud_space: bool = False
115
- ) -> None:
114
+ def update_progress(self, percentage: int, message: str = "", is_base_studio: bool = False) -> None:
116
115
  """Update progress bar with current percentage and message."""
117
116
  if not self.progress or self.task_id is None:
118
117
  return
119
118
 
120
119
  if self.operation_type == "switch":
121
- display_message = get_switching_progress_message(percentage, is_base_studio, is_new_cloud_space)
120
+ display_message = get_switching_progress_message(percentage, is_base_studio)
122
121
  else:
123
122
  display_message = message or f"{self.operation_type.capitalize()}ing Studio..."
124
123
 
@@ -153,7 +152,7 @@ class StudioProgressTracker:
153
152
  message_stability_delay = 3.0 # Seconds to wait before changing message
154
153
 
155
154
  # Show initial progress immediately
156
- self.update_progress(5, StartupPhase.ALLOCATING_MACHINE.value)
155
+ self.update_progress(5, StartupPhase.STARTING_STUDIO.value)
157
156
 
158
157
  while True:
159
158
  try:
@@ -163,7 +162,7 @@ class StudioProgressTracker:
163
162
  # Default fallback progress based on time
164
163
  time_based_progress = min(95, int((elapsed / timeout) * 100))
165
164
  current_progress = max(last_progress, time_based_progress)
166
- current_message = StartupPhase.ALLOCATING_MACHINE.value
165
+ current_message = StartupPhase.STARTING_STUDIO.value
167
166
 
168
167
  # Check if we have detailed status information
169
168
  if hasattr(status, "in_use") and status.in_use:
@@ -178,7 +177,7 @@ class StudioProgressTracker:
178
177
  hasattr(startup_status, "top_up_restore_finished")
179
178
  and startup_status.top_up_restore_finished
180
179
  ):
181
- self.complete(StartupPhase.COMPLETED.value)
180
+ self.complete(StartupPhase.DONE.value)
182
181
  break
183
182
 
184
183
  # Check other phases in descending priority
@@ -186,14 +185,14 @@ class StudioProgressTracker:
186
185
  hasattr(startup_status, "initial_restore_finished")
187
186
  and startup_status.initial_restore_finished
188
187
  ):
189
- current_progress = max(current_progress, 85)
190
- current_message = StartupPhase.FINALIZING_SETUP.value
188
+ current_progress = max(current_progress, 80)
189
+ current_message = StartupPhase.PREPARING_STUDIO.value
191
190
  elif hasattr(startup_status, "container_ready") and startup_status.container_ready:
192
- current_progress = max(current_progress, 70)
193
- current_message = StartupPhase.RESTORING_STATE.value
194
- elif hasattr(startup_status, "machine_ready") and startup_status.machine_ready:
195
191
  current_progress = max(current_progress, 60)
196
- current_message = StartupPhase.SETTING_UP_ENVIRONMENT.value
192
+ current_message = StartupPhase.SETTING_UP_MACHINE.value
193
+ elif hasattr(startup_status, "machine_ready") and startup_status.machine_ready:
194
+ current_progress = max(current_progress, 30)
195
+ current_message = StartupPhase.GETTING_MACHINE.value
197
196
 
198
197
  # Check general phase information
199
198
  if hasattr(in_use, "phase") and in_use.phase:
@@ -201,7 +200,7 @@ class StudioProgressTracker:
201
200
 
202
201
  if phase == "CLOUD_SPACE_INSTANCE_STATE_RUNNING":
203
202
  current_progress = max(current_progress, 80)
204
- current_message = StartupPhase.STUDIO_STARTING.value
203
+ current_message = StartupPhase.SETTING_UP_MACHINE.value
205
204
  elif phase == "CLOUD_SPACE_INSTANCE_STATE_PENDING":
206
205
  # Track time in pending phase for smoother progress
207
206
  if "pending" not in phase_start_times:
@@ -211,7 +210,7 @@ class StudioProgressTracker:
211
210
  # Progress more smoothly through pending phase (10-60%)
212
211
  pending_progress = 10 + min(50, int((pending_elapsed / 60) * 50))
213
212
  current_progress = max(current_progress, pending_progress)
214
- current_message = StartupPhase.ALLOCATING_MACHINE.value
213
+ current_message = StartupPhase.GETTING_MACHINE.value
215
214
 
216
215
  # Check for requested machine status (pre-allocation)
217
216
  elif hasattr(status, "requested") and status.requested:
@@ -222,7 +221,7 @@ class StudioProgressTracker:
222
221
  # Progress through allocation phase (5-30%)
223
222
  allocation_progress = 5 + min(25, int((allocation_elapsed / 30) * 25))
224
223
  current_progress = max(current_progress, allocation_progress)
225
- current_message = StartupPhase.ALLOCATING_MACHINE.value
224
+ current_message = StartupPhase.GETTING_MACHINE.value
226
225
 
227
226
  # Ensure progress never decreases and moves smoothly
228
227
  if current_progress > last_progress:
@@ -262,12 +261,12 @@ class StudioProgressTracker:
262
261
 
263
262
  # Only update message if enough time has passed
264
263
  current_time = time.time()
265
- should_update_message = StartupPhase.ALLOCATING_MACHINE.value != self._last_message and (
264
+ should_update_message = StartupPhase.GETTING_MACHINE.value != self._last_message and (
266
265
  current_time - last_message_time >= message_stability_delay or last_message_time == 0
267
266
  )
268
267
 
269
268
  if should_update_message:
270
- self.update_progress(fallback_progress, StartupPhase.ALLOCATING_MACHINE.value)
269
+ self.update_progress(fallback_progress, StartupPhase.GETTING_MACHINE.value)
271
270
  last_message_time = current_time
272
271
  else:
273
272
  # Update progress but keep existing message
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: lightning_sdk
3
- Version: 2025.9.30
3
+ Version: 2025.10.8
4
4
  Summary: SDK to develop using Lightning AI Studios
5
5
  Author-email: Lightning-AI <justus@lightning.ai>
6
6
  License: MIT License
@@ -1,8 +1,8 @@
1
1
  docs/source/conf.py,sha256=r8yX20eC-4mHhMTd0SbQb5TlSWHhO6wnJ0VJ_FBFpag,13249
2
- lightning_sdk/__init__.py,sha256=1ZQWzgRpFgZg43Tyx8qq56e3b0ELkQ-sE12C9T7wg_o,1235
2
+ lightning_sdk/__init__.py,sha256=MLcIAWdrZoDrz0HJ7KhaLRbFrSg8Qs_XFS8p2vr2cwQ,1235
3
3
  lightning_sdk/agents.py,sha256=ly6Ma1j0ZgGPFyvPvMN28JWiB9dATIstFa5XM8pMi6I,1577
4
4
  lightning_sdk/ai_hub.py,sha256=iI1vNhgcz_Ff1c3rN1ogN7dK-r-HXRj6NMtS2cA14UA,6925
5
- lightning_sdk/base_studio.py,sha256=_Pwwl37R9GRd7t-f2kO5aQXiLNrP4sUtUNht2ZkP8LE,3678
5
+ lightning_sdk/base_studio.py,sha256=FIjHUiZtRf42m5YeqYYaJxcdgOtJVAJLZU42gzmldO8,4109
6
6
  lightning_sdk/constants.py,sha256=ztl1PTUBULnqTf3DyKUSJaV_O20hNtUYT6XvAYIrmIk,749
7
7
  lightning_sdk/helpers.py,sha256=KWMWnORHItIIA3PGn71YPs-7RjzGi8IXa2kQ5Qo4U8M,2459
8
8
  lightning_sdk/lit_container.py,sha256=8ys49TXI9MO89jLTA7MwDrKrssTRARAIF9OVmolDLq0,5273
@@ -14,7 +14,7 @@ lightning_sdk/plugin.py,sha256=f3P5-xZY6x-MX0Fs2z_Q2erSxPSiHZARO0BVkCezHw4,15192
14
14
  lightning_sdk/sandbox.py,sha256=_NvnWotEXW2rBiVFZZ4krKXxVjuAqfNh04qELSM0-Pg,5786
15
15
  lightning_sdk/serve.py,sha256=uW7zLhQ3X90ifetpxzTb8FNxifv5vIs7qZlgfEjVKzk,11794
16
16
  lightning_sdk/status.py,sha256=lLGAuSvXBoXQFEEsEYwdCi0RcSNatUn5OPjJVjDtoM0,386
17
- lightning_sdk/studio.py,sha256=_mPNeH_IBoWC-bwEiGXCeWizl1D2YXdWMUVUH3QyFsw,33145
17
+ lightning_sdk/studio.py,sha256=pij-pT-Cd11j5-LMAHxORLRP9a-Gcow87ZkTPcluu38,34350
18
18
  lightning_sdk/teamspace.py,sha256=iYaEaZH8N8Yj4XXj6M4zIH0tsjPVrAf_BzAAv6Jt9VI,26556
19
19
  lightning_sdk/user.py,sha256=TSYh38rxoi7qKOfrK2JYh_Nknya2Kbz2ngDIY85fFOY,1778
20
20
  lightning_sdk/api/__init__.py,sha256=xrp_RNECJGQtL5rZHF69WOzEuEIbWSLtjWAJAz4R5K4,500
@@ -35,8 +35,10 @@ lightning_sdk/api/teamspace_api.py,sha256=UZ6Ka6TL7zggs5qhQB4wl47VRusVTMpth1-toG
35
35
  lightning_sdk/api/user_api.py,sha256=o9gZmtvqNfj4kdSpo2fyyRuFAP7bM5w7mx0Oj9__Ads,4702
36
36
  lightning_sdk/api/utils.py,sha256=1NJgW4HCfzMqgnAqbMA7RRr2v3iM3KTuPIUQK5klDeQ,27127
37
37
  lightning_sdk/cli/__init__.py,sha256=lksw08t_ZIuHOH47LCIqSVHeZ8cUXI2aJWHYhyujYHM,32
38
- lightning_sdk/cli/entrypoint.py,sha256=kvzeMKNuMoGSfsvg8v9QsBuMHR4U3pcWqBEHfvQQzhM,4595
39
- lightning_sdk/cli/groups.py,sha256=OPH6PHN82w5ERX00wzmstqU0NjEHmViVLAmy8-y8jQk,1131
38
+ lightning_sdk/cli/entrypoint.py,sha256=wMAyzfKRqZ6EZ73SxpuJDk3u8J3o-I5vF0VnESJh2QE,4646
39
+ lightning_sdk/cli/groups.py,sha256=3l9_U-8T613RJ2fHZy_-iJbtpq4fH8auWH7eFGSdZmg,1373
40
+ lightning_sdk/cli/base_studio/__init__.py,sha256=npBEmXWNVuiawm98TFVNHZf038kD4aW3gzVHsjbUPNI,272
41
+ lightning_sdk/cli/base_studio/list.py,sha256=DCXzSIY4IB0lHEstP0PA-8e7VqQQSUXL6OVv8VvwBKM,1244
40
42
  lightning_sdk/cli/config/__init__.py,sha256=_aJ7uZZOdjomZ-ABmz8Cu3fLmnSujnQK4hKAEKM4aF4,395
41
43
  lightning_sdk/cli/config/get.py,sha256=JXGRNVGxELS6Es2RESgStzTKnjg5rr-M6yOmx2ty9A8,1606
42
44
  lightning_sdk/cli/config/set.py,sha256=9p8q8nUWGWNcAcflQ3hS_tj-59jY7cmkzccuiZl92h4,3117
@@ -73,11 +75,11 @@ lightning_sdk/cli/legacy/deploy/devbox.py,sha256=SdOtYdGefzldn-WCa6sz31zysyf2nP5
73
75
  lightning_sdk/cli/legacy/deploy/serve.py,sha256=bxN04mHlArjHtxR_QkCM0S_RPmL3WIPsgnPgFadSXwI,15478
74
76
  lightning_sdk/cli/mmt/__init__.py,sha256=CLgr-ZHHLS6Db_JGqpxbn4G2pYrKi4Qn-uhi8e0kFNc,145
75
77
  lightning_sdk/cli/studio/__init__.py,sha256=jqUuEniYlf3fwJni-MGLII3oprpsBSQEFqzBJwPIxMk,925
76
- lightning_sdk/cli/studio/connect.py,sha256=iHZiY0VvlBO4t7-tE3R5kkDsCQo6BSQ0CztOK9_uOKI,3038
78
+ lightning_sdk/cli/studio/connect.py,sha256=yb4Ga4lLjrg5C0sLPwAoEkkTpT7rv165AwqBOsYBIWc,6488
77
79
  lightning_sdk/cli/studio/create.py,sha256=kTxLtvTFjo-SgkdzOIngdXu_6yYLHwVbkfMhHjpIFkI,2505
78
80
  lightning_sdk/cli/studio/delete.py,sha256=Gi1KLCrSgjTiyoQe9NRvF4ZpsNyeM07KHnBxfdJg2FE,1483
79
81
  lightning_sdk/cli/studio/list.py,sha256=ieLiHRIfAb7E5e_r48BeTC03ib8RY6maeF7nkmrqaek,3301
80
- lightning_sdk/cli/studio/ssh.py,sha256=ly7FyaVYTfjQiB8T1iKJ_OoeEZTZCYhzVkgIEavYUoY,2299
82
+ lightning_sdk/cli/studio/ssh.py,sha256=zHpVUe2QiFvM91Pj5kzNCnacavwrptxNsPpxyxpKLno,2194
81
83
  lightning_sdk/cli/studio/start.py,sha256=GHfV_1I1XwV8UHOWBP2WVfd1B2BHyw1lQ9HMv6Eomes,3236
82
84
  lightning_sdk/cli/studio/stop.py,sha256=hVXIyrv68rvrwIQgyEPUKt3Th8S1iipOVflopPnoVOY,1398
83
85
  lightning_sdk/cli/studio/switch.py,sha256=x6F1biDtfR0dk1t7X3NSkmfj5jkzGJe6-nvj3SonQlQ,2069
@@ -88,7 +90,7 @@ lightning_sdk/cli/utils/owner_selection.py,sha256=PbKbFoGblxi60lF_kZxnlqxjCmI9r9
88
90
  lightning_sdk/cli/utils/resolve.py,sha256=RWQ8MgXEJCdsoY50y8Pa9B-A6NA_SvxOqL5pVKfvSr8,917
89
91
  lightning_sdk/cli/utils/richt_print.py,sha256=psSY65nLAQbxK_K0w93Qq9857aUUbm77FLS1sc3Oecg,1262
90
92
  lightning_sdk/cli/utils/save_to_config.py,sha256=mMjjfA4Yv1e7MtE03iADTXcTSAuUL0Ws9OZObx6ufo0,1164
91
- lightning_sdk/cli/utils/ssh_connection.py,sha256=REB4rRsDrYrjACBK1PWlmcehLEututqjCLeGTHca1Lk,1677
93
+ lightning_sdk/cli/utils/ssh_connection.py,sha256=cfppkWOLRLVSE-LhkENRANgdPNpT3U6beg8Lyx1xoxo,1984
92
94
  lightning_sdk/cli/utils/studio_selection.py,sha256=z0SMd82sOt_mfTiQZgVY27MxLkMAjM__SWhlOlfXJT0,4347
93
95
  lightning_sdk/cli/utils/teamspace_selection.py,sha256=n-FFfWLTV63SxoK6Fi2VQikBIfKYqMNTew9b9v-1UQk,4945
94
96
  lightning_sdk/cli/vm/__init__.py,sha256=tAJpluJVDOF3fUQyRwko5ZYIpqglczMyRv9Kz4OVaIk,711
@@ -1205,11 +1207,11 @@ lightning_sdk/utils/config.py,sha256=LfCAbko_f-1TgFoYD1KT68tZHEIgn3M2y1HUqTMruUo
1205
1207
  lightning_sdk/utils/dynamic.py,sha256=glUTO1JC9APtQ6Gr9SO02a3zr56-sPAXM5C3NrTpgyQ,1959
1206
1208
  lightning_sdk/utils/enum.py,sha256=h2JRzqoBcSlUdanFHmkj_j5DleBHAu1esQYUsdNI-hU,4106
1207
1209
  lightning_sdk/utils/names.py,sha256=1EuXbIh7wldkDp1FG10oz9vIOyWrpGWeFFVy-DQBgzA,18162
1208
- lightning_sdk/utils/progress.py,sha256=IXfEcUF-rL5jIw0Hir6eSxN7VBZfR--1O2LaEhGAU70,12698
1210
+ lightning_sdk/utils/progress.py,sha256=bLWw39fzq29PMWoFXaPIVfoS3Ug245950oWOFJ2ZaiU,12596
1209
1211
  lightning_sdk/utils/resolve.py,sha256=4TyEnIgIrvkSvYk5i5PmcIogD_5Y9pBhiphRLfLMttc,10477
1210
- lightning_sdk-2025.9.30.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
1211
- lightning_sdk-2025.9.30.dist-info/METADATA,sha256=Ju9urPSHGtE4gfSRZGGJqjYVK7MA-EjocNrUfLKtr8A,4130
1212
- lightning_sdk-2025.9.30.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
1213
- lightning_sdk-2025.9.30.dist-info/entry_points.txt,sha256=OoZa4Fc8NMs6GSN0cdA1J8e6couzAcL82CbM1yo4f_M,122
1214
- lightning_sdk-2025.9.30.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
1215
- lightning_sdk-2025.9.30.dist-info/RECORD,,
1212
+ lightning_sdk-2025.10.8.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
1213
+ lightning_sdk-2025.10.8.dist-info/METADATA,sha256=N8b-kMsJedFH9daVWxNta30bWOraumpFMRESf8W-Sh4,4130
1214
+ lightning_sdk-2025.10.8.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
1215
+ lightning_sdk-2025.10.8.dist-info/entry_points.txt,sha256=OoZa4Fc8NMs6GSN0cdA1J8e6couzAcL82CbM1yo4f_M,122
1216
+ lightning_sdk-2025.10.8.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
1217
+ lightning_sdk-2025.10.8.dist-info/RECORD,,