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

lightning_sdk/job/v2.py

@@ -4,9 +4,10 @@ from lightning_sdk.api.job_api import JobApiV2
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.job.base import _BaseJob
 from lightning_sdk.status import Status
+from lightning_sdk.utils.resolve import _get_org_id
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
@@ -37,18 +38,22 @@ class _JobV2(_BaseJob):
 
     def _submit(
         self,
-        machine: "Machine",
+        machine: Union["Machine", str],
         command: Optional[str] = None,
         studio: Optional["Studio"] = None,
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        artifacts_local: Optional[str] = None,
-        artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
+        artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
+        reuse_snapshot: bool = True,
     ) -> "_JobV2":
         """Submit a new job to the Lightning AI platform.
 
@@ -70,17 +75,21 @@ class _JobV2(_BaseJob):
             artifacts_local: 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.
-                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.
+            path_mappings: Dictionary of path mappings. The keys are the path inside the container whereas the value
+                represents the data-connection name and the path inside that connection.
+                Should be of form
+                    {
+                        "<CONTAINER_PATH_1>": "<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1>",
+                        "<CONTAINER_PATH_2>": "<CONNECTION_NAME_2>"
+                    }
+                If the path inside the connection is omitted it's assumed to be the root path of that connection.
                 Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -99,6 +108,13 @@ class _JobV2(_BaseJob):
             if image is None:
                 raise ValueError("either image or studio must be provided")
 
+        cloud_account = self._cloud_account_api.resolve_cloud_account(
+            self._teamspace.id,
+            cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
+            default_cloud_account=self._teamspace.default_cloud_account,
+        )
+
         submitted = self._job_api.submit_job(
             name=self.name,
             command=command,
@@ -114,6 +130,9 @@ class _JobV2(_BaseJob):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             entrypoint=entrypoint,
+            path_mappings=path_mappings,
+            max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         self._job = submitted
         self._name = submitted.name
@@ -140,13 +159,31 @@ class _JobV2(_BaseJob):
     @property
     def status(self) -> "Status":
         """The current status of the job."""
-        return self._job_api._job_state_to_external(self._latest_job.state)
+        try:
+            return self._job_api._job_state_to_external(self._latest_job.state)
+        except Exception:
+            raise RuntimeError(
+                f"Job {self._name} does not exist in Teamspace {self.teamspace.name}. Did you delete it?"
+            ) from None
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """The machine type the job is running on."""
         # only fetch the job it it hasn't been fetched yet as machine cannot change over time
-        return self._job_api._get_job_machine_from_spec(self._guaranteed_job.spec)
+
+        return self._job_api._get_job_machine_from_spec(
+            self._guaranteed_job.spec,
+            self.teamspace.id,
+            _get_org_id(self.teamspace),
+        )
+
+    @property
+    def public_ip(self) -> Optional[str]:
+        """Get the public IP of the machine the job is running on."""
+        try:
+            return self._job.public_ip_address
+        except AttributeError:
+            return None
 
     @property
     def artifact_path(self) -> Optional[str]:

lightning_sdk/job/work.py

@@ -1,6 +1,7 @@
-from typing import TYPE_CHECKING, Any, Optional, Protocol
+from typing import TYPE_CHECKING, Any, Optional, Protocol, Union
 
 from lightning_sdk.api.job_api import JobApiV1
+from lightning_sdk.utils.resolve import _get_org_id
 
 if TYPE_CHECKING:
     from lightning_sdk.job.base import MachineDict
@@ -51,8 +52,11 @@ class Work:
         return self._job._name_filter(self._guaranteed_work.name)
 
     @property
-    def machine(self) -> "Machine":
-        return self._job_api.get_machine_from_work(self._guaranteed_work)
+    def machine(self) -> Union["Machine", str]:
+        return self._job_api.get_machine_from_work(
+            self._guaranteed_work,
+            org_id=_get_org_id(self._teamspace),
+        )
 
     @property
     def artifact_path(self) -> Optional[str]:

lightning_sdk/lightning_cloud/login.py

@@ -5,7 +5,7 @@ import logging
 import os
 import pathlib
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Literal
 from urllib.parse import urlencode
 
 import webbrowser
@@ -17,14 +17,26 @@ from starlette.background import BackgroundTask
 from starlette.responses import RedirectResponse
 
 from lightning_sdk.lightning_cloud import env
+from lightning_sdk.lightning_cloud.openapi import ApiClient, Configuration
+from lightning_sdk.lightning_cloud.openapi.api import \
+    AuthServiceApi
+from lightning_sdk.lightning_cloud.openapi.models.v1_guest_login_request import \
+    V1GuestLoginRequest
+from lightning_sdk.lightning_cloud.openapi.models.v1_token_login_request import \
+    V1TokenLoginRequest
+from lightning_sdk.lightning_cloud.openapi.models.v1_refresh_request import \
+    V1RefreshRequest
 
 logger = logging.getLogger(__name__)
 
+# Authentication override types
+AuthOverride = Literal["auth_token", "api_key", "guest"]
 
 class Keys(Enum):
     # USERNAME = "LIGHTNING_USERNAME"
     USER_ID = "LIGHTNING_USER_ID"
     API_KEY = "LIGHTNING_API_KEY"
+    AUTH_TOKEN = "LIGHTNING_AUTH_TOKEN"
 
     @property
     def suffix(self):
@@ -36,21 +48,18 @@ class Auth:
     # username: Optional[str] = None
     user_id: Optional[str] = None
     api_key: Optional[str] = None
+    auth_token: Optional[str] = None
 
     secrets_file = pathlib.Path(env.LIGHTNING_CREDENTIAL_PATH)
 
     def __post_init__(self):
         for key in Keys:
-            setattr(self, key.suffix, os.environ.get(key.value, None))
+            # Only set from environment if not already set
+            if getattr(self, key.suffix) is None:
+                setattr(self, key.suffix, os.environ.get(key.value, None))
 
-        self._with_env_var = bool(
-            self.user_id and self.api_key)  # used by authenticate method
-        if self.api_key and not self.user_id:
-            raise ValueError(
-                f"{Keys.USER_ID.value} is missing from env variables. "
-                "To use env vars for authentication both "
-                f"{Keys.USER_ID.value} and {Keys.API_KEY.value} should be set."
-            )
+        # used by authenticate method
+        self._with_env_var = bool(self.api_key or self.auth_token)
 
     def load(self) -> bool:
         """Load credentials from disk and update properties with credentials.
@@ -72,6 +81,7 @@ class Auth:
              token: str = "",
              user_id: str = "",
              api_key: str = "",
+             auth_token: str = "",
              # username: str = "",
         ) -> None:
         """save credentials to disk."""
@@ -82,6 +92,7 @@ class Auth:
                     # f"{Keys.USERNAME.suffix}": username,
                     f"{Keys.USER_ID.suffix}": user_id,
                     f"{Keys.API_KEY.suffix}": api_key,
+                    f"{Keys.AUTH_TOKEN.suffix}": auth_token,
                 },
                 f,
             )
@@ -89,6 +100,7 @@ class Auth:
         # self.username = username
         self.user_id = user_id
         self.api_key = api_key
+        self.auth_token = auth_token
         logger.debug("credentials saved successfully")
 
     @classmethod
@@ -100,15 +112,78 @@ class Auth:
             os.environ.pop(key.value, None)
         logger.debug("credentials removed successfully")
 
-    @property
-    def auth_header(self) -> Optional[str]:
-        """authentication header used by lightning-cloud client."""
-        if self.api_key:
+    def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[str]:
+        """Get authentication header with optional override.
+        
+        By default, uses the automatic priority selection (auth_token > api_key).
+        You can override this for specific cases where you need a different auth method.
+        
+        Parameters
+        ----------
+        override : AuthOverride, optional
+            Override the default authentication method:
+            - "auth_token": Force use of JWT auth token (Bearer)
+            - "api_key": Force use of API key (Basic)
+            - "guest": Force use of guest credentials (Basic)
+            - None: Use automatic selection (default)
+            
+        Returns
+        -------
+        Optional[str]
+            The authorization header for the specified method.
+            
+        Raises
+        ------
+        ValueError
+            If the specified override is not available or invalid.
+        """
+        if override == "auth_token":
+            if not self.auth_token:
+                raise ValueError(
+                    "Auth token override requested but no JWT token available. "
+                    "Please use token_login() method first."
+                )
+            return f"Bearer {self.auth_token}"
+        
+        elif override == "api_key":
+            if not self.api_key or not self.user_id:
+                raise ValueError(
+                    "API key override requested but no API key or user ID available. "
+                    "Please set LIGHTNING_API_KEY and LIGHTNING_USER_ID environment variables "
+                    "or use authenticate() method."
+                )
             token = f"{self.user_id}:{self.api_key}"
             return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}"  # noqa E501
-        raise AttributeError(
-            "Authentication Failed, no authentication header available. "
-            "This is most likely a bug in the LightningCloud Framework")
+        
+        elif override == "guest":
+            if not self.api_key or not self.user_id:
+                raise ValueError(
+                    "Guest override requested but no guest credentials available. "
+                    "Please call guest_login() method first."
+                )
+            token = f"{self.user_id}:{self.api_key}"
+            return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}"  # noqa E501
+        
+        elif override is None:
+            # Use the original automatic selection logic (default behavior)
+            if self.auth_token:
+                return f"Bearer {self.auth_token}"
+            elif self.api_key:
+                token = f"{self.user_id}:{self.api_key}"
+                return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}"  # noqa E501
+            else:
+                raise ValueError(
+                    "No authentication credentials available. Please authenticate first using "
+                    "token_login(), guest_login(), or authenticate() methods."
+                )
+        
+        else:
+            raise ValueError(f"Invalid authentication override: {override}")
+
+    @property
+    def auth_header(self) -> Optional[str]:
+        """authentication header used by lightning-cloud client (automatic selection)."""
+        return self.get_auth_header()
 
     def _run_server(self) -> None:
         """start a server to complete authentication."""
@@ -131,6 +206,8 @@ class Auth:
             self._run_server()
             return self.auth_header
 
+        elif self.auth_token:
+            return self.auth_header
         elif self.user_id and self.api_key:
             return self.auth_header
 
@@ -138,6 +215,233 @@ class Auth:
             "We couldn't find any credentials linked to your account. Please try logging in using the CLI command `lightning login`"
         )
 
+    def guest_login(self) -> Optional[str]:
+        """Performs guest user authentication.
+        This method sends a request to the guest login endpoint to get temporary
+        credentials, saves them, and returns the authorization header.
+        Useful to log experiments as a non signed in user, using a guest account
+        in the background.
+        Returns
+        -------
+        Optional[str]
+            The authorization header to use for subsequent requests.
+        Raises
+        ------
+        RuntimeError
+            If the guest login request fails.
+        ValueError
+            If the response from the server is invalid.
+        """
+
+        config = Configuration()
+        config.host = env.LIGHTNING_CLOUD_URL
+        api_client = ApiClient(configuration=config)
+        auth_api = AuthServiceApi(api_client)
+
+        logger.debug(f"Attempting guest login to {config.host}")
+
+        try:
+            # The body is an empty object for a guest login.
+            body = V1GuestLoginRequest()
+            credentials = auth_api.auth_service_guest_login(body)
+
+        except requests.RequestException as e:
+            logger.error(f"Guest login request failed: {e}")
+            raise RuntimeError(
+                "Failed to connect to the guest login endpoint. "
+                "Please check your network connection and the server status."
+            ) from e
+
+        # attributes based on the `V1GuestLoginResponse` model.
+        user = getattr(credentials, "user", None)
+        user_id = getattr(user, "id", None) if user else None
+        api_key = getattr(user, "api_key", None) if user else None
+
+        if not all([user_id, api_key]):
+            logger.error(
+                f"Incomplete credentials received from guest login: {credentials}"
+            )
+            raise ValueError(
+                "The guest login response did not contain the required 'user_id' and 'api_key' fields."
+            )
+
+        self.save(user_id=user_id, api_key=api_key)
+        logger.info("Successfully authenticated as a guest user.")
+
+        return self.auth_header
+
+    def token_login(self, token_key: str, save_token: bool = True) -> Optional[str]:
+        """Performs token-based authentication.
+        
+        This method sends a request to the token login endpoint to authenticate
+        using an auth token key, optionally saves the returned JWT token, and 
+        returns the authorization header.
+        
+        Parameters
+        ----------
+        token_key : str
+            The auth token key to use for authentication.
+        save_token : bool, optional
+            Whether to save the JWT token for future use. Defaults to True.
+            If False, the token is only used for the current session.
+            
+        Returns
+        -------
+        Optional[str]
+            The authorization header to use for subsequent requests.
+            
+        Raises
+        ------
+        RuntimeError
+            If the token login request fails.
+        ValueError
+            If the response from the server is invalid.
+        """
+        config = Configuration()
+        config.host = env.LIGHTNING_CLOUD_URL
+        api_client = ApiClient(configuration=config)
+        auth_api = AuthServiceApi(api_client)
+
+        logger.debug(f"Attempting token login to {config.host}")
+
+        try:
+            body = V1TokenLoginRequest(token_key=token_key)
+            response = auth_api.auth_service_token_login(body)
+
+        except requests.RequestException as e:
+            logger.error(f"Token login request failed: {e}")
+            raise RuntimeError(
+                "Failed to connect to the token login endpoint. "
+                "Please check your network connection and the server status."
+            ) from e
+
+        # Extract the JWT token from the response
+        jwt_token = getattr(response, "token", None)
+
+        if not jwt_token:
+            logger.error(
+                f"No token received from token login response: {response}"
+            )
+            raise ValueError(
+                "The token login response did not contain a valid JWT token."
+            )
+
+        # Set the JWT token in memory
+        self.auth_token = jwt_token
+        
+        # Optionally save the JWT token to disk
+        if save_token:
+            self.save(auth_token=jwt_token)
+            logger.info("Successfully authenticated using auth token and saved to disk.")
+        else:
+            logger.info("Successfully authenticated using auth token (not saved to disk).")
+
+        return self.auth_header
+
+    def refresh_token(self, duration: int = 43200) -> Optional[str]:
+        """Refreshes the current JWT token.
+        
+        This method sends a request to the refresh endpoint to get a new JWT token
+        with the specified duration, saves the new token, and returns the updated
+        authorization header.
+        
+        Parameters
+        ----------
+        duration : int, optional
+            Duration in seconds for the new token. Can range from 900 seconds (15 minutes)
+            up to a maximum of 129,600 seconds (36 hours), with a default of 43,200 seconds (12 hours).
+            
+        Returns
+        -------
+        Optional[str]
+            The updated authorization header with the new JWT token.
+            
+        Raises
+        ------
+        RuntimeError
+            If the refresh request fails.
+        ValueError
+            If no valid JWT token is available to refresh, or if the response is invalid.
+        """
+        if not self.auth_token:
+            raise ValueError(
+                "No JWT token available to refresh. Please authenticate first using "
+                "token_login() or authenticate() methods."
+            )
+        
+        config = Configuration()
+        config.host = env.LIGHTNING_CLOUD_URL
+        # Set the current auth token as the authorization header
+        config.api_key_prefix['Authorization'] = 'Bearer'
+        config.api_key['Authorization'] = self.auth_token
+        
+        api_client = ApiClient(configuration=config)
+        auth_api = AuthServiceApi(api_client)
+
+        logger.debug(f"Attempting to refresh JWT token with duration {duration} seconds")
+
+        try:
+            body = V1RefreshRequest(duration=str(duration))
+            response = auth_api.auth_service_refresh(body)
+
+        except requests.RequestException as e:
+            logger.error(f"Token refresh request failed: {e}")
+            raise RuntimeError(
+                "Failed to connect to the refresh endpoint. "
+                "Please check your network connection and the server status."
+            ) from e
+
+        # Extract the new JWT token from the response
+        new_jwt_token = getattr(response, "token", None)
+
+        if not new_jwt_token:
+            logger.error(
+                f"No token received from refresh response: {response}"
+            )
+            raise ValueError(
+                "The refresh response did not contain a valid JWT token."
+            )
+
+        # Save the new JWT token
+        self.save(auth_token=new_jwt_token)
+        logger.info("Successfully refreshed JWT token.")
+
+        return self.auth_header
+
+    def create_api_client(self, override: Optional[AuthOverride] = None) -> 'ApiClient':
+        """Create an API client with optional authentication override.
+        
+        This is a convenience method for creating API clients that use a specific
+        authentication method instead of the default automatic selection.
+        
+        Parameters
+        ----------
+        override : AuthOverride, optional
+            Override the default authentication method for this API client.
+            See get_auth_header() for available options.
+            
+        Returns
+        -------
+        ApiClient
+            Configured API client with the specified authentication method.
+        """
+        from lightning_sdk.lightning_cloud.openapi import ApiClient, Configuration
+        
+        config = Configuration()
+        config.host = env.LIGHTNING_CLOUD_URL
+        
+        # Get the auth header for the specified override
+        auth_header = self.get_auth_header(override)
+        
+        # Create the API client
+        client = ApiClient(configuration=config)
+        
+        # Set the Authorization header directly in default_headers
+        if auth_header:
+            client.set_default_header('Authorization', auth_header)
+        
+        return client
+
 
 class AuthServer:
 
@@ -168,7 +472,10 @@ class AuthServer:
         ok = webbrowser.open(url)
         if not ok:
             # can't open a browser, authentication failed
-            raise RuntimeError("Failed to authenticate to Lightning. When running without access to a browser, `LIGHTNING_USER_ID` and `LIGHTNING_API_KEY` should be exported.")
+            deployment_id = os.environ.get("LIGHTNING_DEPLOYMENT_ID", None)
+            if deployment_id is not None and deployment_id != "":
+                raise RuntimeError("Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the 'Environment' section of the deployment settings.")
+            raise RuntimeError("Failed to authenticate to Lightning. When running without access to a browser, 'LIGHTNING_USER_ID' and 'LIGHTNING_API_KEY' should be exported.")
 
         @app.get("/login-complete")
         async def save_token(request: Request,