lightning-sdk 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_server_alert_phase.py

@@ -0,0 +1,104 @@
+# coding: utf-8
+
+"""
+    external/v1/auth_service.proto
+
+    No description provided (generated by Swagger Codegen https://github.com/swagger-api/swagger-codegen)  # noqa: E501
+
+    OpenAPI spec version: version not set
+    
+    Generated by: https://github.com/swagger-api/swagger-codegen.git
+
+    NOTE
+    ----
+    standard swagger-codegen-cli for this python client has been modified
+    by custom templates. The purpose of these templates is to include
+    typing information in the API and Model code. Please refer to the
+    main grid repository for more info
+"""
+
+import pprint
+import re  # noqa: F401
+
+from typing import TYPE_CHECKING
+
+import six
+
+if TYPE_CHECKING:
+    from datetime import datetime
+    from lightning_sdk.lightning_cloud.openapi.models import *
+
+class V1ServerAlertPhase(object):
+    """NOTE: This class is auto generated by the swagger code generator program.
+
+    Do not edit the class manually.
+    """
+
+    """
+    allowed enum values
+    """
+    UNSPECIFIED = "SERVER_ALERT_PHASE_UNSPECIFIED"
+    STARTUP = "SERVER_ALERT_PHASE_STARTUP"
+    SHUTDOWN = "SERVER_ALERT_PHASE_SHUTDOWN"
+    """
+    Attributes:
+      swagger_types (dict): The key is attribute name
+                            and the value is attribute type.
+      attribute_map (dict): The key is attribute name
+                            and the value is json key in definition.
+    """
+    swagger_types = {
+    }
+
+    attribute_map = {
+    }
+
+    def __init__(self):  # noqa: E501
+        """V1ServerAlertPhase - a model defined in Swagger"""  # noqa: E501
+        self.discriminator = None
+
+    def to_dict(self) -> dict:
+        """Returns the model properties as a dict"""
+        result = {}
+
+        for attr, _ in six.iteritems(self.swagger_types):
+            value = getattr(self, attr)
+            if isinstance(value, list):
+                result[attr] = list(map(
+                    lambda x: x.to_dict() if hasattr(x, "to_dict") else x,
+                    value
+                ))
+            elif hasattr(value, "to_dict"):
+                result[attr] = value.to_dict()
+            elif isinstance(value, dict):
+                result[attr] = dict(map(
+                    lambda item: (item[0], item[1].to_dict())
+                    if hasattr(item[1], "to_dict") else item,
+                    value.items()
+                ))
+            else:
+                result[attr] = value
+        if issubclass(V1ServerAlertPhase, dict):
+            for key, value in self.items():
+                result[key] = value
+
+        return result
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model"""
+        return pprint.pformat(self.to_dict())
+
+    def __repr__(self) -> str:
+        """For `print` and `pprint`"""
+        return self.to_str()
+
+    def __eq__(self, other: 'V1ServerAlertPhase') -> bool:
+        """Returns true if both objects are equal"""
+        if not isinstance(other, V1ServerAlertPhase):
+            return False
+
+        return self.__dict__ == other.__dict__
+
+    def __ne__(self, other: 'V1ServerAlertPhase') -> bool:
+        """Returns true if both objects are not equal"""
+        return not self == other

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -79,6 +79,7 @@ class V1UserFeatures(object):
         'lambda_labs': 'bool',
         'landing_studios': 'bool',
         'lit_logger': 'bool',
+        'manage_default_studio_machine': 'bool',
         'marketplace': 'bool',
         'mmt_fault_tolerance': 'bool',
         'mmt_strategy_selector': 'bool',
@@ -107,6 +108,7 @@ class V1UserFeatures(object):
         'runnable_public_studio_page': 'bool',
         'security_docs': 'bool',
         'show_dev_admin': 'bool',
+        'single_wallet': 'bool',
         'slurm': 'bool',
         'slurm_machine_selector': 'bool',
         'stop_ide_container_on_shutdown': 'bool',
@@ -159,6 +161,7 @@ class V1UserFeatures(object):
         'lambda_labs': 'lambdaLabs',
         'landing_studios': 'landingStudios',
         'lit_logger': 'litLogger',
+        'manage_default_studio_machine': 'manageDefaultStudioMachine',
         'marketplace': 'marketplace',
         'mmt_fault_tolerance': 'mmtFaultTolerance',
         'mmt_strategy_selector': 'mmtStrategySelector',
@@ -187,6 +190,7 @@ class V1UserFeatures(object):
         'runnable_public_studio_page': 'runnablePublicStudioPage',
         'security_docs': 'securityDocs',
         'show_dev_admin': 'showDevAdmin',
+        'single_wallet': 'singleWallet',
         'slurm': 'slurm',
         'slurm_machine_selector': 'slurmMachineSelector',
         'stop_ide_container_on_shutdown': 'stopIdeContainerOnShutdown',
@@ -200,7 +204,7 @@ class V1UserFeatures(object):
         'vultr': 'vultr'
     }
 
-    def __init__(self, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byoc_litcr: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloud_space_environment_templates: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_reservations: 'bool' =None, dgx_cloud: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_crypto_crackdown: 'bool' =None, enable_storage_limits: 'bool' =None, enterprise_compute_admin: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, filestore: 'bool' =None, inactive_notify_delete: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, lambda_labs: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, org_admin_alerts: 'bool' =None, org_level_member_permissions: 'bool' =None, pipelines: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_lightning_apps: 'bool' =None, plugin_lightning_apps_distributed: 'bool' =None, plugin_milvus: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_react: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, plugin_weviate: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, stop_ide_container_on_shutdown: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, studios_dashboard: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None):  # noqa: E501
+    def __init__(self, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byoc_litcr: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloud_space_environment_templates: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_reservations: 'bool' =None, dgx_cloud: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_crypto_crackdown: 'bool' =None, enable_storage_limits: 'bool' =None, enterprise_compute_admin: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, filestore: 'bool' =None, inactive_notify_delete: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, lambda_labs: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, manage_default_studio_machine: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, multicloud_saas: 'bool' =None, multiple_studio_versions: 'bool' =None, org_admin_alerts: 'bool' =None, org_level_member_permissions: 'bool' =None, pipelines: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_lightning_apps: 'bool' =None, plugin_lightning_apps_distributed: 'bool' =None, plugin_milvus: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_react: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, plugin_weviate: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, stop_ide_container_on_shutdown: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, studios_dashboard: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None):  # noqa: E501
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
         self._affiliate_links = None
         self._agents_v2 = None
@@ -240,6 +244,7 @@ class V1UserFeatures(object):
         self._lambda_labs = None
         self._landing_studios = None
         self._lit_logger = None
+        self._manage_default_studio_machine = None
         self._marketplace = None
         self._mmt_fault_tolerance = None
         self._mmt_strategy_selector = None
@@ -268,6 +273,7 @@ class V1UserFeatures(object):
         self._runnable_public_studio_page = None
         self._security_docs = None
         self._show_dev_admin = None
+        self._single_wallet = None
         self._slurm = None
         self._slurm_machine_selector = None
         self._stop_ide_container_on_shutdown = None
@@ -356,6 +362,8 @@ class V1UserFeatures(object):
             self.landing_studios = landing_studios
         if lit_logger is not None:
             self.lit_logger = lit_logger
+        if manage_default_studio_machine is not None:
+            self.manage_default_studio_machine = manage_default_studio_machine
         if marketplace is not None:
             self.marketplace = marketplace
         if mmt_fault_tolerance is not None:
@@ -412,6 +420,8 @@ class V1UserFeatures(object):
             self.security_docs = security_docs
         if show_dev_admin is not None:
             self.show_dev_admin = show_dev_admin
+        if single_wallet is not None:
+            self.single_wallet = single_wallet
         if slurm is not None:
             self.slurm = slurm
         if slurm_machine_selector is not None:
@@ -1233,6 +1243,27 @@ class V1UserFeatures(object):
 
         self._lit_logger = lit_logger
 
+    @property
+    def manage_default_studio_machine(self) -> 'bool':
+        """Gets the manage_default_studio_machine of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The manage_default_studio_machine of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._manage_default_studio_machine
+
+    @manage_default_studio_machine.setter
+    def manage_default_studio_machine(self, manage_default_studio_machine: 'bool'):
+        """Sets the manage_default_studio_machine of this V1UserFeatures.
+
+
+        :param manage_default_studio_machine: The manage_default_studio_machine of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._manage_default_studio_machine = manage_default_studio_machine
+
     @property
     def marketplace(self) -> 'bool':
         """Gets the marketplace of this V1UserFeatures.  # noqa: E501
@@ -1821,6 +1852,27 @@ class V1UserFeatures(object):
 
         self._show_dev_admin = show_dev_admin
 
+    @property
+    def single_wallet(self) -> 'bool':
+        """Gets the single_wallet of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The single_wallet of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._single_wallet
+
+    @single_wallet.setter
+    def single_wallet(self, single_wallet: 'bool'):
+        """Sets the single_wallet of this V1UserFeatures.
+
+
+        :param single_wallet: The single_wallet of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._single_wallet = single_wallet
+
     @property
     def slurm(self) -> 'bool':
         """Gets the slurm of this V1UserFeatures.  # noqa: E501

lightning_sdk/lightning_cloud/rest_client.py

@@ -2,41 +2,39 @@ import functools
 import logging
 import time
 from functools import wraps
-from typing import Any, Callable, Optional
+from typing import Callable, Optional, Any
 
 import urllib3
-
 from lightning_sdk.lightning_cloud import env
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     ApiClient,
-    AssistantsServiceApi,
     AuthServiceApi,
-    BillingServiceApi,
-    CloudSpaceEnvironmentTemplateServiceApi,
     CloudSpaceServiceApi,
     ClusterServiceApi,
     Configuration,
     DataConnectionServiceApi,
-    DatasetServiceApi,
-    DeploymentTemplatesServiceApi,
-    EndpointServiceApi,
-    JobsServiceApi,
     LightningappInstanceServiceApi,
     LightningappV2ServiceApi,
     LightningworkServiceApi,
-    LitLoggerServiceApi,
-    LitRegistryServiceApi,
-    ModelsStoreApi,
-    OrganizationsServiceApi,
-    PipelinesServiceApi,
     ProjectsServiceApi,
-    SchedulesServiceApi,
     SecretServiceApi,
-    SlurmJobsUserServiceApi,
     SSHPublicKeyServiceApi,
-    StorageServiceApi,
+    DatasetServiceApi,
+    OrganizationsServiceApi,
     UserServiceApi,
+    BillingServiceApi,
+    EndpointServiceApi,
+    SlurmJobsUserServiceApi,
+    LitLoggerServiceApi,
+    JobsServiceApi,
+    AssistantsServiceApi,
+    StorageServiceApi,
+    DeploymentTemplatesServiceApi,
+    ModelsStoreApi,
+    LitRegistryServiceApi,
+    PipelinesServiceApi,
+    SchedulesServiceApi,
 )
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.source_code.logs_socket_api import LightningLogsSocketAPI
@@ -73,34 +71,34 @@ def create_swagger_client(check_context=True):
 
 
 class GridRestClient(
-    LightningLogsSocketAPI,
-    LightningappInstanceServiceApi,
-    LightningappV2ServiceApi,
-    AuthServiceApi,
-    CloudSpaceServiceApi,
-    ClusterServiceApi,
-    ProjectsServiceApi,
-    LightningworkServiceApi,
-    SecretServiceApi,
-    SSHPublicKeyServiceApi,
-    DataConnectionServiceApi,
-    DatasetServiceApi,
-    OrganizationsServiceApi,
-    UserServiceApi,
-    BillingServiceApi,
-    EndpointServiceApi,
-    SlurmJobsUserServiceApi,
-    LitLoggerServiceApi,
-    JobsServiceApi,
-    AssistantsServiceApi,
-    StorageServiceApi,
-    DeploymentTemplatesServiceApi,
-    ModelsStoreApi,
-    LitRegistryServiceApi,
-    PipelinesServiceApi,
-    SchedulesServiceApi,
-    CloudSpaceEnvironmentTemplateServiceApi,
+        LightningLogsSocketAPI,
+        LightningappInstanceServiceApi,
+        LightningappV2ServiceApi,
+        AuthServiceApi,
+        CloudSpaceServiceApi,
+        ClusterServiceApi,
+        ProjectsServiceApi,
+        LightningworkServiceApi,
+        SecretServiceApi,
+        SSHPublicKeyServiceApi,
+        DataConnectionServiceApi,
+        DatasetServiceApi,
+        OrganizationsServiceApi,
+        UserServiceApi,
+        BillingServiceApi,
+        EndpointServiceApi,
+        SlurmJobsUserServiceApi,
+        LitLoggerServiceApi,
+        JobsServiceApi,
+        AssistantsServiceApi,
+        StorageServiceApi,
+        DeploymentTemplatesServiceApi,
+        ModelsStoreApi,
+        LitRegistryServiceApi,
+        PipelinesServiceApi,
+        SchedulesServiceApi,
 ):
+
     def __init__(self, api_client: Optional[ApiClient] = None):
         api_client = api_client if api_client else create_swagger_client()
         api_client.request = request_auth_warning_wrapper(api_client.request)

lightning_sdk/lit_container.py

@@ -64,7 +64,14 @@ class LitContainer:
         return self._api.delete_container(project_id, container)
 
     def upload_container(
-        self, container: str, teamspace: str, org: Optional[str] = None, user: Optional[str] = None, tag: str = "latest"
+        self,
+        container: str,
+        teamspace: str,
+        org: Optional[str] = None,
+        user: Optional[str] = None,
+        tag: str = "latest",
+        cloud_account: Optional[str] = None,
+        platform: Optional[str] = "linux/amd64",
     ) -> None:
         """Upload a container to the docker registry.
 
@@ -74,18 +81,26 @@ class LitContainer:
             org: The organization which contains the container.
             user: The user which contains the container.
             tag: The tag to use for the container.
+            cloud_account: The cloud account where the container is stored.
+            platform: The platform the container is meant to run on.
         """
         try:
             teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
         except Exception as e:
             raise ValueError(f"Could not resolve teamspace: {e}") from e
 
-        resp = self._api.upload_container(container, teamspace, tag)
+        resp = self._api.upload_container(container, teamspace, tag, cloud_account, platform=platform)
         for line in resp:
             print(line)
 
     def download_container(
-        self, container: str, teamspace: str, org: Optional[str] = None, user: Optional[str] = None, tag: str = "latest"
+        self,
+        container: str,
+        teamspace: str,
+        org: Optional[str] = None,
+        user: Optional[str] = None,
+        tag: str = "latest",
+        cloud_account: Optional[str] = None,
     ) -> None:
         """Download a container from the docker registry.
 
@@ -101,4 +116,4 @@ class LitContainer:
         except Exception as e:
             raise ValueError(f"Could not resolve teamspace: {e}") from e
 
-        return self._api.download_container(container, teamspace, tag)
+        return self._api.download_container(container, teamspace, tag, cloud_account)

lightning_sdk/models.py

@@ -144,7 +144,7 @@ def download_model(
 
 def upload_model(
     name: str,
-    path: Union[Path, str] = ".",
+    path: Union[str, Path, List[Union[str, Path]]] = ".",
     cloud_account: Optional[str] = None,
     progress_bar: bool = True,
 ) -> UploadedModelInfo:

lightning_sdk/serve.py

@@ -2,18 +2,18 @@ import os
 import shlex
 import subprocess
 from pathlib import Path
-from typing import Generator, Optional
+from typing import Generator, List, Optional, Union
 from urllib.parse import urlencode
 
 import docker
 from rich.console import Console
-from rich.progress import Progress
 
 from lightning_sdk import Deployment, Machine, Teamspace
-from lightning_sdk.api.deployment_api import AutoScaleConfig
+from lightning_sdk.api.deployment_api import AutoScaleConfig, DeploymentApi, Env, Secret
 from lightning_sdk.api.lit_container_api import LitContainerApi
-from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.api.utils import _get_cloud_url, _get_registry_url
 from lightning_sdk.lightning_cloud import env
+from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
 from lightning_sdk.lightning_cloud.login import Auth, AuthServer
 
 _DOCKER_NOT_RUNNING_MSG = (
@@ -35,16 +35,18 @@ class _Auth(Auth):
         _AuthServer().login_with_browser(self)
 
 
+def authenticate() -> None:
+    auth = _Auth()
+    auth.authenticate()
+
+
 class _LitServeDeployer:
-    def __init__(self) -> None:
+    def __init__(self, name: Optional[str], teamspace: Optional[Teamspace]) -> None:
+        self.name = name
+        self.teamspace = teamspace
         self._console = Console()
         self._client = None
 
-    @staticmethod
-    def authenticate() -> None:
-        auth = _Auth()
-        auth.authenticate()
-
     @property
     def client(self) -> docker.DockerClient:
         if self._client is None:
@@ -56,6 +58,10 @@ class _LitServeDeployer:
                 raise RuntimeError(_DOCKER_NOT_RUNNING_MSG) from None
         return self._client
 
+    @property
+    def created(self) -> bool:
+        return DeploymentApi().get_deployment_by_name(self.name, self.teamspace.id) is not None
+
     def dockerize_api(
         self,
         server_filename: str,
@@ -176,45 +182,79 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
 
         return log_generator()
 
-    def build_container(self, path: str, repository: str, tag: str, console: Console, progress: Progress) -> None:
-        build_task = progress.add_task("Building Docker image", total=None)
+    def build_container(self, path: str, repository: str, tag: str) -> Generator[str, None, None]:
         build_logs = self._docker_build_with_logs(path, repository, tag=tag)
 
         for line in build_logs:
             if "error" in line:
-                progress.stop()
-                console.print(f"\n[red]{line}[/red]")
                 raise RuntimeError(f"Failed to build image: {line}")
             else:
-                console.print(
-                    line.strip(),
-                )
-                progress.update(build_task, description="Building Docker image")
-
-        progress.update(build_task, description="[green]Build completed![/green]")
+                yield line.strip()
 
     def push_container(
-        self, repository: str, tag: str, teamspace: Teamspace, lit_cr: LitContainerApi, progress: Progress
-    ) -> dict:
-        console = self._console
-        push_task = progress.add_task("Pushing to registry", total=None)
-        console.print("\nPushing image...", style="bold blue")
+        self,
+        repository: str,
+        tag: str,
+        teamspace: Teamspace,
+        lit_cr: LitContainerApi,
+        cloud_account: str,
+    ) -> Generator[dict, None, dict]:
         lit_cr.authenticate()
-        push_status = lit_cr.upload_container(repository, teamspace, tag=tag)
-        last_status = {}
+        push_status = lit_cr.upload_container(
+            repository, teamspace, tag=tag, cloud_account=cloud_account, platform=None
+        )
         for line in push_status:
-            last_status = line
             if "error" in line:
-                progress.stop()
-                console.print(f"\n[red]{line}[/red]")
                 raise RuntimeError(f"Failed to push image: {line}")
             if "status" in line:
-                console.print(line["status"].strip())
-                progress.update(push_task, description="Pushing to registry")
-        progress.update(push_task, description="[green]Push completed![/green]")
-        return last_status
+                yield {"status": line["status"].strip()}
 
-    def _run_on_cloud(
+        registry_url = _get_registry_url()
+        container_basename = repository.split("/")[-1]
+        repository = (
+            f"{registry_url}/lit-container{f'-{cloud_account}' if cloud_account is not None else ''}/"
+            f"{teamspace.owner.name}/{teamspace.name}/{container_basename}"
+        )
+        yield {
+            "finish": True,
+            "status": "Container pushed successfully",
+            "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container_basename}"
+            f"{f'?clusterId={cloud_account}' if cloud_account is not None else ''}",
+            "image": repository,
+        }
+
+    def _update_deployment(
+        self,
+        deployment: Deployment,
+        machine: Optional[Machine] = None,
+        image: Optional[str] = None,
+        entrypoint: Optional[str] = None,
+        command: Optional[str] = None,
+        env: Optional[List[Union[Env, Secret]]] = None,
+        min_replica: Optional[int] = 0,
+        max_replica: Optional[int] = 1,
+        spot: Optional[bool] = None,
+        replicas: Optional[int] = 1,
+        cloud_account: Optional[str] = None,
+        port: Optional[int] = 8000,
+        include_credentials: Optional[bool] = True,
+    ) -> None:
+        return deployment.update(
+            machine=machine,
+            image=image,
+            entrypoint=entrypoint,
+            command=command,
+            env=env,
+            max_replicas=max_replica,
+            min_replicas=min_replica,
+            replicas=replicas,
+            spot=spot,
+            cloud_account=cloud_account,
+            ports=[port],
+            include_credentials=include_credentials,
+        )
+
+    def run_on_cloud(
         self,
         deployment_name: str,
         teamspace: Teamspace,
@@ -229,18 +269,44 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
         port: Optional[int] = 8000,
         include_credentials: Optional[bool] = True,
     ) -> dict:
+        """Run a deployment on the cloud. If the deployment already exists, it will be updated.
+
+        Args:
+            deployment_name: The name of the deployment.
+            teamspace: The teamspace to run the deployment on.
+            image: The image to run the deployment on.
+            metric: The metric to use for autoscaling. Defaults to None.
+            machine: The machine to run the deployment on. Defaults to None.
+            min_replica: The minimum number of replicas to run. Defaults to 0.
+            max_replica: The maximum number of replicas to run. Defaults to 1.
+            spot: Whether to run the deployment on spot instances. Defaults to None.
+            replicas: The number of replicas to run. Defaults to 1.
+            cloud_account: The cloud account to run the deployment on. Defaults to None.
+            port: The port to run the deployment on. Defaults to 8000.
+            include_credentials: Whether to include credentials in the deployment. Defaults to True.
+
+        Returns:
+            dict: The deployment and the URL of the deployment.
+        """
         url = f"{_get_cloud_url()}/{teamspace.owner.name}/{teamspace.name}/jobs/{deployment_name}"
         machine = machine or Machine.CPU
         metric = metric or ("CPU" if machine.is_cpu() else "GPU")
         deployment = Deployment(deployment_name, teamspace)
         if deployment.is_started:
-            raise RuntimeError(
-                f"Deployment with name {deployment_name} already running: {url} \n"
-                "To update the deployment, use the Deployment API:\n"
-                "from lightning_sdk import Deployment\n"
-                f"deployment = Deployment('{deployment_name}', teamspace)\n"
-                "deployment.update(...)"
+            self._console.print(f"Deployment with name {deployment_name} already running. Updating the deployment.")
+            self._update_deployment(
+                deployment,
+                machine,
+                image,
+                min_replica=min_replica,
+                max_replica=max_replica,
+                spot=spot,
+                replicas=replicas,
+                cloud_account=cloud_account,
+                port=port,
+                include_credentials=include_credentials,
             )
+            return {"deployment": deployment, "url": url, "updated": True}
         autoscale = AutoScaleConfig(min_replicas=min_replica, max_replicas=max_replica, metric=metric, threshold=0.95)
         deployment.start(
             machine=machine,