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

lightning_sdk/cli/teamspace_menu.py

@@ -0,0 +1,94 @@
+from typing import Dict, List, Optional
+
+from simple_term_menu import TerminalMenu
+
+from lightning_sdk.api import OrgApi
+from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _get_authed_user
+
+
+class _TeamspacesMenu:
+    def _get_teamspace_from_interactive_menu(self, possible_teamspaces: Dict[str, Dict[str, str]]) -> Dict[str, str]:
+        teamspace_ids = sorted(possible_teamspaces.keys())
+        terminal_menu = self._prepare_terminal_menu_teamspaces([possible_teamspaces[k] for k in teamspace_ids])
+        terminal_menu.show()
+
+        selected_id = teamspace_ids[terminal_menu.chosen_menu_index]
+        return possible_teamspaces[selected_id]
+
+    def _get_teamspace_from_name(
+        self, teamspace: str, possible_teamspaces: Dict[str, Dict[str, str]]
+    ) -> Dict[str, str]:
+        owner, name = teamspace.split("/", maxsplit=1)
+        for _, ts in possible_teamspaces.items():
+            if ts["name"] == name and (ts["user"] == owner or ts["org"] == owner):
+                return ts
+
+        print("Could not find Teamspace {teamspace}, please select it from the list:")
+        return self._get_teamspace_from_interactive_menu(possible_teamspaces)
+
+    @staticmethod
+    def _prepare_terminal_menu_teamspaces(
+        possible_teamspaces: List[Dict[str, str]], title: Optional[str] = None
+    ) -> TerminalMenu:
+        if title is None:
+            title = "Please select a Teamspace of the following:"
+
+        return TerminalMenu(
+            [f"{t['user'] or t['org']}/{t['name']}" for t in possible_teamspaces], title=title, clear_menu_on_exit=True
+        )
+
+    @staticmethod
+    def _get_possible_teamspaces(user: User, is_owner: bool = True) -> Dict[str, Dict[str, str]]:
+        org_api = OrgApi()
+        user_api = user._user_api
+
+        user_api._get_organizations_for_authed_user()
+        memberships = user_api._get_all_teamspace_memberships(user_id=user.id)
+
+        teamspaces = {}
+        # get all teamspace memberships
+        for membership in memberships:
+            teamspace_id = membership.project_id
+            teamspace_name = membership.name
+
+            # get organization if necessary
+            if membership.owner_type == "organization":
+                org_name = org_api._get_org_by_id(membership.owner_id).name
+                user_name = None
+            else:
+                org_name = None
+
+                # don't do a request if not necessary
+                if membership.owner_id == user.id:
+                    user_name = user.name
+                else:
+                    user_name = user_api._get_user_by_id(membership.owner_id).username
+
+            teamspaces[teamspace_id] = {"user": user_name, "org": org_name, "name": teamspace_name}
+
+        return teamspaces
+
+    def _resolve_teamspace(self, teamspace: Optional[str]) -> Teamspace:
+        try:
+            user = _get_authed_user()
+
+            possible_teamspaces = self._get_possible_teamspaces(user)
+            if teamspace is None:
+                teamspace_dict = self._get_teamspace_from_interactive_menu(possible_teamspaces=possible_teamspaces)
+            else:
+                teamspace_dict = self._get_teamspace_from_name(
+                    teamspace=teamspace, possible_teamspaces=possible_teamspaces
+                )
+
+            return Teamspace(**teamspace_dict)
+        except KeyboardInterrupt:
+            raise KeyboardInterrupt from None
+
+        except Exception as e:
+            raise StudioCliError(
+                f"Could not find the given Teamspace {teamspace}. "
+                "Please contact Lightning AI directly to resolve this issue."
+            ) from e

lightning_sdk/job/base.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union
 
+from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace
 
 if TYPE_CHECKING:
@@ -12,6 +13,19 @@ if TYPE_CHECKING:
     from lightning_sdk.user import User
 
 
+class MachineDict(TypedDict):
+    name: str
+    status: "Status"
+    machine: "Machine"
+
+
+class JobDict(MachineDict):
+    command: str
+    teamspace: str
+    studio: Optional[str]
+    image: Optional[str]
+
+
 class _BaseJob(ABC):
     """Base interface to all job types."""
 
@@ -65,6 +79,7 @@ class _BaseJob(ABC):
         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
     ) -> "_BaseJob":
         """Run async workloads using a docker image or a compute environment from your studio.
@@ -98,6 +113,10 @@ class _BaseJob(ABC):
                 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
 
@@ -145,6 +164,9 @@ class _BaseJob(ABC):
                     "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(
@@ -174,6 +196,7 @@ class _BaseJob(ABC):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )
 
     @abstractmethod
@@ -190,6 +213,7 @@ class _BaseJob(ABC):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> "_BaseJob":
         """Submit a new job to the Lightning AI platform.
 
@@ -218,6 +242,9 @@ class _BaseJob(ABC):
                 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.
         """
 
     @abstractmethod
@@ -278,10 +305,51 @@ class _BaseJob(ABC):
     def logs(self) -> str:
         """The logs of the job."""
 
+    @property
+    @abstractmethod
+    def image(self) -> Optional[str]:
+        """The image used to submit the job."""
+
+    @property
+    @abstractmethod
+    def studio(self) -> Optional["Studio"]:
+        """The studio used to submit the job."""
+
+    @property
+    @abstractmethod
+    def command(self) -> str:
+        """The command the job is running."""
+
+    def dict(self) -> JobDict:
+        """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,
+        }
+
+    def json(self) -> str:
+        """JSON representation of this job."""
+        import json
+
+        return json.dumps(self.dict(), indent=4, sort_keys=True, default=str)
+
     @property
     def link(self) -> str:
         """A link to view the current job in the UI."""
-        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=jobs&job_name={self.name}"
+        studio_name = self._job_api.get_studio_name(self._guaranteed_job)
+        if not studio_name:
+            raise RuntimeError("Cannot extract studio name from job")
+        return (
+            f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/studios/"
+            f"{studio_name}/app?app_id=jobs&job_name={self.name}"
+        )
 
     @property
     def _guaranteed_job(self) -> Any:

lightning_sdk/job/job.py

@@ -50,15 +50,40 @@ class Job(_BaseJob):
             user: the name of the user owning the :param`teamspace`
                 in case it is owned directly by a user instead of an org.
         """
-        internal_job_cls = _JobV2 if _has_jobs_v2() and not self._force_v1 else _JobV1
-
-        self._internal_job = internal_job_cls(
-            name=name,
-            teamspace=teamspace,
-            org=org,
-            user=user,
-            _fetch_job=_fetch_job,
-        )
+        from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+
+        if _has_jobs_v2() and not self._force_v1:
+            # try with v2 and fall back to v1
+            try:
+                job = _JobV2(
+                    name=name,
+                    teamspace=teamspace,
+                    org=org,
+                    user=user,
+                    _fetch_job=_fetch_job,
+                )
+            except ApiException as e:
+                try:
+                    job = _JobV1(
+                        name=name,
+                        teamspace=teamspace,
+                        org=org,
+                        user=user,
+                        _fetch_job=_fetch_job,
+                    )
+                except ApiException:
+                    raise e from e
+
+        else:
+            job = _JobV1(
+                name=name,
+                teamspace=teamspace,
+                org=org,
+                user=user,
+                _fetch_job=_fetch_job,
+            )
+
+        self._internal_job = job
 
     @classmethod
     def run(
@@ -78,6 +103,7 @@ class Job(_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
     ) -> "Job":
         """Run async workloads using a docker image or a compute environment from your studio.
@@ -111,6 +137,10 @@ class Job(_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.
         """
         ret_val = super().run(
             name=name,
@@ -128,6 +158,7 @@ class Job(_BaseJob):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
             cluster=cluster,
         )
         # required for typing with "Job"
@@ -149,6 +180,7 @@ class Job(_BaseJob):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> "Job":
         """Submit a new job to the Lightning AI platform.
 
@@ -177,6 +209,9 @@ class Job(_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.
         """
         self._job = self._internal_job._submit(
             machine=machine,
@@ -245,6 +280,21 @@ class Job(_BaseJob):
         """The teamspace the job is part of."""
         return self._internal_job._teamspace
 
+    @property
+    def studio(self) -> Optional["Studio"]:
+        """The studio used to submit the job."""
+        return self._internal_job.studio
+
+    @property
+    def image(self) -> Optional[str]:
+        """The image used to submit the job."""
+        return self._internal_job.image
+
+    @property
+    def command(self) -> str:
+        """The command the job is running."""
+        return self._internal_job.command
+
     @property
     def logs(self) -> str:
         """The logs of the job."""

lightning_sdk/job/v1.py

@@ -100,6 +100,7 @@ class _JobV1(_BaseJob):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> "_JobV1":
         """Submit a job to run on a machine.
 
@@ -115,6 +116,10 @@ class _JobV1(_BaseJob):
             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).
+            entrypoint: The entrypoint of your docker container (not supported).
+                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.
 
         Returns:
             The submitted job.
@@ -132,6 +137,10 @@ class _JobV1(_BaseJob):
             raise ValueError("Environment variables are not supported for submitting jobs")
         if command is None:
             raise ValueError("Command is required for submitting jobs")
+
+        if entrypoint != "sh -c":
+            raise ValueError("Specifying the entrypoint is not yet supported with jobs")
+
         # TODO: add support for empty names (will give an empty string)
         _submitted = self._job_api.submit_job(
             name=self._name,
@@ -215,6 +224,25 @@ class _JobV1(_BaseJob):
         """The logs of the job."""
         return self.work.logs
 
+    @property
+    def image(self) -> Optional[str]:
+        """The image used to submit the job."""
+        # jobsv1 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
     def _id(self) -> str:

lightning_sdk/job/v2.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from lightning_sdk.api.job_api import JobApiV2
+from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.job.base import _BaseJob
 
 if TYPE_CHECKING:
@@ -47,6 +48,7 @@ class _JobV2(_BaseJob):
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> "_JobV2":
         """Submit a new job to the Lightning AI platform.
 
@@ -75,6 +77,10 @@ class _JobV2(_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.
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -107,6 +113,7 @@ class _JobV2(_BaseJob):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )
         self._job = submitted
         self._name = submitted.name
@@ -178,13 +185,42 @@ class _JobV2(_BaseJob):
 
     @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"
-            )
+        mmt_name = self._job_api.get_mmt_name(self._guaranteed_job)
+
+        if self._job_api.get_image_name(self._guaranteed_job):
+            if mmt_name:
+                # don't go via the studio unless we use studio env
+                return (
+                    f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/"
+                    f"jobs/{mmt_name}?app_id=mmt&machine_name={self.name}"
+                )
+            return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=jobs"
 
+        # TODO: MMT env with studio
         return super().link
 
+    @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)
+
     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)

lightning_sdk/job/work.py

@@ -3,9 +3,10 @@ from typing import TYPE_CHECKING, Any, Optional, Protocol
 from lightning_sdk.api.job_api import JobApiV1
 
 if TYPE_CHECKING:
+    from lightning_sdk.job.base import MachineDict
+    from lightning_sdk.machine import Machine
     from lightning_sdk.status import Status
     from lightning_sdk.teamspace import Teamspace
-from lightning_sdk.machine import Machine
 
 
 class _WorkHolder(Protocol):
@@ -70,3 +71,11 @@ class Work:
             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)
+
+    def dict(self) -> "MachineDict":
+        """Dict representation of the work."""
+        return {
+            "name": self.name,
+            "status": self.status,
+            "machine": self.machine,
+        }

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -341,6 +341,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningapp_instanc
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningapp_instance_response import V1DeleteLightningappInstanceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningwork_response import V1DeleteLightningworkResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_page_response import V1DeleteLitPageResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_repository_response import V1DeleteLitRepositoryResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_logger_artifact_response import V1DeleteLoggerArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_managed_endpoint_response import V1DeleteManagedEndpointResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_metrics_stream_response import V1DeleteMetricsStreamResponse
@@ -793,8 +794,10 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_vultr_direct_v1 import V1Vu
 from lightning_sdk.lightning_cloud.openapi.models.v1_work import V1Work
 from lightning_sdk.lightning_cloud.openapi.models.validate import Validate
 from lightning_sdk.lightning_cloud.openapi.models.validateautojoindomain_domain_body import ValidateautojoindomainDomainBody
+from lightning_sdk.lightning_cloud.openapi.models.version_default_body import VersionDefaultBody
 from lightning_sdk.lightning_cloud.openapi.models.version_uploads_body import VersionUploadsBody
 from lightning_sdk.lightning_cloud.openapi.models.versions_id_body import VersionsIdBody
+from lightning_sdk.lightning_cloud.openapi.models.versions_version_body import VersionsVersionBody
 from lightning_sdk.lightning_cloud.openapi.models.works_id_body import WorksIdBody
 from lightning_sdk.lightning_cloud.openapi.models.v1_image_spec import V1ImageSpec as Gridv1ImageSpec
 from lightning_sdk.lightning_cloud.openapi.models.clusters_id_body import ClustersIdBody as Body