lightning-sdk 2025.7.10__py3-none-any.whl → 2025.7.22__py3-none-any.whl

lightning_sdk/cli/start.py

@@ -4,10 +4,10 @@ import click
 
 from lightning_sdk import Machine, Studio
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
-from lightning_sdk.studio import Provider
+from lightning_sdk.machine import CloudProvider
 
 _MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
-_PROVIDER_VALUES = tuple([provider.value for provider in Provider])
+_PROVIDER_VALUES = tuple([provider.value for provider in CloudProvider])
 
 
 @click.group("start")
@@ -35,14 +35,28 @@ def start() -> None:
     type=click.Choice(_MACHINE_VALUES),
     help="The machine type to start the studio on.",
 )
+@click.option(
+    "--cloud-provider",
+    default=None,
+    type=click.Choice(_PROVIDER_VALUES),
+    help=("The provider to create the studio on. If --cloud-account is specified, this option is prioritized."),
+)
 @click.option(
     "--provider",
     default=None,
-    show_default=True,
     type=click.Choice(_PROVIDER_VALUES),
-    help="The provider to start the studio on.",
+    help=(
+        "Deprecated. Use --cloud-provider instead. The provider to create the studio on. "
+        "If --cloud-account is specified, this option is prioritized."
+    ),
 )
-def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU", provider: Optional[str] = None) -> None:
+def studio(
+    name: str,
+    teamspace: Optional[str] = None,
+    machine: str = "CPU",
+    cloud_provider: Optional[str] = None,
+    provider: Optional[str] = None,
+) -> None:
     """Start a studio on a given machine.
 
     Example:
@@ -59,10 +73,26 @@ def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU", pro
         owner, teamspace = None, None
 
     try:
-        studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False, provider=provider)
+        studio = Studio(
+            name=name,
+            teamspace=teamspace,
+            org=owner,
+            user=None,
+            create_ok=False,
+            cloud_provider=cloud_provider,
+            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)
+            studio = Studio(
+                name=name,
+                teamspace=teamspace,
+                org=None,
+                user=owner,
+                create_ok=False,
+                cloud_provider=cloud_provider,
+                provider=provider,
+            )
         except (RuntimeError, ValueError, ApiException) as second_error:
             raise first_error from second_error
 

lightning_sdk/deployment/deployment.py

@@ -130,6 +130,7 @@ class Deployment:
         include_credentials: Optional[bool] = None,
         from_onboarding: Optional[bool] = None,
         from_litserve: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
     ) -> None:
         """The Lightning AI Deployment.
 
@@ -157,6 +158,10 @@ class Deployment:
             quantity: The number of machines per replica to deploy.
             include_credentials: Whether to include the environment variables for the SDK to authenticate
             from_onboarding: Whether the deployment is from onboarding.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
 
         Note:
             Since a teamspace can either be owned by an org or by a user directly,
@@ -227,6 +232,7 @@ class Deployment:
                     quantity=quantity,
                     include_credentials=include_credentials if include_credentials is not None else True,
                     cloudspace_id=cloudspace_id,
+                    max_runtime=max_runtime,
                 ),
                 strategy=to_strategy(release_strategy),
             ),
@@ -262,6 +268,7 @@ class Deployment:
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
         quantity: Optional[int] = None,
         include_credentials: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
     ) -> None:
         cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
 
@@ -288,6 +295,7 @@ class Deployment:
             release_strategy=release_strategy,
             quantity=quantity,
             include_credentials=include_credentials,
+            max_runtime=max_runtime,
         )
 
     def stop(self) -> None:

lightning_sdk/job/base.py

@@ -2,11 +2,12 @@ import warnings
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union
 
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, in_studio
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.status import Status
     from lightning_sdk.studio import Studio
@@ -64,6 +65,7 @@ class _BaseJob(ABC):
             self._update_internal_job()
 
         self._prevent_refetch_latest = False
+        self._cloud_account_api = CloudAccountApi()
 
     @classmethod
     def run(
@@ -77,6 +79,7 @@ class _BaseJob(ABC):
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
@@ -85,6 +88,7 @@ class _BaseJob(ABC):
         artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseJob":
         """Run async workloads using a docker image or a compute environment from your studio.
@@ -101,7 +105,11 @@ class _BaseJob(ABC):
             user: The user owning the teamspace (if any). Defaults to the current user.
             cloud_account: The cloud account to run the job on.
                 Defaults to the studio cloud account if running with studio compute env.
-                If not provided will fall back to the teamspaces default cloud account.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             env: Environment variables to set inside the job.
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             image_credentials: The credentials used to pull the image. Required if the image is private.
@@ -122,6 +130,10 @@ class _BaseJob(ABC):
                 just runs the provided command in a standard shell.
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
         """
         from lightning_sdk.lightning_cloud.openapi.rest import ApiException
         from lightning_sdk.studio import Studio
@@ -199,6 +211,7 @@ class _BaseJob(ABC):
         return inst._submit(
             machine=machine,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             command=command,
             studio=studio,
             image=image,
@@ -210,6 +223,7 @@ class _BaseJob(ABC):
             artifacts_remote=artifacts_remote,
             entrypoint=entrypoint,
             path_mappings=path_mappings,
+            max_runtime=max_runtime,
         )
 
     @abstractmethod
@@ -222,12 +236,14 @@ class _BaseJob(ABC):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
     ) -> "_BaseJob":
         """Submit a new job to the Lightning AI platform.
 
@@ -241,7 +257,11 @@ class _BaseJob(ABC):
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             cloud_account: The cloud account to run the job on.
                 Defaults to the studio cloud account if running with studio compute env.
-                If not provided will fall back to the teamspaces default cloud account.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
@@ -259,6 +279,10 @@ class _BaseJob(ABC):
             entrypoint: The entrypoint of your docker container. Defaults to sh -c.
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
         """
 
     @abstractmethod
@@ -275,12 +299,13 @@ class _BaseJob(ABC):
         Caution: This also deletes all artifacts and snapshots associated with the job.
         """
 
-    def wait(self, interval: float = 5.0, timeout: Optional[float] = None) -> None:
+    def wait(self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False) -> None:
         """Waits for the job to be either completed, manually stopped or failed.
 
         Args:
             interval: The number of seconds to spend in-between status checks.
             timeout: The maximum number of seconds to wait before raising an error. If None, waits forever.
+            stop_on_timeout: Whether to stop the job if it didn't finish within the timeout.
         """
         import time
 
@@ -292,16 +317,21 @@ class _BaseJob(ABC):
                 break
 
             if timeout is not None and time.time() - start > timeout:
+                if stop_on_timeout:
+                    self.stop()  # ensure the job is stopped if it didn't finish
                 raise TimeoutError("Job didn't finish within the provided timeout.")
 
             time.sleep(interval)
 
-    async def async_wait(self, interval: float = 5.0, timeout: Optional[float] = None) -> None:
+    async def async_wait(
+        self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False
+    ) -> None:
         """Waits for the job to be either completed, manually stopped or failed.
 
         Args:
             interval: The number of seconds to spend in-between status checks.
             timeout: The maximum number of seconds to wait before raising an error. If None, waits forever.
+            stop_on_timeout: Whether to stop the job if it didn't finish within the timeout.
         """
         import asyncio
 
@@ -313,6 +343,8 @@ class _BaseJob(ABC):
                 break
 
             if timeout is not None and asyncio.get_event_loop().time() - start > timeout:
+                if stop_on_timeout:
+                    self.stop()  # ensure the job is stopped if it didn't finish
                 raise TimeoutError("Job didn't finish within the provided timeout.")
 
             await asyncio.sleep(interval)

lightning_sdk/job/job.py

@@ -8,7 +8,7 @@ from lightning_sdk.utils.resolve import _setup_logger
 _logger = _setup_logger(__name__)
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.status import Status
     from lightning_sdk.studio import Studio
@@ -86,12 +86,14 @@ class Job(_BaseJob):
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in terms of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in terms of path_mappings
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
@@ -108,9 +110,13 @@ class Job(_BaseJob):
         teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
         org: The organization owning the teamspace (if any). Defaults to the current organization.
         user: The user owning the teamspace (if any). Defaults to the current user.
-        cloud_account: The cloud acocunt to run the job on.
+        cloud_account: The cloud account to run the job on.
             Defaults to the studio cloud account if running with studio compute env.
-        If not provided will fall back to the teamspaces default cloud account.
+            If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+            will fall back to the teamspaces default cloud account.
+        cloud_account_provider: The provider to select the cloud-account from.
+            If set, must be in agreement with the provider from the cloud_account (if specified).
+            If not specified, falls backto the teamspace default cloud account.
         env: Environment variables to set inside the job.
         interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
         image_credentials: The credentials used to pull the image. Required if the image is private.
@@ -130,6 +136,10 @@ class Job(_BaseJob):
                 }
             If the path inside the connection is omitted it's assumed to be the root path of that connection.
             Only applicable when submitting docker jobs.
+        max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
         """
         ret_val = super().run(
             name=name,
@@ -141,6 +151,7 @@ class Job(_BaseJob):
             org=org,
             user=user,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
@@ -149,6 +160,7 @@ class Job(_BaseJob):
             artifacts_remote=artifacts_remote,
             entrypoint=entrypoint,
             path_mappings=path_mappings,
+            max_runtime=max_runtime,
             cluster=cluster,
         )
         # required for typing with "Job"
@@ -166,12 +178,14 @@ class Job(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
         artifacts_local: Optional[str] = None,  # deprecated in terms of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in terms of path_mappings
+        max_runtime: Optional[int] = None,
     ) -> "Job":
         """Submit a new job to the Lightning AI platform.
 
@@ -185,7 +199,11 @@ class Job(_BaseJob):
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             cloud_account: The cloud account to run the job on.
                 Defaults to the studio cloud account if running with studio compute env.
-                If not provided will fall back to the teamspaces default cloud account.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
@@ -202,10 +220,15 @@ class Job(_BaseJob):
                     }
                 If the path inside the connection is omitted it's assumed to be the root path of that connection.
                 Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
         """
         self._job = self._internal_job._submit(
             machine=machine,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             command=command,
             studio=studio,
             image=image,
@@ -217,6 +240,7 @@ class Job(_BaseJob):
             path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            max_runtime=max_runtime,
         )
         return self
 

lightning_sdk/job/v1.py

@@ -5,7 +5,7 @@ from lightning_sdk.job.base import _BaseJob
 from lightning_sdk.status import Status
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
@@ -51,6 +51,7 @@ class _JobV1(_BaseJob):
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[str] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseJob":
@@ -80,12 +81,14 @@ class _JobV1(_BaseJob):
             org=org,
             user=user,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             env=None,
             interruptible=interruptible,
             image_credentials=None,
             cloud_account_auth=False,
             cluster=cluster,
             path_mappings=None,
+            max_runtime=None,
         )
 
     def _submit(
@@ -97,12 +100,14 @@ class _JobV1(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
     ) -> "_JobV1":
         """Submit a job to run on a machine.
 
@@ -121,6 +126,10 @@ class _JobV1(_BaseJob):
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
             path_mappings: The mappings from data connection inside your container (not supported)
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
 
         Returns:
             The submitted job.

lightning_sdk/job/v2.py

@@ -6,7 +6,7 @@ from lightning_sdk.job.base import _BaseJob
 from lightning_sdk.status import Status
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
@@ -44,10 +44,12 @@ class _JobV2(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> "_JobV2":
@@ -80,6 +82,10 @@ class _JobV2(_BaseJob):
                     }
                 If the path inside the connection is omitted it's assumed to be the root path of that connection.
                 Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -98,6 +104,13 @@ class _JobV2(_BaseJob):
             if image is None:
                 raise ValueError("either image or studio must be provided")
 
+        cloud_account = self._cloud_account_api.resolve_cloud_account(
+            self._teamspace.id,
+            cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
+            default_cloud_account=self._teamspace.default_cloud_account,
+        )
+
         submitted = self._job_api.submit_job(
             name=self.name,
             command=command,
@@ -114,6 +127,7 @@ class _JobV2(_BaseJob):
             artifacts_remote=artifacts_remote,
             entrypoint=entrypoint,
             path_mappings=path_mappings,
+            max_runtime=max_runtime,
         )
         self._job = submitted
         self._name = submitted.name

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -102,6 +102,7 @@ from lightning_sdk.lightning_cloud.openapi.models.clusters_id_body import Cluste
 from lightning_sdk.lightning_cloud.openapi.models.clusters_id_body1 import ClustersIdBody1
 from lightning_sdk.lightning_cloud.openapi.models.command_argument_command_argument_type import CommandArgumentCommandArgumentType
 from lightning_sdk.lightning_cloud.openapi.models.conversations_id_body import ConversationsIdBody
+from lightning_sdk.lightning_cloud.openapi.models.conversations_id_body1 import ConversationsIdBody1
 from lightning_sdk.lightning_cloud.openapi.models.create import Create
 from lightning_sdk.lightning_cloud.openapi.models.create_checkout_session_request_wallet_type import CreateCheckoutSessionRequestWalletType
 from lightning_sdk.lightning_cloud.openapi.models.create_deployment_request_defines_a_spec_for_the_job_that_allows_for_autoscaling_jobs import CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs
@@ -165,6 +166,7 @@ from lightning_sdk.lightning_cloud.openapi.models.litloggermetrics_id_body impor
 from lightning_sdk.lightning_cloud.openapi.models.litpages_id_body import LitpagesIdBody
 from lightning_sdk.lightning_cloud.openapi.models.litregistry_lit_repo_name_body import LitregistryLitRepoNameBody
 from lightning_sdk.lightning_cloud.openapi.models.loggermetrics_id_body import LoggermetricsIdBody
+from lightning_sdk.lightning_cloud.openapi.models.messages_id_body import MessagesIdBody
 from lightning_sdk.lightning_cloud.openapi.models.metrics_stream_id_loggerartifacts_body import MetricsStreamIdLoggerartifactsBody
 from lightning_sdk.lightning_cloud.openapi.models.metricsstream_create_body import MetricsstreamCreateBody
 from lightning_sdk.lightning_cloud.openapi.models.metricsstream_delete_body import MetricsstreamDeleteBody
@@ -208,6 +210,7 @@ from lightning_sdk.lightning_cloud.openapi.models.project_id_storage_body import
 from lightning_sdk.lightning_cloud.openapi.models.projects_id_body import ProjectsIdBody
 from lightning_sdk.lightning_cloud.openapi.models.projects_project_id_body import ProjectsProjectIdBody
 from lightning_sdk.lightning_cloud.openapi.models.protobuf_any import ProtobufAny
+from lightning_sdk.lightning_cloud.openapi.models.protobuf_null_value import ProtobufNullValue
 from lightning_sdk.lightning_cloud.openapi.models.query_query_id_body import QueryQueryIdBody
 from lightning_sdk.lightning_cloud.openapi.models.rpc_status import RpcStatus
 from lightning_sdk.lightning_cloud.openapi.models.schedules_id_body import SchedulesIdBody
@@ -239,6 +242,7 @@ from lightning_sdk.lightning_cloud.openapi.models.usagerestrictions_id_body impo
 from lightning_sdk.lightning_cloud.openapi.models.user_id_affiliatelinks_body import UserIdAffiliatelinksBody
 from lightning_sdk.lightning_cloud.openapi.models.user_id_membershiprolebindings_body import UserIdMembershiprolebindingsBody
 from lightning_sdk.lightning_cloud.openapi.models.user_id_membershiprolebindings_body1 import UserIdMembershiprolebindingsBody1
+from lightning_sdk.lightning_cloud.openapi.models.user_id_upgradetrigger_body import UserIdUpgradetriggerBody
 from lightning_sdk.lightning_cloud.openapi.models.v1_aws_cluster_credentials import V1AWSClusterCredentials
 from lightning_sdk.lightning_cloud.openapi.models.v1_aws_direct_v1 import V1AWSDirectV1
 from lightning_sdk.lightning_cloud.openapi.models.v1_aws_direct_v1_status import V1AWSDirectV1Status
@@ -251,6 +255,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_agent_job import V1AgentJob
 from lightning_sdk.lightning_cloud.openapi.models.v1_agent_job_artifact import V1AgentJobArtifact
 from lightning_sdk.lightning_cloud.openapi.models.v1_agent_upload_multipart_url import V1AgentUploadMultipartUrl
 from lightning_sdk.lightning_cloud.openapi.models.v1_agent_upload_part_response import V1AgentUploadPartResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_ai_pod_v1 import V1AiPodV1
 from lightning_sdk.lightning_cloud.openapi.models.v1_alert_method import V1AlertMethod
 from lightning_sdk.lightning_cloud.openapi.models.v1_alerts_config import V1AlertsConfig
 from lightning_sdk.lightning_cloud.openapi.models.v1_api_pricing_spec import V1ApiPricingSpec
@@ -263,6 +268,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_assistant import V1Assistan
 from lightning_sdk.lightning_cloud.openapi.models.v1_assistant_knowledge_item_status import V1AssistantKnowledgeItemStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_assistant_knowledge_status import V1AssistantKnowledgeStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_assistant_model_status import V1AssistantModelStatus
+from lightning_sdk.lightning_cloud.openapi.models.v1_assistant_session_daily_aggregated import V1AssistantSessionDailyAggregated
 from lightning_sdk.lightning_cloud.openapi.models.v1_author import V1Author
 from lightning_sdk.lightning_cloud.openapi.models.v1_auto_join_domain_validation import V1AutoJoinDomainValidation
 from lightning_sdk.lightning_cloud.openapi.models.v1_auto_join_org_response import V1AutoJoinOrgResponse
@@ -358,6 +364,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_count_metrics_streams_respo
 from lightning_sdk.lightning_cloud.openapi.models.v1_create_agent_multipart_upload_response import V1CreateAgentMultipartUploadResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_create_billing_portal_session_request import V1CreateBillingPortalSessionRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_create_billing_portal_session_response import V1CreateBillingPortalSessionResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_create_billing_upgrade_trigger_record_response import V1CreateBillingUpgradeTriggerRecordResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_create_blog_post_request import V1CreateBlogPostRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_create_checkout_session_request import V1CreateCheckoutSessionRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_create_checkout_session_response import V1CreateCheckoutSessionResponse
@@ -526,6 +533,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_work import V1Fi
 from lightning_sdk.lightning_cloud.openapi.models.v1_find_capacity_block_offering_response import V1FindCapacityBlockOfferingResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_flowserver import V1Flowserver
 from lightning_sdk.lightning_cloud.openapi.models.v1_folder_index_status import V1FolderIndexStatus
+from lightning_sdk.lightning_cloud.openapi.models.v1_function_tool import V1FunctionTool
 from lightning_sdk.lightning_cloud.openapi.models.v1_gcp_direct_vpc import V1GCPDirectVPC
 from lightning_sdk.lightning_cloud.openapi.models.v1_gcs_folder_data_connection import V1GCSFolderDataConnection
 from lightning_sdk.lightning_cloud.openapi.models.v1_gpu_system_metrics import V1GPUSystemMetrics
@@ -539,6 +547,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_get_affiliate_link_response
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_agent_job_env_response import V1GetAgentJobEnvResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_agent_job_logs_metadata_response import V1GetAgentJobLogsMetadataResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_artifacts_page_response import V1GetArtifactsPageResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_get_assistant_session_daily_aggregated_response import V1GetAssistantSessionDailyAggregatedResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_cloud_space_cold_start_metrics_stats_response import V1GetCloudSpaceColdStartMetricsStatsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_cloud_space_instance_open_ports_response import V1GetCloudSpaceInstanceOpenPortsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_cloud_space_instance_status_response import V1GetCloudSpaceInstanceStatusResponse
@@ -629,6 +638,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_lightningwork_spec import V
 from lightning_sdk.lightning_cloud.openapi.models.v1_lightningwork_state import V1LightningworkState
 from lightning_sdk.lightning_cloud.openapi.models.v1_lightningwork_status import V1LightningworkStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_lightningwork_status_reason import V1LightningworkStatusReason
+from lightning_sdk.lightning_cloud.openapi.models.v1_like_status import V1LikeStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_affiliate_links_response import V1ListAffiliateLinksResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_agent_job_artifacts_response import V1ListAgentJobArtifactsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_agent_jobs_response import V1ListAgentJobsResponse
@@ -699,7 +709,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_model_versions_respons
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_models_response import V1ListModelsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_multi_machine_job_events_response import V1ListMultiMachineJobEventsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_multi_machine_jobs_response import V1ListMultiMachineJobsResponse
-from lightning_sdk.lightning_cloud.openapi.models.v1_list_new_features_for_user_response import V1ListNewFeaturesForUserResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_list_notification_dialogs_response import V1ListNotificationDialogsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_org_members_response import V1ListOrgMembersResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_org_membership_role_binding_response import V1ListOrgMembershipRoleBindingResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_org_memberships_response import V1ListOrgMembershipsResponse
@@ -719,6 +729,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_project_memberships_re
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_project_roles_response import V1ListProjectRolesResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_published_cloud_spaces_response import V1ListPublishedCloudSpacesResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_published_deployment_templates_response import V1ListPublishedDeploymentTemplatesResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_list_published_managed_endpoint_models_response import V1ListPublishedManagedEndpointModelsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_quests_response import V1ListQuestsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_slurm_cluster_users_response import V1ListSLURMClusterUsersResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_ssh_public_keys_response import V1ListSSHPublicKeysResponse
@@ -918,6 +929,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_system_metrics_list import
 from lightning_sdk.lightning_cloud.openapi.models.v1_telemetry import V1Telemetry
 from lightning_sdk.lightning_cloud.openapi.models.v1_timestamp_code_telemetry import V1TimestampCodeTelemetry
 from lightning_sdk.lightning_cloud.openapi.models.v1_token_usage import V1TokenUsage
+from lightning_sdk.lightning_cloud.openapi.models.v1_tool import V1Tool
 from lightning_sdk.lightning_cloud.openapi.models.v1_transaction import V1Transaction
 from lightning_sdk.lightning_cloud.openapi.models.v1_transfer_cloud_space_response import V1TransferCloudSpaceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_transfer_org_balance_response import V1TransferOrgBalanceResponse
@@ -935,6 +947,8 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_update_cloud_space_visibili
 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
+from lightning_sdk.lightning_cloud.openapi.models.v1_update_conversation_like_response import V1UpdateConversationLikeResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_update_conversation_message_like_response import V1UpdateConversationMessageLikeResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_deployment_visibility_response import V1UpdateDeploymentVisibilityResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_index_response import V1UpdateIndexResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_job_visibility_response import V1UpdateJobVisibilityResponse