lightning-sdk 0.1.57__py3-none-any.whl → 0.1.58__py3-none-any.whl

lightning_sdk/__init__.py

@@ -29,6 +29,6 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.57"
+__version__ = "0.1.58"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/lit_container_api.py

@@ -1,4 +1,5 @@
-from typing import Any, Callable, Generator, List
+import time
+from typing import Any, Callable, Dict, Generator, Iterator, List
 
 import docker
 import requests
@@ -12,11 +13,11 @@ from lightning_sdk.teamspace import Teamspace
 
 class LCRAuthFailedError(Exception):
     def __init__(self) -> None:
-        super().__init__(
-            "Failed to authenticate with Lightning Container Registry. Please login manually "
-            "using the following command:\n  "
-            "echo $LIGHTNING_API_KEY | docker login litcr.io --username=LIGHTNING_USERNAME --password-stdin"
-        )
+        super().__init__("Failed to authenticate with Lightning Container Registry. Please try again.")
+
+
+class DockerPushError(Exception):
+    pass
 
 
 def retry_on_lcr_auth_failure(func: Callable) -> Callable:
@@ -24,7 +25,7 @@ def retry_on_lcr_auth_failure(func: Callable) -> Callable:
         try:
             return func(self, *args, **kwargs)
         except LCRAuthFailedError:
-            self.authenticate()
+            self.authenticate(reauth=True)
             return func(self, *args, **kwargs)
 
     return wrapper
@@ -37,15 +38,35 @@ class LitContainerApi:
         try:
             self._docker_client = docker.from_env()
             self._docker_client.ping()
+            self._docker_auth_config = {}
         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 authenticate(self, reauth: bool = False) -> bool:
+        try:
+            authed_user = self._client.auth_service_get_user()
+            username = authed_user.username
+            api_key = authed_user.api_key
+            registry = _get_registry_url()
+            resp = self._docker_client.login(username, password=api_key, registry=registry, reauth=reauth)
+
+            if (
+                resp.get("username", None) == username
+                and resp.get("password", None) == api_key
+                and resp.get("serveraddress", None) == registry
+            ):
+                self._docker_auth_config = {"username": username, "password": api_key}
+                return True
+
+            # This is a new 200 response auth attempt from the client.
+            if "Status" in resp and resp["Status"] == "Login Succeeded":
+                self._docker_auth_config = {"username": username, "password": api_key}
+                return True
+
+            return False
+        except Exception as e:
+            print(f"Authentication error: {e} resp: {resp}")
+            return False
 
     def list_containers(self, project_id: str) -> List:
         project = self._client.lit_registry_service_get_lit_project_registry(project_id)
@@ -61,7 +82,15 @@ class LitContainerApi:
         try:
             self._docker_client.images.get(container)
         except docker.errors.ImageNotFound:
-            raise ValueError(f"Container {container} does not exist") from None
+            try:
+                self._docker_client.images.pull(container, tag)
+                self._docker_client.images.get(container)
+            except docker.errors.APIError as e:
+                raise ValueError(f"Could not pull container {container}") from e
+            except docker.errors.ImageNotFound as e:
+                raise ValueError(f"Container {container} does not exist") from e
+            except Exception as e:
+                raise ValueError(f"Unable to upload {container}") from e
 
         registry_url = _get_registry_url()
         container_basename = container.split("/")[-1]
@@ -69,22 +98,55 @@ 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}")
-        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"] or "unauth" in line["error"]):
-                raise LCRAuthFailedError()
-            yield line
+        yield from self._push_with_retry(repository)
         yield {
             "finish": True,
             "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container_basename}",
         }
 
+    def _push_with_retry(self, repository: str, max_retries: int = 3) -> Iterator[Dict[str, Any]]:
+        def is_auth_error(error_msg: str) -> bool:
+            auth_errors = ["unauthorized", "authentication required", "unauth"]
+            return any(err in error_msg.lower() for err in auth_errors)
+
+        def is_timeout_error(error_msg: str) -> bool:
+            timeout_errors = ["proxyconnect tcp", "i/o timeout"]
+            return any(err in error_msg.lower() for err in timeout_errors)
+
+        for attempt in range(max_retries):
+            try:
+                if attempt > 0:
+                    # This is important, if we don't set reauth here then we just keep using the
+                    # same authentication context that we know just failed.
+                    self.authenticate(reauth=True)
+                    time.sleep(2)
+
+                lines = self._docker_client.api.push(repository, stream=True, decode=True, auth_config=self._docker_auth_config)
+                for line in lines:
+                    if isinstance(line, dict) and "error" in line:
+                        error = line["error"]
+                        if is_auth_error(error) or is_timeout_error(error):
+                            if attempt < max_retries - 1:
+                                break
+                            raise DockerPushError(f"Max retries reached: {error}")
+                        raise DockerPushError(f"Push error: {error}")
+                    yield line
+                else:
+                    return
+
+            except docker.errors.APIError as e:
+                if (is_auth_error(str(e)) or is_timeout_error(str(e))) and attempt < max_retries - 1:
+                    continue
+                raise DockerPushError(f"Push failed: {e}") from e
+
+        raise DockerPushError("Max push retries reached")
+
     @retry_on_lcr_auth_failure
     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)
+            self._docker_client.images.pull(repository, tag=tag, auth_config=self._docker_auth_config)
         except requests.exceptions.HTTPError as e:
             if "unauthorized" in e.response.text:
                 raise LCRAuthFailedError() from e

lightning_sdk/cli/upload.py

@@ -154,7 +154,9 @@ class _Uploads(_StudiosMenu, _TeamspacesMenu):
         console.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)
+        """Upload a container to Lightning AI's container registry."""
+        menu = _TeamspacesMenu()
+        teamspace = menu._resolve_teamspace(teamspace)
         api = LitContainerApi()
         console = Console()
         with Progress(
@@ -166,11 +168,19 @@ class _Uploads(_StudiosMenu, _TeamspacesMenu):
         ) as progress:
             push_task = progress.add_task("Pushing Docker image", total=None)
             try:
+                console.print("Authenticating with Lightning Container Registry...")
+                try:
+                    api.authenticate()
+                    console.print("Authenticated with Lightning Container Registry", style="green")
+                except Exception:
+                    # let the push with retry take control of auth moving forward
+                    pass
+
                 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():
+                console.print("Re-authenticating with Lightning Container Registry...")
+                if not api.authenticate(reauth=True):
                     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)

{lightning_sdk-0.1.57.dist-info → lightning_sdk-0.1.58.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.1.57
+Version: 0.1.58
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License

{lightning_sdk-0.1.57.dist-info → lightning_sdk-0.1.58.dist-info}/RECORD

@@ -1,5 +1,5 @@
 docs/source/conf.py,sha256=r8yX20eC-4mHhMTd0SbQb5TlSWHhO6wnJ0VJ_FBFpag,13249
-lightning_sdk/__init__.py,sha256=OA3km-3Y0trAWAJMgGucXNd2lF73ByuAOu0VJ59EFak,1039
+lightning_sdk/__init__.py,sha256=ruIP3Znv7pOMFZGza2nk5WQAKSi4E64V4RBeB28YTFc,1039
 lightning_sdk/agents.py,sha256=ly6Ma1j0ZgGPFyvPvMN28JWiB9dATIstFa5XM8pMi6I,1577
 lightning_sdk/ai_hub.py,sha256=dfw5CS93ehGmPlZCbbE2Pu3v0ANI6ZY05-ULbcHFqs4,7189
 lightning_sdk/constants.py,sha256=ztl1PTUBULnqTf3DyKUSJaV_O20hNtUYT6XvAYIrmIk,749
@@ -19,7 +19,7 @@ lightning_sdk/api/agents_api.py,sha256=G47TbFo9kYqnBMqdw2RW-lfS1VAUBSXDmzs6fpIEM
 lightning_sdk/api/ai_hub_api.py,sha256=Yr9VvxueIqPUUeMqExbjbWD5CE_yeWjq0dXQQv47erg,6779
 lightning_sdk/api/deployment_api.py,sha256=SwpTh8elvCeEE-Ilf5pJc4nM_xUXBUNWV5ejH4Y5glI,21903
 lightning_sdk/api/job_api.py,sha256=pHJMs209qMfML0uO2jV-6v5QfA8bb-whMxQfmiInb5U,12404
-lightning_sdk/api/lit_container_api.py,sha256=QJ2-6jnKztjbg-Em6KW82z-_hz8XHmzEToVvic3WcLo,4231
+lightning_sdk/api/lit_container_api.py,sha256=jaE4uXD9Y5ThhXhDw2qcmwQ8umtvsbrtYqRsnqIIPvg,7018
 lightning_sdk/api/mmt_api.py,sha256=UXUJaHG4R-8jDWkYh6hFYr0uJSiAYiXFiWTJn4_eVHQ,7515
 lightning_sdk/api/org_api.py,sha256=Ze3z_ATVrukobujV5YdC42DKj45Vuwl7X52q_Vr-o3U,803
 lightning_sdk/api/studio_api.py,sha256=N0XbvNvFha-et94bTlQLLjj6OPTXvrpIWgnSMUXi4UI,26327
@@ -48,7 +48,7 @@ lightning_sdk/cli/stop.py,sha256=KzykI4-mVfla0yMRvSuBAxoo06hq3gG342sVuSFpBA4,264
 lightning_sdk/cli/studios_menu.py,sha256=TA9rO6_fFHGMz0Nt4rJ6iV80X5pZE4xShrSiyXoU-oQ,4129
 lightning_sdk/cli/switch.py,sha256=K78yxn9voyVCZR46lV29ZD6A539EPRmvg6XYgb2K6mc,1876
 lightning_sdk/cli/teamspace_menu.py,sha256=C3g3spTKgtMwoK7pnooy0MBPz4AKhFjcObkvZyZ4v04,3797
-lightning_sdk/cli/upload.py,sha256=fy7Bz1TnAzxhjuYKrJQoClSA_hzXf02b0d1FxTvsv9o,11589
+lightning_sdk/cli/upload.py,sha256=qcEdBYFFHc_LfSNir6wl1_GZG3PkijFxYatapwv-VyM,12094
 lightning_sdk/deployment/__init__.py,sha256=BLu7_cVLp97TYxe6qe-J1zKUSZXAVcvCjgcA7plV2k4,497
 lightning_sdk/deployment/deployment.py,sha256=Fn8iJVimaWB0Ngz53oJ03kBUzVe5_-3oHAbvEqFTf0s,15681
 lightning_sdk/job/__init__.py,sha256=1MxjQ6rHkyUHCypSW9RuXuVMVH11WiqhIXcU2LCFMwE,64
@@ -896,9 +896,9 @@ lightning_sdk/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 lightning_sdk/utils/dynamic.py,sha256=glUTO1JC9APtQ6Gr9SO02a3zr56-sPAXM5C3NrTpgyQ,1959
 lightning_sdk/utils/enum.py,sha256=h2JRzqoBcSlUdanFHmkj_j5DleBHAu1esQYUsdNI-hU,4106
 lightning_sdk/utils/resolve.py,sha256=idhRk4zg0jBOlySwKwjhsGgRorpsNVHOKOADOxMUYY0,6062
-lightning_sdk-0.1.57.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
-lightning_sdk-0.1.57.dist-info/METADATA,sha256=zK1E4W0HFXVjevGaq4MI9eqcvAYTOcToZAOFatRGOlE,4054
-lightning_sdk-0.1.57.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-lightning_sdk-0.1.57.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
-lightning_sdk-0.1.57.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
-lightning_sdk-0.1.57.dist-info/RECORD,,
+lightning_sdk-0.1.58.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
+lightning_sdk-0.1.58.dist-info/METADATA,sha256=YJJ6KEn8kLe7WF-r72Dhg5loqInqGLlJALO231wf1lo,4054
+lightning_sdk-0.1.58.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+lightning_sdk-0.1.58.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
+lightning_sdk-0.1.58.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
+lightning_sdk-0.1.58.dist-info/RECORD,,