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

lightning_sdk/_mmt/mmt.py

@@ -1,161 +0,0 @@
-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

@@ -1,69 +0,0 @@
-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

@@ -1,141 +0,0 @@
-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

lightning_sdk/cli/mmt.py

@@ -1,137 +0,0 @@
-from typing import Dict, Optional
-
-from fire import Fire
-
-from lightning_sdk._mmt import MMT
-from lightning_sdk.api.studio_api import _cloud_url
-from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.machine import Machine
-from lightning_sdk.teamspace import Teamspace
-
-_MACHINE_VALUES = tuple([machine.value for machine in Machine])
-
-
-class MMTCLI:
-    """Command line interface (CLI) to interact with/manage Lightning AI MMT."""
-
-    def __init__(self) -> None:
-        # Need to set the docstring here for f-strings to work.
-        # Sadly this is the only way to really show options as f-strings are not allowed as docstrings directly
-        # and fire does not show values for literals, just that it is a literal.
-        docstr = f"""Run async workloads on multiple machines using a docker image.
-
-        Args:
-            name: The name of the job. Needs to be unique within the teamspace.
-            num_machines: The number of Machines to run on. Defaults to 2 Machines
-            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}. Defaults to CPU
-            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.
-            image: The docker image to run the job with. Mutually exclusive with studio.
-            teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
-            org: The organization owning the teamspace (if any). Defaults to the current organization.
-            user: The user owning the teamspace (if any). Defaults to the current user.
-            cluster: The cluster to run the job on. Defaults to the studio cluster if running with studio compute env.
-                If not provided will fall back to the teamspaces default cluster.
-            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.
-            cluster_auth: Whether to authenticate with the cluster to pull the image.
-                Required if the registry is part of a cluster 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.
-        """
-        # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
-        # might need to switch to explicit cli definition
-        self.run.__func__.__doc__ = docstr
-
-    def login(self) -> None:
-        """Login to Lightning AI Studios."""
-        auth = Auth()
-        auth.clear()
-
-        try:
-            auth.authenticate()
-        except ConnectionError:
-            raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None
-
-    def logout(self) -> None:
-        """Logout from Lightning AI Studios."""
-        auth = Auth()
-        auth.clear()
-
-    # TODO: sadly, fire displays both Optional[type] and Union[type, None] as Optional[Optional]
-    # see https://github.com/google/python-fire/pull/513
-    # might need to move to different cli library
-    def run(
-        self,
-        name: Optional[str] = None,
-        num_machines: int = 2,
-        machine: Optional[str] = None,
-        command: Optional[str] = None,
-        studio: Optional[str] = None,
-        image: Optional[str] = None,
-        teamspace: Optional[str] = None,
-        org: Optional[str] = None,
-        user: Optional[str] = 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,
-    ) -> None:
-        if name is None:
-            from datetime import datetime
-
-            timestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-            name = f"mmt-{timestr}"
-
-        if machine is None:
-            # TODO: infer from studio
-            machine = "CPU"
-        machine_enum = Machine(machine.upper())
-
-        teamspace = Teamspace(name=teamspace, org=org, user=user)
-        if cluster is None:
-            cluster = teamspace.default_cluster
-
-        if image is None:
-            raise RuntimeError("Currently only docker images are specified")
-        MMT.run(
-            name=name,
-            num_machines=num_machines,
-            machine=machine_enum,
-            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,
-        )
-
-
-def main_cli() -> None:
-    """CLI entrypoint."""
-    Fire(MMTCLI(), name="_mmt")
-
-
-if __name__ == "__main__":
-    main_cli()