lightning-sdk 0.1.46__py3-none-any.whl → 0.1.48__py3-none-any.whl

lightning_sdk/mmt/base.py

@@ -1,7 +1,8 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Dict, Optional, Protocol, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Tuple, Union
 
 if TYPE_CHECKING:
+    from lightning_sdk.job.base import MachineDict
     from lightning_sdk.machine import Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.status import Status
@@ -41,6 +42,10 @@ class MMTMachine(Protocol):
         """The logs of the given machine."""
         ...
 
+    def dict(self) -> "MachineDict":
+        """Dict representation of the given machine."""
+        ...
+
 
 class _BaseMMT(_BaseJob):
     """Base interface to all job types."""
@@ -64,6 +69,7 @@ class _BaseMMT(_BaseJob):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseMMT":
         """Run async workloads using a docker image across multiple machines.
@@ -98,6 +104,10 @@ class _BaseMMT(_BaseJob):
                 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.
         """
         from lightning_sdk.studio import Studio
 
@@ -148,6 +158,9 @@ class _BaseMMT(_BaseJob):
                     "Other jobs will automatically persist artifacts to the teamspace distributed filesystem."
                 )
 
+            if entrypoint != "sh -c":
+                raise ValueError("Specifying the entrypoint has no effect for jobs with Studio envs.")
+
         else:
             if studio is not None:
                 raise RuntimeError(
@@ -178,6 +191,7 @@ class _BaseMMT(_BaseJob):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )
         return inst
 
@@ -196,6 +210,7 @@ class _BaseMMT(_BaseJob):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> None:
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -225,6 +240,9 @@ class _BaseMMT(_BaseJob):
                 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.
         """
 
     @property
@@ -283,6 +301,26 @@ class _BaseMMT(_BaseJob):
         """Logs of the rank 0 machine."""
         return self.machines[0].logs
 
+    def dict(
+        self
+    ) -> Dict[str, Union[str, "Studio", "Status", "Machine", None, List[Dict[str, Union[str, "Status", "Machine"]]]]]:
+        """Dict representation of this job."""
+        studio = self.studio
+
+        return {
+            "name": self.name,
+            "teamspace": f"{self.teamspace.owner.name}/{self.teamspace.name}",
+            "studio": studio.name if studio else None,
+            "image": self.image,
+            "command": self.command,
+            "status": self.status,
+            "machine": self.machine,
+            "machines": [
+                {"name": d["name"], "status": d["status"], "machine": d["machine"]}
+                for d in (x.dict() for x in self.machines)
+            ],
+        }
+
     @abstractmethod
     def _update_internal_job(self) -> None:
         pass

lightning_sdk/mmt/mmt.py

@@ -1,4 +1,3 @@
-from contextlib import suppress
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
@@ -58,15 +57,40 @@ class MMT(_BaseMMT):
             user: the name of the user owning the :param`teamspace`
                 in case it is owned directly by a user instead of an org.
         """
-        internal_mmt_cls = _MMTV2 if _has_mmt_v2() and not self._force_v1 else _MMTV1
-
-        self._internal_mmt = internal_mmt_cls(
-            name=name,
-            teamspace=teamspace,
-            org=org,
-            user=user,
-            _fetch_job=_fetch_job,
-        )
+        from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+
+        if _has_mmt_v2() and not self._force_v1:
+            # try with v2 and fall back to v1
+            try:
+                mmt = _MMTV2(
+                    name=name,
+                    teamspace=teamspace,
+                    org=org,
+                    user=user,
+                    _fetch_job=_fetch_job,
+                )
+            except ApiException as e:
+                try:
+                    mmt = _MMTV1(
+                        name=name,
+                        teamspace=teamspace,
+                        org=org,
+                        user=user,
+                        _fetch_job=_fetch_job,
+                    )
+                except ApiException:
+                    raise e from e
+
+        else:
+            mmt = _MMTV1(
+                name=name,
+                teamspace=teamspace,
+                org=org,
+                user=user,
+                _fetch_job=_fetch_job,
+            )
+
+        self._internal_mmt = mmt
 
     @classmethod
     def run(
@@ -87,6 +111,7 @@ class MMT(_BaseMMT):
         cloud_account_auth: bool = False,
         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.
@@ -121,6 +146,9 @@ class MMT(_BaseMMT):
                 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.
         """
         ret_val = super().run(
             name=name,
@@ -144,10 +172,7 @@ class MMT(_BaseMMT):
         # required for typing with "MMT"
         assert isinstance(ret_val, cls)
 
-        msg = "Multi-Machine Job was successfully launched."
-
-        with suppress(NotImplementedError):
-            msg += f" View it at {ret_val.link}"
+        msg = f"Multi-Machine Job was successfully launched. View it at {ret_val.link}"
 
         _logger.info(msg)
         return ret_val
@@ -166,6 +191,7 @@ class MMT(_BaseMMT):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> "MMT":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -195,6 +221,10 @@ class MMT(_BaseMMT):
                 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.
         """
         self._job = self._internal_mmt._submit(
             num_machines=num_machines,
@@ -270,6 +300,21 @@ class MMT(_BaseMMT):
     def link(self) -> str:
         return self._internal_mmt.link
 
+    @property
+    def studio(self) -> Optional["Studio"]:
+        """The studio used to submit the MMT."""
+        return self._internal_mmt.studio
+
+    @property
+    def image(self) -> Optional[str]:
+        """The image used to submit the MMT."""
+        return self._internal_mmt.image
+
+    @property
+    def command(self) -> str:
+        """The command the MMT is running."""
+        return self._internal_mmt.command
+
     def __getattr__(self, key: str) -> Any:
         """Forward the attribute lookup to the internal job implementation."""
         try:

lightning_sdk/mmt/v1.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 
 from lightning_sdk.api.mmt_api import MMTApiV1
+from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.job.v1 import _internal_status_to_external_status
 from lightning_sdk.job.work import Work
 
@@ -53,6 +54,7 @@ class _MMTV1(_BaseMMT):
         cloud_account_auth: bool = False,
         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.
 
@@ -82,10 +84,14 @@ class _MMTV1(_BaseMMT):
                 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.
         """
         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:
+        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:
@@ -174,7 +180,29 @@ class _MMTV1(_BaseMMT):
 
     @property
     def link(self) -> str:
-        return f"https://lightning.ai/{self.teamspace.owner.name}/{self.teamspace.name}/studios/{self._job_api.get_studio_name(self._guaranteed_job)}/app?app_id=mmt&app_tab=Runs&job_name={self.name}"
+        return (
+            f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/studios/{self._job_api.get_studio_name(self._guaranteed_job)}/"
+            f"app?app_id=mmt&app_tab=Runs&job_name={self.name}"
+        )
+
+    @property
+    def image(self) -> Optional[str]:
+        """The image used to submit the job."""
+        # mmtv1 don't support images, so return None here
+        return None
+
+    @property
+    def studio(self) -> Optional["Studio"]:
+        """The studio used to submit the job."""
+        from lightning_sdk.studio import Studio
+
+        studio_name = self._job_api.get_studio_name(self._guaranteed_job)
+        return Studio(studio_name, teamspace=self.teamspace)
+
+    @property
+    def command(self) -> str:
+        """The command the job is running."""
+        return self._job_api.get_command(self._guaranteed_job)
 
     # the following and functions are solely to make the Work class function
     @property

lightning_sdk/mmt/v2.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 from lightning_sdk.api.mmt_api import MMTApiV2
+from lightning_sdk.api.utils import _get_cloud_url
 
 if TYPE_CHECKING:
     from lightning_sdk.job.job import Job
@@ -52,6 +53,7 @@ class _MMTV2(_BaseMMT):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> "_MMTV2":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -81,6 +83,10 @@ class _MMTV2(_BaseMMT):
                 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.
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -113,6 +119,7 @@ class _MMTV2(_BaseMMT):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )
         self._job = submitted
         self._name = submitted.name
@@ -189,5 +196,27 @@ class _MMTV2(_BaseMMT):
 
     @property
     def link(self) -> str:
-        # TODO: Since we don't have a UI for this yet, we can't have a link
-        raise NotImplementedError
+        # TODO: MMT env with studio -> go to studio plugin
+        return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=mmt"
+
+    @property
+    def image(self) -> Optional[str]:
+        """The image used to submit the job."""
+        return self._job_api.get_image_name(self._guaranteed_job)
+
+    @property
+    def studio(self) -> Optional["Studio"]:
+        """The studio used to submit the job."""
+        from lightning_sdk.studio import Studio
+
+        studio_name = self._job_api.get_studio_name(self._guaranteed_job)
+
+        # if job was submitted with image, studio will be None
+        if not studio_name:
+            return None
+        return Studio(studio_name, teamspace=self.teamspace)
+
+    @property
+    def command(self) -> str:
+        """The command the job is running."""
+        return self._job_api.get_command(self._guaranteed_job)

lightning_sdk/status.py

@@ -4,10 +4,14 @@ from enum import Enum
 class Status(Enum):
     """Enum holding all possible studio status types."""
 
-    NotCreated = 1
-    Pending = 2
-    Running = 3
-    Stopping = 4
-    Stopped = 5
-    Completed = 6
-    Failed = 7
+    NotCreated = "NotCreated"
+    Pending = "Pending"
+    Running = "Running"
+    Stopping = "Stopping"
+    Stopped = "Stopped"
+    Completed = "Completed"
+    Failed = "Failed"
+
+    def __str__(self) -> str:
+        """String representation of the enum."""
+        return self.value

lightning_sdk/teamspace.py

@@ -1,6 +1,6 @@
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 from lightning_sdk.agents import Agent
 from lightning_sdk.api import TeamspaceApi
@@ -17,6 +17,8 @@ from lightning_sdk.utils.resolve import (
 )
 
 if TYPE_CHECKING:
+    from lightning_sdk.job import Job
+    from lightning_sdk.mmt import MMT
     from lightning_sdk.studio import Studio
 
 
@@ -126,6 +128,54 @@ class Teamspace:
         )
         return self.cloud_accounts
 
+    @property
+    def jobs(self) -> Tuple["Job", ...]:
+        from lightning_sdk.job import Job
+        from lightning_sdk.plugin import forced_v1
+
+        jobsv1, jobsv2 = self._teamspace_api.list_jobs(teamspace_id=self.id)
+
+        jobs = []
+
+        for j1 in jobsv1:
+            with forced_v1(Job):
+                # _fetch_job = False to prevent refetching on init since we already got it
+                job = Job(name=j1.name, teamspace=self, _fetch_job=False)
+            job._internal_job._job = j1
+            jobs.append(job)
+
+        for j2 in jobsv2:
+            # _fetch_job = False to prevent refetching on init since we already got it
+            job = Job(name=j2.name, teamspace=self, _fetch_job=False)
+            job._internal_job._job = j2
+            jobs.append(job)
+
+        return tuple(jobs)
+
+    @property
+    def multi_machine_jobs(self) -> Tuple["MMT", ...]:
+        from lightning_sdk.mmt import MMT
+        from lightning_sdk.plugin import forced_v1
+
+        mmtsv1, mmtsv2 = self._teamspace_api.list_mmts(teamspace_id=self.id)
+
+        mmts = []
+
+        for m1 in mmtsv1:
+            with forced_v1(MMT):
+                # _fetch_job = False to prevent refetching on init since we already got it
+                mmt = MMT(name=m1.name, teamspace=self, _fetch_job=False)
+            mmt._internal_mmt._job = m1
+            mmts.append(mmt)
+
+        for m2 in mmtsv2:
+            # _fetch_job = False to prevent refetching on init since we already got it
+            mmt = MMT(name=m2.name, teamspace=self, _fetch_job=False)
+            mmt._internal_mmt._job = m2
+            mmts.append(mmt)
+
+        return tuple(mmts)
+
     def __eq__(self, other: "Teamspace") -> bool:
         """Checks whether the provided other object is equal to this one."""
         return (

{lightning_sdk-0.1.46.dist-info → lightning_sdk-0.1.48.dist-info}/METADATA

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