lightning-sdk 0.2.21rc1__py3-none-any.whl → 0.2.23__py3-none-any.whl

lightning_sdk/__init__.py

@@ -31,6 +31,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "0.2.21.rc1"
+__version__ = "0.2.23"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/license_api.py

@@ -10,8 +10,8 @@ LICENSE_CODE = os.environ.get("LICENSE_CODE", "d9s79g79ss")
 LICENSE_SIGNING_URL = f"{env.LIGHTNING_CLOUD_URL}?settings=licenses"
 
 
-def generate_url_user_settings(redirect_to: str = LICENSE_SIGNING_URL) -> str:
-    params = urlencode({"redirectTo": redirect_to, "okbhrt": LICENSE_CODE})
+def generate_url_user_settings(name: str, redirect_to: str = LICENSE_SIGNING_URL) -> str:
+    params = urlencode({"redirectTo": redirect_to, "okbhrt": LICENSE_CODE, "licenseName": name})
     return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
 
 

lightning_sdk/api/llm_api.py

@@ -2,7 +2,9 @@ import asyncio
 import base64
 import json
 import os
-from typing import AsyncGenerator, Dict, Generator, List, Optional, Union
+import threading
+import warnings
+from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
 
 from pip._vendor.urllib3 import HTTPResponse
 
@@ -28,35 +30,40 @@ class LLMApi:
         result = self._client.assistants_service_list_assistants(user_id=user_id)
         return result.assistants
 
+    def _parse_stream_line(self, decoded_line: str) -> Optional[V1ConversationResponseChunk]:
+        try:
+            payload = json.loads(decoded_line)
+            result_data = payload.get("result", {})
+
+            choices = []
+            for choice in result_data.get("choices", []):
+                delta = choice.get("delta", {})
+                choices.append(
+                    V1ResponseChoice(
+                        delta=V1ResponseChoiceDelta(**delta),
+                        finish_reason=choice.get("finishReason"),
+                        index=choice.get("index"),
+                    )
+                )
+
+            return V1ConversationResponseChunk(
+                choices=choices,
+                conversation_id=result_data.get("conversationId"),
+                executable=result_data.get("executable"),
+                id=result_data.get("id"),
+                throughput=result_data.get("throughput"),
+            )
+        except json.JSONDecodeError:
+            warnings.warn("Error decoding JSON:", decoded_line)
+            return None
+
     def _stream_chat_response(self, result: HTTPResponse) -> Generator[V1ConversationResponseChunk, None, None]:
         for line in result.stream():
             decoded_lines = line.decode("utf-8").strip()
             for decoded_line in decoded_lines.splitlines():
-                try:
-                    payload = json.loads(decoded_line)
-                    result_data = payload.get("result", {})
-
-                    choices = []
-                    for choice in result_data.get("choices", []):
-                        delta = choice.get("delta", {})
-                        choices.append(
-                            V1ResponseChoice(
-                                delta=V1ResponseChoiceDelta(**delta),
-                                finish_reason=choice.get("finishReason"),
-                                index=choice.get("index"),
-                            )
-                        )
-
-                    yield V1ConversationResponseChunk(
-                        choices=choices,
-                        conversation_id=result_data.get("conversationId"),
-                        executable=result_data.get("executable"),
-                        id=result_data.get("id"),
-                        throughput=result_data.get("throughput"),
-                    )
-
-                except json.JSONDecodeError:
-                    print("Error decoding JSON:", decoded_line)
+                chunk = self._parse_stream_line(decoded_line)
+                if chunk:
+                    yield chunk
 
     def _encode_image_bytes_to_data_url(self, image: str) -> str:
         with open(image, "rb") as image_file:
@@ -92,6 +99,7 @@ class LLMApi:
             "stream": stream,
             "metadata": metadata or {},
             "internal_conversation": is_internal_conversation,
+            "system_prompt": system_prompt,
         }
         if images:
             for image in images:
@@ -139,6 +147,7 @@ class LLMApi:
             "stream": stream,
             "metadata": metadata or {},
             "internal_conversation": is_internal_conversation,
+            "system_prompt": system_prompt,
         }
         if images:
             for image in images:
@@ -160,7 +169,40 @@ class LLMApi:
             result = await asyncio.to_thread(thread.get)
             return result.result
 
-        raise NotImplementedError("Streaming is not supported in this client.")
+        conversation_thread = await asyncio.to_thread(
+            self._client.assistants_service_start_conversation,
+            body,
+            assistant_id,
+            async_req=True,
+            _preload_content=False,
+        )
+
+        return self.stream_response(conversation_thread)
+
+    async def stream_response(self, thread: Any) -> AsyncGenerator[V1ConversationResponseChunk, None]:
+        loop = asyncio.get_event_loop()
+        response = await asyncio.to_thread(thread.get)
+
+        queue = asyncio.Queue()
+
+        def enqueue() -> None:
+            try:
+                for line in response:
+                    decoded_lines = line.decode("utf-8").strip()
+                    for decoded_line in decoded_lines.splitlines():
+                        chunk = self._parse_stream_line(decoded_line)
+                        if chunk:
+                            asyncio.run_coroutine_threadsafe(queue.put(chunk), loop)
+            finally:
+                asyncio.run_coroutine_threadsafe(queue.put(None), loop)
+
+        threading.Thread(target=enqueue, daemon=True).start()
+
+        while True:
+            item = await queue.get()
+            if item is None:
+                break
+            yield item
 
     def list_conversations(self, assistant_id: str) -> List[str]:
         result = self._client.assistants_service_list_conversations(assistant_id)

lightning_sdk/api/pipeline_api.py

@@ -1,15 +1,20 @@
-from typing import List
+from typing import TYPE_CHECKING, List, Optional
 
 from lightning_sdk.lightning_cloud.openapi.models import (
     ProjectIdPipelinesBody,
+    ProjectIdSchedulesBody,
     V1DeletePipelineResponse,
     V1Pipeline,
     V1PipelineStep,
+    V1ScheduleResourceType,
     V1SharedFilesystem,
 )
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
+if TYPE_CHECKING:
+    from lightning_sdk.pipeline.schedule import Schedule
+
 
 class PipelineApi:
     """Internal API client for Pipeline requests (mainly http requests)."""
@@ -17,13 +22,23 @@ class PipelineApi:
     def __init__(self) -> None:
         self._client = LightningClient(retry=False, max_tries=0)
 
-    def get_pipeline_by_id(self, project_id: str, pipeline_id: str) -> V1Pipeline:
-        try:
-            return self._client.jobs_service_get_deployment(project_id=project_id, id=pipeline_id)
-        except ApiException as ex:
-            if "Reason: Not Found" in str(ex):
-                return None
-            raise ex
+    def get_pipeline_by_id(self, project_id: str, pipeline_id_or_name: str) -> Optional[V1Pipeline]:
+        if pipeline_id_or_name.startswith("pip_"):
+            try:
+                return self._client.pipelines_service_get_pipeline(project_id=project_id, id=pipeline_id_or_name)
+            except ApiException as ex:
+                if "not found" in str(ex):
+                    return None
+                raise ex
+        else:
+            try:
+                return self._client.pipelines_service_get_pipeline_by_name(
+                    project_id=project_id, name=pipeline_id_or_name
+                )
+            except ApiException as ex:
+                if "not found" in str(ex):
+                    return None
+                raise ex
 
     def create_pipeline(
         self,
@@ -31,15 +46,40 @@ class PipelineApi:
         project_id: str,
         steps: List["V1PipelineStep"],
         shared_filesystem: bool,
+        schedules: List["Schedule"],
+        parent_pipeline_id: Optional[str],
     ) -> V1Pipeline:
         body = ProjectIdPipelinesBody(
             name=name,
             steps=steps,
             shared_filesystem=V1SharedFilesystem(
                 enabled=shared_filesystem,
+                s3_folder=True,
             ),
+            parent_pipeline_id=parent_pipeline_id or "",
         )
-        return self._client.pipelines_service_create_pipeline(body, project_id)
+
+        pipeline = self._client.pipelines_service_create_pipeline(body, project_id)
+
+        # Delete the previous schedules
+        if parent_pipeline_id is not None:
+            current_schedules = self._client.schedules_service_list_schedules(project_id).schedules
+            for schedule in current_schedules:
+                self._client.schedules_service_delete_schedule(project_id, schedule.id)
+
+        if len(schedules):
+            for schedule in schedules:
+                body = ProjectIdSchedulesBody(
+                    cron_expression=schedule.cron_expression,
+                    display_name=schedule.name,
+                    resource_id=pipeline.id,
+                    parent_resource_id=parent_pipeline_id or "",
+                    resource_type=V1ScheduleResourceType.PIPELINE,
+                )
+
+                self._client.schedules_service_create_schedule(body, project_id)
+
+        return pipeline
 
     def stop(self, pipeline: V1Pipeline) -> V1Pipeline:
         body = pipeline

lightning_sdk/api/studio_api.py

@@ -123,6 +123,7 @@ class StudioApi:
         teamspace_id: str,
         cloud_account: Optional[str] = None,
         source: Optional[V1CloudSpaceSourceType] = None,
+        disable_secrets: bool = False,
     ) -> V1CloudSpace:
         """Create a Studio with a given name in a given Teamspace on a possibly given cloud_account."""
         body = ProjectIdCloudspacesBody(
@@ -131,6 +132,7 @@ class StudioApi:
             display_name=name,
             seed_files=[V1CloudSpaceSeedFile(path="main.py", contents="print('Hello, Lightning World!')\n")],
             source=source,
+            disable_secrets=disable_secrets,
         )
         studio = self._client.cloud_space_service_create_cloud_space(body, teamspace_id)
 

lightning_sdk/cli/configure.py

@@ -17,33 +17,10 @@ def configure() -> None:
     """Configure access to resources on the Lightning AI platform."""
 
 
-@configure.command(name="ssh")
-@click.option(
-    "--name",
-    default=None,
-    help=(
-        "The name of the studio to obtain SSH config. "
-        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
-    ),
-)
-@click.option(
-    "--teamspace",
-    default=None,
-    help=(
-        "The teamspace the studio is part of. "
-        "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
-        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
-    ),
-)
-@click.option(
-    "--overwrite",
-    is_flag=True,
-    flag_value=True,
-    default=False,
-    help="Whether to overwrite the SSH key and config if they already exist.",
-)
-def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
-    """Get SSH config entry for a studio."""
+def _configure_ssh_internal(
+    name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False
+) -> None:
+    """Internal function to configure SSH without Click decorators."""
     auth = Auth()
     auth.authenticate()
     console = Console()
@@ -78,6 +55,36 @@ def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite:
         console.print(f"SSH config updated at {config_path}")
 
 
+@configure.command(name="ssh")
+@click.option(
+    "--name",
+    default=None,
+    help=(
+        "The name of the studio to obtain SSH config. "
+        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
+    ),
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace the studio is part of. "
+        "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
+        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
+    ),
+)
+@click.option(
+    "--overwrite",
+    is_flag=True,
+    flag_value=True,
+    default=False,
+    help="Whether to overwrite the SSH key and config if they already exist.",
+)
+def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
+    """Get SSH config entry for a studio."""
+    _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=overwrite)
+
+
 def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
     """Download a file from a URL."""
     import requests

lightning_sdk/cli/connect.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 import click
 
-from lightning_sdk.cli.configure import ssh as configure_ssh
+from lightning_sdk.cli.configure import _configure_ssh_internal
 from lightning_sdk.cli.studios_menu import _StudiosMenu
 
 
@@ -22,7 +22,7 @@ def connect() -> None:
 )
 def studio(name: Optional[str], teamspace: Optional[str]) -> None:
     """Connect to a studio via SSH."""
-    configure_ssh(name=name, teamspace=teamspace, overwrite=False)
+    _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=False)
 
     menu = _StudiosMenu()
     studio = menu._get_studio(name=name, teamspace=teamspace)

lightning_sdk/cli/download.py

@@ -241,8 +241,8 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
 
 
 @download.command(name="licenses")
-def licenses() -> None:
-    """Download licenses for all products/packages.
+def download_licenses() -> None:
+    """Download licenses for all user's products/packages.
 
     Example:
       lightning download licenses
@@ -261,3 +261,39 @@ def licenses() -> None:
     with licenses_file.open("w") as fp:
         json.dump(licenses_short, fp, indent=4)
     Console().print(f"Licenses downloaded to {licenses_file}", style="green")
+
+
+@download.command(name="license")
+@click.argument("name")
+def download_license(name: str) -> None:
+    """Download license for specific products/packages.
+
+    Example:
+      lightning download license NAME
+
+    NAME: The name of the product/package to download the license for.
+    """
+    user = _get_authed_user()
+    api = LicenseApi()
+    licenses = api.list_user_licenses(user.id)
+    licenses_short = {ll.product_name: ll.license_key for ll in licenses if ll.is_valid}
+
+    if name not in licenses_short:
+        Console().print(f"Missing valid license for {name}", style="red")
+        return
+
+    user_home = Path.home()
+    lit_dir = user_home / ".lightning"
+    lit_dir.mkdir(parents=True, exist_ok=True)
+    licenses_file = lit_dir / "licenses.json"
+
+    licenses_loaded = {}
+    if licenses_file.exists():
+        with licenses_file.open("r") as fp:
+            licenses_loaded = json.load(fp)
+
+    licenses_loaded[name] = licenses_short[name]
+
+    with licenses_file.open("w") as fp:
+        json.dump(licenses_loaded, fp, indent=4)
+    Console().print(f"Updated license for {name} in {licenses_file}", style="green")

lightning_sdk/cli/entrypoint.py

@@ -4,8 +4,9 @@ from types import TracebackType
 from typing import Type
 
 import click
-from rich.console import Console
+from rich.console import Console, Group
 from rich.panel import Panel
+from rich.syntax import Syntax
 from rich.text import Text
 
 from lightning_sdk import __version__
@@ -34,23 +35,24 @@ from lightning_sdk.lightning_cloud.login import Auth
 
 def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:
     """CLI won't show tracebacks, just print the exception message."""
-    # if debug mode, print the traceback using rich
     console = Console()
-    if value.args:
-        message = str(value.args[0]) if value.args[0] else str(value)
-    else:
-        message = str(value) or "An unknown error occurred"
 
-    error_content = Text()
-    error_content.append(f"{exception_type.__name__}: ", style="bold red")
-    error_content.append(message, style="white")
+    message = str(value.args[0]) if value.args else str(value) or "An unknown error occurred"
+
+    error_text = Text()
+    error_text.append(f"{exception_type.__name__}: ", style="bold red")
+    error_text.append(message, style="white")
+
+    renderables = [error_text]
 
     if _LIGHTNING_DEBUG:
-        error_content.append("\n\nFull traceback:\n", style="bold yellow")
-        tb_lines = traceback.format_exception(exception_type, value, tb)
-        error_content.append("".join(tb_lines), style="dim white")
+        tb_text = "".join(traceback.format_exception(exception_type, value, tb))
+        renderables.append(Text("\n\nFull traceback:\n", style="bold yellow"))
+        renderables.append(Syntax(tb_text, "python", theme="monokai", line_numbers=False, word_wrap=True))
+    else:
+        renderables.append(Text("\n\nTo see the full traceback, set the LIGHTNING_DEBUG environment variable to 1."))
 
-    console.print(Panel(error_content, title="⚡ Lightning CLI Error", border_style="red"))
+    console.print(Panel(Group(*renderables), title="⚡ Lightning CLI Error", border_style="red"))
 
 
 @click.group(name="lightning", help="Command line interface (CLI) to interact with/manage Lightning AI Studios.")

lightning_sdk/cli/start.py

@@ -60,8 +60,11 @@ def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU", pro
 
     try:
         studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False, provider=provider)
-    except (RuntimeError, ValueError, ApiException):
-        studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False, provider=provider)
+    except (RuntimeError, ValueError, ApiException) as first_error:
+        try:
+            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False, provider=provider)
+        except (RuntimeError, ValueError, ApiException) as second_error:
+            raise first_error from second_error
 
     try:
         resolved_machine = getattr(Machine, machine.upper(), Machine(machine, machine))

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -89,6 +89,7 @@ from lightning_sdk.lightning_cloud.openapi.models.cloud_space_id_versions_body i
 from lightning_sdk.lightning_cloud.openapi.models.cloudspace_id_metric_body import CloudspaceIdMetricBody
 from lightning_sdk.lightning_cloud.openapi.models.cloudspace_id_runs_body import CloudspaceIdRunsBody
 from lightning_sdk.lightning_cloud.openapi.models.cloudspace_id_systemmetrics_body import CloudspaceIdSystemmetricsBody
+from lightning_sdk.lightning_cloud.openapi.models.cloudspace_id_visibility_body import CloudspaceIdVisibilityBody
 from lightning_sdk.lightning_cloud.openapi.models.cloudspaces_id_body import CloudspacesIdBody
 from lightning_sdk.lightning_cloud.openapi.models.cluster_id_capacityblock_body import ClusterIdCapacityblockBody
 from lightning_sdk.lightning_cloud.openapi.models.cluster_id_capacityreservations_body import ClusterIdCapacityreservationsBody
@@ -275,6 +276,8 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_cpu_system_metrics import V
 from lightning_sdk.lightning_cloud.openapi.models.v1_cancel_cloud_space_instance_switch_response import V1CancelCloudSpaceInstanceSwitchResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_cancellation_metadata import V1CancellationMetadata
 from lightning_sdk.lightning_cloud.openapi.models.v1_capacity_block_offering import V1CapacityBlockOffering
+from lightning_sdk.lightning_cloud.openapi.models.v1_check_cluster_name_availability_request import V1CheckClusterNameAvailabilityRequest
+from lightning_sdk.lightning_cloud.openapi.models.v1_check_cluster_name_availability_response import V1CheckClusterNameAvailabilityResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_check_external_service_status_response import V1CheckExternalServiceStatusResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_check_snowflake_connection_response import V1CheckSnowflakeConnectionResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_checkbox import V1Checkbox
@@ -420,6 +423,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningwork_respon
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_dataset_response import V1DeleteLitDatasetResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_dataset_version_response import V1DeleteLitDatasetVersionResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_page_response import V1DeleteLitPageResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_registry_repository_image_artifact_version_by_digest_response import V1DeleteLitRegistryRepositoryImageArtifactVersionByDigestResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_repository_response import V1DeleteLitRepositoryResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_logger_artifact_response import V1DeleteLoggerArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_managed_endpoint_response import V1DeleteManagedEndpointResponse
@@ -498,6 +502,8 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_execute_in_cloud_space_sess
 from lightning_sdk.lightning_cloud.openapi.models.v1_execute_snowflake_query_response import V1ExecuteSnowflakeQueryResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_experiment import V1Experiment
 from lightning_sdk.lightning_cloud.openapi.models.v1_export_snowflake_query_response import V1ExportSnowflakeQueryResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_external_cluster import V1ExternalCluster
+from lightning_sdk.lightning_cloud.openapi.models.v1_external_cluster_spec import V1ExternalClusterSpec
 from lightning_sdk.lightning_cloud.openapi.models.v1_external_search_user import V1ExternalSearchUser
 from lightning_sdk.lightning_cloud.openapi.models.v1_file_endpoint import V1FileEndpoint
 from lightning_sdk.lightning_cloud.openapi.models.v1_filestore_data_connection import V1FilestoreDataConnection
@@ -718,6 +724,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_lit_page_type import V1LitP
 from lightning_sdk.lightning_cloud.openapi.models.v1_lit_registry_artifact import V1LitRegistryArtifact
 from lightning_sdk.lightning_cloud.openapi.models.v1_lit_registry_project import V1LitRegistryProject
 from lightning_sdk.lightning_cloud.openapi.models.v1_lit_repository import V1LitRepository
+from lightning_sdk.lightning_cloud.openapi.models.v1_lite_published_cloud_space_response import V1LitePublishedCloudSpaceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_locked_resource import V1LockedResource
 from lightning_sdk.lightning_cloud.openapi.models.v1_log_page import V1LogPage
 from lightning_sdk.lightning_cloud.openapi.models.v1_logger_artifact import V1LoggerArtifact
@@ -910,6 +917,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_update_agent_status_respons
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_billing_subscription_request import V1UpdateBillingSubscriptionRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_cloud_space_collab_response import V1UpdateCloudSpaceCollabResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_cloud_space_publication_response import V1UpdateCloudSpacePublicationResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_update_cloud_space_visibility_response import V1UpdateCloudSpaceVisibilityResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_cluster_accelerators_request import V1UpdateClusterAcceleratorsRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_cluster_accelerators_response import V1UpdateClusterAcceleratorsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_cluster_availability_request import V1UpdateClusterAvailabilityRequest
@@ -958,6 +966,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_validate_managed_model_resp
 from lightning_sdk.lightning_cloud.openapi.models.v1_verify_verification_response import V1VerifyVerificationResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_voltage_park_direct_v1 import V1VoltageParkDirectV1
 from lightning_sdk.lightning_cloud.openapi.models.v1_volume import V1Volume
+from lightning_sdk.lightning_cloud.openapi.models.v1_volume_state import V1VolumeState
 from lightning_sdk.lightning_cloud.openapi.models.v1_vultr_direct_v1 import V1VultrDirectV1
 from lightning_sdk.lightning_cloud.openapi.models.v1_weka_data_connection import V1WekaDataConnection
 from lightning_sdk.lightning_cloud.openapi.models.v1_work import V1Work