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

lightning_sdk/job/v1.py

@@ -0,0 +1,241 @@
+from typing import TYPE_CHECKING, Dict, Optional, Union
+
+from lightning_sdk.api.job_api import JobApiV1
+from lightning_sdk.job.base import _BaseJob
+from lightning_sdk.status import Status
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import Machine
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.studio import Studio
+    from lightning_sdk.teamspace import Teamspace
+    from lightning_sdk.user import User
+
+from functools import cached_property
+
+from lightning_sdk.job.work import Work
+
+
+class _JobV1(_BaseJob):
+    """Implementation to run async workloads from your Studio."""
+
+    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:
+        """Fetch already existing jobs.
+
+        Args:
+            name: the name of the job
+            teamspace: the teamspace the job is part of
+            org: the name of the organization owning the :param`teamspace` in case it is owned by an org
+            user: the name of the user owning the :param`teamspace`
+                in case it is owned directly by a user instead of an org
+        """
+        self._job_api = JobApiV1()
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
+
+    @classmethod
+    def run(
+        cls,
+        name: str,
+        machine: "Machine",
+        command: str,
+        studio: "Studio",
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cloud_account: Optional[str] = None,
+        interruptible: bool = False,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
+    ) -> "_BaseJob":
+        """Start a new async workload from your studio.
+
+        Args:
+            name: the name of the job
+            machine: the machine to run the workload on
+            command: the command to execute
+            studio: the studio the job belongs to
+            teamspace: the teamspace the job is part of
+            org: the organization owning the teamspace (if applicable)
+            user: the user owning the teamspace (if applicable)
+            cloud_account: the cloud account to run the workload on
+            interruptible: whether the workload can be interrupted
+
+        Returns:
+            the created job
+        """
+        return super().run(
+            name=name,
+            machine=machine,
+            command=command,
+            studio=studio,
+            image=None,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            cloud_account=cloud_account,
+            env=None,
+            interruptible=interruptible,
+            image_credentials=None,
+            cloud_account_auth=False,
+            cluster=cluster,
+        )
+
+    def _submit(
+        self,
+        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,
+    ) -> "_JobV1":
+        """Submit a job to run on a machine.
+
+        Args:
+            machine: The machine to run the job on.
+            command: The command to execute.
+            studio: The studio the job belongs to.
+            image: The image to use for the job (not supported).
+            env: The environment variables for the job (not supported).
+            interruptible: Whether the job can be interrupted.
+            cloud_account: The cloud account to run the job on.
+            image_credentials: The image credentials for the job (not supported).
+            cloud_account_auth: Whether to use cloud account authentication for the job (not supported).
+            artifacts_local: The local path for persisting artifacts (not supported).
+            artifacts_remote: The remote path for persisting artifacts (not supported).
+
+        Returns:
+            The submitted job.
+
+        """
+        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 jobs")
+        # TODO: add support for empty names (will give an empty string)
+        _submitted = self._job_api.submit_job(
+            name=self._name,
+            command=command,
+            studio_id=studio._studio.id,
+            teamspace_id=self._teamspace.id,
+            cloud_account=cloud_account or "",
+            machine=machine,
+            interruptible=interruptible,
+        )
+        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 status(self) -> "Status":
+        """Returns the status of the job."""
+        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"Job {self._name} does not exist in Teamspace {self.teamspace.name}. Did you delete it?"
+            ) from None
+
+    def stop(self) -> None:
+        """Stops the job. is blocking until the ob is stopped."""
+        if self.status in (Status.Stopped, Status.Failed):
+            return None
+
+        return self._job_api.stop_job(self._job.id, self.teamspace.id)
+
+    def delete(self) -> None:
+        """Deletes the job.
+
+        Caution: this also deletes all artifacts created by the job.
+        """
+        self._job_api.delete_job(self._job.id, self.teamspace.id)
+
+    @cached_property
+    def work(self) -> Work:
+        """Get the work associated with the job."""
+        _work = self._job_api.list_works(self._job.id, self.teamspace.id)
+        if len(_work) == 0:
+            raise ValueError("No works found for job")
+        return Work(_work[0].id, self, self.teamspace)
+
+    @property
+    def machine(self) -> "Machine":
+        """Get the machine the job is running on."""
+        return self.work.machine
+
+    @property
+    def name(self) -> str:
+        """The name of the job."""
+        return self._job.name
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        """The path to the artifacts of the job in the distributed teamspace filesystem."""
+        return self.work.artifact_path
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        """The path to the snapshot of the job in the distributed teamspace filesystem."""
+        return f"/teamspace/jobs/{self.name}/snapshot"
+
+    @property
+    def share_path(self) -> Optional[str]:
+        """The path to the share of the job in the distributed teamspace filesystem."""
+        return f"/teamspace/jobs/{self.name}/share"
+
+    @property
+    def logs(self) -> str:
+        """The logs of the job."""
+        return self.work.logs
+
+    # 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.", "")
+
+
+def _internal_status_to_external_status(internal_status: str) -> "Status":
+    """Converts internal status strings from HTTP requests to external enums."""
+    return {
+        # don't get a status if no instance alive
+        None: Status.Stopped,
+        # TODO: should we have deleted in here?
+        "LIGHTNINGAPP_INSTANCE_STATE_UNSPECIFIED": Status.Pending,
+        "LIGHTNINGAPP_INSTANCE_STATE_IMAGE_BUILDING": Status.Pending,
+        "LIGHTNINGAPP_INSTANCE_STATE_NOT_STARTED": Status.Pending,
+        "LIGHTNINGAPP_INSTANCE_STATE_PENDING": Status.Pending,
+        "LIGHTNINGAPP_INSTANCE_STATE_RUNNING": Status.Running,
+        "LIGHTNINGAPP_INSTANCE_STATE_FAILED": Status.Failed,
+        "LIGHTNINGAPP_INSTANCE_STATE_STOPPED": Status.Stopped,
+        "LIGHTNINGAPP_INSTANCE_STATE_COMPLETED": Status.Completed,
+    }[internal_status]

lightning_sdk/job/v2.py

@@ -0,0 +1,193 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+from lightning_sdk.api.job_api import JobApiV2
+from lightning_sdk.job.base import _BaseJob
+
+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
+
+
+class _JobV2(_BaseJob):
+    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:
+        """Fetch already existing jobs.
+
+        Args:
+            name: the name of the job
+            teamspace: the teamspace the job is part of
+            org: the name of the organization owning the :param`teamspace` in case it is owned by an org
+            user: the name of the user owning the :param`teamspace`
+                in case it is owned directly by a user instead of an org.
+        """
+        self._job_api = JobApiV2()
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
+
+    def _submit(
+        self,
+        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,
+    ) -> "_JobV2":
+        """Submit a new job to the Lightning AI platform.
+
+        Args:
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            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.
+            env: Environment variables to set inside the job.
+            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.
+            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.
+        """
+        # 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,
+            command=command,
+            cloud_account=cloud_account,
+            teamspace_id=self._teamspace.id,
+            studio_id=studio_id,
+            image=image,
+            machine=machine,
+            interruptible=interruptible,
+            env=env,
+            image_credentials=image_credentials,
+            cloud_account_auth=cloud_account_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+        self._job = submitted
+        self._name = submitted.name
+        return self
+
+    def stop(self) -> None:
+        """Stop the job. If the job is already stopped, this is a no-op. This is blocking until the job is stopped."""
+        self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
+
+    def delete(self) -> None:
+        """Delete the job.
+
+        Caution: This also deletes all artifacts created by the job.
+        """
+        self._job_api.delete_job(
+            job_id=self._guaranteed_job.id,
+            teamspace_id=self._teamspace.id,
+            cloudspace_id=self._guaranteed_job.spec.cloudspace_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 status(self) -> "Status":
+        """The current status of the job."""
+        return self._job_api._job_state_to_external(self._latest_job.state)
+
+    @property
+    def machine(self) -> "Machine":
+        """The machine type the job is running on."""
+        # only fetch the job it it hasn't been fetched yet as machine cannot change over time
+        return self._job_api._get_job_machine_from_spec(self._guaranteed_job.spec)
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        """The path to the artifacts of the job within the distributed teamspace filesystem."""
+        if self._guaranteed_job.spec.image != "":
+            if self._guaranteed_job.spec.artifacts_destination != "":
+                splits = self._guaranteed_job.spec.artifacts_destination.split(":")
+                return f"/teamspace/{splits[0]}_connections/{splits[1]}/{splits[2]}"
+            return None
+
+        return f"/teamspace/jobs/{self._guaranteed_job.name}/artifacts"
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        """The path to the snapshot of the Studio used to create the job within the distributed teamspace filesystem."""
+        if self._guaranteed_job.spec.image != "":
+            return None
+        return f"/teamspace/jobs/{self._guaranteed_job.name}/snapshot"
+
+    @property
+    def share_path(self) -> Optional[str]:
+        """The path to the share of the job within the distributed teamspace filesystem."""
+        raise NotImplementedError("Not implemented yet")
+
+    @property
+    def logs(self) -> str:
+        from lightning_sdk.status import Status
+
+        if self.status not in (Status.Failed, Status.Completed, Status.Stopped):
+            raise RuntimeError("Getting jobs logs while the job is pending or running is not supported yet!")
+
+        return self._job_api.get_logs_finished(job_id=self._guaranteed_job.id, teamspace_id=self.teamspace.id)
+
+    @property
+    def link(self) -> str:
+        if self._guaranteed_job.spec.image:
+            return (
+                f"https://lightning.ai/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=jobs"
+            )
+
+        return super().link
+
+    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)

lightning_sdk/job/work.py

@@ -0,0 +1,72 @@
+from typing import TYPE_CHECKING, Any, Optional, Protocol
+
+from lightning_sdk.api.job_api import JobApiV1
+
+if TYPE_CHECKING:
+    from lightning_sdk.status import Status
+    from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.machine import Machine
+
+
+class _WorkHolder(Protocol):
+    @property
+    def _id(self) -> str:
+        ...
+
+    @property
+    def name(self) -> str:
+        ...
+
+    def _name_filter(self, name: str) -> str:
+        ...
+
+
+class Work:
+    def __init__(self, work_id: str, job: _WorkHolder, teamspace: "Teamspace") -> None:
+        self._id = work_id
+        self._job = job
+        self._teamspace = teamspace
+        self._job_api = JobApiV1()
+        self._work = None
+
+    @property
+    def _latest_work(self) -> Any:
+        self._work = self._job_api.get_work(work_id=self._id, job_id=self._job._id, teamspace_id=self._teamspace.id)
+        return self._work
+
+    @property
+    def _guaranteed_work(self) -> Any:
+        if self._work is None:
+            return self._latest_work
+
+        return self._work
+
+    @property
+    def id(self) -> str:
+        return self._guaranteed_work.id
+
+    @property
+    def name(self) -> str:
+        return self._job._name_filter(self._guaranteed_work.name)
+
+    @property
+    def machine(self) -> "Machine":
+        return self._job_api.get_machine_from_work(self._guaranteed_work)
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        return f"/teamspace/jobs/{self._job.name}/{self.name}"
+
+    @property
+    def status(self) -> "Status":
+        return self._job_api.get_status_from_work(self._latest_work)
+
+    @property
+    def logs(self) -> str:
+        """The logs of the work."""
+        from lightning_sdk.status import Status
+
+        if self.status not in (Status.Failed, Status.Completed, Status.Stopped):
+            raise RuntimeError("Getting jobs logs while the job is pending or running is not supported yet!")
+
+        return self._job_api.get_logs_finished(job_id=self._job._id, work_id=self._id, teamspace_id=self._teamspace.id)

lightning_sdk/lightning_cloud/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.64"
+__version__ = "0.5.70"

lightning_sdk/lightning_cloud/cli/__main__.py

@@ -1,16 +1,9 @@
-import click
 
 
-@click.group()
-def main():
-    pass
-
-
-@main.command()
-def login():
+def login() -> None:
     """Authorize the CLI to access Grid AI resources for a particular user.
-    Use login command to force authenticate,
-    a web browser will open to complete the authentication.
+
+    Use login command to force authenticate, a web browser will open to complete the authentication.
     """
     from lightning_sdk.lightning_cloud.login import Auth  # local to avoid circular import
 
@@ -18,10 +11,19 @@ def login():
     auth.clear()
     auth._run_server()
 
-
-@main.command()
-def logout():
+def logout() -> None:
     """Logout from LightningCloud"""
     from lightning_sdk.lightning_cloud.login import Auth  # local to avoid circular import
 
     Auth.clear()
+
+
+def main() -> None:
+    """CLI entrypoint."""
+    from fire import Fire
+
+    Fire({"login": login, "logout": logout})
+
+
+if __name__ == "__main__":
+    main()

lightning_sdk/lightning_cloud/env.py

@@ -34,6 +34,7 @@ DOT_IGNORE_FILENAME = ".lightningignore"
 
 LEEWAY = 100
 IS_DEV_ENV = True
+LIGHTNING_CLOUD_PROJECT_ID = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
 
 
 def reset_global_variables() -> None:

lightning_sdk/lightning_cloud/login.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass
 from typing import Optional
 from urllib.parse import urlencode
 
-import click
+import webbrowser
 
 import requests
 import uvicorn
@@ -22,7 +22,7 @@ logger = logging.getLogger(__name__)
 
 
 class Keys(Enum):
-    USERNAME = "LIGHTNING_USERNAME"
+    # USERNAME = "LIGHTNING_USERNAME"
     USER_ID = "LIGHTNING_USER_ID"
     API_KEY = "LIGHTNING_API_KEY"
 
@@ -33,7 +33,7 @@ class Keys(Enum):
 
 @dataclass
 class Auth:
-    username: Optional[str] = None
+    # username: Optional[str] = None
     user_id: Optional[str] = None
     api_key: Optional[str] = None
 
@@ -72,20 +72,21 @@ class Auth:
              token: str = "",
              user_id: str = "",
              api_key: str = "",
-             username: str = "") -> None:
+             # username: str = "",
+        ) -> None:
         """save credentials to disk."""
         self.secrets_file.parent.mkdir(exist_ok=True, parents=True)
         with self.secrets_file.open("w") as f:
             json.dump(
                 {
-                    f"{Keys.USERNAME.suffix}": username,
+                    # f"{Keys.USERNAME.suffix}": username,
                     f"{Keys.USER_ID.suffix}": user_id,
                     f"{Keys.API_KEY.suffix}": api_key,
                 },
                 f,
             )
 
-        self.username = username
+        # self.username = username
         self.user_id = user_id
         self.api_key = api_key
         logger.debug("credentials saved successfully")
@@ -164,7 +165,10 @@ class AuthServer:
             )
 
         logger.info(f"login started for lightning.ai, opening {url}")
-        click.launch(url)
+        ok = webbrowser.open(url)
+        if not ok:
+            # can't open a browser, authentication failed
+            raise RuntimeError("Failed to authenticate to Lightning. When running without access to a browser, `LIGHTNING_USER_ID` and `LIGHTNING_API_KEY` should be exported.")
 
         @app.get("/login-complete")
         async def save_token(request: Request,
@@ -173,7 +177,7 @@ class AuthServer:
                              user_id: str = Query("", alias="userID")):
             if token:
                 auth.save(token=token,
-                          username=user_id,
+                          # username=user_id,
                           user_id=user_id,
                           api_key=key)
                 logger.info("Authentication Successful")