lightning-sdk 0.1.51__py3-none-any.whl → 0.1.53__py3-none-any.whl

lightning_sdk/__init__.py

@@ -29,5 +29,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.51"
+__version__ = "0.1.53"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/ai_hub.py

@@ -1,11 +1,10 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 from urllib.parse import quote
 
-from lightning_sdk.api import AIHubApi, UserApi
-from lightning_sdk.lightning_cloud import login
-from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
+from lightning_sdk.api import AIHubApi
+from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_org, _resolve_teamspace
+from lightning_sdk.utils.resolve import _resolve_teamspace
 
 if TYPE_CHECKING:
     from lightning_sdk import Organization, Teamspace
@@ -99,26 +98,6 @@ class AIHub:
             results.append(result)
         return results
 
-    def _authenticate(
-        self,
-        teamspace: Optional[Union[str, "Teamspace"]] = None,
-        org: Optional[Union[str, "Organization"]] = None,
-        user: Optional[Union[str, "User"]] = None,
-    ) -> "Teamspace":
-        if self._auth is None:
-            self._auth = login.Auth()
-        try:
-            self._auth.authenticate()
-            user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
-        except ConnectionError as e:
-            raise e
-
-        org = _resolve_org(org)
-        teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user if org is None else None)
-        if teamspace is None:
-            raise ValueError("You need to pass a teamspace or an org for your deployment.")
-        return teamspace
-
     def run(
         self,
         api_id: str,
@@ -127,6 +106,7 @@ class AIHub:
         cloud_account: Optional[str] = None,
         teamspace: Optional[Union[str, "Teamspace"]] = None,
         org: Optional[Union[str, "Organization"]] = None,
+        user: Optional[Union[str, "User"]] = None,
     ) -> Dict[str, Union[str, bool]]:
         """Deploy an API from the AI Hub.
 
@@ -146,7 +126,8 @@ class AIHub:
             cloud_account: The cloud account 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.
+            org: The organization for deployment. Don't pass user with this. Defaults to None.
+            user: The user for deployment. Don't pass org with this. Defaults to None.
 
         Returns:
             A dictionary containing the name of the deployed API,
@@ -156,7 +137,13 @@ class AIHub:
             ValueError: If a teamspace or organization is not provided.
             ConnectionError: If there is an issue with logging in.
         """
-        teamspace = self._authenticate(teamspace, org)
+        if user is not None and org is not None:
+            raise ValueError("User and org are mutually exclusive. Please only specify the one owns the teamspace.")
+
+        teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        if teamspace is None:
+            raise ValueError("You need to pass a teamspace or an org for your deployment.")
+
         teamspace_id = teamspace.id
 
         api_arguments = api_arguments or {}
@@ -167,13 +154,15 @@ class AIHub:
             name=name,
             api_arguments=api_arguments,
         )
+
         url = (
             quote(
-                f"{LIGHTNING_CLOUD_URL}/{teamspace._org.name}/{teamspace.name}/jobs/{deployment.name}",
+                f"{_get_cloud_url}/{teamspace.owner.name}/{teamspace.name}/jobs/{deployment.name}",
                 safe=":/()",
             )
             + "?app_id=deployment"
         )
+
         print("Deployment available at:", url)
 
         return {

lightning_sdk/api/ai_hub_api.py

@@ -116,6 +116,12 @@ class AIHubApi:
         name = name or template.name
         template.spec_v2.endpoint.id = None
 
+        # These are needed to ensure templates with a max replicas of 0 will start on creation
+        if template.spec_v2.autoscaling.max_replicas == "0":
+            template.spec_v2.autoscaling.max_replicas = "1"
+        if not template.spec_v2.autoscaling.enabled:
+            template.spec_v2.autoscaling.enabled = True
+
         AIHubApi._set_parameters(template.spec_v2.job, template.parameter_spec.parameters, api_arguments)
         return self._client.jobs_service_create_deployment(
             project_id=project_id,
@@ -124,7 +130,7 @@ class AIHubApi:
                 cluster_id=cloud_account,
                 endpoint=template.spec_v2.endpoint,
                 name=name,
-                replicas=0,
+                replicas=1,
                 spec=template.spec_v2.job,
             ),
         )

lightning_sdk/cli/run.py

@@ -166,16 +166,7 @@ class _Run:
             cloud_account = resolved_teamspace.default_cloud_account
         machine_enum = Machine(machine.upper())
 
-        path_mappings_dict = {}
-        for mapping in path_mappings.split(","):
-            splits = str(mapping).split(":", 1)
-            if len(splits) != 2:
-                raise RuntimeError(
-                    "Mapping needs to be of form <CONTAINER_PATH>:<CONNECTION_NAME>[:<PATH_WITHIN_CONNECTION>], "
-                    f"but got {mapping}"
-                )
-
-            path_mappings_dict[splits[0]] = splits[1]
+        path_mappings_dict = self._resolve_path_mapping(path_mappings=path_mappings)
 
         Job.run(
             name=name,
@@ -242,16 +233,7 @@ class _Run:
         if image is None:
             raise RuntimeError("Image needs to be specified to run a multi-machine job")
 
-        path_mappings_dict = {}
-        for mapping in path_mappings.split(","):
-            splits = str(mapping).split(":", 1)
-            if len(splits) != 2:
-                raise RuntimeError(
-                    "Mapping needs to be of form <CONTAINER_PATH>:<CONNECTION_NAME>[:<PATH_WITHIN_CONNECTION>], "
-                    f"but got {mapping}"
-                )
-
-            path_mappings_dict[splits[0]] = splits[1]
+        path_mappings_dict = self._resolve_path_mapping(path_mappings=path_mappings)
 
         MMT.run(
             name=name,
@@ -273,3 +255,26 @@ class _Run:
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
+
+    @staticmethod
+    def _resolve_path_mapping(path_mappings: str) -> Dict[str, str]:
+        path_mappings = path_mappings.strip()
+
+        if not path_mappings:
+            return {}
+
+        path_mappings_dict = {}
+        for mapping in path_mappings.split(","):
+            if not mapping.strip():
+                continue
+
+            splits = str(mapping).split(":", 1)
+            if len(splits) != 2:
+                raise RuntimeError(
+                    "Mapping needs to be of form <CONTAINER_PATH>:<CONNECTION_NAME>[:<PATH_WITHIN_CONNECTION>], "
+                    f"but got {mapping}"
+                )
+
+            path_mappings_dict[splits[0].strip()] = splits[1].strip()
+
+        return path_mappings_dict

lightning_sdk/job/base.py

@@ -275,6 +275,22 @@ class _BaseJob(ABC):
         Caution: This also deletes all artifacts and snapshots associated with the job.
         """
 
+    def wait(self, interval: float = 5.0) -> 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.
+        """
+        import time
+
+        from lightning_sdk.status import Status
+
+        while True:
+            if self.status in (Status.Completed, Status.Stopped, Status.Failed):
+                break
+
+            time.sleep(interval)
+
     @property
     @abstractmethod
     def status(self) -> "Status":

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -80,6 +80,7 @@ from lightning_sdk.lightning_cloud.openapi.models.cloudspaces_id_body import Clo
 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
 from lightning_sdk.lightning_cloud.openapi.models.cluster_id_proxies_body import ClusterIdProxiesBody
+from lightning_sdk.lightning_cloud.openapi.models.cluster_id_usagerestrictions_body import ClusterIdUsagerestrictionsBody
 from lightning_sdk.lightning_cloud.openapi.models.clusters_id_body import ClustersIdBody
 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
@@ -105,6 +106,7 @@ from lightning_sdk.lightning_cloud.openapi.models.id_artifacts_body import IdArt
 from lightning_sdk.lightning_cloud.openapi.models.id_codeconfig_body import IdCodeconfigBody
 from lightning_sdk.lightning_cloud.openapi.models.id_collaborate_body import IdCollaborateBody
 from lightning_sdk.lightning_cloud.openapi.models.id_complete_body import IdCompleteBody
+from lightning_sdk.lightning_cloud.openapi.models.id_contactowner_body import IdContactownerBody
 from lightning_sdk.lightning_cloud.openapi.models.id_engage_body import IdEngageBody
 from lightning_sdk.lightning_cloud.openapi.models.id_engage_body1 import IdEngageBody1
 from lightning_sdk.lightning_cloud.openapi.models.id_execute_body import IdExecuteBody
@@ -191,6 +193,7 @@ from lightning_sdk.lightning_cloud.openapi.models.update import Update
 from lightning_sdk.lightning_cloud.openapi.models.upload_id_complete_body import UploadIdCompleteBody
 from lightning_sdk.lightning_cloud.openapi.models.upload_id_parts_body import UploadIdPartsBody
 from lightning_sdk.lightning_cloud.openapi.models.uploads_upload_id_body import UploadsUploadIdBody
+from lightning_sdk.lightning_cloud.openapi.models.usagerestrictions_id_body import UsagerestrictionsIdBody
 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
@@ -235,11 +238,14 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_capacity_block_offering imp
 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
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_provider import V1CloudProvider
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space import V1CloudSpace
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_app import V1CloudSpaceApp
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_app_action import V1CloudSpaceAppAction
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_app_action_type import V1CloudSpaceAppActionType
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_app_owner_type import V1CloudSpaceAppOwnerType
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_artifact_event import V1CloudSpaceArtifactEvent
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_artifact_event_type import V1CloudSpaceArtifactEventType
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_code_version import V1CloudSpaceCodeVersion
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_code_version_status import V1CloudSpaceCodeVersionStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_engagement_response import V1CloudSpaceEngagementResponse
@@ -269,6 +275,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_cluster_state import V1Clus
 from lightning_sdk.lightning_cloud.openapi.models.v1_cluster_status import V1ClusterStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_cluster_tagging_options import V1ClusterTaggingOptions
 from lightning_sdk.lightning_cloud.openapi.models.v1_cluster_type import V1ClusterType
+from lightning_sdk.lightning_cloud.openapi.models.v1_cluster_usage_restriction import V1ClusterUsageRestriction
 from lightning_sdk.lightning_cloud.openapi.models.v1_collab_action import V1CollabAction
 from lightning_sdk.lightning_cloud.openapi.models.v1_collab_session import V1CollabSession
 from lightning_sdk.lightning_cloud.openapi.models.v1_command_argument import V1CommandArgument
@@ -281,6 +288,8 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_complete_upload_project_art
 from lightning_sdk.lightning_cloud.openapi.models.v1_complete_upload_service_execution_artifact_response import V1CompleteUploadServiceExecutionArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_completed_part import V1CompletedPart
 from lightning_sdk.lightning_cloud.openapi.models.v1_compute_config import V1ComputeConfig
+from lightning_sdk.lightning_cloud.openapi.models.v1_contact_assistant_owner_reason import V1ContactAssistantOwnerReason
+from lightning_sdk.lightning_cloud.openapi.models.v1_contact_assistant_owner_response import V1ContactAssistantOwnerResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation import V1Conversation
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
 from lightning_sdk.lightning_cloud.openapi.models.v1_count_metrics_streams_response import V1CountMetricsStreamsResponse
@@ -328,6 +337,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_delete_cluster_capacity_res
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_cluster_encryption_key_response import V1DeleteClusterEncryptionKeyResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_cluster_proxy_response import V1DeleteClusterProxyResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_cluster_response import V1DeleteClusterResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_delete_cluster_usage_restriction_response import V1DeleteClusterUsageRestrictionResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_conversation_response import V1DeleteConversationResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_data_connection_response import V1DeleteDataConnectionResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_dataset_response import V1DeleteDatasetResponse
@@ -417,6 +427,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_app import V1Fil
 from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_cloud_space import V1FilesystemCloudSpace
 from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_dataset import V1FilesystemDataset
 from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_job import V1FilesystemJob
+from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_mmt import V1FilesystemMMT
 from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_slurm_job import V1FilesystemSlurmJob
 from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_snowflake_connection import V1FilesystemSnowflakeConnection
 from lightning_sdk.lightning_cloud.openapi.models.v1_filesystem_work import V1FilesystemWork
@@ -526,6 +537,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_cluster_accelerators_r
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_cluster_availabilities_response import V1ListClusterAvailabilitiesResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_cluster_capacity_reservations_response import V1ListClusterCapacityReservationsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_cluster_proxies_response import V1ListClusterProxiesResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_list_cluster_usage_restrictions_response import V1ListClusterUsageRestrictionsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_clusters_response import V1ListClustersResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_conversations_response import V1ListConversationsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_data_connection_artifacts_response import V1ListDataConnectionArtifactsResponse
@@ -544,6 +556,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_filesystem_apps_respon
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_filesystem_cloud_spaces_response import V1ListFilesystemCloudSpacesResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_filesystem_datasets_response import V1ListFilesystemDatasetsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_filesystem_jobs_response import V1ListFilesystemJobsResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_list_filesystem_mm_ts_response import V1ListFilesystemMMTsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_filesystem_slurm_jobs_response import V1ListFilesystemSlurmJobsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_filesystem_snowflake_response import V1ListFilesystemSnowflakeResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_gallery_components_response import V1ListGalleryComponentsResponse
@@ -652,6 +665,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_path_telemetry import V1Pat
 from lightning_sdk.lightning_cloud.openapi.models.v1_phase_type import V1PhaseType
 from lightning_sdk.lightning_cloud.openapi.models.v1_plugin import V1Plugin
 from lightning_sdk.lightning_cloud.openapi.models.v1_plugins_list_response import V1PluginsListResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_post_cloud_space_artifact_events_response import V1PostCloudSpaceArtifactEventsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_presigned_url import V1PresignedUrl
 from lightning_sdk.lightning_cloud.openapi.models.v1_profiler_capture import V1ProfilerCapture
 from lightning_sdk.lightning_cloud.openapi.models.v1_profiler_enabled_response import V1ProfilerEnabledResponse

lightning_sdk/lightning_cloud/openapi/api/assistants_service_api.py

@@ -43,6 +43,111 @@ class AssistantsServiceApi(object):
             api_client = ApiClient()
         self.api_client = api_client
 
+    def assistants_service_contact_assistant_owner(self, body: 'IdContactownerBody', id: 'str', **kwargs) -> 'V1ContactAssistantOwnerResponse':  # noqa: E501
+        """assistants_service_contact_assistant_owner  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.assistants_service_contact_assistant_owner(body, id, async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :param IdContactownerBody body: (required)
+        :param str id: (required)
+        :return: V1ContactAssistantOwnerResponse
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+        kwargs['_return_http_data_only'] = True
+        if kwargs.get('async_req'):
+            return self.assistants_service_contact_assistant_owner_with_http_info(body, id, **kwargs)  # noqa: E501
+        else:
+            (data) = self.assistants_service_contact_assistant_owner_with_http_info(body, id, **kwargs)  # noqa: E501
+            return data
+
+    def assistants_service_contact_assistant_owner_with_http_info(self, body: 'IdContactownerBody', id: 'str', **kwargs) -> 'V1ContactAssistantOwnerResponse':  # noqa: E501
+        """assistants_service_contact_assistant_owner  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.assistants_service_contact_assistant_owner_with_http_info(body, id, async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :param IdContactownerBody body: (required)
+        :param str id: (required)
+        :return: V1ContactAssistantOwnerResponse
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+
+        all_params = ['body', 'id']  # noqa: E501
+        all_params.append('async_req')
+        all_params.append('_return_http_data_only')
+        all_params.append('_preload_content')
+        all_params.append('_request_timeout')
+
+        params = locals()
+        for key, val in six.iteritems(params['kwargs']):
+            if key not in all_params:
+                raise TypeError(
+                    "Got an unexpected keyword argument '%s'"
+                    " to method assistants_service_contact_assistant_owner" % key
+                )
+            params[key] = val
+        del params['kwargs']
+        # verify the required parameter 'body' is set
+        if ('body' not in params or
+                params['body'] is None):
+            raise ValueError("Missing the required parameter `body` when calling `assistants_service_contact_assistant_owner`")  # noqa: E501
+        # verify the required parameter 'id' is set
+        if ('id' not in params or
+                params['id'] is None):
+            raise ValueError("Missing the required parameter `id` when calling `assistants_service_contact_assistant_owner`")  # noqa: E501
+
+        collection_formats = {}
+
+        path_params = {}
+        if 'id' in params:
+            path_params['id'] = params['id']  # noqa: E501
+
+        query_params = []
+
+        header_params = {}
+
+        form_params = []
+        local_var_files = {}
+
+        body_params = None
+        if 'body' in params:
+            body_params = params['body']
+        # HTTP header `Accept`
+        header_params['Accept'] = self.api_client.select_header_accept(
+            ['application/json'])  # noqa: E501
+
+        # HTTP header `Content-Type`
+        header_params['Content-Type'] = self.api_client.select_header_content_type(  # noqa: E501
+            ['application/json'])  # noqa: E501
+
+        # Authentication setting
+        auth_settings = []  # noqa: E501
+
+        return self.api_client.call_api(
+            '/v1/agents/{id}/contact-owner', 'POST',
+            path_params,
+            query_params,
+            header_params,
+            body=body_params,
+            post_params=form_params,
+            files=local_var_files,
+            response_type='V1ContactAssistantOwnerResponse',  # noqa: E501
+            auth_settings=auth_settings,
+            async_req=params.get('async_req'),
+            _return_http_data_only=params.get('_return_http_data_only'),
+            _preload_content=params.get('_preload_content', True),
+            _request_timeout=params.get('_request_timeout'),
+            collection_formats=collection_formats)
+
     def assistants_service_create_assistant(self, body: 'ProjectIdAgentsBody', project_id: 'str', **kwargs) -> 'V1Assistant':  # noqa: E501
         """assistants_service_create_assistant  # noqa: E501