lightning-sdk 0.1.47__py3-none-any.whl → 0.1.49__py3-none-any.whl

lightning_sdk/__init__.py

@@ -4,6 +4,7 @@ from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__  # noq
 from lightning_sdk.helpers import _check_version_and_prompt_upgrade
 from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
+from lightning_sdk.mmt import MMT
 from lightning_sdk.organization import Organization
 from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin, SlurmJobsPlugin
 from lightning_sdk.status import Status
@@ -15,6 +16,7 @@ __all__ = [
     "Job",
     "JobsPlugin",
     "Machine",
+    "MMT",
     "MultiMachineTrainingPlugin",
     "Organization",
     "Plugin",
@@ -27,5 +29,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.47"
+__version__ = "0.1.49"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/api/job_api.py

@@ -180,8 +180,13 @@ class JobApiV1:
 
         raise RuntimeError("Could not extract command from app")
 
+    def get_total_cost(self, job: Externalv1LightningappInstance) -> float:
+        status: V1LightningappInstanceStatus = job.status
+        return status.total_cost
+
 
 class JobApiV2:
+    # these are stages the job can be in.
     v2_job_state_pending = "pending"
     v2_job_state_running = "running"
     v2_job_state_stopped = "stopped"
@@ -189,6 +194,9 @@ class JobApiV2:
     v2_job_state_failed = "failed"
     v2_job_state_stopping = "stopping"
 
+    # this is the user action to stop the job.
+    v2_job_state_stop = "stop"
+
     def __init__(self) -> None:
         self._cloud_url = _cloud_url()
         self._client = LightningClient(max_tries=7)
@@ -262,7 +270,7 @@ class JobApiV2:
             return
 
         if current_state != Status.Stopping:
-            update_body = JobsIdBody1(cloudspace_id=current_job.spec.cloudspace_id, state=self.v2_job_state_stopped)
+            update_body = JobsIdBody1(cloudspace_id=current_job.spec.cloudspace_id, state=self.v2_job_state_stop)
             self._client.jobs_service_update_job(body=update_body, project_id=teamspace_id, id=job_id)
 
         while True:
@@ -302,6 +310,13 @@ class JobApiV2:
     def get_command(self, job: V1Job) -> str:
         return job.spec.command
 
+    def get_mmt_name(self, job: V1Job) -> str:
+        if job.multi_machine_job_id:
+            splits = job.name.rsplit("-", 1)
+            if len(splits) == 2:
+                return splits[0]
+        return ""
+
     def _job_state_to_external(self, state: str) -> "Status":
         from lightning_sdk.status import Status
 
@@ -326,3 +341,6 @@ class JobApiV2:
         return _COMPUTE_NAME_TO_MACHINE.get(
             instance_type, _COMPUTE_NAME_TO_MACHINE.get(instance_name, instance_type or instance_name)
         )
+
+    def get_total_cost(self, job: V1Job) -> float:
+        return job.total_cost

lightning_sdk/api/lit_container_api.py

@@ -0,0 +1,42 @@
+from typing import Generator, List
+
+from lightning_sdk.api.utils import _get_registry_url
+from lightning_sdk.lightning_cloud.openapi.models import V1DeleteLitRepositoryResponse
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.teamspace import Teamspace
+
+
+class LitContainerApi:
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=3)
+
+    def list_containers(self, project_id: str) -> List:
+        project = self._client.lit_registry_service_get_lit_project_registry(project_id)
+        return project.repositories
+
+    def delete_container(self, project_id: str, container: str) -> V1DeleteLitRepositoryResponse:
+        try:
+            return self._client.lit_registry_service_delete_lit_repository(project_id, container)
+        except Exception as ex:
+            raise ValueError(f"Could not delete container {container} from project {project_id}") from ex
+
+    def upload_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
+        import docker
+
+        try:
+            client = docker.from_env()
+            client.ping()
+        except docker.errors.DockerException as e:
+            raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
+
+        try:
+            client.images.get(container)
+        except docker.errors.ImageNotFound:
+            raise ValueError(f"Container {container} does not exist") from None
+
+        registry_url = _get_registry_url()
+        repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
+        tagged = client.api.tag(container, repository, tag)
+        if not tagged:
+            raise ValueError(f"Could not tag container {container} with {repository}:{tag}")
+        return client.api.push(repository, stream=True, decode=True)

lightning_sdk/api/mmt_api.py

@@ -203,3 +203,9 @@ class MMTApiV2:
         return _COMPUTE_NAME_TO_MACHINE.get(
             instance_type, _COMPUTE_NAME_TO_MACHINE.get(instance_name, instance_type or instance_name)
         )
+
+    def get_total_cost(self, job: V1MultiMachineJob) -> float:
+        return job.total_cost
+
+    def get_num_machines(self, job: V1MultiMachineJob) -> int:
+        return job.machines

lightning_sdk/api/utils.py

@@ -353,7 +353,8 @@ def _machine_to_compute_name(machine: Union[Machine, str]) -> str:
 
 _COMPUTE_NAME_TO_MACHINE: Dict[str, Machine] = {v: k for k, v in _MACHINE_TO_COMPUTE_NAME.items()}
 
-_DEFAULT_CLOUD_URL = "https://lightning.ai:443"
+_DEFAULT_CLOUD_URL = "https://lightning.ai"
+_DEFAULT_REGISTRY_URL = "litcr.io"
 
 
 def _get_cloud_url() -> str:
@@ -362,6 +363,12 @@ def _get_cloud_url() -> str:
     return cloud_url
 
 
+def _get_registry_url() -> str:
+    registry_url = os.environ.get("LIGHTNING_REGISTRY_URL", _DEFAULT_REGISTRY_URL)
+    os.environ["LIGHTNING_REGISTRY_URL"] = registry_url
+    return registry_url
+
+
 def _sanitize_studio_remote_path(path: str, studio_id: str) -> str:
     return f"/cloudspaces/{studio_id}/code/content/{path.replace('/teamspace/studios/this_studio/', '')}"
 

lightning_sdk/cli/delete.py

@@ -0,0 +1,58 @@
+from typing import Optional
+
+from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
+from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.lit_container import LitContainer
+
+
+class _Delete(_JobAndMMTAction, _TeamspacesMenu):
+    """Delete resources on the Lightning AI platform."""
+
+    def container(self, container: str, teamspace: Optional[str] = None) -> None:
+        """Delete a docker container.
+
+        Args:
+            container: The name of the container to delete.
+            teamspace: The teamspace to delete the container from. Should be specified as {owner}/{name}
+                If not provided, can be selected in an interactive menu.
+        """
+        api = LitContainer()
+        resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
+        try:
+            api.delete_container(container, resolved_teamspace.name, resolved_teamspace.owner.name)
+            print(f"Container {container} deleted successfully.")
+        except Exception as e:
+            raise StudioCliError(
+                f"Could not delete container {container} from project {resolved_teamspace.name}: {e}"
+            ) from None
+
+    def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Delete a job.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        job = super().job(name=name, teamspace=teamspace)
+
+        job.delete()
+        print(f"Successfully deleted {job.name}!")
+
+    def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Delete a multi-machine job.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        mmt = super().mmt(name=name, teamspace=teamspace)
+
+        mmt.delete()
+        print(f"Successfully deleted {mmt.name}!")

lightning_sdk/cli/entrypoint.py

@@ -1,13 +1,20 @@
+import sys
+from types import TracebackType
+from typing import Type
+
 from fire import Fire
 from lightning_utilities.core.imports import RequirementCache
 
 from lightning_sdk.api.studio_api import _cloud_url
 from lightning_sdk.cli.ai_hub import _AIHub
+from lightning_sdk.cli.delete import _Delete
 from lightning_sdk.cli.download import _Downloads
+from lightning_sdk.cli.inspect import _Inspect
 from lightning_sdk.cli.legacy import _LegacyLightningCLI
 from lightning_sdk.cli.list import _List
 from lightning_sdk.cli.run import _Run
 from lightning_sdk.cli.serve import _Docker, _LitServe
+from lightning_sdk.cli.stop import _Stop
 from lightning_sdk.cli.upload import _Uploads
 from lightning_sdk.lightning_cloud.login import Auth
 
@@ -25,6 +32,11 @@ class StudioCLI:
         self.serve = _LitServe()
         self.dockerize = _Docker()
         self.list = _List()
+        self.delete = _Delete()
+        self.inspect = _Inspect()
+        self.stop = _Stop()
+
+        sys.excepthook = _notify_exception
 
     def login(self) -> None:
         """Login to Lightning AI Studios."""
@@ -42,6 +54,11 @@ class StudioCLI:
         auth.clear()
 
 
+def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:  # No
+    """CLI won't show tracebacks, just print the exception message."""
+    print(value)
+
+
 def main_cli() -> None:
     """CLI entrypoint."""
     Fire(StudioCLI(), name="lightning")

lightning_sdk/cli/inspect.py

@@ -0,0 +1,31 @@
+from typing import Optional
+
+from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
+
+
+class _Inspect(_JobAndMMTAction):
+    """Inspect resources of the Lightning AI platform to get additional details as JSON."""
+
+    def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Inspect a job for further details as JSON.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        print(super().job(name=name, teamspace=teamspace).json())
+
+    def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Inspect a multi-machine job for further details as JSON.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        print(super().mmt(name=name, teamspace=teamspace).json())

lightning_sdk/cli/job_and_mmt_action.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+from lightning_sdk.cli.jobs_menu import _JobsMenu
+from lightning_sdk.cli.mmts_menu import _MMTsMenu
+from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.job import Job
+from lightning_sdk.mmt import MMT
+
+
+class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
+    """Inspect resources of the Lightning AI platform to get additional details as JSON."""
+
+    def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> Job:
+        """Fetch a job for further processing.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        resolved_teamspace = self._resolve_teamspace(teamspace)
+        return self._resolve_job(name, teamspace=resolved_teamspace)
+
+    def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> MMT:
+        """Fetch a multi-machine job for further processing.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        resolved_teamspace = self._resolve_teamspace(teamspace)
+        return self._resolve_mmt(name, teamspace=resolved_teamspace)

lightning_sdk/cli/jobs_menu.py

@@ -0,0 +1,57 @@
+from typing import Dict, List, Optional
+
+from simple_term_menu import TerminalMenu
+
+from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.job import Job
+from lightning_sdk.teamspace import Teamspace
+
+
+class _JobsMenu:
+    def _get_job_from_interactive_menu(self, possible_jobs: Dict[str, Job]) -> Job:
+        job_ids = sorted(possible_jobs.keys())
+        terminal_menu = self._prepare_terminal_menu_jobs([possible_jobs[k] for k in job_ids])
+        terminal_menu.show()
+
+        return possible_jobs[terminal_menu.chosen_menu_entry]
+
+    def _get_job_from_name(self, job: str, possible_jobs: Dict[str, Job]) -> Job:
+        for _, j in possible_jobs.items():
+            if j.name == job:
+                return j
+
+        print("Could not find Job {job}, please select it from the list:")
+        return self._get_job_from_interactive_menu(possible_jobs)
+
+    @staticmethod
+    def _prepare_terminal_menu_jobs(possible_jobs: List[Job], title: Optional[str] = None) -> TerminalMenu:
+        if title is None:
+            title = "Please select a Job of the following:"
+
+        return TerminalMenu([j.name for j in possible_jobs], title=title, clear_menu_on_exit=True)
+
+    @staticmethod
+    def _get_possible_jobs(teamspace: Teamspace) -> Dict[str, Job]:
+        jobs = {}
+        for j in teamspace.jobs:
+            jobs[j.name] = j
+
+        return jobs
+
+    def _resolve_job(self, job: Optional[str], teamspace: Teamspace) -> Job:
+        try:
+            possible_jobs = self._get_possible_jobs(teamspace)
+            if job is None:
+                resolved_job = self._get_job_from_interactive_menu(possible_jobs)
+            else:
+                resolved_job = self._get_job_from_name(job=job, possible_jobs=possible_jobs)
+
+            return resolved_job
+        except KeyboardInterrupt:
+            raise KeyboardInterrupt from None
+
+        except Exception as e:
+            raise StudioCliError(
+                f"Could not find the given Job {job} in Teamspace {teamspace.name}. "
+                "Please contact Lightning AI directly to resolve this issue."
+            ) from e

lightning_sdk/cli/list.py

@@ -4,7 +4,7 @@ from rich.console import Console
 from rich.table import Table
 
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
-from lightning_sdk.lit_registry import LitRegistry
+from lightning_sdk.lit_container import LitContainer
 
 
 class _List(_TeamspacesMenu):
@@ -20,7 +20,36 @@ class _List(_TeamspacesMenu):
         """
         resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
 
-        print("Available Jobs:\n" + "\n".join([j.name for j in resolved_teamspace.jobs]))
+        jobs = resolved_teamspace.jobs
+
+        table = Table(
+            pad_edge=True,
+        )
+        table.add_column("Name")
+        table.add_column("Teamspace")
+        table.add_column("Studio")
+        table.add_column("Image")
+        table.add_column("Status")
+        table.add_column("Machine")
+        table.add_column("Total Cost")
+        for j in jobs:
+            # we know we just fetched these, so no need to refetch
+            j._prevent_refetch_latest = True
+            j._internal_job._prevent_refetch_latest = True
+
+            studio = j.studio
+            table.add_row(
+                j.name,
+                f"{j.teamspace.owner.name}/{j.teamspace.name}",
+                studio.name if studio else None,
+                j.image,
+                str(j.status),
+                str(j.machine),
+                f"{j.total_cost:.3f}",
+            )
+
+        console = Console()
+        console.print(table)
 
     def mmts(self, teamspace: Optional[str] = None) -> None:
         """List multi-machine jobs for a given teamspace.
@@ -32,7 +61,36 @@ class _List(_TeamspacesMenu):
         """
         resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
 
-        print("Available MMTs:\n" + "\n".join([j.name for j in resolved_teamspace.multi_machine_jobs]))
+        jobs = resolved_teamspace.multi_machine_jobs
+
+        table = Table(pad_edge=True)
+        table.add_column("Name")
+        table.add_column("Teamspace")
+        table.add_column("Studio")
+        table.add_column("Image")
+        table.add_column("Status")
+        table.add_column("Machine")
+        table.add_column("Num Machines")
+        table.add_column("Total Cost")
+        for j in jobs:
+            # we know we just fetched these, so no need to refetch
+            j._prevent_refetch_latest = True
+            j._internal_job._prevent_refetch_latest = True
+
+            studio = j.studio
+            table.add_row(
+                j.name,
+                f"{j.teamspace.owner.name}/{j.teamspace.name}",
+                studio.name if studio else None,
+                j.image,
+                str(j.status),
+                str(j.machine),
+                str(j.num_machines),
+                str(j.total_cost),
+            )
+
+        console = Console()
+        console.print(table)
 
     def containers(self, teamspace: Optional[str] = None) -> None:
         """Display the list of available containers.
@@ -41,7 +99,7 @@ class _List(_TeamspacesMenu):
             teamspace: The teamspace to list containers from. Should be specified as {owner}/{name}
                 If not provided, can be selected in an interactive menu.
         """
-        api = LitRegistry()
+        api = LitContainer()
         resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
         result = api.list_containers(teamspace=resolved_teamspace.name, org=resolved_teamspace.owner.name)
         table = Table(pad_edge=True, box=None)

lightning_sdk/cli/mmts_menu.py

@@ -0,0 +1,57 @@
+from typing import Dict, List, Optional
+
+from simple_term_menu import TerminalMenu
+
+from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.mmt import MMT
+from lightning_sdk.teamspace import Teamspace
+
+
+class _MMTsMenu:
+    def _get_mmt_from_interactive_menu(self, possible_mmts: Dict[str, MMT]) -> MMT:
+        job_ids = sorted(possible_mmts.keys())
+        terminal_menu = self._prepare_terminal_menu_mmts([possible_mmts[k] for k in job_ids])
+        terminal_menu.show()
+
+        return possible_mmts[terminal_menu.chosen_menu_entry]
+
+    def _get_mmt_from_name(self, mmt: str, possible_mmts: Dict[str, MMT]) -> MMT:
+        for _, j in possible_mmts.items():
+            if j.name == mmt:
+                return j
+
+        print("Could not find Multi-Machine Job {mmt}, please select it from the list:")
+        return self._get_mmt_from_interactive_menu(possible_mmts)
+
+    @staticmethod
+    def _prepare_terminal_menu_mmts(possible_mmts: List[MMT], title: Optional[str] = None) -> TerminalMenu:
+        if title is None:
+            title = "Please select a Multi-Machine Job of the following:"
+
+        return TerminalMenu([m.name for m in possible_mmts], title=title, clear_menu_on_exit=True)
+
+    @staticmethod
+    def _get_possible_mmts(teamspace: Teamspace) -> Dict[str, MMT]:
+        jobs = {}
+        for j in teamspace.multi_machine_jobs:
+            jobs[j.name] = j
+
+        return jobs
+
+    def _resolve_mmt(self, mmt: Optional[str], teamspace: Teamspace) -> MMT:
+        try:
+            possible_mmts = self._get_possible_mmts(teamspace)
+            if mmt is None:
+                resolved_mmt = self._get_mmt_from_interactive_menu(possible_mmts)
+            else:
+                resolved_mmt = self._get_mmt_from_name(mmt=mmt, possible_mmts=possible_mmts)
+
+            return resolved_mmt
+        except KeyboardInterrupt:
+            raise KeyboardInterrupt from None
+
+        except Exception as e:
+            raise StudioCliError(
+                f"Could not find the given Multi-Machine-Job {mmt} in Teamspace {teamspace.name}. "
+                "Please contact Lightning AI directly to resolve this issue."
+            ) from e

lightning_sdk/cli/run.py

@@ -53,6 +53,10 @@ class _Run:
                 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.
         """
         # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
         # might need to switch to explicit cli definition
@@ -93,6 +97,10 @@ class _Run:
                 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.
         """
         # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
         # might need to switch to explicit cli definition
@@ -103,7 +111,7 @@ class _Run:
     # might need to move to different cli library
     def job(
         self,
-        name: str,
+        name: Optional[str] = None,
         machine: Optional[str] = None,
         command: Optional[str] = None,
         studio: Optional[str] = None,
@@ -118,17 +126,25 @@ class _Run:
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> None:
+        if not name:
+            from datetime import datetime
+
+            timestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+            name = f"job-{timestr}"
+
         if machine is None:
             # TODO: infer from studio
             machine = "CPU"
-        machine_enum = Machine(machine.upper())
+        machine_enum = Machine[machine.upper()]
 
         resolved_teamspace = Teamspace(name=teamspace, org=org, user=user)
 
         if cloud_account is None:
             cloud_account = resolved_teamspace.default_cloud_account
         machine_enum = Machine(machine.upper())
+
         Job.run(
             name=name,
             machine=machine_enum,
@@ -145,6 +161,7 @@ class _Run:
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )
 
     # TODO: sadly, fire displays both Optional[type] and Union[type, None] as Optional[Optional]
@@ -167,6 +184,7 @@ class _Run:
         cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        entrypoint: str = "sh -c",
     ) -> None:
         if name is None:
             from datetime import datetime
@@ -177,7 +195,7 @@ class _Run:
         if machine is None:
             # TODO: infer from studio
             machine = "CPU"
-        machine_enum = Machine(machine.upper())
+        machine_enum = Machine[machine.upper()]
 
         resolved_teamspace = Teamspace(name=teamspace, org=org, user=user)
         if cloud_account is None:
@@ -203,4 +221,5 @@ class _Run:
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )

lightning_sdk/cli/stop.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
+
+
+class _Stop(_JobAndMMTAction):
+    """Stop resources on the Lightning AI platform."""
+
+    def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Stop a job.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        job = super().job(name=name, teamspace=teamspace)
+
+        job.stop()
+        print(f"Successfully stopped {job.name}!")
+
+    def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Stop a multi-machine job.
+
+        Args:
+            name: the name of the job. If not specified can be selected interactively.
+            teamspace: the name of the teamspace the job lives in.
+                Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace).
+                If not specified can be selected interactively.
+
+        """
+        mmt = super().mmt(name=name, teamspace=teamspace)
+
+        mmt.stop()
+        print(f"Successfully stopped {mmt.name}!")