lightning-sdk 0.1.37__py3-none-any.whl → 0.1.38__py3-none-any.whl

lightning_sdk/__init__.py

@@ -27,5 +27,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.37"
+__version__ = "0.1.38"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/ai_hub.py

@@ -47,17 +47,7 @@ class AIHub:
             including its name, description, creation and update timestamps,
             parameters, tags, job specifications, and autoscaling settings.
         """
-        template = self._api.api_info(api_id)
-
-        api_arguments = [
-            {
-                "name": param.name,
-                "short_description": param.short_description,
-                "required": param.required,
-                "default": param.input.default_value,
-            }
-            for param in template.parameter_spec.parameters
-        ]
+        template, api_arguments = self._api.api_info(api_id)
 
         return {
             "name": template.name,
@@ -129,33 +119,33 @@ class AIHub:
             raise ValueError("You need to pass a teamspace or an org for your deployment.")
         return teamspace
 
-    def deploy(
+    def run(
         self,
         api_id: str,
-        cluster: Optional[str] = None,
+        api_arguments: Optional[Dict[str, Any]] = None,
         name: Optional[str] = None,
+        cluster: Optional[str] = None,
         teamspace: Optional[Union[str, "Teamspace"]] = None,
         org: Optional[Union[str, "Organization"]] = None,
-        api_arguments: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Union[str, bool]]:
         """Deploy an API from the AI Hub.
 
         Example:
             from lightning_sdk import AIHub
             hub = AIHub()
-            deployment = hub.deploy("temp_xxxx")
+            deployment = hub.run("temp_xxxx")
 
             # Using API arguments
-            api_arugments = {"model": "unitary/toxic-bert", "batch_size" 10, "token": "lit_xxxx"}
-            deployment = hub.deploy("temp_xxxx", api_arugments=api_arugments)
+            api_arugments = {"model": "unitary/toxic-bert", "batch_size": 10, "token": "lit_xxxx"}
+            deployment = hub.run("temp_xxxx", api_arugments=api_arugments)
 
         Args:
-            api_id: The ID of the API you want to deploy.
-            cluster: The cluster where you want to deploy the API, such as "lightning-public-prod". Defaults to None.
+            api_id: The ID of the AIHub template you want to run.
+            api_arguments: Additional API argument, such as model name, or batch size.
             name: Name for the deployed API. Defaults to None.
+            cluster: The cluster where you want to run the template, such as "lightning-public-prod". Defaults to None.
             teamspace: The team or group for deployment. Defaults to None.
             org: The organization for deployment. Defaults to None.
-            api_arguments: Additional API argument, such as model name, or batch size.
 
         Returns:
             A dictionary containing the name of the deployed API,
@@ -169,14 +159,22 @@ class AIHub:
         teamspace_id = teamspace.id
 
         api_arguments = api_arguments or {}
-        deployment = self._api.deploy_api(
+        deployment = self._api.run_api(
             template_id=api_id, cluster_id=cluster, project_id=teamspace_id, name=name, api_arguments=api_arguments
         )
-        url = quote(f"{LIGHTNING_CLOUD_URL}/{teamspace._org.name}/{teamspace.name}/jobs/{deployment.name}", safe=":/()")
+        url = (
+            quote(
+                f"{LIGHTNING_CLOUD_URL}/{teamspace._org.name}/{teamspace.name}/jobs/{deployment.name}",
+                safe=":/()",
+            )
+            + "?app_id=deployment"
+        )
         print("Deployment available at:", url)
+
         return {
             "id": deployment.id,
             "name": deployment.name,
-            "base_url": deployment.status.urls[0],
+            "deployment_url": url,
+            "api_endpoint": deployment.status.urls[0],
             "interruptible": deployment.spec.spot,
         }

lightning_sdk/api/ai_hub_api.py

@@ -1,5 +1,5 @@
 import traceback
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import backoff
 
@@ -22,15 +22,40 @@ class AIHubApi:
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=3)
 
-    def api_info(self, api_id: str) -> "V1DeploymentTemplate":
+    def api_info(self, api_id: str) -> Tuple[V1DeploymentTemplate, List[Dict[str, str]]]:
         try:
-            return self._client.deployment_templates_service_get_deployment_template(api_id)
+            template = self._client.deployment_templates_service_get_deployment_template(api_id)
         except Exception as e:
             stack_trace = traceback.format_exc()
             if "record not found" in stack_trace:
                 raise ValueError(f"api_id={api_id} not found.") from None
             raise e
 
+        api_arguments = []
+        for param in template.parameter_spec.parameters:
+            default = None
+            if param.type == V1DeploymentTemplateParameterType.INPUT and param.input:
+                default = param.input.default_value
+            if param.type == V1DeploymentTemplateParameterType.SELECT and param.select:
+                default = param.select.options[0]
+            if param.type == V1DeploymentTemplateParameterType.CHECKBOX and param.checkbox:
+                default = (
+                    (param.checkbox.true_value or "True")
+                    if param.checkbox.is_checked
+                    else (param.checkbox.false_value or "False")
+                )
+
+            api_arguments.append(
+                {
+                    "name": param.name,
+                    "short_description": param.short_description,
+                    "required": param.required,
+                    "type": param.type,
+                    "default": default,
+                }
+            )
+        return template, api_arguments
+
     @backoff.on_predicate(backoff.expo, lambda x: not x, max_tries=5)
     def list_apis(self, search_query: str) -> List[V1DeploymentTemplateGalleryResponse]:
         kwargs = {"show_globally_visible": True}
@@ -84,7 +109,7 @@ class AIHubApi:
 
         return job
 
-    def deploy_api(
+    def run_api(
         self, template_id: str, project_id: str, cluster_id: str, name: Optional[str], api_arguments: Dict[str, str]
     ) -> V1Deployment:
         template = self._client.deployment_templates_service_get_deployment_template(template_id)

lightning_sdk/api/job_api.py

@@ -141,13 +141,15 @@ class JobApiV2:
         self,
         name: str,
         command: Optional[str],
-        cluster_id: str,
+        cluster_id: Optional[str],
         teamspace_id: str,
         studio_id: Optional[str],
         image: Optional[str],
         machine: Machine,
         interruptible: bool,
         env: Optional[Dict[str, str]],
+        image_credentials: Optional[str],
+        cluster_auth: bool,
     ) -> V1Job:
         env_vars = []
         if env is not None:
@@ -160,13 +162,15 @@ class JobApiV2:
 
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
-            cluster_id=cluster_id,
+            cluster_id=cluster_id or "",
             command=command or "",
             env=env_vars,
             image=image or "",
             instance_name=instance_name,
             run_id=run_id,
             spot=interruptible,
+            image_cluster_credentials=cluster_auth,
+            image_secret_ref=image_credentials or "",
         )
         body = ProjectIdJobsBody(name=name, spec=spec)
 

lightning_sdk/api/teamspace_api.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Dict, List, Optional
 
-from lightning_sdk.api.utils import _download_model_files, _DummyBody, _ModelFileUploader
+from lightning_sdk.api.utils import _download_model_files, _DummyBody, _get_model_version, _ModelFileUploader
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     ModelIdVersionsBody,
@@ -159,6 +159,16 @@ class TeamspaceApi:
 
         return self._client.assistants_service_create_assistant(body=body, project_id=teamspace_id)
 
+    # lazy property which is only created when needed
+    @property
+    def models(self) -> ModelsStoreApi:
+        if not self._models:
+            self._models = ModelsStoreApi(self._client.api_client)
+        return self._models
+
+    def get_model_version(self, name: str, version: str, teamspace_id: str) -> V1ModelVersionArchive:
+        return _get_model_version(client=self._client, name=name, version=version, teamspace_id=teamspace_id)
+
     def create_model(
         self,
         name: str,
@@ -167,17 +177,15 @@ class TeamspaceApi:
         teamspace_id: str,
         cluster_id: str,
     ) -> V1ModelVersionArchive:
-        if not self._models:
-            self._models = ModelsStoreApi(self._client.api_client)
         # ask if such model already exists by listing models with specific name
-        models = self._models.models_store_list_models(project_id=teamspace_id, name=name).models
+        models = self.models.models_store_list_models(project_id=teamspace_id, name=name).models
         if len(models) == 0:
-            return self._models.models_store_create_model(
+            return self.models.models_store_create_model(
                 body=ProjectIdModelsBody(cluster_id=cluster_id, metadata=metadata, name=name, private=private),
                 project_id=teamspace_id,
             )
         assert len(models) == 1, "Multiple models with the same name found"
-        return self._models.models_store_create_model_version(
+        return self.models.models_store_create_model_version(
             body=ModelIdVersionsBody(cluster_id=cluster_id),
             project_id=teamspace_id,
             model_id=models[0].id,
@@ -185,18 +193,16 @@ class TeamspaceApi:
 
     def delete_model(self, name: str, version: Optional[str], teamspace_id: str) -> None:
         """Delete a model or a version from the model store."""
-        if not self._models:
-            self._models = ModelsStoreApi(self._client.api_client)
-        models = self._models.models_store_list_models(project_id=teamspace_id, name=name).models
+        models = self.models.models_store_list_models(project_id=teamspace_id, name=name).models
         assert len(models) == 1, "Multiple models with the same name found"
         model_id = models[0].id
         # decide if delete only version of whole model
         if version:
             if version == "latest":
                 version = models[0].latest_version
-            self._models.models_store_delete_model_version(project_id=teamspace_id, model_id=model_id, version=version)
+            self.models.models_store_delete_model_version(project_id=teamspace_id, model_id=model_id, version=version)
         else:
-            self._models.models_store_delete_model(project_id=teamspace_id, model_id=model_id)
+            self.models.models_store_delete_model(project_id=teamspace_id, model_id=model_id)
 
     def upload_model_file(
         self,
@@ -242,9 +248,7 @@ class TeamspaceApi:
             )
 
     def complete_model_upload(self, model_id: str, version: str, teamspace_id: str) -> None:
-        if not self._models:
-            self._models = ModelsStoreApi(self._client.api_client)
-        self._models.models_store_complete_model_upload(
+        self.models.models_store_complete_model_upload(
             body=_DummyBody(),
             project_id=teamspace_id,
             model_id=model_id,

lightning_sdk/api/utils.py

@@ -28,6 +28,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1UploadProjectArtifactResponse,
     VersionUploadsBody,
 )
+from lightning_sdk.lightning_cloud.openapi.models.v1_model_version_archive import V1ModelVersionArchive
 
 try:
     from lightning_sdk.lightning_cloud.openapi import AppsIdBody1 as AppsIdBody
@@ -492,6 +493,24 @@ class _FileDownloader:
         os.rename(tmp_filename, self.local_path)
 
 
+def _get_model_version(client: LightningClient, teamspace_id: str, name: str, version: str) -> V1ModelVersionArchive:
+    api = ModelsStoreApi(client.api_client)
+    models = api.models_store_list_models(project_id=teamspace_id, name=name).models
+    if not models:
+        raise ValueError(f"Model `{name}` does not exist")
+    elif len(models) > 1:
+        raise ValueError("Multiple models with the same name found")
+    if version == "latest":
+        return models[0].latest_version
+    versions = api.models_store_list_model_versions(project_id=teamspace_id, model_id=models[0].id).versions
+    if not versions:
+        raise ValueError(f"Model `{name}` does not have any versions")
+    for ver in versions:
+        if ver.version == version:
+            return ver
+    raise ValueError(f"Model `{name}` does not have version `{version}`")
+
+
 def _download_model_files(
     client: LightningClient,
     teamspace_id: str,

lightning_sdk/cli/ai_hub.py

@@ -46,4 +46,4 @@ class _AIHub(_StudiosMenu):
           teamspace: Teamspace to deploy the API to. Defaults to user's default teamspace.
           org: Organization to deploy the API to. Defaults to user's default organization.
         """
-        return self._hub.deploy(api_id, cluster=cluster, name=name, teamspace=teamspace, org=org)
+        return self._hub.run(api_id, cluster=cluster, name=name, teamspace=teamspace, org=org)

lightning_sdk/cli/run.py

@@ -36,6 +36,10 @@ class _Run:
                 If not provided will fall back to the teamspaces default cluster.
             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.
+                This should be the name of the respective credentials secret created on the Lightning AI platform.
+            cluster_auth: Whether to authenticate with the cluster to pull the image.
+                Required if the registry is part of a cluster provider (e.g. ECR).
         """
         self.job.__func__.__doc__ = docstr
 
@@ -55,6 +59,8 @@ class _Run:
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> None:
         machine_enum = Machine(machine.upper())
         Job.run(
@@ -69,4 +75,6 @@ class _Run:
             cluster=cluster,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
         )

lightning_sdk/job/base.py

@@ -50,6 +50,8 @@ class _BaseJob(ABC):
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> "_BaseJob":
         from lightning_sdk.studio import Studio
 
@@ -80,6 +82,13 @@ class _BaseJob(ABC):
                     "Studio cluster does not match provided cluster. "
                     "Can only run jobs with Studio envs in the same cluster."
                 )
+
+            if image_credentials is not None:
+                raise ValueError("image_credentials is only supported when using a custom image")
+
+            if cluster_auth:
+                raise ValueError("cluster_auth is only supported when using a custom image")
+
         else:
             if studio is not None:
                 raise RuntimeError(
@@ -95,6 +104,8 @@ class _BaseJob(ABC):
             image=image,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
         )
         return inst
 
@@ -108,6 +119,8 @@ class _BaseJob(ABC):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> None:
         """Submits a job and updates the internal _job attribute as well as the _name attribute."""
 

lightning_sdk/job/job.py

@@ -58,6 +58,8 @@ class Job(_BaseJob):
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> "Job":
         ret_val = super().run(
             name=name,
@@ -71,6 +73,8 @@ class Job(_BaseJob):
             cluster=cluster,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
         )
         # required for typing with "Job"
         assert isinstance(ret_val, cls)
@@ -85,6 +89,8 @@ class Job(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> None:
         return self._internal_job._submit(
             machine=machine,
@@ -94,6 +100,8 @@ class Job(_BaseJob):
             image=image,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
         )
 
     def stop(self) -> None:

lightning_sdk/job/v1.py

@@ -54,6 +54,8 @@ class _JobV1(_BaseJob):
             cluster=cluster,
             env=None,
             interruptible=interruptible,
+            image_credentials=None,
+            cluster_auth=False,
         )
 
     def _submit(
@@ -65,11 +67,13 @@ class _JobV1(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> None:
         if studio is None:
             raise ValueError("Studio is required for submitting jobs")
 
-        if image is not None:
+        if image is not None or image_credentials is not None or cluster_auth:
             raise ValueError("Image is not supported for submitting jobs")
 
         if env is not None:

lightning_sdk/job/v2.py

@@ -34,6 +34,8 @@ class _JobV2(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> None:
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -62,6 +64,8 @@ class _JobV2(_BaseJob):
             machine=machine,
             interruptible=interruptible,
             env=env,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
         )
         self._job = submitted
         self._name = submitted.name

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -138,6 +138,7 @@ from lightning_sdk.lightning_cloud.openapi.models.metricsstream_delete_body impo
 from lightning_sdk.lightning_cloud.openapi.models.metricsstream_id_body import MetricsstreamIdBody
 from lightning_sdk.lightning_cloud.openapi.models.model_id_versions_body import ModelIdVersionsBody
 from lightning_sdk.lightning_cloud.openapi.models.models_model_id_body import ModelsModelIdBody
+from lightning_sdk.lightning_cloud.openapi.models.multimachinejobs_id_body import MultimachinejobsIdBody
 from lightning_sdk.lightning_cloud.openapi.models.multipartuploads_upload_id_body import MultipartuploadsUploadIdBody
 from lightning_sdk.lightning_cloud.openapi.models.org_id_memberships_body import OrgIdMembershipsBody
 from lightning_sdk.lightning_cloud.openapi.models.org_id_roles_body import OrgIdRolesBody
@@ -157,6 +158,7 @@ from lightning_sdk.lightning_cloud.openapi.models.project_id_invite_body import
 from lightning_sdk.lightning_cloud.openapi.models.project_id_jobs_body import ProjectIdJobsBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_memberships_body import ProjectIdMembershipsBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_models_body import ProjectIdModelsBody
+from lightning_sdk.lightning_cloud.openapi.models.project_id_multimachinejobs_body import ProjectIdMultimachinejobsBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_projectclustersbindings_body import ProjectIdProjectclustersbindingsBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_secrets_body import ProjectIdSecretsBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_serviceexecution_body import ProjectIdServiceexecutionBody
@@ -200,6 +202,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_api_pricing_spec import V1ApiPricingSpec
 from lightning_sdk.lightning_cloud.openapi.models.v1_app_type import V1AppType
 from lightning_sdk.lightning_cloud.openapi.models.v1_append_logger_metrics_response import V1AppendLoggerMetricsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_approve_auto_join_domain_response import V1ApproveAutoJoinDomainResponse
@@ -340,6 +343,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_delete_managed_model_respon
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_metrics_stream_response import V1DeleteMetricsStreamResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_model_response import V1DeleteModelResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_model_version_response import V1DeleteModelVersionResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_delete_multi_machine_job_response import V1DeleteMultiMachineJobResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_org_membership_response import V1DeleteOrgMembershipResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_org_membership_role_binding_response import V1DeleteOrgMembershipRoleBindingResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_org_role_response import V1DeleteOrgRoleResponse
@@ -364,6 +368,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_deployment import V1Deploym
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_event import V1DeploymentEvent
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_event_type import V1DeploymentEventType
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_metrics import V1DeploymentMetrics
+from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_performance import V1DeploymentPerformance
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_release import V1DeploymentRelease
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_spec import V1DeploymentSpec
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_state import V1DeploymentState
@@ -388,9 +393,8 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_drive_status import V1Drive
 from lightning_sdk.lightning_cloud.openapi.models.v1_drive_type import V1DriveType
 from lightning_sdk.lightning_cloud.openapi.models.v1_drive_type_spec import V1DriveTypeSpec
 from lightning_sdk.lightning_cloud.openapi.models.v1_drive_type_status import V1DriveTypeStatus
-from lightning_sdk.lightning_cloud.openapi.models.v1_efs_folder_data_connection import V1EFSFolderDataConnection
 from lightning_sdk.lightning_cloud.openapi.models.v1_ebs import V1Ebs
-from lightning_sdk.lightning_cloud.openapi.models.v1_efs_data_connection import V1EfsDataConnection
+from lightning_sdk.lightning_cloud.openapi.models.v1_efs_config import V1EfsConfig
 from lightning_sdk.lightning_cloud.openapi.models.v1_endpoint import V1Endpoint
 from lightning_sdk.lightning_cloud.openapi.models.v1_endpoint_auth import V1EndpointAuth
 from lightning_sdk.lightning_cloud.openapi.models.v1_endpoint_prewarm import V1EndpointPrewarm
@@ -478,6 +482,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_job_timing import V1JobTimi
 from lightning_sdk.lightning_cloud.openapi.models.v1_joinable_organization import V1JoinableOrganization
 from lightning_sdk.lightning_cloud.openapi.models.v1_keep_alive_cloud_space_instance_response import V1KeepAliveCloudSpaceInstanceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_knowledge_configuration import V1KnowledgeConfiguration
+from lightning_sdk.lightning_cloud.openapi.models.v1_lambda_labs_direct_v1 import V1LambdaLabsDirectV1
 from lightning_sdk.lightning_cloud.openapi.models.v1_lightning_app_user import V1LightningAppUser
 from lightning_sdk.lightning_cloud.openapi.models.v1_lightning_auth import V1LightningAuth
 from lightning_sdk.lightning_cloud.openapi.models.v1_lightning_basic_auth import V1LightningBasicAuth
@@ -552,6 +557,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_memberships_response i
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_metrics_streams_response import V1ListMetricsStreamsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_model_versions_response import V1ListModelVersionsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_models_response import V1ListModelsResponse
+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_org_members_response import V1ListOrgMembersResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_org_membership_role_binding_response import V1ListOrgMembershipRoleBindingResponse
@@ -587,6 +593,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_login_response import V1Log
 from lightning_sdk.lightning_cloud.openapi.models.v1_logout_request import V1LogoutRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_logout_response import V1LogoutResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_logs_response import V1LogsResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_machines_selector import V1MachinesSelector
 from lightning_sdk.lightning_cloud.openapi.models.v1_magic_link_login_request import V1MagicLinkLoginRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_magic_link_login_response import V1MagicLinkLoginResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_managed_endpoint import V1ManagedEndpoint
@@ -595,6 +602,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_membership import V1Members
 from lightning_sdk.lightning_cloud.openapi.models.v1_message import V1Message
 from lightning_sdk.lightning_cloud.openapi.models.v1_message_author import V1MessageAuthor
 from lightning_sdk.lightning_cloud.openapi.models.v1_message_content import V1MessageContent
+from lightning_sdk.lightning_cloud.openapi.models.v1_message_content_type import V1MessageContentType
 from lightning_sdk.lightning_cloud.openapi.models.v1_metadata import V1Metadata
 from lightning_sdk.lightning_cloud.openapi.models.v1_metric_value import V1MetricValue
 from lightning_sdk.lightning_cloud.openapi.models.v1_metrics import V1Metrics
@@ -604,6 +612,9 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_metrics_tracker import V1Me
 from lightning_sdk.lightning_cloud.openapi.models.v1_model import V1Model
 from lightning_sdk.lightning_cloud.openapi.models.v1_model_version_archive import V1ModelVersionArchive
 from lightning_sdk.lightning_cloud.openapi.models.v1_mount_target import V1MountTarget
+from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job import V1MultiMachineJob
+from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_state import V1MultiMachineJobState
+from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_status import V1MultiMachineJobStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_named_get_logger_metrics import V1NamedGetLoggerMetrics
 from lightning_sdk.lightning_cloud.openapi.models.v1_network_config import V1NetworkConfig
 from lightning_sdk.lightning_cloud.openapi.models.v1_new_feature import V1NewFeature

lightning_sdk/lightning_cloud/openapi/api/cluster_service_api.py

@@ -2197,6 +2197,7 @@ class ClusterServiceApi(object):
         >>> result = thread.get()
 
         :param async_req bool
+        :param bool include_pricing:
         :return: V1ListDefaultClusterAcceleratorsResponse
                  If the method is called asynchronously,
                  returns the request thread.
@@ -2217,12 +2218,13 @@ class ClusterServiceApi(object):
         >>> result = thread.get()
 
         :param async_req bool
+        :param bool include_pricing:
         :return: V1ListDefaultClusterAcceleratorsResponse
                  If the method is called asynchronously,
                  returns the request thread.
         """
 
-        all_params = []  # noqa: E501
+        all_params = ['include_pricing']  # noqa: E501
         all_params.append('async_req')
         all_params.append('_return_http_data_only')
         all_params.append('_preload_content')
@@ -2243,6 +2245,8 @@ class ClusterServiceApi(object):
         path_params = {}
 
         query_params = []
+        if 'include_pricing' in params:
+            query_params.append(('includePricing', params['include_pricing']))  # noqa: E501
 
         header_params = {}