lightning-sdk 0.1.40__py3-none-any.whl → 0.1.42__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_validate_deployment_image_response.py

@@ -0,0 +1,97 @@
+# coding: utf-8
+
+"""
+    external/v1/auth_service.proto
+
+    No description provided (generated by Swagger Codegen https://github.com/swagger-api/swagger-codegen)  # noqa: E501
+
+    OpenAPI spec version: version not set
+    
+    Generated by: https://github.com/swagger-api/swagger-codegen.git
+
+    NOTE
+    ----
+    standard swagger-codegen-cli for this python client has been modified
+    by custom templates. The purpose of these templates is to include
+    typing information in the API and Model code. Please refer to the
+    main grid repository for more info
+"""
+
+import pprint
+import re  # noqa: F401
+
+from typing import TYPE_CHECKING
+
+import six
+
+if TYPE_CHECKING:
+    from datetime import datetime
+    from lightning_sdk.lightning_cloud.openapi.models import *
+
+class V1ValidateDeploymentImageResponse(object):
+    """NOTE: This class is auto generated by the swagger code generator program.
+
+    Do not edit the class manually.
+    """
+    """
+    Attributes:
+      swagger_types (dict): The key is attribute name
+                            and the value is attribute type.
+      attribute_map (dict): The key is attribute name
+                            and the value is json key in definition.
+    """
+    swagger_types = {
+    }
+
+    attribute_map = {
+    }
+
+    def __init__(self):  # noqa: E501
+        """V1ValidateDeploymentImageResponse - a model defined in Swagger"""  # noqa: E501
+        self.discriminator = None
+
+    def to_dict(self) -> dict:
+        """Returns the model properties as a dict"""
+        result = {}
+
+        for attr, _ in six.iteritems(self.swagger_types):
+            value = getattr(self, attr)
+            if isinstance(value, list):
+                result[attr] = list(map(
+                    lambda x: x.to_dict() if hasattr(x, "to_dict") else x,
+                    value
+                ))
+            elif hasattr(value, "to_dict"):
+                result[attr] = value.to_dict()
+            elif isinstance(value, dict):
+                result[attr] = dict(map(
+                    lambda item: (item[0], item[1].to_dict())
+                    if hasattr(item[1], "to_dict") else item,
+                    value.items()
+                ))
+            else:
+                result[attr] = value
+        if issubclass(V1ValidateDeploymentImageResponse, dict):
+            for key, value in self.items():
+                result[key] = value
+
+        return result
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model"""
+        return pprint.pformat(self.to_dict())
+
+    def __repr__(self) -> str:
+        """For `print` and `pprint`"""
+        return self.to_str()
+
+    def __eq__(self, other: 'V1ValidateDeploymentImageResponse') -> bool:
+        """Returns true if both objects are equal"""
+        if not isinstance(other, V1ValidateDeploymentImageResponse):
+            return False
+
+        return self.__dict__ == other.__dict__
+
+    def __ne__(self, other: 'V1ValidateDeploymentImageResponse') -> bool:
+        """Returns true if both objects are not equal"""
+        return not self == other

lightning_sdk/lightning_cloud/rest_client.py

@@ -31,6 +31,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     AssistantsServiceApi,
     StorageServiceApi,
     DeploymentTemplatesServiceApi,
+    ModelsStoreApi,
 )
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.source_code.logs_socket_api import LightningLogsSocketAPI
@@ -89,6 +90,7 @@ class GridRestClient(
         AssistantsServiceApi,
         StorageServiceApi,
         DeploymentTemplatesServiceApi,
+        ModelsStoreApi,
 ):
 
     def __init__(self, api_client: Optional[ApiClient] = None):

lightning_sdk/mmt/__init__.py

@@ -0,0 +1,3 @@
+from lightning_sdk.mmt.mmt import MMT
+
+__all__ = ["MMT"]

lightning_sdk/{_mmt → mmt}/base.py

@@ -11,6 +11,7 @@ if TYPE_CHECKING:
 
 from lightning_sdk.job.base import _BaseJob
 from lightning_sdk.job.job import Job
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
 
 
 class _BaseMMT(_BaseJob):
@@ -26,16 +27,19 @@ class _BaseMMT(_BaseJob):
         teamspace: Union[str, "Teamspace", None] = None,
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseMMT":
         from lightning_sdk.studio import Studio
 
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+
         if num_machines <= 1:
             raise ValueError("Multi-Machine training cannot be run with less than 2 Machines")
 
@@ -44,7 +48,9 @@ class _BaseMMT(_BaseJob):
 
         if image is None:
             if not isinstance(studio, Studio):
-                studio = Studio(name=studio, teamspace=teamspace, org=org, user=user, cluster=cluster, create_ok=False)
+                studio = Studio(
+                    name=studio, teamspace=teamspace, org=org, user=user, cloud_account=cloud_account, create_ok=False
+                )
 
             # studio is a Studio instance at this point
             if teamspace is None:
@@ -58,20 +64,20 @@ class _BaseMMT(_BaseJob):
                         "Can only run jobs with Studio envs in the teamspace of that Studio."
                     )
 
-            if cluster is None:
-                cluster = studio.cluster
+            if cloud_account is None:
+                cloud_account = studio.cloud_account
 
-            if cluster != studio.cluster:
+            if cloud_account != studio.cloud_account:
                 raise ValueError(
-                    "Studio cluster does not match provided cluster. "
-                    "Can only run jobs with Studio envs in the same cluster."
+                    "Studio cloud_account does not match provided cloud_account. "
+                    "Can only run jobs with Studio envs in the same cloud_account."
                 )
 
             if image_credentials is not None:
                 raise ValueError("image_credentials is only supported when using a custom image")
 
-            if cluster_auth:
-                raise ValueError("cluster_auth is only supported when using a custom image")
+            if cloud_account_auth:
+                raise ValueError("cloud_account_auth is only supported when using a custom image")
 
             if artifacts_local is not None or artifacts_remote is not None:
                 raise ValueError(
@@ -99,14 +105,14 @@ class _BaseMMT(_BaseJob):
         inst._submit(
             num_machines=num_machines,
             machine=machine,
-            cluster=cluster,
+            cloud_account=cloud_account,
             command=command,
             studio=studio,
             image=image,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
@@ -122,9 +128,9 @@ class _BaseMMT(_BaseJob):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> None:

lightning_sdk/{_mmt → mmt}/mmt.py

@@ -1,12 +1,11 @@
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Tuple, Union
 
-from lightning_sdk._mmt.base import _BaseMMT
-from lightning_sdk._mmt.v1 import _MMTV1
-from lightning_sdk._mmt.v2 import _MMTV2
 from lightning_sdk.job.job import _has_jobs_v2
+from lightning_sdk.mmt.base import _BaseMMT
+from lightning_sdk.mmt.v1 import _MMTV1
+from lightning_sdk.mmt.v2 import _MMTV2
 
 if TYPE_CHECKING:
-    from lightning_sdk.job import Job
     from lightning_sdk.machine import Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.status import Status
@@ -15,7 +14,31 @@ if TYPE_CHECKING:
     from lightning_sdk.user import User
 
 
+class MMTMachine(Protocol):
+    """A single machine in multi-machine training."""
+
+    @property
+    def name(self) -> str:
+        ...
+
+    @property
+    def machine(self) -> "Machine":
+        ...
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        ...
+
+    @property
+    def status(self) -> "Status":
+        ...
+
+
 class MMT(_BaseMMT):
+    _force_v1: (
+        bool
+    ) = False  # required for studio plugin still working correctly as v2 currently does not support the studio env
+
     def __init__(
         self,
         name: str,
@@ -25,7 +48,7 @@ class MMT(_BaseMMT):
         *,
         _fetch_job: bool = True,
     ) -> None:
-        internal_mmt_cls = _MMTV2 if _has_jobs_v2() else _MMTV1
+        internal_mmt_cls = _MMTV2 if _has_jobs_v2() and not self._force_v1 else _MMTV1
 
         self._internal_mmt = internal_mmt_cls(
             name=name,
@@ -47,13 +70,14 @@ class MMT(_BaseMMT):
         teamspace: Union[str, "Teamspace", None] = None,
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "MMT":
         ret_val = super().run(
             name=name,
@@ -65,13 +89,14 @@ class MMT(_BaseMMT):
             teamspace=teamspace,
             org=org,
             user=user,
-            cluster=cluster,
+            cloud_account=cloud_account,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            cluster=cluster,  # deprecated in favor of cloud_account
         )
         # required for typing with "Job"
         assert isinstance(ret_val, cls)
@@ -86,23 +111,23 @@ class MMT(_BaseMMT):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> "MMT":
         self._job = self._internal_mmt._submit(
             num_machines=num_machines,
             machine=machine,
-            cluster=cluster,
+            cloud_account=cloud_account,
             command=command,
             studio=studio,
             image=image,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
@@ -119,7 +144,7 @@ class MMT(_BaseMMT):
         return self._internal_mmt.status
 
     @property
-    def machines(self) -> Tuple["Job", ...]:
+    def machines(self) -> Tuple[MMTMachine, ...]:
         return self._internal_mmt.machines
 
     @property
@@ -150,8 +175,8 @@ class MMT(_BaseMMT):
         return self._internal_mmt._teamspace
 
     @property
-    def cluster(self) -> Optional[str]:
-        return self._internal_mmt.cluster
+    def cloud_account(self) -> Optional[str]:
+        return self._internal_mmt.cloud_account
 
     def __getattr__(self, key: str) -> Any:
         """Forward the attribute lookup to the internal job implementation."""
@@ -159,3 +184,7 @@ class MMT(_BaseMMT):
             return getattr(super(), key)
         except AttributeError:
             return getattr(self._internal_mmt, key)
+
+    @property
+    def _guaranteed_job(self) -> Any:
+        return self._internal_mmt._guaranteed_job

lightning_sdk/mmt/v1.py

@@ -0,0 +1,129 @@
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+
+from lightning_sdk.api.mmt_api import MMTApiV1
+from lightning_sdk.job.v1 import _internal_status_to_external_status
+from lightning_sdk.job.work import Work
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import Machine
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.status import Status
+    from lightning_sdk.studio import Studio
+    from lightning_sdk.teamspace import Teamspace
+    from lightning_sdk.user import User
+
+from lightning_sdk.mmt.base import _BaseMMT
+
+
+class _MMTV1(_BaseMMT):
+    def __init__(
+        self,
+        name: str,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        *,
+        _fetch_job: bool = True,
+    ) -> None:
+        self._job_api = MMTApiV1()
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
+
+    def _submit(
+        self,
+        num_machines: int,
+        machine: "Machine",
+        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,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> "_MMTV1":
+        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:
+            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
+
+    def _update_internal_job(self) -> None:
+        try:
+            self._job = self._job_api.get_job(self._name, self.teamspace.id)
+        except ValueError as e:
+            raise ValueError(f"Job {self._name} does not exist in Teamspace {self.teamspace.name}") from e
+
+    @property
+    def machines(self) -> Tuple["Work", ...]:
+        works = self._job_api.list_works(self._guaranteed_job.id, self.teamspace.id)
+
+        return tuple(Work(w.id, self, self.teamspace) for w in works)
+
+    def stop(self) -> None:
+        self._job_api.stop_job(self._guaranteed_job.id, self.teamspace.id)
+
+    def delete(self) -> None:
+        self._job_api.delete_job(self._guaranteed_job.id, self.teamspace.id)
+
+    @property
+    def status(self) -> "Status":
+        try:
+            status = self._job_api.get_job_status(self._job.id, self.teamspace.id)
+            return _internal_status_to_external_status(status)
+        except Exception:
+            raise RuntimeError(
+                f"MMT {self._name} does not exist in Teamspace {self.teamspace.name}. Did you delete it?"
+            ) from None
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        return f"/teamspace/jobs/{self.name}"
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        return f"/teamspace/jobs/{self.name}/snapshot"
+
+    @property
+    def machine(self) -> "Machine":
+        return self.machines[0].machine
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def teamspace(self) -> "Teamspace":
+        return self._teamspace
+
+    # the following and functions are solely to make the Work class function
+    @property
+    def _id(self) -> str:
+        return self._guaranteed_job.id
+
+    def _name_filter(self, name: str) -> str:
+        return name.replace("root.", "")

lightning_sdk/{_mmt → mmt}/v2.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
-from lightning_sdk.api.mmt_api import MMTApi
+from lightning_sdk.api.mmt_api import MMTApiV2
 
 if TYPE_CHECKING:
     from lightning_sdk.job.job import Job
@@ -11,7 +11,7 @@ if TYPE_CHECKING:
     from lightning_sdk.teamspace import Teamspace
     from lightning_sdk.user import User
 
-from lightning_sdk._mmt.base import _BaseMMT
+from lightning_sdk.mmt.base import _BaseMMT
 
 
 class _MMTV2(_BaseMMT):
@@ -24,7 +24,7 @@ class _MMTV2(_BaseMMT):
         *,
         _fetch_job: bool = True,
     ) -> None:
-        self._job_api = MMTApi()
+        self._job_api = MMTApiV2()
         super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
 
     def _submit(
@@ -36,9 +36,9 @@ class _MMTV2(_BaseMMT):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> "_MMTV2":
@@ -62,7 +62,7 @@ class _MMTV2(_BaseMMT):
             name=self.name,
             num_machines=num_machines,
             command=command,
-            cluster_id=cluster,
+            cloud_account=cloud_account,
             teamspace_id=self._teamspace.id,
             studio_id=studio_id,
             image=image,
@@ -70,7 +70,7 @@ class _MMTV2(_BaseMMT):
             interruptible=interruptible,
             env=env,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
@@ -80,7 +80,12 @@ class _MMTV2(_BaseMMT):
 
     @property
     def machines(self) -> Tuple["Job", ...]:
-        raise NotImplementedError
+        from lightning_sdk.job import Job
+
+        return tuple(
+            Job(name=j.name, teamspace=self.teamspace)
+            for j in self._job_api.list_mmt_subjobs(self._guaranteed_job.id, self.teamspace.id)
+        )
 
     def stop(self) -> None:
         self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
@@ -97,28 +102,18 @@ class _MMTV2(_BaseMMT):
         self._update_internal_job()
         return self._job
 
-    @property
-    def _guaranteed_job(self) -> Any:
-        """Guarantees that the job was fetched at some point before returning it.
-
-        Doesn't guarantee to have the lastest version of the job. Use _latest_job for that.
-        """
-        if getattr(self, "_job", None) is None:
-            self._update_internal_job()
-
-        return self._job
-
     @property
     def status(self) -> "Status":
-        # TODO: Should this rather be a list of states from the individual machines?
-        return self._job_api._job_state_to_external(self._latest_job.desired_state)
+        return self._job_api._job_state_to_external(self._latest_job.state)
 
     @property
     def artifact_path(self) -> Optional[str]:
+        # TODO: Since grouping for those is not done yet on the BE, we cannot yet have a unified link here
         raise NotImplementedError
 
     @property
     def snapshot_path(self) -> Optional[str]:
+        # TODO: Since grouping for those is not done yet on the BE, we cannot yet have a unified link here
         raise NotImplementedError
 
     @property