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

lightning_sdk/job/job.py

@@ -28,17 +28,20 @@ class Job(_BaseJob):
     def __init__(
         self,
         name: str,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
-        cluster: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         *,
         _fetch_job: bool = True,
     ) -> None:
         internal_job_cls = _JobV2 if _has_jobs_v2() else _JobV1
 
         self._internal_job = internal_job_cls(
-            name=name, teamspace=teamspace, org=org, user=user, cluster=cluster, _fetch_job=_fetch_job
+            name=name,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            _fetch_job=_fetch_job,
         )
 
     @classmethod
@@ -47,14 +50,16 @@ class Job(_BaseJob):
         name: str,
         machine: "Machine",
         command: Optional[str] = None,
-        studio: Optional["Studio"] = None,
-        image: Optional[str] = None,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Union[str, None] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         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,
@@ -68,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)
@@ -81,9 +88,20 @@ class Job(_BaseJob):
         image: Optional[str] = None,
         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, command=command, studio=studio, image=image, env=env, interruptible=interruptible
+            machine=machine,
+            cluster=cluster,
+            command=command,
+            studio=studio,
+            image=image,
+            env=env,
+            interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
         )
 
     def stop(self) -> None:

lightning_sdk/job/v1.py

@@ -20,15 +20,14 @@ class _JobV1(_BaseJob):
     def __init__(
         self,
         name: str,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
-        cluster: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         *,
         _fetch_job: bool = True,
     ) -> None:
         self._job_api = JobApiV1()
-        super().__init__(name=name, teamspace=teamspace, org=org, user=user, cluster=cluster, _fetch_job=_fetch_job)
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
 
     @classmethod
     def run(
@@ -37,9 +36,9 @@ class _JobV1(_BaseJob):
         machine: "Machine",
         command: str,
         studio: "Studio",
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         cluster: Optional[str] = None,
         interruptible: bool = False,
     ) -> "_BaseJob":
@@ -55,6 +54,8 @@ class _JobV1(_BaseJob):
             cluster=cluster,
             env=None,
             interruptible=interruptible,
+            image_credentials=None,
+            cluster_auth=False,
         )
 
     def _submit(
@@ -65,11 +66,14 @@ class _JobV1(_BaseJob):
         image: Optional[str] = None,
         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:
@@ -85,7 +89,7 @@ class _JobV1(_BaseJob):
             command=command,
             studio_id=studio._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._cluster,
+            cluster_id=cluster,
             machine=machine,
             interruptible=interruptible,
         )

lightning_sdk/job/v2.py

@@ -16,15 +16,14 @@ class _JobV2(_BaseJob):
     def __init__(
         self,
         name: str,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
-        cluster: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         *,
         _fetch_job: bool = True,
     ) -> None:
         self._job_api = JobApiV2()
-        super().__init__(name=name, teamspace=teamspace, org=org, user=user, cluster=cluster, _fetch_job=_fetch_job)
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
 
     def _submit(
         self,
@@ -34,6 +33,9 @@ class _JobV2(_BaseJob):
         image: Optional[str] = None,
         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
@@ -55,13 +57,15 @@ class _JobV2(_BaseJob):
         submitted = self._job_api.submit_job(
             name=self.name,
             command=command,
-            cluster_id=self._cluster,
+            cluster_id=cluster,
             teamspace_id=self._teamspace.id,
             studio_id=studio_id,
             image=image,
             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 = {}