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

lightning_sdk/mmt/base.py

@@ -1,3 +1,4 @@
+import warnings
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Tuple, Union
 
@@ -11,7 +12,7 @@ if TYPE_CHECKING:
     from lightning_sdk.user import User
 
 from lightning_sdk.job.base import _BaseJob
-from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, in_studio
 
 
 class MMTMachine(Protocol):
@@ -67,9 +68,10 @@ class _BaseMMT(_BaseJob):
         interruptible: bool = False,
         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,
+        artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseMMT":
         """Run async workloads using a docker image across multiple machines.
@@ -94,21 +96,21 @@ class _BaseMMT(_BaseJob):
                 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.
-                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.
                 Only applicable when submitting docker jobs.
+            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.
         """
+        from lightning_sdk.lightning_cloud.openapi.rest import ApiException
         from lightning_sdk.studio import Studio
 
         cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
@@ -167,6 +169,13 @@ class _BaseMMT(_BaseJob):
                     "image and studio are mutually exclusive as both define the environment to run the job in"
                 )
 
+            if cloud_account is None and in_studio():
+                try:
+                    resolve_studio = Studio(teamspace=teamspace, user=user, org=org)
+                    cloud_account = resolve_studio.cloud_account
+                except (ValueError, ApiException):
+                    warnings.warn("Could not infer cloud account from studio. Using teamspace default.")
+
             # they either need to specified both or none of them
             if bool(artifacts_local) != bool(artifacts_remote):
                 raise ValueError("Artifact persistence requires both artifacts_local and artifacts_remote to be set")
@@ -189,9 +198,10 @@ class _BaseMMT(_BaseJob):
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )
         return inst
 
@@ -208,9 +218,10 @@ class _BaseMMT(_BaseJob):
         cloud_account: Optional[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,
+        artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> None:
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -230,19 +241,18 @@ class _BaseMMT(_BaseJob):
                 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.
-                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.
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
+            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.
         """
 
     @property

lightning_sdk/mmt/mmt.py

@@ -109,9 +109,10 @@ class MMT(_BaseMMT):
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "MMT":
         """Run async workloads using a docker image across multiple machines.
@@ -119,7 +120,7 @@ class MMT(_BaseMMT):
         Args:
             name: The name of the job. Needs to be unique within the teamspace.
             machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
-            num_machine: The number of machines to run on.
+            num_machines: The number of machines to run on.
             command: The command to run inside your job. Required if using a studio. Optional if using an image.
                 If not provided for images, will run the container entrypoint and default command.
             studio: The studio env to run the job with. Mutually exclusive with image.
@@ -136,19 +137,18 @@ class MMT(_BaseMMT):
                 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.
-                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.
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
+            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.
         """
         ret_val = super().run(
             name=name,
@@ -165,9 +165,10 @@ class MMT(_BaseMMT):
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
             cluster=cluster,  # deprecated in favor of cloud_account
         )
         # required for typing with "MMT"
@@ -190,9 +191,10 @@ class MMT(_BaseMMT):
         cloud_account: Optional[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,
+        artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> "MMT":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -212,20 +214,19 @@ class MMT(_BaseMMT):
                 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.
-                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.
                 Only applicable when submitting docker jobs.
+            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.
         """
         self._job = self._internal_mmt._submit(
             num_machines=num_machines,
@@ -238,9 +239,10 @@ class MMT(_BaseMMT):
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )
         return self
 

lightning_sdk/mmt/v1.py

@@ -52,9 +52,10 @@ class _MMTV1(_BaseMMT):
         cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
     ) -> "_MMTV1":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -88,6 +89,8 @@ class _MMTV1(_BaseMMT):
                 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: The mappings from data connection inside your container (not supported)
+
         """
         if studio is None:
             raise ValueError("Studio is required for submitting jobs")

lightning_sdk/mmt/v2.py

@@ -51,9 +51,10 @@ class _MMTV2(_BaseMMT):
         cloud_account: Optional[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,
+        artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> "_MMTV2":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -73,20 +74,19 @@ class _MMTV2(_BaseMMT):
                 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.
-                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.
                 Only applicable when submitting docker jobs.
+            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.
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -117,9 +117,10 @@ class _MMTV2(_BaseMMT):
             env=env,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )
         self._job = submitted
         self._name = submitted.name

lightning_sdk/models.py

@@ -121,10 +121,11 @@ def download_model(
             progress_bar=progress_bar,
         )
     except ApiException as e:
-        # if we get an error, check if the teamspace actually exists (and print the list)
-        # TODO: ideally this would match a specific error about teamspace not being found
-        _ = _get_teamspace(name=teamspace_name, organization=teamspace_owner_name)
-        raise e
+        if e.status == 404:
+            raise RuntimeError(
+                f"Model '{name}' not found. Either the model doesn't exist or you don't have access to it."
+            ) from None
+        raise RuntimeError(f"Error downloading model. Status code: {e.status}.") from None
 
 
 def upload_model(

lightning_sdk/studio.py

@@ -1,6 +1,6 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
 from lightning_sdk.api.studio_api import StudioApi
 from lightning_sdk.api.utils import _machine_to_compute_name
@@ -14,6 +14,8 @@ from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, _setup_logger
 
 if TYPE_CHECKING:
+    from lightning_sdk.job import Job
+    from lightning_sdk.mmt import MMT
     from lightning_sdk.plugin import Plugin
 
 _logger = _setup_logger(__name__)
@@ -281,6 +283,71 @@ class Studio:
             cloud_account=self._studio.cluster_id,
         )
 
+    def run_job(
+        self,
+        name: str,
+        machine: Union["Machine", str],
+        command: str,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+    ) -> "Job":
+        """Run async workloads using the compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            command: The command to run inside your job.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+        """
+        from lightning_sdk.job import Job
+
+        return Job.run(
+            name=name,
+            machine=machine,
+            command=command,
+            studio=self,
+            image=None,
+            teamspace=self.teamspace,
+            cloud_account=self.cloud_account,
+            env=env,
+            interruptible=interruptible,
+        )
+
+    def run_mmt(
+        self,
+        name: str,
+        num_machines: int,
+        machine: Union["Machine", str],
+        command: str,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+    ) -> "MMT":
+        """Run async workloads using the compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            num_machines: The number of machines to run on.
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            command: The command to run inside your job.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+        """
+        from lightning_sdk.mmt import MMT
+
+        return MMT.run(
+            name=name,
+            num_machines=num_machines,
+            machine=machine,
+            command=command,
+            studio=self,
+            image=None,
+            teamspace=self.teamspace,
+            cloud_account=self.cloud_account,
+            env=env,
+            interruptible=interruptible,
+        )
+
     @property
     def auto_sleep(self) -> bool:
         """Returns if a Studio has auto-sleep enabled."""

lightning_sdk/utils/resolve.py

@@ -187,3 +187,10 @@ def _parse_model_and_version(name: str) -> Tuple[str, str]:
         "Model version is expected to be in the format `entity/modelname:version` separated by a"
         f" single colon, but got: {name}"
     )
+
+
+def in_studio() -> bool:
+    """Returns true if inside a studio, else false."""
+    has_cloudspace_id = bool(os.getenv("LIGHTNING_CLOUD_SPACE_ID", None))
+    is_interactive = os.getenv("LIGHTNING_INTERACTIVE", "false") == "true"
+    return has_cloudspace_id and is_interactive

{lightning_sdk-0.1.50.dist-info → lightning_sdk-0.1.53.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.1.50
+Version: 0.1.53
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License
@@ -45,9 +45,9 @@ Requires-Dist: tqdm
 Requires-Dist: fire
 Requires-Dist: simple-term-menu
 Requires-Dist: lightning-utilities
+Requires-Dist: docker
 Provides-Extra: serve
 Requires-Dist: litserve>=0.2.5; extra == "serve"
-Requires-Dist: docker; extra == "serve"
 
 # Lightning SDK