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

lightning_sdk/mmt/mmt.py

@@ -1,15 +1,13 @@
-from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
-from lightning_sdk.api.user_api import UserApi
-from lightning_sdk.job.job import _has_jobs_v2
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.mmt.base import MMTMachine, _BaseMMT
 from lightning_sdk.mmt.v1 import _MMTV1
 from lightning_sdk.mmt.v2 import _MMTV2
 from lightning_sdk.utils.resolve import _setup_logger
 
 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.status import Status
     from lightning_sdk.studio import Studio
@@ -19,25 +17,12 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
-@lru_cache(maxsize=None)
-def _has_mmt_v2() -> bool:
-    # users need both mmtv2 and jobsv2 flags in order for mmtv2 to work correctly
-    if not _has_jobs_v2():
-        return False
-
-    api = UserApi()
-    try:
-        return api._get_feature_flags().mmt_v2
-    except Exception:
-        return False
-
-
 class MMT(_BaseMMT):
     """Class to submit and manage multi-machine jobs on the Lightning AI Platform."""
 
-    _force_v1: (
-        bool
-    ) = False  # required for studio plugin still working correctly as v2 currently does not support the studio env
+    _force_v1: bool = (
+        False  # required for studio plugin still working correctly as v2 currently does not support the studio env
+    )
 
     def __init__(
         self,
@@ -59,7 +44,7 @@ class MMT(_BaseMMT):
         """
         from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 
-        if _has_mmt_v2() and not self._force_v1:
+        if not self._force_v1:
             # try with v2 and fall back to v1
             try:
                 mmt = _MMTV2(
@@ -91,13 +76,14 @@ class MMT(_BaseMMT):
             )
 
         self._internal_mmt = mmt
+        self._cloud_account_api = CloudAccountApi()
 
     @classmethod
     def run(
         cls,
         name: str,
         num_machines: int,
-        machine: "Machine",
+        machine: Union["Machine", str],
         command: Optional[str] = None,
         studio: Union["Studio", str, None] = None,
         image: Union[str, None] = None,
@@ -105,21 +91,25 @@ class MMT(_BaseMMT):
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = 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
+        reuse_snapshot: bool = True,
     ) -> "MMT":
         """Run async workloads using a docker image across multiple machines.
 
         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.
@@ -129,26 +119,31 @@ class MMT(_BaseMMT):
             user: The user owning the teamspace (if any). Defaults to the current user.
             cloud_account: The cloud account to run the job on.
                 Defaults to the studio cloud account if running with studio compute env.
-                If not provided will fall back to the teamspaces default cloud account.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             env: Environment variables to set inside the job.
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 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.
+            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.
         """
         ret_val = super().run(
             name=name,
@@ -161,13 +156,18 @@ class MMT(_BaseMMT):
             org=org,
             user=user,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             env=env,
             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,
             cluster=cluster,  # deprecated in favor of cloud_account
+            max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         # required for typing with "MMT"
         assert isinstance(ret_val, cls)
@@ -180,18 +180,22 @@ class MMT(_BaseMMT):
     def _submit(
         self,
         num_machines: int,
-        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,
     ) -> "MMT":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -206,30 +210,40 @@ class MMT(_BaseMMT):
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             cloud_account: The cloud account to run the job on.
                 Defaults to the studio cloud account if running with studio compute env.
-                If not provided will fall back to the teamspaces default cloud account.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 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.
+            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.
         """
         self._job = self._internal_mmt._submit(
             num_machines=num_machines,
             machine=machine,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             command=command,
             studio=studio,
             image=image,
@@ -237,8 +251,12 @@ 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,
+            max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         return self
 
@@ -264,7 +282,7 @@ class MMT(_BaseMMT):
         return self._internal_mmt.machines
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
         return self._internal_mmt.machine
 

lightning_sdk/mmt/v1.py

@@ -7,7 +7,7 @@ from lightning_sdk.job.work import Work
 from lightning_sdk.status import Status
 
 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
@@ -43,18 +43,22 @@ class _MMTV1(_BaseMMT):
     def _submit(
         self,
         num_machines: int,
-        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,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
+        reuse_snapshot: bool = True,
     ) -> "_MMTV1":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -88,35 +92,15 @@ 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)
+            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.
         """
-        if studio is None:
-            raise ValueError("Studio is required for submitting jobs")
-        if image is not None or image_credentials is not None or cloud_account_auth or entrypoint != "sh -c":
-            raise ValueError("Image is not supported for submitting jobs")
-
-        if artifacts_local is not None or artifacts_remote is not None:
-            raise ValueError("Specifying how to persist artifacts is not yet supported with jobs")
-
-        if env is not None:
-            raise ValueError("Environment variables are not supported for submitting jobs")
-        if command is None:
-            raise ValueError("Command is required for submitting multi-machine jobs")
-
-        _submitted = self._job_api.submit_job(
-            name=self._name,
-            num_machines=num_machines,
-            command=command,
-            studio_id=studio._studio.id,
-            teamspace_id=self._teamspace.id,
-            cloud_account=cloud_account or "",
-            machine=machine,
-            interruptible=interruptible,
-            strategy="parallel",
-        )
-
-        self._name = _submitted.name
-        self._job = _submitted
-        return self
+        raise NotImplementedError("Cannot submit new mmts with MMTV1!")
 
     def _update_internal_job(self) -> None:
         try:
@@ -166,7 +150,7 @@ class _MMTV1(_BaseMMT):
         return f"/teamspace/jobs/{self.name}/snapshot"
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
         return self.machines[0].machine
 

lightning_sdk/mmt/v2.py

@@ -3,10 +3,11 @@ from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 from lightning_sdk.api.mmt_api import MMTApiV2
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.status import Status
+from lightning_sdk.utils.resolve import _get_org_id
 
 if TYPE_CHECKING:
     from lightning_sdk.job.job import Job
-    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
@@ -42,18 +43,22 @@ class _MMTV2(_BaseMMT):
     def _submit(
         self,
         num_machines: int,
-        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,
     ) -> "_MMTV2":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -68,25 +73,34 @@ class _MMTV2(_BaseMMT):
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             cloud_account: The cloud account to run the job on.
                 Defaults to the studio cloud account if running with studio compute env.
-                If not provided will fall back to the teamspaces default cloud account.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 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.
+            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
@@ -104,6 +118,14 @@ class _MMTV2(_BaseMMT):
             studio_id = None
             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,
             num_machines=num_machines,
@@ -117,9 +139,12 @@ 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,
+            max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         self._job = submitted
         self._name = submitted.name
@@ -169,9 +194,13 @@ class _MMTV2(_BaseMMT):
         raise NotImplementedError
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
-        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),
+        )
 
     def _update_internal_job(self) -> None:
         if getattr(self, "_job", None) is None:

lightning_sdk/models.py

@@ -4,10 +4,11 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 from lightning_sdk.api import OrgApi, TeamspaceApi, UserApi
+from lightning_sdk.lightning_cloud.openapi import V1ModelVersionArchive
 from lightning_sdk.lightning_cloud.openapi.models import V1Membership, V1OwnerType
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _get_authed_user
+from lightning_sdk.utils.resolve import _get_authed_user, _resolve_teamspace
 
 if TYPE_CHECKING:
     from lightning_sdk.teamspace import Teamspace
@@ -72,24 +73,37 @@ def _get_teamspace(name: str, organization: str) -> "Teamspace":
     raise RuntimeError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: {os.linesep}\t{options}")
 
 
-def _parse_model_name_and_version(name: str) -> Tuple[str, str, str, str]:
+def _extend_model_name_with_teamspace(name: str) -> str:
+    """Extend the model name with the teamspace if it can be determined from env. variables."""
+    if "/" in name:
+        return name
+    # do some magic if you run studio
+    teamspace = _resolve_teamspace(None, None, None)
+    if not teamspace:
+        raise ValueError(
+            f"Model name must be in the format `organization/teamspace/model_name` but you provided '{name}'."
+        )
+    return f"{teamspace.owner.name}/{teamspace.name}/{name}"
+
+
+def _parse_org_teamspace_model_version(name: str) -> Tuple[str, str, str, Optional[str]]:
     """Parse the name argument into its components."""
     try:
         org_name, teamspace_name, model_name = name.split("/")
-        parts = model_name.split(":")
-        if len(parts) == 1:
-            return org_name, teamspace_name, parts[0], "latest"
-        if len(parts) == 2:
-            return org_name, teamspace_name, parts[0], parts[1]
-        # The rest of the validation for name and version happens in the backend
-        raise ValueError(
-            "Model version is expected to be in the format `entity/modelname:version` separated by a"
-            f" single colon, but got: {name}"
-        )
     except ValueError as err:
         raise ValueError(
-            f"Model name must be in the format 'organization/teamspace/model' but you provided '{name}'."
+            f"Model name must be in the format `organization/teamspace/model_name` but you provided '{name}'."
         ) from err
+    parts = model_name.split(":")
+    if len(parts) == 1:
+        return org_name, teamspace_name, parts[0], None
+    if len(parts) == 2:
+        return org_name, teamspace_name, parts[0], parts[1]
+    # The rest of the validation for name and version happens in the backend
+    raise ValueError(
+        "Model version is expected to be in the format `organization/teamspace/model_name:version`"
+        f" separated by a single colon, but got: {name}"
+    )
 
 
 def download_model(
@@ -105,33 +119,33 @@ def download_model(
         download_dir: The directory where the Model should be downloaded.
         progress_bar: Whether to show a progress bar when downloading.
     """
-    teamspace_owner_name, teamspace_name, model_name, version = _parse_model_name_and_version(name)
-
-    download_dir = Path(download_dir)
-
+    name = _extend_model_name_with_teamspace(name)
+    teamspace_owner_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
     api = TeamspaceApi()
 
     try:
         return api.download_model_files(
             name=model_name,
             version=version,
-            download_dir=download_dir,
+            download_dir=Path(download_dir).expanduser().resolve(),
             teamspace_name=teamspace_name,
             teamspace_owner_name=teamspace_owner_name,
             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(
     name: str,
-    path: Union[Path, str] = ".",
+    path: Union[str, Path, List[Union[str, Path]]] = ".",
     cloud_account: Optional[str] = None,
     progress_bar: bool = True,
+    metadata: Optional[Dict[str, Any]] = None,
 ) -> UploadedModelInfo:
     """Upload a Model.
 
@@ -142,12 +156,48 @@ def upload_model(
         cloud_account: The name of the cloud account to store the Model in.
             If not provided, the default cloud account for the Teamspace will be used.
         progress_bar: Whether to show a progress bar for the upload.
+        metadata: Metadata to attach to the uploaded model.
+            If not provided, an empty dictionary will be used.
     """
-    org_name, teamspace_name, model_name, _ = _parse_model_name_and_version(name)
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
     teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
     return teamspace.upload_model(
         path=path,
         name=model_name,
+        version=version,
         cloud_account=cloud_account,
         progress_bar=progress_bar,
+        metadata=metadata,
     )
+
+
+def delete_model(
+    name: str,
+) -> None:
+    """Delete a model or a version of model.
+
+    Args:
+        name: The name of the model you want to delete or with specified version it deltes only that version.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME> for full model deletion
+            or <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>:<VERSION> for version deletion.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    teamspace.delete_model(name=f"{model_name}:{version}" if version else model_name)
+
+
+def list_model_versions(
+    name: str,
+) -> List[V1ModelVersionArchive]:
+    """List all versions of a model.
+
+    Args:
+        name: The name of the model you want to list versions for.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, _ = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    return teamspace.list_model_versions(name=model_name)

lightning_sdk/organization.py

@@ -40,6 +40,10 @@ class Organization(Owner):
         """The organization's ID."""
         return self._org.id
 
+    @property
+    def default_cloud_account(self) -> Optional[str]:
+        return self._org.preferred_cluster or None
+
     def __repr__(self) -> str:
         """Returns reader friendly representation."""
         return f"Organization(name={self.name})"

lightning_sdk/owner.py

@@ -2,12 +2,13 @@ from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, List
 
 from lightning_sdk.api import TeamspaceApi
+from lightning_sdk.utils.logging import TrackCallsABCMeta
 
 if TYPE_CHECKING:
     from lightning_sdk.teamspace import Teamspace
 
 
-class Owner(ABC):
+class Owner(ABC, metaclass=TrackCallsABCMeta):
     """Represents an owner of teamspaces and studios."""
 
     def __init__(self) -> None: