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

lightning_sdk/__init__.py

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

lightning_sdk/api/job_api.py

@@ -171,6 +171,15 @@ class JobApiV1:
         data = urlopen(resp.url).read().decode("utf-8")
         return remove_datetime_prefix(str(data))
 
+    def get_command(self, job: Externalv1LightningappInstance) -> str:
+        env = job.spec.env
+
+        for e in env:
+            if e.name == "COMMAND":
+                return e.value
+
+        raise RuntimeError("Could not extract command from app")
+
 
 class JobApiV2:
     v2_job_state_pending = "pending"
@@ -199,6 +208,7 @@ class JobApiV2:
         cloud_account_auth: bool,
         artifacts_local: Optional[str],
         artifacts_remote: Optional[str],
+        entrypoint: str,
     ) -> V1Job:
         env_vars = []
         if env is not None:
@@ -213,6 +223,7 @@ class JobApiV2:
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
             command=command or "",
+            entrypoint=entrypoint,
             env=env_vars,
             image=image or "",
             instance_name=instance_name,
@@ -276,11 +287,20 @@ class JobApiV2:
         data = urlopen(resp.url).read().decode("utf-8")
         return remove_datetime_prefix(str(data))
 
-    def get_studio_name(self, job: V1Job) -> str:
-        cs: V1CloudSpace = self._client.cloud_space_service_get_cloud_space(
-            project_id=job.project_id, id=job.spec.cloudspace_id
-        )
-        return cs.name
+    def get_studio_name(self, job: V1Job) -> Optional[str]:
+        if job.spec.cloudspace_id:
+            cs: V1CloudSpace = self._client.cloud_space_service_get_cloud_space(
+                project_id=job.project_id, id=job.spec.cloudspace_id
+            )
+            return cs.name
+
+        return None
+
+    def get_image_name(self, job: V1Job) -> Optional[str]:
+        return job.spec.image or None
+
+    def get_command(self, job: V1Job) -> str:
+        return job.spec.command
 
     def _job_state_to_external(self, state: str) -> "Status":
         from lightning_sdk.status import Status

lightning_sdk/api/lit_registry_api.py

@@ -0,0 +1,12 @@
+from typing import List
+
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+class LitRegistryApi:
+    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

lightning_sdk/api/mmt_api.py

@@ -16,6 +16,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     Externalv1LightningappInstance,
     MultimachinejobsIdBody,
     ProjectIdMultimachinejobsBody,
+    V1CloudSpace,
     V1EnvVar,
     V1Job,
     V1JobSpec,
@@ -86,6 +87,7 @@ class MMTApiV2:
         cloud_account_auth: bool,
         artifacts_local: Optional[str],
         artifacts_remote: Optional[str],
+        entrypoint: str,
     ) -> V1MultiMachineJob:
         env_vars = []
         if env is not None:
@@ -100,7 +102,7 @@ class MMTApiV2:
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
             command=command or "",
-            entrypoint="sh -c",
+            entrypoint=entrypoint,
             env=env_vars,
             image=image or "",
             instance_name=instance_name,
@@ -179,6 +181,21 @@ class MMTApiV2:
             return Status.Failed
         return Status.Pending
 
+    def get_studio_name(self, job: V1MultiMachineJob) -> Optional[str]:
+        if job.spec.cloudspace_id:
+            cs: V1CloudSpace = self._client.cloud_space_service_get_cloud_space(
+                project_id=job.project_id, id=job.spec.cloudspace_id
+            )
+            return cs.name
+
+        return None
+
+    def get_image_name(self, job: V1MultiMachineJob) -> Optional[str]:
+        return job.spec.image or None
+
+    def get_command(self, job: V1MultiMachineJob) -> str:
+        return job.spec.command
+
     def _get_job_machine_from_spec(self, spec: V1JobSpec) -> "Machine":
         instance_name = spec.instance_name
         instance_type = spec.instance_type

lightning_sdk/api/teamspace_api.py

@@ -1,12 +1,13 @@
 import os
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 from tqdm.auto import tqdm
 
 from lightning_sdk.api.utils import _download_model_files, _DummyBody, _get_model_version, _ModelFileUploader
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
+    Externalv1LightningappInstance,
     ModelIdVersionsBody,
     ModelsStoreApi,
     ProjectIdAgentsBody,
@@ -14,7 +15,9 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1Assistant,
     V1CloudSpace,
     V1Endpoint,
+    V1Job,
     V1ModelVersionArchive,
+    V1MultiMachineJob,
     V1Project,
     V1ProjectClusterBinding,
     V1PromptSuggestion,
@@ -277,3 +280,18 @@ class TeamspaceApi:
             download_dir=download_dir,
             progress_bar=progress_bar,
         )
+
+    def list_jobs(self, teamspace_id: str) -> Tuple[List[Externalv1LightningappInstance], List[V1Job]]:
+        apps = self._client.lightningapp_instance_service_list_lightningapp_instances(
+            project_id=teamspace_id, source_app="job_run_plugin"
+        ).lightningapps
+        jobs = self._client.jobs_service_list_jobs(project_id=teamspace_id, standalone=True).jobs
+
+        return apps, jobs
+
+    def list_mmts(self, teamspace_id: str) -> Tuple[List[Externalv1LightningappInstance], List[V1MultiMachineJob]]:
+        apps = self._client.lightningapp_instance_service_list_lightningapp_instances(
+            project_id=teamspace_id, source_app="distributed_plugin"
+        ).lightningapps
+        jobs = self._client.jobs_service_list_multi_machine_jobs(project_id=teamspace_id).multi_machine_jobs
+        return apps, jobs

lightning_sdk/cli/entrypoint.py

@@ -5,6 +5,7 @@ from lightning_sdk.api.studio_api import _cloud_url
 from lightning_sdk.cli.ai_hub import _AIHub
 from lightning_sdk.cli.download import _Downloads
 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.upload import _Uploads
@@ -23,6 +24,7 @@ class StudioCLI:
         self.run = _Run(legacy_run=_LegacyLightningCLI() if _LIGHTNING_AVAILABLE else None)
         self.serve = _LitServe()
         self.dockerize = _Docker()
+        self.list = _List()
 
     def login(self) -> None:
         """Login to Lightning AI Studios."""

lightning_sdk/cli/list.py

@@ -0,0 +1,54 @@
+from typing import Optional
+
+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
+
+
+class _List(_TeamspacesMenu):
+    """List resources on the Lightning AI platform."""
+
+    def jobs(self, teamspace: Optional[str] = None) -> None:
+        """List jobs for a given teamspace.
+
+        Args:
+            teamspace: the teamspace to list jobs from. Should be specified as {owner}/{name}
+                If not provided, can be selected in an interactive menu.
+
+        """
+        resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
+
+        print("Available Jobs:\n" + "\n".join([j.name for j in resolved_teamspace.jobs]))
+
+    def mmts(self, teamspace: Optional[str] = None) -> None:
+        """List multi-machine jobs for a given teamspace.
+
+        Args:
+            teamspace: the teamspace to list jobs from. Should be specified as {owner}/{name}
+                If not provided, can be selected in an interactive menu.
+
+        """
+        resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
+
+        print("Available MMTs:\n" + "\n".join([j.name for j in resolved_teamspace.multi_machine_jobs]))
+
+    def containers(self, teamspace: Optional[str] = None) -> None:
+        """Display the list of available containers.
+
+        Args:
+            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()
+        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)
+        table.add_column("REPOSITORY")
+        table.add_column("IMAGE ID")
+        table.add_column("CREATED")
+        for repo in result:
+            table.add_row(repo["REPOSITORY"], repo["IMAGE ID"], repo["CREATED"])
+        console = Console()
+        console.print(table)

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,5 +1,5 @@
 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.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace
 
@@ -12,6 +12,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 +78,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 +112,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 +163,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 +195,7 @@ class _BaseJob(ABC):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )
 
     @abstractmethod
@@ -190,6 +212,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 +241,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 +304,48 @@ 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"https://lightning.ai/{self.teamspace.owner.name}/{self.teamspace.name}/studios/{studio_name}/app?app_id=jobs&job_name={self.name}"
 
     @property
     def _guaranteed_job(self) -> Any:

lightning_sdk/job/job.py

@@ -78,6 +78,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 +112,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 +133,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 +155,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 +184,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 +255,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

@@ -47,6 +47,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 +76,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 +112,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 +184,35 @@ class _JobV2(_BaseJob):
 
     @property
     def link(self) -> str:
-        if self._guaranteed_job.spec.image:
+        if self._job_api.get_image_name(self._guaranteed_job):
             return (
                 f"https://lightning.ai/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=jobs"
             )
 
         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,
+        }