fix(sdk): fail promptly when Studio startup terminates

python/lightning_sdk/api/studio_api.py

@@ -6,7 +6,6 @@
 from threading import Event, Thread
 from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
-import backoff
 import requests
 from tqdm import tqdm
 
@@ -44,6 +43,7 @@
     V1Assistant,
     V1CloudSpace,
     V1CloudSpaceInstanceConfig,
+    V1CloudSpaceInstanceState,
     V1CloudSpaceSeedFile,
     V1CloudSpaceSourceType,
     V1CloudSpaceState,
@@ -229,24 +229,12 @@ def get_studio_status(self, studio_id: str, teamspace_id: str) -> V1GetCloudSpac
             id=studio_id,
         )
 
-    @backoff.on_exception(backoff.expo, AttributeError, max_tries=10)
     def _check_code_status_top_up_restore_finished(self, studio_id: str, teamspace_id: str) -> bool:
-        """Retries checking the top_up_restore_finished value of the code status when there's an AttributeError.
-
-        Args:
-            studio_id: Studio (cloud space) ID.
-            teamspace_id: ID of the owning teamspace.
-
-        Returns:
-            ``True`` if the top-up restore has finished, ``False`` otherwise.
-        """
-        if (
-            self.get_studio_status(studio_id, teamspace_id) is None
-            or self.get_studio_status(studio_id, teamspace_id).in_use is None
-        ):
+        """Check readiness using a single Studio status snapshot."""
+        status = self.get_studio_status(studio_id, teamspace_id)
+        if status is None or status.in_use is None or status.in_use.startup_status is None:
             return False
-        startup_status = self.get_studio_status(studio_id, teamspace_id).in_use.startup_status
-        return startup_status and startup_status.top_up_restore_finished
+        return bool(status.in_use.startup_status.top_up_restore_finished)
 
     def start_studio(
         self,
@@ -264,7 +252,22 @@ def start_studio(
             machine: Machine type to start the Studio on.
             interruptible: Whether to use a spot (interruptible) instance.
             max_runtime: Optional maximum runtime in seconds before the Studio is stopped.
+
+        Raises:
+            RuntimeError: If the started instance fails or stops before becoming ready.
         """
+        terminal_phases = {
+            V1CloudSpaceInstanceState.FAILED,
+            V1CloudSpaceInstanceState.STOPPING,
+            V1CloudSpaceInstanceState.STOPPED,
+            V1CloudSpaceInstanceState.DELETED,
+        }
+        previous_status = self.get_studio_status(studio_id, teamspace_id)
+        previous_terminal_ids = {
+            instance.cloud_space_instance_id
+            for instance in ((previous_status.requested, previous_status.in_use) if previous_status else ())
+            if instance is not None and instance.phase in terminal_phases
+        }
         # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
         optional_kwargs_compute_body = {}
 
@@ -282,14 +285,54 @@ def start_studio(
             studio_id,
         )
 
+        instance_id = None
+        status_message = None
         while True:
-            if self._check_code_status_top_up_restore_finished(studio_id, teamspace_id):
-                break
+            status = self.get_studio_status(studio_id, teamspace_id)
+            if status is None:
+                time.sleep(1)
+                continue
+            instances = [instance for instance in (status.requested, status.in_use) if instance is not None]
+            if instance_id is None:
+                # The start response has no operation ID. Ignore old stopped instances until the new one appears.
+                instance = next(
+                    (
+                        candidate
+                        for candidate in instances
+                        if candidate.cloud_space_instance_id
+                        and candidate.cloud_space_instance_id not in previous_terminal_ids
+                    ),
+                    None,
+                )
+                if instance is not None:
+                    instance_id = instance.cloud_space_instance_id
+            else:
+                instance = next(
+                    (candidate for candidate in instances if candidate.cloud_space_instance_id == instance_id), None
+                )
+                if instance is None:
+                    reason = status_message or "instance disappeared before becoming ready"
+                    raise RuntimeError(f"Studio {studio_id} failed to start (instance {instance_id}): {reason}")
+
+            if instance is not None:
+                status_message = instance.status_message or status_message
+                if instance.phase in terminal_phases:
+                    reason = status_message or "instance stopped before becoming ready"
+                    raise RuntimeError(
+                        f"Studio {studio_id} failed to start (instance {instance_id}, {instance.phase}): {reason}"
+                    )
+                in_use = status.in_use
+                startup_status = in_use.startup_status if in_use is not None else None
+                if (
+                    in_use is not None
+                    and in_use.cloud_space_instance_id == instance_id
+                    and startup_status
+                    and startup_status.top_up_restore_finished
+                ):
+                    break
             time.sleep(1)
 
         if _LIGHTNING_DEBUG:
-            code_status = self.get_studio_status(studio_id, teamspace_id)
-            instance_id = code_status.in_use.cloud_space_instance_id
             print(f"Studio started | {teamspace_id=} {studio_id=} {instance_id=}")
 
     def start_studio_async(

python/lightning_sdk/studio.py

@@ -358,6 +358,7 @@ def start(
         Raises:
             RuntimeError: If the Studio is already running on a different machine or is not stopped.
             RuntimeError: If the requested machine has no available capacity.
+            RuntimeError: If the started instance fails or stops before becoming ready.
 
         Warns:
             UserWarning: If the requested machine is not a known machine type for the selected

python/tests/api/test_studio_api.py

@@ -39,6 +39,7 @@
     V1Resources,
     V1SearchUser,
     V1SearchUsersResponse,
+    V1StartCloudSpaceInstanceResponse,
     V1UserRequestedComputeConfig,
 )
 from lightning_sdk.machine import Machine
@@ -281,16 +282,19 @@ def side_effect(*args, **kwargs):
 )
 @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock())
 def test_start_studio(mock_start_instance, mock_get_status):
+    mock_start_instance.return_value = V1StartCloudSpaceInstanceResponse()
     mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse(
         requested=Externalv1CloudSpaceInstanceStatus(
+            cloud_space_instance_id="instance-new",
             startup_status=V1CloudSpaceInstanceStartupStatus(
                 initial_restore_finished=True, top_up_restore_finished=True
-            )
+            ),
         ),
         in_use=Externalv1CloudSpaceInstanceStatus(
+            cloud_space_instance_id="instance-new",
             startup_status=V1CloudSpaceInstanceStartupStatus(
                 initial_restore_finished=True, top_up_restore_finished=True
-            )
+            ),
         ),
     )
 
@@ -708,6 +712,7 @@ def test_duplicate_user(
     )
     mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse(
         in_use=Externalv1CloudSpaceInstanceStatus(
+            cloud_space_instance_id="instance-new",
             startup_status=V1CloudSpaceInstanceStartupStatus(
                 initial_restore_finished=True, top_up_restore_finished=True
             ),
@@ -771,6 +776,7 @@ def test_duplicate_org(
     )
     mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse(
         in_use=Externalv1CloudSpaceInstanceStatus(
+            cloud_space_instance_id="instance-new",
             startup_status=V1CloudSpaceInstanceStartupStatus(
                 initial_restore_finished=True, top_up_restore_finished=True
             ),
@@ -832,6 +838,7 @@ def test_duplicate_user_with_new_name(
     )
     mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse(
         in_use=Externalv1CloudSpaceInstanceStatus(
+            cloud_space_instance_id="instance-new",
             startup_status=V1CloudSpaceInstanceStartupStatus(
                 initial_restore_finished=True, top_up_restore_finished=True
             ),
@@ -901,6 +908,7 @@ def test_duplicate_org_with_new_name(
     )
     mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse(
         in_use=Externalv1CloudSpaceInstanceStatus(
+            cloud_space_instance_id="instance-new",
             startup_status=V1CloudSpaceInstanceStartupStatus(
                 initial_restore_finished=True, top_up_restore_finished=True
             ),

python/tests/api/test_studio_start_failure.py

@@ -0,0 +1,135 @@
+from unittest import mock
+
+import pytest
+
+from lightning_sdk.api.studio_api import StudioApi
+from lightning_sdk.lightning_cloud.openapi import (
+    Externalv1CloudSpaceInstanceStatus,
+    V1CloudSpaceInstanceStartupStatus,
+    V1StartCloudSpaceInstanceResponse,
+)
+from lightning_sdk.lightning_cloud.openapi import (
+    V1CloudSpaceInstanceState as State,
+)
+from lightning_sdk.lightning_cloud.openapi import (
+    V1GetCloudSpaceInstanceStatusResponse as Status,
+)
+from lightning_sdk.machine import Machine
+
+
+def instance(instance_id="new", phase=State.PENDING, message=None, ready=False):
+    return Externalv1CloudSpaceInstanceStatus(
+        cloud_space_instance_id=instance_id,
+        phase=phase,
+        status_message=message,
+        startup_status=V1CloudSpaceInstanceStartupStatus(top_up_restore_finished=ready),
+    )
+
+
+@pytest.fixture()
+def start_api(mocker):
+    mocker.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock())
+    api = StudioApi()
+    api._client = mock.MagicMock()
+    api._client.cloud_space_service_start_cloud_space_instance.return_value = V1StartCloudSpaceInstanceResponse()
+    sleep = mocker.patch("lightning_sdk.api.studio_api.time.sleep")
+    return api, sleep
+
+
+@pytest.mark.parametrize("phase", [State.FAILED, State.STOPPING, State.STOPPED, State.DELETED])
+@pytest.mark.parametrize("first_poll", [True, False])
+def test_start_raises_on_terminal_instance(start_api, phase, first_poll):
+    api, sleep = start_api
+    statuses = [Status(in_use=instance("old", State.STOPPED))]
+    if not first_poll:
+        statuses.append(Status(in_use=instance()))
+    statuses.append(Status(in_use=instance(phase=phase, message="workload exited (status code: 2)")))
+    api._client.cloud_space_service_get_cloud_space_instance_status.side_effect = statuses
+
+    with pytest.raises(RuntimeError, match=f"instance new, {phase}.*workload exited.*2"):
+        api.start_studio("studio", "teamspace", Machine.CPU)
+
+    assert sleep.call_count == (0 if first_poll else 1)
+    assert api._client.cloud_space_service_get_cloud_space_instance_status.call_count == len(statuses)
+
+
+def test_start_waits_for_new_instance_and_in_use_readiness(start_api):
+    api, sleep = start_api
+    old = instance("old", State.DELETED, "previous failure", ready=True)
+    statuses = [
+        Status(in_use=old),
+        Status(in_use=old),
+        Status(),
+        Status(in_use=old, requested=instance(ready=True)),
+        Status(in_use=instance(phase=State.RUNNING, ready=True)),
+    ]
+    api._client.cloud_space_service_get_cloud_space_instance_status.side_effect = statuses
+
+    api.start_studio("studio", "teamspace", Machine.CPU)
+
+    assert sleep.call_count == 3
+    assert api._client.cloud_space_service_get_cloud_space_instance_status.call_count == len(statuses)
+
+
+@pytest.mark.parametrize("final_status", [Status(), Status(in_use=instance("other", State.RUNNING, ready=True))])
+def test_start_raises_if_observed_instance_disappears(start_api, final_status):
+    api, sleep = start_api
+    api._client.cloud_space_service_get_cloud_space_instance_status.side_effect = [
+        Status(),
+        Status(requested=instance(message="last observed diagnostic")),
+        final_status,
+    ]
+
+    with pytest.raises(RuntimeError, match="instance new.*last observed diagnostic"):
+        api.start_studio("studio", "teamspace", Machine.CPU)
+    sleep.assert_called_once_with(1)
+
+
+def test_start_reports_terminal_phase_without_message(start_api):
+    api, _ = start_api
+    api._client.cloud_space_service_get_cloud_space_instance_status.side_effect = [
+        Status(),
+        Status(in_use=instance(phase=State.STOPPING, ready=True)),
+    ]
+    with pytest.raises(RuntimeError, match="STOPPING.*stopped before becoming ready"):
+        api.start_studio("studio", "teamspace", Machine.CPU)
+
+
+def test_start_accepts_healthy_running_instance(start_api):
+    api, sleep = start_api
+    api._client.cloud_space_service_get_cloud_space_instance_status.return_value = Status(
+        in_use=instance(phase=State.RUNNING, ready=True)
+    )
+    api.start_studio("studio", "teamspace", Machine.CPU)
+    sleep.assert_not_called()
+
+
+def test_start_accepts_same_instance_in_requested_and_in_use(start_api):
+    api, sleep = start_api
+    api._client.cloud_space_service_get_cloud_space_instance_status.return_value = Status(
+        requested=instance(), in_use=instance(phase=State.RUNNING, ready=True)
+    )
+    api.start_studio("studio", "teamspace", Machine.CPU)
+    sleep.assert_not_called()
+
+
+def test_start_waits_through_missing_status(start_api):
+    api, sleep = start_api
+    api._client.cloud_space_service_get_cloud_space_instance_status.side_effect = [
+        None,
+        None,
+        Status(in_use=instance()),
+        None,
+        Status(in_use=instance(phase=State.RUNNING, ready=True)),
+    ]
+    api.start_studio("studio", "teamspace", Machine.CPU)
+    assert sleep.call_count == 3
+
+
+@pytest.mark.parametrize("status", [None, Status(), Status(in_use=Externalv1CloudSpaceInstanceStatus())])
+def test_readiness_helper_handles_missing_status(start_api, status):
+    api, _ = start_api
+    getter = api._client.cloud_space_service_get_cloud_space_instance_status
+    getter.return_value = status
+    assert not api._check_code_status_top_up_restore_finished("studio", "teamspace")
+    getter.assert_called_once()

python/tests/conftest.py

@@ -423,14 +423,16 @@ def internal_studio_api_mocker_start_studio(mocker):
         "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_get_cloud_space_instance_status",
         return_value=V1GetCloudSpaceInstanceStatusResponse(
             requested=Externalv1CloudSpaceInstanceStatus(
+                cloud_space_instance_id="instance-new",
                 startup_status=V1CloudSpaceInstanceStartupStatus(
                     initial_restore_finished=True, top_up_restore_finished=True
-                )
+                ),
             ),
             in_use=Externalv1CloudSpaceInstanceStatus(
+                cloud_space_instance_id="instance-new",
                 startup_status=V1CloudSpaceInstanceStartupStatus(
                     initial_restore_finished=True, top_up_restore_finished=True
-                )
+                ),
             ),
         ),
         autospec=True,