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

lightning_sdk/__init__.py

@@ -29,5 +29,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.50"
+__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/api/job_api.py

@@ -7,10 +7,9 @@ from lightning_sdk.api.utils import (
     _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,
@@ -214,9 +213,10 @@ class JobApiV2:
         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:
@@ -227,6 +227,12 @@ class JobApiV2:
 
         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,23 +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)
 
-        import docker
-
         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
@@ -29,8 +42,6 @@ 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:
             self._docker_client.images.get(container)
         except docker.errors.ImageNotFound:
@@ -41,11 +52,17 @@ class LitContainerApi:
         tagged = self._docker_client.api.tag(container, repository, tag)
         if not tagged:
             raise ValueError(f"Could not tag container {container} with {repository}:{tag}")
-        return self._docker_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]:
-        import docker
-
         registry_url = _get_registry_url()
         repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
         try:

lightning_sdk/api/mmt_api.py

@@ -7,10 +7,9 @@ from lightning_sdk.api.utils import (
     _COMPUTE_NAME_TO_MACHINE,
     _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,
@@ -85,9 +84,10 @@ class MMTApiV2:
         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:
@@ -98,6 +98,12 @@ class MMTApiV2:
 
         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/run.py

@@ -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,14 +140,15 @@ 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:
@@ -149,6 +166,8 @@ class _Run:
             cloud_account = resolved_teamspace.default_cloud_account
         machine_enum = Machine(machine.upper())
 
+        path_mappings_dict = self._resolve_path_mapping(path_mappings=path_mappings)
+
         Job.run(
             name=name,
             machine=machine_enum,
@@ -163,9 +182,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]
@@ -186,14 +206,15 @@ 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:
@@ -212,6 +233,8 @@ class _Run:
         if image is None:
             raise RuntimeError("Image needs to be specified to run a multi-machine job")
 
+        path_mappings_dict = self._resolve_path_mapping(path_mappings=path_mappings)
+
         MMT.run(
             name=name,
             num_machines=num_machines,
@@ -227,7 +250,31 @@ 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,
         )
+
+    @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/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