lightning-sdk 0.1.39__py3-none-any.whl → 0.1.41__py3-none-any.whl

lightning_sdk/__init__.py

@@ -27,5 +27,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.39"
+__version__ = "0.1.41"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/_mmt/__init__.py

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

lightning_sdk/_mmt/base.py

@@ -0,0 +1,180 @@
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+
+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.job.base import _BaseJob
+from lightning_sdk.job.job import Job
+
+
+class _BaseMMT(_BaseJob):
+    @classmethod
+    def run(
+        cls,
+        name: str,
+        machine: "Machine",
+        num_machines: int,
+        command: Optional[str] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cluster: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> "_BaseMMT":
+        from lightning_sdk.studio import Studio
+
+        if num_machines <= 1:
+            raise ValueError("Multi-Machine training cannot be run with less than 2 Machines")
+
+        if not name:
+            raise ValueError("A job needs to have a name!")
+
+        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 is a Studio instance at this point
+            if teamspace is None:
+                teamspace = studio.teamspace
+            else:
+                teamspace_name = teamspace if isinstance(teamspace, str) else teamspace.name
+
+                if studio.teamspace.name != teamspace_name:
+                    raise ValueError(
+                        "Studio teamspace does not match provided teamspace. "
+                        "Can only run jobs with Studio envs in the teamspace of that Studio."
+                    )
+
+            if cluster is None:
+                cluster = studio.cluster
+
+            if cluster != studio.cluster:
+                raise ValueError(
+                    "Studio cluster does not match provided cluster. "
+                    "Can only run jobs with Studio envs in the same cluster."
+                )
+
+            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 artifacts_local is not None or artifacts_remote is not None:
+                raise ValueError(
+                    "Specifying artifacts persistence is supported for docker images only. "
+                    "Other jobs will automatically persist artifacts to the teamspace distributed filesystem."
+                )
+
+        else:
+            if studio is not None:
+                raise RuntimeError(
+                    "image and studio are mutually exclusive as both define the environment to run the job in"
+                )
+
+            # 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")
+
+            if artifacts_remote and len(artifacts_remote.split(":")) != 3:
+                raise ValueError(
+                    "Artifact persistence requires exactly three arguments separated by colon of kind "
+                    f"<CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>, got {artifacts_local}"
+                )
+
+        inst = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False)
+        inst._submit(
+            num_machines=num_machines,
+            machine=machine,
+            cluster=cluster,
+            command=command,
+            studio=studio,
+            image=image,
+            env=env,
+            interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+        return inst
+
+    @abstractmethod
+    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,
+        cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> None:
+        """Submits a job and updates the internal _job attribute as well as the _name attribute."""
+
+    @property
+    @abstractmethod
+    def machines(self) -> Tuple["Job", ...]:
+        pass
+
+    @property
+    @abstractmethod
+    def machine(self) -> "Machine":
+        pass
+
+    @abstractmethod
+    def stop(self) -> None:
+        pass
+
+    @abstractmethod
+    def delete(self) -> None:
+        pass
+
+    @property
+    @abstractmethod
+    def status(self) -> "Status":
+        pass
+
+    @property
+    @abstractmethod
+    def artifact_path(self) -> Optional[str]:
+        pass
+
+    @property
+    @abstractmethod
+    def snapshot_path(self) -> Optional[str]:
+        pass
+
+    @property
+    def share_path(self) -> Optional[str]:
+        return None
+
+    @abstractmethod
+    def _update_internal_job(self) -> None:
+        pass
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def teamspace(self) -> "Teamspace":
+        return self._teamspace

lightning_sdk/_mmt/mmt.py

@@ -0,0 +1,161 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, 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
+
+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
+    from lightning_sdk.studio import Studio
+    from lightning_sdk.teamspace import Teamspace
+    from lightning_sdk.user import User
+
+
+class MMT(_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:
+        internal_mmt_cls = _MMTV2 if _has_jobs_v2() else _MMTV1
+
+        self._internal_mmt = internal_mmt_cls(
+            name=name,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            _fetch_job=_fetch_job,
+        )
+
+    @classmethod
+    def run(
+        cls,
+        name: str,
+        num_machines: int,
+        machine: "Machine",
+        command: Optional[str] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Union[str, None] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cluster: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> "MMT":
+        ret_val = super().run(
+            name=name,
+            num_machines=num_machines,
+            machine=machine,
+            command=command,
+            studio=studio,
+            image=image,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            cluster=cluster,
+            env=env,
+            interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+        # required for typing with "Job"
+        assert isinstance(ret_val, cls)
+        return ret_val
+
+    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,
+        cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_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,
+            command=command,
+            studio=studio,
+            image=image,
+            env=env,
+            interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+        return self
+
+    def stop(self) -> None:
+        return self._internal_mmt.stop()
+
+    def delete(self) -> None:
+        return self._internal_mmt.delete()
+
+    @property
+    def status(self) -> "Status":
+        return self._internal_mmt.status
+
+    @property
+    def machines(self) -> Tuple["Job", ...]:
+        return self._internal_mmt.machines
+
+    @property
+    def machine(self) -> "Machine":
+        return self._internal_mmt.machine
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        return self._internal_mmt.artifact_path
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        return self._internal_mmt.snapshot_path
+
+    @property
+    def share_path(self) -> Optional[str]:
+        return None
+
+    def _update_internal_job(self) -> None:
+        return self._internal_mmt._update_internal_job()
+
+    @property
+    def name(self) -> str:
+        return self._internal_mmt.name
+
+    @property
+    def teamspace(self) -> "Teamspace":
+        return self._internal_mmt._teamspace
+
+    @property
+    def cluster(self) -> Optional[str]:
+        return self._internal_mmt.cluster
+
+    def __getattr__(self, key: str) -> Any:
+        """Forward the attribute lookup to the internal job implementation."""
+        try:
+            return getattr(super(), key)
+        except AttributeError:
+            return getattr(self._internal_mmt, key)

lightning_sdk/_mmt/v1.py

@@ -0,0 +1,69 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
+
+if TYPE_CHECKING:
+    from lightning_sdk.job.job import Job
+    from lightning_sdk.machine import Machine
+    from lightning_sdk.status import Status
+    from lightning_sdk.studio import Studio
+    from lightning_sdk.teamspace import Teamspace
+
+from lightning_sdk._mmt.base import _BaseMMT
+
+
+class _MMTV1(_BaseMMT):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError
+
+    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,
+        cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> None:
+        raise NotImplementedError
+
+    @property
+    def machines(self) -> Tuple["Job", ...]:
+        raise NotImplementedError
+
+    def stop(self) -> None:
+        raise NotImplementedError
+
+    def delete(self) -> None:
+        raise NotImplementedError
+
+    @property
+    def status(self) -> "Status":
+        raise NotImplementedError
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        raise NotImplementedError
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        raise NotImplementedError
+
+    @property
+    def machine(self) -> "Machine":
+        raise NotImplementedError
+
+    def _update_internal_job(self) -> None:
+        raise NotImplementedError
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def teamspace(self) -> "Teamspace":
+        return self._teamspace

lightning_sdk/_mmt/v2.py

@@ -0,0 +1,141 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+
+from lightning_sdk.api.mmt_api import MMTApi
+
+if TYPE_CHECKING:
+    from lightning_sdk.job.job import Job
+    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 _MMTV2(_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 = MMTApi()
+        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,
+        cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> "_MMTV2":
+        # Command is required if Studio is provided to know what to run
+        # Image is mutually exclusive with Studio
+        # Command is optional for Image
+        # Either image or studio must be provided
+        if studio is not None:
+            studio_id = studio._studio.id
+            if image is not None:
+                raise ValueError(
+                    "image and studio are mutually exclusive as both define the environment to run the job in"
+                )
+            if command is None:
+                raise ValueError("command is required when using a studio")
+        else:
+            studio_id = None
+            if image is None:
+                raise ValueError("either image or studio must be provided")
+        submitted = self._job_api.submit_job(
+            name=self.name,
+            num_machines=num_machines,
+            command=command,
+            cluster_id=cluster,
+            teamspace_id=self._teamspace.id,
+            studio_id=studio_id,
+            image=image,
+            machine=machine,
+            interruptible=interruptible,
+            env=env,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+        self._job = submitted
+        self._name = submitted.name
+        return self
+
+    @property
+    def machines(self) -> Tuple["Job", ...]:
+        raise NotImplementedError
+
+    def stop(self) -> None:
+        self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
+
+    def delete(self) -> None:
+        self._job_api.delete_job(
+            job_id=self._guaranteed_job.id,
+            teamspace_id=self._teamspace.id,
+        )
+
+    @property
+    def _latest_job(self) -> Any:
+        """Guarantees to fetch the latest version of a job before returning it."""
+        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)
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        raise NotImplementedError
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        raise NotImplementedError
+
+    @property
+    def machine(self) -> "Machine":
+        return self._job_api._get_job_machine_from_spec(self._guaranteed_job.spec)
+
+    def _update_internal_job(self) -> None:
+        if getattr(self, "_job", None) is None:
+            self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id)
+            return
+
+        self._job = self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id)
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def teamspace(self) -> "Teamspace":
+        return self._teamspace