diff --git a/python/lightning_sdk/api/job_api.py b/python/lightning_sdk/api/job_api.py index 527f489d..cfdf0928 100644 --- a/python/lightning_sdk/api/job_api.py +++ b/python/lightning_sdk/api/job_api.py @@ -496,6 +496,8 @@ def stream_logs( return # idle timeout on a snapshot: treat the stream as finished except WebSocketConnectionClosedException: break # fall through to reconnect logic + if isinstance(message, bytes): + message = message.decode("utf-8", errors="replace") if message == "": break yield from _decode_log_messages(message, timestamps=timestamps) diff --git a/python/lightning_sdk/api/logs_api.py b/python/lightning_sdk/api/logs_api.py index cde0602b..fa6b9904 100644 --- a/python/lightning_sdk/api/logs_api.py +++ b/python/lightning_sdk/api/logs_api.py @@ -541,6 +541,8 @@ def follow( continue except (WebSocketConnectionClosedException, OSError): break # fall through to the reconnect decision + if isinstance(message, bytes): + message = message.decode("utf-8", errors="replace") if message == "": break entries = parse_log_entries(message) diff --git a/python/lightning_sdk/studio.py b/python/lightning_sdk/studio.py index 2e44741b..55b890c1 100644 --- a/python/lightning_sdk/studio.py +++ b/python/lightning_sdk/studio.py @@ -12,7 +12,7 @@ from lightning_sdk.api.utils import AccessibleResource, raise_access_error_if_not_allowed from lightning_sdk.base_studio import BaseStudio from lightning_sdk.constants import _LIGHTNING_DEBUG -from lightning_sdk.exceptions import NotSupportedError, OutOfCapacityError +from lightning_sdk.exceptions import OutOfCapacityError from lightning_sdk.lightning_cloud.openapi import V1ClusterType, V1Endpoint from lightning_sdk.machine import DEFAULT_MACHINE, CloudProvider, Machine from lightning_sdk.organization import Organization @@ -357,7 +357,11 @@ def start( Raises: RuntimeError: If the Studio is already running on a different machine or is not stopped. - RuntimeError: If the requested machine is not supported or has no available capacity. + RuntimeError: If the requested machine has no available capacity. + + Warns: + UserWarning: If the requested machine is not a known machine type for the selected + cloud account. It is treated as a custom instance type and passed through as-is. """ # Check to see if we're inside a studio and if its running current_studio_machine = None @@ -410,9 +414,10 @@ def start( if not self._studio_api.machine_is_supported( new_machine, self._teamspace.id, self.cloud_account, _get_org_id(self._teamspace) ): - raise NotSupportedError( - "Requested machine is not supported in the selected cloud account. " - "Try a different machine or cloud account by setting the `machine` or `cloud_account` argument." + warnings.warn( + f"Machine {new_machine} is a custom instance type that hasn't been vetted by Lightning. " + "It may not be available in the selected cloud account and startup may fail. " + "Continue at your own risk." ) if not self._studio_api.machine_has_capacity( diff --git a/python/tests/core/studio/test_studio_start.py b/python/tests/core/studio/test_studio_start.py index 7abbfbf7..5f4f480f 100644 --- a/python/tests/core/studio/test_studio_start.py +++ b/python/tests/core/studio/test_studio_start.py @@ -1,6 +1,8 @@ import os from unittest import mock +import pytest + from lightning_sdk.lightning_cloud.openapi import ( CloudSpaceServiceCreateCloudSpaceBody, CloudSpaceServiceStartCloudSpaceInstanceBody, @@ -41,6 +43,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): return _list_cloudspaces_side_effect +@pytest.mark.parametrize("machine_supported", [True, False]) @mock.patch( "lightning_sdk.lightning_cloud.openapi.api.cluster_service_api.ClusterServiceApi.cluster_service_list_default_cluster_accelerators", autospec=True, @@ -99,6 +102,7 @@ def test_studio_start( mock_list_project_clusters, mock_list_clusters, mock_list_accelerators, + machine_supported, ): # Setup state from internal_studio_start_mocker status = {"st-abc": None} @@ -207,7 +211,7 @@ def _create_lightning_run_side_effect(body, project_id, cloudspace_id, **kwargs) mock_list_cloudspaces.side_effect = list_cloudspaces_side_effect(existing_studios) mock_create_cloudspace.side_effect = _create_cloudspace_side_effect mock_create_lightning_run.side_effect = _create_lightning_run_side_effect - mock_machine_is_supported.return_value = True + mock_machine_is_supported.return_value = machine_supported # Setup teamspace and org mocks mock_get_teamspace.return_value = V1Project( @@ -225,7 +229,12 @@ def _create_lightning_run_side_effect(body, project_id, cloudspace_id, **kwargs) assert studio.machine is None assert studio.teamspace.start_studios_on_interruptible is True - studio.start() + if machine_supported: + studio.start() + else: + # unvetted machines emit a warning but still start + with pytest.warns(UserWarning, match="custom instance type"): + studio.start() assert studio.status == Status.Running assert studio.interruptible is True