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

lightning_sdk/__init__.py

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

lightning_sdk/api/job_api.py

@@ -1,16 +1,15 @@
 import time
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 from urllib.request import urlopen
 
 from lightning_sdk.api.utils import (
     _COMPUTE_NAME_TO_MACHINE,
-    _MACHINE_TO_COMPUTE_NAME,
     _create_app,
+    _machine_to_compute_name,
     remove_datetime_prefix,
+    resolve_path_mappings,
 )
-from lightning_sdk.api.utils import (
-    _get_cloud_url as _cloud_url,
-)
+from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
 from lightning_sdk.lightning_cloud.openapi import (
     AppinstancesIdBody,
@@ -120,7 +119,7 @@ class JobApiV1:
         studio_id: str,
         teamspace_id: str,
         cloud_account: str,
-        machine: Machine,
+        machine: Union[Machine, str],
         interruptible: bool,
     ) -> Externalv1LightningappInstance:
         """Creates an arbitrary app."""
@@ -130,7 +129,7 @@ class JobApiV1:
             teamspace_id=teamspace_id,
             cloud_account=cloud_account,
             plugin_type="job",
-            compute=_MACHINE_TO_COMPUTE_NAME[machine],
+            compute=_machine_to_compute_name(machine),
             name=name,
             entrypoint=command,
             interruptible=interruptible,
@@ -209,24 +208,31 @@ class JobApiV2:
         teamspace_id: str,
         studio_id: Optional[str],
         image: Optional[str],
-        machine: Machine,
+        machine: Union[Machine, str],
         interruptible: bool,
         env: Optional[Dict[str, str]],
         image_credentials: Optional[str],
         cloud_account_auth: bool,
-        artifacts_local: Optional[str],
-        artifacts_remote: Optional[str],
         entrypoint: str,
+        path_mappings: Optional[Dict[str, str]],
+        artifacts_local: Optional[str],  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
     ) -> V1Job:
         env_vars = []
         if env is not None:
             for k, v in env.items():
                 env_vars.append(V1EnvVar(name=k, value=v))
 
-        instance_name = _MACHINE_TO_COMPUTE_NAME[machine]
+        instance_name = _machine_to_compute_name(machine)
 
         run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if studio_id is not None else ""
 
+        path_mappings_list = resolve_path_mappings(
+            mappings=path_mappings or {},
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
@@ -239,8 +245,7 @@ class JobApiV2:
             spot=interruptible,
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
-            artifacts_source=artifacts_local or "",
-            artifacts_destination=artifacts_remote or "",
+            path_mappings=path_mappings_list,
         )
         body = ProjectIdJobsBody(name=name, spec=spec)
 

lightning_sdk/api/lit_container_api.py

@@ -1,15 +1,36 @@
 from typing import Generator, List
 
+import docker
+
 from lightning_sdk.api.utils import _get_registry_url
+from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
 from lightning_sdk.lightning_cloud.openapi.models import V1DeleteLitRepositoryResponse
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.teamspace import Teamspace
 
 
+class LCRAuthFailedError(Exception):
+    def __init__(self) -> None:
+        super().__init__("Failed to authenticate with Lightning Container Registry")
+
+
 class LitContainerApi:
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=3)
 
+        try:
+            self._docker_client = docker.from_env()
+            self._docker_client.ping()
+        except docker.errors.DockerException as e:
+            raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
+
+    def authenticate(self) -> bool:
+        authed_user = self._client.auth_service_get_user()
+        username = authed_user.username
+        api_key = authed_user.api_key
+        resp = self._docker_client.login(username, password=api_key, registry=_get_registry_url())
+        return resp["Status"] == "Login Succeeded"
+
     def list_containers(self, project_id: str) -> List:
         project = self._client.lit_registry_service_get_lit_project_registry(project_id)
         return project.repositories
@@ -21,22 +42,31 @@ class LitContainerApi:
             raise ValueError(f"Could not delete container {container} from project {project_id}") from ex
 
     def upload_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
-        import docker
-
-        try:
-            client = docker.from_env()
-            client.ping()
-        except docker.errors.DockerException as e:
-            raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
-
         try:
-            client.images.get(container)
+            self._docker_client.images.get(container)
         except docker.errors.ImageNotFound:
             raise ValueError(f"Container {container} does not exist") from None
 
         registry_url = _get_registry_url()
         repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
-        tagged = client.api.tag(container, repository, tag)
+        tagged = self._docker_client.api.tag(container, repository, tag)
         if not tagged:
             raise ValueError(f"Could not tag container {container} with {repository}:{tag}")
-        return client.api.push(repository, stream=True, decode=True)
+        lines = self._docker_client.api.push(repository, stream=True, decode=True)
+        for line in lines:
+            if "errorDetail" in line and "authorization failed" in line["error"]:
+                raise LCRAuthFailedError()
+            yield line
+        yield {
+            "finish": True,
+            "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container}",
+        }
+
+    def download_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
+        registry_url = _get_registry_url()
+        repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
+        try:
+            self._docker_client.images.pull(repository, tag=tag)
+        except docker.errors.APIError as e:
+            raise ValueError(f"Could not pull container {container} from {repository}:{tag}") from e
+        return self._docker_client.api.tag(repository, container, tag)

lightning_sdk/api/mmt_api.py

@@ -1,16 +1,15 @@
 import json
 import time
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from lightning_sdk.api.job_api import JobApiV1
 from lightning_sdk.api.utils import (
     _COMPUTE_NAME_TO_MACHINE,
-    _MACHINE_TO_COMPUTE_NAME,
     _create_app,
+    _machine_to_compute_name,
+    resolve_path_mappings,
 )
-from lightning_sdk.api.utils import (
-    _get_cloud_url as _cloud_url,
-)
+from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
 from lightning_sdk.lightning_cloud.openapi import (
     Externalv1LightningappInstance,
@@ -43,13 +42,13 @@ class MMTApiV1(JobApiV1):
         cloud_account: Optional[str],
         teamspace_id: str,
         studio_id: str,
-        machine: Machine,
+        machine: Union[Machine, str],
         interruptible: bool,
         strategy: str,
     ) -> Externalv1LightningappInstance:
         """Creates a multi-machine job with given commands."""
         distributed_args = {
-            "cloud_compute": _MACHINE_TO_COMPUTE_NAME[machine],
+            "cloud_compute": _machine_to_compute_name(machine),
             "num_instances": num_machines,
             "strategy": strategy,
         }
@@ -80,24 +79,31 @@ class MMTApiV2:
         teamspace_id: str,
         studio_id: Optional[str],
         image: Optional[str],
-        machine: Machine,
+        machine: Union[Machine, str],
         interruptible: bool,
         env: Optional[Dict[str, str]],
         image_credentials: Optional[str],
         cloud_account_auth: bool,
-        artifacts_local: Optional[str],
-        artifacts_remote: Optional[str],
         entrypoint: str,
+        path_mappings: Optional[Dict[str, str]],
+        artifacts_local: Optional[str],  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
     ) -> V1MultiMachineJob:
         env_vars = []
         if env is not None:
             for k, v in env.items():
                 env_vars.append(V1EnvVar(name=k, value=v))
 
-        instance_name = _MACHINE_TO_COMPUTE_NAME[machine]
+        instance_name = _machine_to_compute_name(machine)
 
         run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if studio_id is not None else ""
 
+        path_mappings_list = resolve_path_mappings(
+            mappings=path_mappings or {},
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
@@ -110,8 +116,7 @@ class MMTApiV2:
             spot=interruptible,
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
-            artifacts_source=artifacts_local or "",
-            artifacts_destination=artifacts_remote or "",
+            path_mappings=path_mappings_list,
         )
         body = ProjectIdMultimachinejobsBody(
             name=name, spec=spec, cluster_id=cloud_account or "", machines=num_machines

lightning_sdk/api/utils.py

@@ -24,6 +24,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     UploadsUploadIdBody,
     V1CompletedPart,
     V1CompleteUpload,
+    V1PathMapping,
     V1PresignedUrl,
     V1SignedUrl,
     V1UploadProjectArtifactPartsResponse,
@@ -614,3 +615,54 @@ def remove_datetime_prefix(text: str) -> str:
     # lines looks something like
     # '[2025-01-08T14:15:03.797142418Z] ⚡  ~ echo Hello\n[2025-01-08T14:15:03.803077717Z] Hello\n'
     return re.sub(r"^\[.*?\] ", "", text, flags=re.MULTILINE)
+
+
+def resolve_path_mappings(
+    mappings: Dict[str, str],
+    artifacts_local: Optional[str],
+    artifacts_remote: Optional[str],
+) -> List[V1PathMapping]:
+    path_mappings_list = []
+    for k, v in mappings.items():
+        splitted = str(v).rsplit(":", 1)
+        connection_name: str
+        connection_path: str
+        if len(splitted) == 1:
+            connection_name = splitted[0]
+            connection_path = ""
+        else:
+            connection_name, connection_path = splitted
+
+        path_mappings_list.append(
+            V1PathMapping(
+                connection_name=connection_name,
+                connection_path=connection_path,
+                container_path=k,
+            )
+        )
+
+    if artifacts_remote:
+        splitted = str(artifacts_remote).rsplit(":", 2)
+        if len(splitted) not in (2, 3):
+            raise RuntimeError(
+                f"Artifacts remote need to be of format efs:connection_name[:path] but got {artifacts_remote}"
+            )
+        else:
+            if not artifacts_local:
+                raise RuntimeError("If Artifacts remote is specified, artifacts local should be specified as well")
+
+            if len(splitted) == 2:
+                _, connection_name = splitted
+                connection_path = ""
+            else:
+                _, connection_name, connection_path = splitted
+
+            path_mappings_list.append(
+                V1PathMapping(
+                    connection_name=connection_name,
+                    connection_path=connection_path,
+                    container_path=artifacts_local,
+                )
+            )
+
+    return path_mappings_list

lightning_sdk/cli/download.py

@@ -3,14 +3,18 @@ import re
 from pathlib import Path
 from typing import Optional
 
+from rich.console import Console
+
+from lightning_sdk.api.lit_container_api import LitContainerApi
 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 download_model
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
 
 
-class _Downloads(_StudiosMenu):
+class _Downloads(_StudiosMenu, _TeamspacesMenu):
     """Download files and folders from Lightning AI."""
 
     def model(self, name: str, download_dir: str = ".") -> None:
@@ -130,3 +134,18 @@ class _Downloads(_StudiosMenu):
                 f"Could not download the file from the given Studio {studio}. "
                 "Please contact Lightning AI directly to resolve this issue."
             ) from e
+
+    def container(self, container: str, teamspace: Optional[str] = None, tag: str = "latest") -> None:
+        """Download a docker container from a teamspace.
+
+        Args:
+          container: The name of the container to download.
+          teamspace: The name of the teamspace to download the container from.
+          tag: The tag of the container to download.
+        """
+        resolved_teamspace = self._resolve_teamspace(teamspace)
+        console = Console()
+        with console.status("Downloading container..."):
+            api = LitContainerApi()
+            api.download_container(container, resolved_teamspace, tag)
+            console.print("Container downloaded successfully", style="green")

lightning_sdk/cli/run.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Dict, Optional
+from typing import TYPE_CHECKING, Dict, Optional, Union
 
 from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
@@ -43,20 +43,28 @@ class _Run:
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
                 Required if the registry is part of a cloud provider (e.g. ECR).
-            artifacts_local: The path of inside the docker container, you want to persist images from.
+            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
+                just runs the provided command in a standard shell.
+                To use the pre-defined entrypoint of the provided image, set this to an empty string.
+                Only applicable when submitting docker jobs.
+            path_mappings: Maps path inside of containers to paths inside data-connections.
+                Should be a comma separated list of form:
+                <MAPPING_1>,<MAPPING_2>,...
+                where each mapping is of the form
+                <CONTAINER_PATH_1>:<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1> and
+                omitting the path inside the connection defaults to the connections root.
+            artifacts_local: Deprecated in favor of path_mappings.
+                The path of inside the docker container, you want to persist images from.
                 CAUTION: When setting this to "/", it will effectively erase your container.
                 Only supported for jobs with a docker image compute environment.
-            artifacts_remote: The remote storage to persist your artifacts to.
+            artifacts_remote: Deprecated in favor of path_mappings.
+                The remote storage to persist your artifacts to.
                 Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
                 PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
                 E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
                 within it.
                 Note that the connection needs to be added to the teamspace already in order for it to be found.
                 Only supported for jobs with a docker image compute environment.
-            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
-                just runs the provided command in a standard shell.
-                To use the pre-defined entrypoint of the provided image, set this to an empty string.
-                Only applicable when submitting docker jobs.
         """
         # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
         # might need to switch to explicit cli definition
@@ -87,20 +95,28 @@ class _Run:
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
                 Required if the registry is part of a cloud provider (e.g. ECR).
-            artifacts_local: The path of inside the docker container, you want to persist images from.
+            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
+                just runs the provided command in a standard shell.
+                To use the pre-defined entrypoint of the provided image, set this to an empty string.
+                Only applicable when submitting docker jobs.
+            path_mappings: Maps path inside of containers to paths inside data-connections.
+                Should be a comma separated list of form:
+                <MAPPING_1>,<MAPPING_2>,...
+                where each mapping is of the form
+                <CONTAINER_PATH_1>:<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1> and
+                omitting the path inside the connection defaults to the connections root.
+            artifacts_local: Deprecated in favor of path_mappings.
+                The path of inside the docker container, you want to persist images from.
                 CAUTION: When setting this to "/", it will effectively erase your container.
                 Only supported for jobs with a docker image compute environment.
-            artifacts_remote: The remote storage to persist your artifacts to.
+            artifacts_remote: Deprecated in favor of path_mappings.
+                The remote storage to persist your artifacts to.
                 Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
                 PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
                 E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
                 within it.
                 Note that the connection needs to be added to the teamspace already in order for it to be found.
                 Only supported for jobs with a docker image compute environment.
-            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
-                just runs the provided command in a standard shell.
-                To use the pre-defined entrypoint of the provided image, set this to an empty string.
-                Only applicable when submitting docker jobs.
         """
         # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
         # might need to switch to explicit cli definition
@@ -124,20 +140,25 @@ class _Run:
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: str = "",
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
     ) -> None:
         if not name:
             from datetime import datetime
 
-            timestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+            timestr = datetime.now().strftime("%b-%d-%H_%M")
             name = f"job-{timestr}"
 
         if machine is None:
             # TODO: infer from studio
             machine = "CPU"
-        machine_enum = Machine[machine.upper()]
+        machine_enum: Union[str, Machine]
+        try:
+            machine_enum = Machine[machine.upper()]
+        except KeyError:
+            machine_enum = machine
 
         resolved_teamspace = Teamspace(name=teamspace, org=org, user=user)
 
@@ -145,6 +166,17 @@ 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]
+
         Job.run(
             name=name,
             machine=machine_enum,
@@ -159,9 +191,10 @@ class _Run:
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings_dict,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )
 
     # TODO: sadly, fire displays both Optional[type] and Union[type, None] as Optional[Optional]
@@ -182,20 +215,25 @@ class _Run:
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: str = "",
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
     ) -> None:
         if name is None:
             from datetime import datetime
 
-            timestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+            timestr = datetime.now().strftime("%b-%d-%H_%M")
             name = f"mmt-{timestr}"
 
         if machine is None:
             # TODO: infer from studio
             machine = "CPU"
-        machine_enum = Machine[machine.upper()]
+        machine_enum: Union[str, Machine]
+        try:
+            machine_enum = Machine[machine.upper()]
+        except KeyError:
+            machine_enum = machine
 
         resolved_teamspace = Teamspace(name=teamspace, org=org, user=user)
         if cloud_account is None:
@@ -204,6 +242,17 @@ 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]
+
         MMT.run(
             name=name,
             num_machines=num_machines,
@@ -219,7 +268,8 @@ class _Run:
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings_dict,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )

lightning_sdk/cli/serve.py

@@ -4,6 +4,7 @@ import warnings
 from pathlib import Path
 from typing import Optional, Union
 
+import docker
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.prompt import Confirm
@@ -89,11 +90,6 @@ class _LitServe:
         tag: str = "litserve-model",
         non_interactive: bool = False,
     ) -> None:
-        try:
-            import docker
-        except ImportError:
-            raise ImportError("docker-py is not installed. Please install it with `pip install docker`") from None
-
         try:
             client = docker.from_env()
             client.ping()

lightning_sdk/cli/upload.py

@@ -2,14 +2,15 @@ import concurrent.futures
 import json
 import os
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, Generator, List, Optional
 
+import rich
 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.lit_container_api import LCRAuthFailedError, 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
@@ -162,21 +163,38 @@ class _Uploads(_StudiosMenu, _TeamspacesMenu):
             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")
+            try:
+                lines = api.upload_container(container, teamspace, tag)
+                self._print_docker_push(lines, console, progress, push_task)
+            except LCRAuthFailedError:
+                console.print("Authenticating with Lightning Container Registry...")
+                if not api.authenticate():
+                    raise StudioCliError("Failed to authenticate with Lightning Container Registry") from None
+                console.print("Authenticated with Lightning Container Registry", style="green")
+                lines = api.upload_container(container, teamspace, tag)
+                self._print_docker_push(lines, console, progress, push_task)
             progress.update(push_task, description="[green]Container pushed![/green]")
 
+    @staticmethod
+    def _print_docker_push(
+        lines: Generator, console: Console, progress: Progress, push_task: rich.progress.TaskID
+    ) -> None:
+        for line in lines:
+            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
+            elif "finish" in line:
+                console.print(f"Container available at [i]{line['url']}[/i]")
+                return
+            else:
+                console.print(line, style="bright_black")
+
     def _start_parallel_upload(
         self, executor: concurrent.futures.ThreadPoolExecutor, studio: Studio, upload_state: Dict[str, str]
     ) -> List[concurrent.futures.Future]:

lightning_sdk/helpers.py

@@ -43,7 +43,7 @@ def _check_version_and_prompt_upgrade(curr_version: str) -> None:
         warnings.warn(
             f"A newer version of {__package_name__} is available ({new_version}). "
             f"Please consider upgrading with `pip install -U {__package_name__}`. "
-            "Not all functionalities of the platform can be guaranteed to work with the current version.",
+            "Not all platform functionality can be guaranteed to work with the current version.",
             UserWarning,
         )
     return