lightning-sdk 0.1.47__py3-none-any.whl → 0.1.49__py3-none-any.whl

lightning_sdk/cli/upload.py

@@ -4,18 +4,22 @@ import os
 from pathlib import Path
 from typing import Dict, List, Optional
 
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from simple_term_menu import TerminalMenu
 from tqdm import tqdm
 
+from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.exceptions import StudioCliError
 from lightning_sdk.cli.studios_menu import _StudiosMenu
+from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.models import upload_model
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
 
 
-class _Uploads(_StudiosMenu):
+class _Uploads(_StudiosMenu, _TeamspacesMenu):
     """Upload files and folders to Lightning AI."""
 
     _studio_upload_status_path = "~/.lightning/studios/uploads"
@@ -146,6 +150,33 @@ class _Uploads(_StudiosMenu):
         )
         print(f"See your file at {studio_url}")
 
+    def container(self, container: str, tag: str = "latest", teamspace: Optional[str] = None) -> None:
+        teamspace = self._resolve_teamspace(teamspace)
+        api = LitContainerApi()
+        console = Console()
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            TimeElapsedColumn(),
+            console=console,
+            transient=False,
+        ) as progress:
+            push_task = progress.add_task("Pushing Docker image", total=None)
+            resp = api.upload_container(container, teamspace, tag)
+            for line in resp:
+                if "status" in line:
+                    console.print(line["status"], style="bright_black")
+                    progress.update(push_task, description="Pushing Docker image")
+                elif "aux" in line:
+                    console.print(line["aux"], style="bright_black")
+                elif "error" in line:
+                    progress.stop()
+                    console.print(f"\n[red]{line}[/red]")
+                    return
+                else:
+                    console.print(line, style="bright_black")
+            progress.update(push_task, description="[green]Container pushed![/green]")
+
     def _start_parallel_upload(
         self, executor: concurrent.futures.ThreadPoolExecutor, studio: Studio, upload_state: Dict[str, str]
     ) -> List[concurrent.futures.Future]:

lightning_sdk/job/base.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union
 
+from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace
 
 if TYPE_CHECKING:
@@ -23,6 +24,7 @@ class JobDict(MachineDict):
     teamspace: str
     studio: Optional[str]
     image: Optional[str]
+    total_cost: float
 
 
 class _BaseJob(ABC):
@@ -60,6 +62,8 @@ class _BaseJob(ABC):
         if _fetch_job:
             self._update_internal_job()
 
+        self._prevent_refetch_latest = False
+
     @classmethod
     def run(
         cls,
@@ -331,6 +335,7 @@ class _BaseJob(ABC):
             "command": self.command,
             "status": self.status,
             "machine": self.machine,
+            "total_cost": self.total_cost,
         }
 
     def json(self) -> str:
@@ -345,7 +350,10 @@ class _BaseJob(ABC):
         studio_name = self._job_api.get_studio_name(self._guaranteed_job)
         if not studio_name:
             raise RuntimeError("Cannot extract studio name from job")
-        return f"https://lightning.ai/{self.teamspace.owner.name}/{self.teamspace.name}/studios/{studio_name}/app?app_id=jobs&job_name={self.name}"
+        return (
+            f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/studios/"
+            f"{studio_name}/app?app_id=jobs&job_name={self.name}"
+        )
 
     @property
     def _guaranteed_job(self) -> Any:
@@ -357,3 +365,18 @@ class _BaseJob(ABC):
             self._update_internal_job()
 
         return self._job
+
+    @property
+    def total_cost(self) -> float:
+        """The number of credits the job was consuming so far."""
+        return self._job_api.get_total_cost(self._latest_job)
+
+    @property
+    def _latest_job(self) -> Any:
+        """Guarantees to fetch the latest version of a job before returning it."""
+        # in some cases we know we just refetched the latest state, no need to refetch again
+        if self._prevent_refetch_latest:
+            return self._guaranteed_job
+
+        self._update_internal_job()
+        return self._job

lightning_sdk/job/job.py

@@ -50,15 +50,40 @@ class Job(_BaseJob):
             user: the name of the user owning the :param`teamspace`
                 in case it is owned directly by a user instead of an org.
         """
-        internal_job_cls = _JobV2 if _has_jobs_v2() and not self._force_v1 else _JobV1
-
-        self._internal_job = internal_job_cls(
-            name=name,
-            teamspace=teamspace,
-            org=org,
-            user=user,
-            _fetch_job=_fetch_job,
-        )
+        from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+
+        if _has_jobs_v2() and not self._force_v1:
+            # try with v2 and fall back to v1
+            try:
+                job = _JobV2(
+                    name=name,
+                    teamspace=teamspace,
+                    org=org,
+                    user=user,
+                    _fetch_job=_fetch_job,
+                )
+            except ApiException as e:
+                try:
+                    job = _JobV1(
+                        name=name,
+                        teamspace=teamspace,
+                        org=org,
+                        user=user,
+                        _fetch_job=_fetch_job,
+                    )
+                except ApiException:
+                    raise e from e
+
+        else:
+            job = _JobV1(
+                name=name,
+                teamspace=teamspace,
+                org=org,
+                user=user,
+                _fetch_job=_fetch_job,
+            )
+
+        self._internal_job = job
 
     @classmethod
     def run(

lightning_sdk/job/v1.py

@@ -174,7 +174,7 @@ class _JobV1(_BaseJob):
 
     def stop(self) -> None:
         """Stops the job. is blocking until the ob is stopped."""
-        if self.status in (Status.Stopped, Status.Failed):
+        if self.status in (Status.Stopped, Status.Completed, Status.Failed):
             return None
 
         return self._job_api.stop_job(self._job.id, self.teamspace.id)

lightning_sdk/job/v2.py

@@ -1,12 +1,13 @@
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Dict, Optional, Union
 
 from lightning_sdk.api.job_api import JobApiV2
+from lightning_sdk.api.utils import _get_cloud_url
 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.organization import Organization
-    from lightning_sdk.status import Status
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
     from lightning_sdk.user import User
@@ -120,6 +121,9 @@ class _JobV2(_BaseJob):
 
     def stop(self) -> None:
         """Stop the job. If the job is already stopped, this is a no-op. This is blocking until the job is stopped."""
+        if self.status in (Status.Stopped, Status.Completed, Status.Failed):
+            return
+
         self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
 
     def delete(self) -> None:
@@ -133,12 +137,6 @@ class _JobV2(_BaseJob):
             cloudspace_id=self._guaranteed_job.spec.cloudspace_id,
         )
 
-    @property
-    def _latest_job(self) -> Any:
-        """Guarantees to fetch the latest version of a job before returning it."""
-        self._update_internal_job()
-        return self._job
-
     @property
     def status(self) -> "Status":
         """The current status of the job."""
@@ -184,11 +182,18 @@ class _JobV2(_BaseJob):
 
     @property
     def link(self) -> str:
+        mmt_name = self._job_api.get_mmt_name(self._guaranteed_job)
+
         if self._job_api.get_image_name(self._guaranteed_job):
-            return (
-                f"https://lightning.ai/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=jobs"
-            )
+            if mmt_name:
+                # don't go via the studio unless we use studio env
+                return (
+                    f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/"
+                    f"jobs/{mmt_name}?app_id=mmt&machine_name={self.name}"
+                )
+            return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=jobs"
 
+        # TODO: MMT env with studio
         return super().link
 
     @property

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -327,7 +327,6 @@ 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_container_response import V1DeleteContainerResponse
 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
@@ -342,6 +341,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningapp_instanc
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningapp_instance_response import V1DeleteLightningappInstanceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningwork_response import V1DeleteLightningworkResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_page_response import V1DeleteLitPageResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_repository_response import V1DeleteLitRepositoryResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_logger_artifact_response import V1DeleteLoggerArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_managed_endpoint_response import V1DeleteManagedEndpointResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_metrics_stream_response import V1DeleteMetricsStreamResponse
@@ -794,8 +794,10 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_vultr_direct_v1 import V1Vu
 from lightning_sdk.lightning_cloud.openapi.models.v1_work import V1Work
 from lightning_sdk.lightning_cloud.openapi.models.validate import Validate
 from lightning_sdk.lightning_cloud.openapi.models.validateautojoindomain_domain_body import ValidateautojoindomainDomainBody
+from lightning_sdk.lightning_cloud.openapi.models.version_default_body import VersionDefaultBody
 from lightning_sdk.lightning_cloud.openapi.models.version_uploads_body import VersionUploadsBody
 from lightning_sdk.lightning_cloud.openapi.models.versions_id_body import VersionsIdBody
+from lightning_sdk.lightning_cloud.openapi.models.versions_version_body import VersionsVersionBody
 from lightning_sdk.lightning_cloud.openapi.models.works_id_body import WorksIdBody
 from lightning_sdk.lightning_cloud.openapi.models.v1_image_spec import V1ImageSpec as Gridv1ImageSpec
 from lightning_sdk.lightning_cloud.openapi.models.clusters_id_body import ClustersIdBody as Body

lightning_sdk/lightning_cloud/openapi/api/lit_registry_service_api.py

@@ -148,45 +148,45 @@ class LitRegistryServiceApi(object):
             _request_timeout=params.get('_request_timeout'),
             collection_formats=collection_formats)
 
-    def lit_registry_service_delete_lit_container(self, project_id: 'str', lit_container_name: 'str', **kwargs) -> 'V1DeleteContainerResponse':  # noqa: E501
-        """lit_registry_service_delete_lit_container  # noqa: E501
+    def lit_registry_service_delete_lit_repository(self, project_id: 'str', lit_repo_name: 'str', **kwargs) -> 'V1DeleteLitRepositoryResponse':  # noqa: E501
+        """lit_registry_service_delete_lit_repository  # noqa: E501
 
         This method makes a synchronous HTTP request by default. To make an
         asynchronous HTTP request, please pass async_req=True
-        >>> thread = api.lit_registry_service_delete_lit_container(project_id, lit_container_name, async_req=True)
+        >>> thread = api.lit_registry_service_delete_lit_repository(project_id, lit_repo_name, async_req=True)
         >>> result = thread.get()
 
         :param async_req bool
         :param str project_id: (required)
-        :param str lit_container_name: (required)
-        :return: V1DeleteContainerResponse
+        :param str lit_repo_name: (required)
+        :return: V1DeleteLitRepositoryResponse
                  If the method is called asynchronously,
                  returns the request thread.
         """
         kwargs['_return_http_data_only'] = True
         if kwargs.get('async_req'):
-            return self.lit_registry_service_delete_lit_container_with_http_info(project_id, lit_container_name, **kwargs)  # noqa: E501
+            return self.lit_registry_service_delete_lit_repository_with_http_info(project_id, lit_repo_name, **kwargs)  # noqa: E501
         else:
-            (data) = self.lit_registry_service_delete_lit_container_with_http_info(project_id, lit_container_name, **kwargs)  # noqa: E501
+            (data) = self.lit_registry_service_delete_lit_repository_with_http_info(project_id, lit_repo_name, **kwargs)  # noqa: E501
             return data
 
-    def lit_registry_service_delete_lit_container_with_http_info(self, project_id: 'str', lit_container_name: 'str', **kwargs) -> 'V1DeleteContainerResponse':  # noqa: E501
-        """lit_registry_service_delete_lit_container  # noqa: E501
+    def lit_registry_service_delete_lit_repository_with_http_info(self, project_id: 'str', lit_repo_name: 'str', **kwargs) -> 'V1DeleteLitRepositoryResponse':  # noqa: E501
+        """lit_registry_service_delete_lit_repository  # noqa: E501
 
         This method makes a synchronous HTTP request by default. To make an
         asynchronous HTTP request, please pass async_req=True
-        >>> thread = api.lit_registry_service_delete_lit_container_with_http_info(project_id, lit_container_name, async_req=True)
+        >>> thread = api.lit_registry_service_delete_lit_repository_with_http_info(project_id, lit_repo_name, async_req=True)
         >>> result = thread.get()
 
         :param async_req bool
         :param str project_id: (required)
-        :param str lit_container_name: (required)
-        :return: V1DeleteContainerResponse
+        :param str lit_repo_name: (required)
+        :return: V1DeleteLitRepositoryResponse
                  If the method is called asynchronously,
                  returns the request thread.
         """
 
-        all_params = ['project_id', 'lit_container_name']  # noqa: E501
+        all_params = ['project_id', 'lit_repo_name']  # noqa: E501
         all_params.append('async_req')
         all_params.append('_return_http_data_only')
         all_params.append('_preload_content')
@@ -197,26 +197,26 @@ class LitRegistryServiceApi(object):
             if key not in all_params:
                 raise TypeError(
                     "Got an unexpected keyword argument '%s'"
-                    " to method lit_registry_service_delete_lit_container" % key
+                    " to method lit_registry_service_delete_lit_repository" % key
                 )
             params[key] = val
         del params['kwargs']
         # verify the required parameter 'project_id' is set
         if ('project_id' not in params or
                 params['project_id'] is None):
-            raise ValueError("Missing the required parameter `project_id` when calling `lit_registry_service_delete_lit_container`")  # noqa: E501
-        # verify the required parameter 'lit_container_name' is set
-        if ('lit_container_name' not in params or
-                params['lit_container_name'] is None):
-            raise ValueError("Missing the required parameter `lit_container_name` when calling `lit_registry_service_delete_lit_container`")  # noqa: E501
+            raise ValueError("Missing the required parameter `project_id` when calling `lit_registry_service_delete_lit_repository`")  # noqa: E501
+        # verify the required parameter 'lit_repo_name' is set
+        if ('lit_repo_name' not in params or
+                params['lit_repo_name'] is None):
+            raise ValueError("Missing the required parameter `lit_repo_name` when calling `lit_registry_service_delete_lit_repository`")  # noqa: E501
 
         collection_formats = {}
 
         path_params = {}
         if 'project_id' in params:
             path_params['projectId'] = params['project_id']  # noqa: E501
-        if 'lit_container_name' in params:
-            path_params['litContainerName'] = params['lit_container_name']  # noqa: E501
+        if 'lit_repo_name' in params:
+            path_params['litRepoName'] = params['lit_repo_name']  # noqa: E501
 
         query_params = []
 
@@ -234,14 +234,14 @@ class LitRegistryServiceApi(object):
         auth_settings = []  # noqa: E501
 
         return self.api_client.call_api(
-            '/v1/projects/{projectId}/lit-registry/{litContainerName}', 'DELETE',
+            '/v1/projects/{projectId}/lit-registry/{litRepoName}', 'DELETE',
             path_params,
             query_params,
             header_params,
             body=body_params,
             post_params=form_params,
             files=local_var_files,
-            response_type='V1DeleteContainerResponse',  # noqa: E501
+            response_type='V1DeleteLitRepositoryResponse',  # noqa: E501
             auth_settings=auth_settings,
             async_req=params.get('async_req'),
             _return_http_data_only=params.get('_return_http_data_only'),
@@ -341,3 +341,104 @@ class LitRegistryServiceApi(object):
             _preload_content=params.get('_preload_content', True),
             _request_timeout=params.get('_request_timeout'),
             collection_formats=collection_formats)
+
+    def lit_registry_service_get_lit_repository(self, project_id: 'str', lit_repo_name: 'str', **kwargs) -> 'V1LitRepository':  # noqa: E501
+        """lit_registry_service_get_lit_repository  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.lit_registry_service_get_lit_repository(project_id, lit_repo_name, async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :param str project_id: (required)
+        :param str lit_repo_name: (required)
+        :return: V1LitRepository
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+        kwargs['_return_http_data_only'] = True
+        if kwargs.get('async_req'):
+            return self.lit_registry_service_get_lit_repository_with_http_info(project_id, lit_repo_name, **kwargs)  # noqa: E501
+        else:
+            (data) = self.lit_registry_service_get_lit_repository_with_http_info(project_id, lit_repo_name, **kwargs)  # noqa: E501
+            return data
+
+    def lit_registry_service_get_lit_repository_with_http_info(self, project_id: 'str', lit_repo_name: 'str', **kwargs) -> 'V1LitRepository':  # noqa: E501
+        """lit_registry_service_get_lit_repository  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.lit_registry_service_get_lit_repository_with_http_info(project_id, lit_repo_name, async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :param str project_id: (required)
+        :param str lit_repo_name: (required)
+        :return: V1LitRepository
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+
+        all_params = ['project_id', 'lit_repo_name']  # 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 lit_registry_service_get_lit_repository" % key
+                )
+            params[key] = val
+        del params['kwargs']
+        # verify the required parameter 'project_id' is set
+        if ('project_id' not in params or
+                params['project_id'] is None):
+            raise ValueError("Missing the required parameter `project_id` when calling `lit_registry_service_get_lit_repository`")  # noqa: E501
+        # verify the required parameter 'lit_repo_name' is set
+        if ('lit_repo_name' not in params or
+                params['lit_repo_name'] is None):
+            raise ValueError("Missing the required parameter `lit_repo_name` when calling `lit_registry_service_get_lit_repository`")  # noqa: E501
+
+        collection_formats = {}
+
+        path_params = {}
+        if 'project_id' in params:
+            path_params['projectId'] = params['project_id']  # noqa: E501
+        if 'lit_repo_name' in params:
+            path_params['litRepoName'] = params['lit_repo_name']  # noqa: E501
+
+        query_params = []
+
+        header_params = {}
+
+        form_params = []
+        local_var_files = {}
+
+        body_params = None
+        # HTTP header `Accept`
+        header_params['Accept'] = self.api_client.select_header_accept(
+            ['application/json'])  # noqa: E501
+
+        # Authentication setting
+        auth_settings = []  # noqa: E501
+
+        return self.api_client.call_api(
+            '/v1/projects/{projectId}/lit-registry/{litRepoName}', 'GET',
+            path_params,
+            query_params,
+            header_params,
+            body=body_params,
+            post_params=form_params,
+            files=local_var_files,
+            response_type='V1LitRepository',  # 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)