lightning-sdk 0.1.35__py3-none-any.whl → 0.1.37__py3-none-any.whl

lightning_sdk/__init__.py

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

lightning_sdk/ai_hub.py

@@ -15,8 +15,17 @@ class AIHub:
     """An interface to interact with the AI Hub.
 
     Example:
-        ai_hub = AIHub()
-        api_list = ai_hub.list_apis()
+        from lightning_sdk import AIHub
+        hub = AIHub()
+
+        # List public API templates
+        api_list = hub.list_apis()
+
+        # Get detailed information about an API template
+        api_info =  hub.api_info("temp_xxxx")
+
+        # Deploy an API template
+        deployment = hub.deploy("temp_xxxx")
     """
 
     def __init__(self) -> None:
@@ -28,7 +37,7 @@ class AIHub:
 
         Example:
             ai_hub = AIHub()
-            api_info = ai_hub.api_info("api_12345")
+            api_info = ai_hub.api_info("temp_xxxx")
 
         Args:
             api_id: The ID of the API for which information is requested.
@@ -70,8 +79,11 @@ class AIHub:
             },
         }
 
-    def list_apis(self, search: Optional[str] = None) -> List[Dict[str, str]]:
-        """Get a list of AI Hub API templates.
+    def list_apis(
+        self,
+        search: Optional[str] = None,
+    ) -> List[Dict[str, str]]:
+        """Get a list of public AI Hub API templates.
 
         Example:
             ai_hub = AIHub()
@@ -93,9 +105,6 @@ class AIHub:
                 "name": template.name,
                 "description": template.description,
                 "creator_username": template.creator_username,
-                "created_on": template.creation_timestamp.strftime("%Y-%m-%d %H:%M:%S")
-                if template.creation_timestamp
-                else None,
             }
             results.append(result)
         return results
@@ -127,14 +136,18 @@ class AIHub:
         name: Optional[str] = None,
         teamspace: Optional[Union[str, "Teamspace"]] = None,
         org: Optional[Union[str, "Organization"]] = None,
-        **kwargs: Dict[str, Any],
+        api_arguments: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Union[str, bool]]:
         """Deploy an API from the AI Hub.
 
         Example:
             from lightning_sdk import AIHub
-            ai_hub = AIHub()
-            deployment = ai_hub.deploy("temp_01jc37n6qpqkdptjpyep0z06hy", batch_size=10)
+            hub = AIHub()
+            deployment = hub.deploy("temp_xxxx")
+
+            # Using API arguments
+            api_arugments = {"model": "unitary/toxic-bert", "batch_size" 10, "token": "lit_xxxx"}
+            deployment = hub.deploy("temp_xxxx", api_arugments=api_arugments)
 
         Args:
             api_id: The ID of the API you want to deploy.
@@ -142,7 +155,7 @@ class AIHub:
             name: Name for the deployed API. Defaults to None.
             teamspace: The team or group for deployment. Defaults to None.
             org: The organization for deployment. Defaults to None.
-            **kwargs: Additional keyword arguments for deployment.
+            api_arguments: Additional API argument, such as model name, or batch size.
 
         Returns:
             A dictionary containing the name of the deployed API,
@@ -155,8 +168,9 @@ class AIHub:
         teamspace = self._authenticate(teamspace, org)
         teamspace_id = teamspace.id
 
+        api_arguments = api_arguments or {}
         deployment = self._api.deploy_api(
-            template_id=api_id, cluster_id=cluster, project_id=teamspace_id, name=name, **kwargs
+            template_id=api_id, cluster_id=cluster, project_id=teamspace_id, name=name, api_arguments=api_arguments
         )
         url = quote(f"{LIGHTNING_CLOUD_URL}/{teamspace._org.name}/{teamspace.name}/jobs/{deployment.name}", safe=":/()")
         print("Deployment available at:", url)

lightning_sdk/api/ai_hub_api.py

@@ -1,6 +1,5 @@
-import re
 import traceback
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 import backoff
 
@@ -8,7 +7,10 @@ from lightning_sdk.lightning_cloud.openapi.models import (
     CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs,
     V1Deployment,
     V1DeploymentTemplate,
-    V1ParameterizationSpec,
+    V1DeploymentTemplateParameter,
+    V1DeploymentTemplateParameterPlacement,
+    V1DeploymentTemplateParameterType,
+    V1JobSpec,
 )
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_template_gallery_response import (
     V1DeploymentTemplateGalleryResponse,
@@ -37,31 +39,59 @@ class AIHubApi:
         ).templates
 
     @staticmethod
-    def _parse_and_update_args(cmd: str, **kwargs: dict) -> list:
-        """Parse the command and update the arguments with the provided kwargs.
+    def _update_parameters(
+        job: V1JobSpec, placements: List[V1DeploymentTemplateParameterPlacement], pattern: str, value: str
+    ) -> None:
+        for placement in placements:
+            if placement == V1DeploymentTemplateParameterPlacement.COMMAND:
+                job.command = job.command.replace(pattern, str(value))
+            if placement == V1DeploymentTemplateParameterPlacement.ENTRYPOINT:
+                job.entrypoint = job.entrypoint.replace(pattern, str(value))
 
-        >>> _parse_and_update_args("--arg1 1 --arg2=2", arg1=3)
-        ['--arg1 3']
-        """
-        keys = [key.lstrip("-") for key in re.findall(r"--\w+", cmd)]
-        arguments = {}
-        for key in keys:
-            if key in kwargs:
-                arguments[key] = kwargs[key]
-        return [f"--{k} {v}" for k, v in arguments.items()]
+            if placement == V1DeploymentTemplateParameterPlacement.ENV:
+                for e in job.env:
+                    if e.value == pattern:
+                        e.value = str(value)
 
     @staticmethod
-    def _resolve_api_arguments(parameter_spec: "V1ParameterizationSpec", **kwargs: dict) -> str:
-        return " ".join(AIHubApi._parse_and_update_args(parameter_spec.command, **kwargs))
+    def _set_parameters(
+        job: V1JobSpec, parameters: List[V1DeploymentTemplateParameter], api_arguments: Dict[str, str]
+    ) -> V1JobSpec:
+        for p in parameters:
+            if p.name not in api_arguments:
+                if p.type == V1DeploymentTemplateParameterType.INPUT and p.input and p.input.default_value:
+                    api_arguments[p.name] = p.input.default_value
+
+                if p.type == V1DeploymentTemplateParameterType.SELECT and p.select and len(p.select.options) > 0:
+                    api_arguments[p.name] = p.select.options[0]
+
+                if p.type == V1DeploymentTemplateParameterType.CHECKBOX and p.checkbox:
+                    api_arguments[p.name] = (
+                        (p.checkbox.true_value or "True")
+                        if p.checkbox.is_checked
+                        else (p.checkbox.false_value or "False")
+                    )
+
+        for p in parameters:
+            name = p.name
+            pattern = f"${{{name}}}"
+            if name in api_arguments:
+                AIHubApi._update_parameters(job, p.placements, pattern, api_arguments[name])
+            elif not p.required:
+                AIHubApi._update_parameters(job, p.placements, pattern, "")
+            else:
+                raise ValueError(f"API reqires argument '{p.name}' but is not provided with api_arguments.")
+
+        return job
 
     def deploy_api(
-        self, template_id: str, project_id: str, cluster_id: str, name: Optional[str], **kwargs: dict
+        self, template_id: str, project_id: str, cluster_id: str, name: Optional[str], api_arguments: Dict[str, str]
     ) -> V1Deployment:
         template = self._client.deployment_templates_service_get_deployment_template(template_id)
         name = name or template.name
         template.spec_v2.endpoint.id = None
-        command = self._resolve_api_arguments(template.parameter_spec, **kwargs)
-        template.spec_v2.job.command = command
+
+        AIHubApi._set_parameters(template.spec_v2.job, template.parameter_spec.parameters, api_arguments)
         return self._client.jobs_service_create_deployment(
             project_id=project_id,
             body=CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs(

lightning_sdk/api/utils.py

@@ -509,6 +509,7 @@ def _download_model_files(
         pbar = tqdm(
             desc=f"Downloading {version}",
             unit="B",
+            total=float(response.size_bytes),
             unit_scale=True,
             unit_divisor=1000,
         )

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.run import _Run
 from lightning_sdk.cli.upload import _Uploads
 from lightning_sdk.lightning_cloud.login import Auth
 
@@ -19,8 +20,7 @@ class StudioCLI:
         self.upload = _Uploads()
         self.aihub = _AIHub()
 
-        if _LIGHTNING_AVAILABLE:
-            self.run = _LegacyLightningCLI()
+        self.run = _Run(legacy_run=_LegacyLightningCLI() if _LIGHTNING_AVAILABLE else None)
 
     def login(self) -> None:
         """Login to Lightning AI Studios."""

lightning_sdk/cli/models.py

@@ -1,8 +1,11 @@
-from typing import Tuple
+import os
+from typing import Any, Dict, List, Tuple
 
 from lightning_sdk.api import OrgApi, UserApi
 from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.lightning_cloud.openapi.models import V1Membership, V1OwnerType
 from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _get_authed_user
 
 
@@ -17,22 +20,49 @@ def _parse_model_name(name: str) -> Tuple[str, str, str]:
     return org_name, teamspace_name, model_name
 
 
+def _get_teamspace_and_path(
+    ts: V1Membership, org_api: OrgApi, user_api: UserApi, authed_user: User
+) -> Tuple[str, Dict[str, Any]]:
+    if ts.owner_type == V1OwnerType.ORGANIZATION:
+        org = org_api._get_org_by_id(ts.owner_id)
+        return f"{org.name}/{ts.name}", {"name": ts.name, "org": org.name}
+
+    if ts.owner_type == V1OwnerType.USER and ts.owner_id != authed_user.id:
+        user = user_api._get_user_by_id(ts.owner_id)  # todo: check also the name
+        return f"{user.username}/{ts.name}", {"name": ts.name, "user": User(name=user.username)}
+
+    if ts.owner_type == V1OwnerType.USER:
+        return f"{authed_user.name}/{ts.name}", {"name": ts.name, "user": authed_user}
+
+    raise StudioCliError(f"Unknown organization type {ts.owner_type}")
+
+
+def _list_teamspaces() -> List[str]:
+    org_api = OrgApi()
+    user_api = UserApi()
+    authed_user = _get_authed_user()
+
+    return [
+        _get_teamspace_and_path(ts, org_api, user_api, authed_user)[0]
+        for ts in user_api._get_all_teamspace_memberships("")
+    ]
+
+
 def _get_teamspace(name: str, organization: str) -> Teamspace:
     """Get a Teamspace object from the SDK."""
     org_api = OrgApi()
-    user = _get_authed_user()
-    teamspaces = {}
-    for ts in UserApi()._get_all_teamspace_memberships(""):
-        if ts.owner_type == "organization":
-            org = org_api._get_org_by_id(ts.owner_id)
-            teamspaces[f"{org.name}/{ts.name}"] = {"name": ts.name, "org": org.name}
-        elif ts.owner_type == "user":  # todo: check also the name
-            teamspaces[f"{user.name}/{ts.name}"] = {"name": ts.name, "user": user}
-        else:
-            raise StudioCliError(f"Unknown organization type {ts.owner_type}")
+    user_api = UserApi()
+    authed_user = _get_authed_user()
 
     requested_teamspace = f"{organization}/{name}".lower()
-    if requested_teamspace not in teamspaces:
-        options = "\n\t".join(teamspaces.keys())
-        raise StudioCliError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: \n\t{options}")
-    return Teamspace(**teamspaces[requested_teamspace])
+
+    for ts in user_api._get_all_teamspace_memberships(""):
+        if ts.name != name:
+            continue
+
+        teamspace_path, teamspace = _get_teamspace_and_path(ts, org_api, user_api, authed_user)
+        if requested_teamspace == teamspace_path:
+            return Teamspace(**teamspace)
+
+    options = f"{os.linesep}\t".join(_list_teamspaces())
+    raise StudioCliError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: {os.linesep}\t{options}")

lightning_sdk/cli/run.py

@@ -0,0 +1,72 @@
+from typing import TYPE_CHECKING, Dict, Optional
+
+from lightning_sdk.job import Job
+from lightning_sdk.machine import Machine
+
+if TYPE_CHECKING:
+    from lightning_sdk.cli.legacy import _LegacyLightningCLI
+
+_MACHINE_VALUES = tuple([machine.value for machine in Machine])
+
+
+class _Run:
+    """Run async workloads on the Lightning AI platform."""
+
+    def __init__(self, legacy_run: Optional["_LegacyLightningCLI"] = None) -> None:
+        if legacy_run is not None:
+            self.app = legacy_run.app
+            self.model = legacy_run.model
+
+        # Need to set the docstring here for f-strings to work.
+        # Sadly this is the only way to really show options as f-strings are not allowed as docstrings directly
+        # and fire does not show values for literals, just that it is a literal.
+        docstr = f"""Run async workloads using a docker image or a compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            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.
+            teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
+            org: The organization owning the teamspace (if any). Defaults to the current organization.
+            user: The user owning the teamspace (if any). Defaults to the current user.
+            cluster: The cluster to run the job on. Defaults to the studio cluster if running with studio compute env.
+                If not provided will fall back to the teamspaces default cluster.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+        """
+        self.job.__func__.__doc__ = docstr
+
+    # TODO: sadly, fire displays both Optional[type] and Union[type, None] as Optional[Optional]
+    # see https://github.com/google/python-fire/pull/513
+    # might need to move to different cli library
+    def job(
+        self,
+        name: str,
+        machine: str,
+        command: Optional[str] = None,
+        studio: Optional[str] = None,
+        image: Optional[str] = None,
+        teamspace: Optional[str] = None,
+        org: Optional[str] = None,
+        user: Optional[str] = None,
+        cluster: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+    ) -> None:
+        machine_enum = Machine(machine.upper())
+        Job.run(
+            name=name,
+            machine=machine_enum,
+            command=command,
+            studio=studio,
+            image=image,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            cluster=cluster,
+            env=env,
+            interruptible=interruptible,
+        )

lightning_sdk/job/base.py

@@ -16,15 +16,20 @@ class _BaseJob(ABC):
     def __init__(
         self,
         name: str,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
-        cluster: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         *,
         _fetch_job: bool = True,
     ) -> None:
-        self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
-        self._cluster = cluster
+        _teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        if _teamspace is None:
+            raise ValueError(
+                "Cannot resolve the teamspace from provided arguments."
+                f" Got teamspace={teamspace}, org={org}, user={user}."
+            )
+        else:
+            self._teamspace = _teamspace
         self._name = name
         self._job = None
 
@@ -37,18 +42,25 @@ class _BaseJob(ABC):
         name: str,
         machine: "Machine",
         command: Optional[str] = None,
-        studio: Optional["Studio"] = None,
+        studio: Union["Studio", str, None] = None,
         image: Optional[str] = None,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
     ) -> "_BaseJob":
+        from lightning_sdk.studio import Studio
+
         if not name:
             raise ValueError("A job needs to have a name!")
-        if studio is not None:
+
+        if image is None:
+            if not isinstance(studio, Studio):
+                studio = Studio(name=studio, teamspace=teamspace, org=org, user=user, cluster=cluster, create_ok=False)
+
+            # studio is a Studio instance at this point
             if teamspace is None:
                 teamspace = studio.teamspace
             else:
@@ -60,11 +72,30 @@ class _BaseJob(ABC):
                         "Can only run jobs with Studio envs in the teamspace of that Studio."
                     )
 
-        # TODO: resolve studio and support string studios
-        # TODO: assertions for studio to be on cluster
-        # TODO: if cluster is not provided use studio cluster if provided, otherwise use default cluster from teamspace
-        inst = cls(name=name, teamspace=teamspace, org=org, user=user, cluster=cluster, _fetch_job=False)
-        inst._submit(machine=machine, command=command, studio=studio, image=image, env=env, interruptible=interruptible)
+            if cluster is None:
+                cluster = studio.cluster
+
+            if cluster != studio.cluster:
+                raise ValueError(
+                    "Studio cluster does not match provided cluster. "
+                    "Can only run jobs with Studio envs in the same cluster."
+                )
+        else:
+            if studio is not None:
+                raise RuntimeError(
+                    "image and studio are mutually exclusive as both define the environment to run the job in"
+                )
+
+        inst = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False)
+        inst._submit(
+            machine=machine,
+            cluster=cluster,
+            command=command,
+            studio=studio,
+            image=image,
+            env=env,
+            interruptible=interruptible,
+        )
         return inst
 
     @abstractmethod
@@ -76,6 +107,7 @@ class _BaseJob(ABC):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        cluster: Optional[str] = None,
     ) -> None:
         """Submits a job and updates the internal _job attribute as well as the _name attribute."""
 
@@ -123,7 +155,3 @@ class _BaseJob(ABC):
     @property
     def teamspace(self) -> "Teamspace":
         return self._teamspace
-
-    @property
-    def cluster(self) -> Optional[str]:
-        return self._cluster

lightning_sdk/job/job.py

@@ -28,17 +28,20 @@ class Job(_BaseJob):
     def __init__(
         self,
         name: str,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
-        cluster: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         *,
         _fetch_job: bool = True,
     ) -> None:
         internal_job_cls = _JobV2 if _has_jobs_v2() else _JobV1
 
         self._internal_job = internal_job_cls(
-            name=name, teamspace=teamspace, org=org, user=user, cluster=cluster, _fetch_job=_fetch_job
+            name=name,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            _fetch_job=_fetch_job,
         )
 
     @classmethod
@@ -47,11 +50,11 @@ class Job(_BaseJob):
         name: str,
         machine: "Machine",
         command: Optional[str] = None,
-        studio: Optional["Studio"] = None,
-        image: Optional[str] = None,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Union[str, None] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
@@ -81,9 +84,16 @@ class Job(_BaseJob):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        cluster: Optional[str] = None,
     ) -> None:
         return self._internal_job._submit(
-            machine=machine, command=command, studio=studio, image=image, env=env, interruptible=interruptible
+            machine=machine,
+            cluster=cluster,
+            command=command,
+            studio=studio,
+            image=image,
+            env=env,
+            interruptible=interruptible,
         )
 
     def stop(self) -> None:

lightning_sdk/job/v1.py

@@ -20,15 +20,14 @@ class _JobV1(_BaseJob):
     def __init__(
         self,
         name: str,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
-        cluster: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         *,
         _fetch_job: bool = True,
     ) -> None:
         self._job_api = JobApiV1()
-        super().__init__(name=name, teamspace=teamspace, org=org, user=user, cluster=cluster, _fetch_job=_fetch_job)
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
 
     @classmethod
     def run(
@@ -37,9 +36,9 @@ class _JobV1(_BaseJob):
         machine: "Machine",
         command: str,
         studio: "Studio",
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         cluster: Optional[str] = None,
         interruptible: bool = False,
     ) -> "_BaseJob":
@@ -65,6 +64,7 @@ class _JobV1(_BaseJob):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        cluster: Optional[str] = None,
     ) -> None:
         if studio is None:
             raise ValueError("Studio is required for submitting jobs")
@@ -85,7 +85,7 @@ class _JobV1(_BaseJob):
             command=command,
             studio_id=studio._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._cluster,
+            cluster_id=cluster,
             machine=machine,
             interruptible=interruptible,
         )

lightning_sdk/job/v2.py

@@ -16,15 +16,14 @@ class _JobV2(_BaseJob):
     def __init__(
         self,
         name: str,
-        teamspace: Union[str, "Teamspace"] = None,
-        org: Union[str, "Organization"] = None,
-        user: Union[str, "User"] = None,
-        cluster: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
         *,
         _fetch_job: bool = True,
     ) -> None:
         self._job_api = JobApiV2()
-        super().__init__(name=name, teamspace=teamspace, org=org, user=user, cluster=cluster, _fetch_job=_fetch_job)
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
 
     def _submit(
         self,
@@ -34,6 +33,7 @@ class _JobV2(_BaseJob):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        cluster: Optional[str] = None,
     ) -> None:
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -55,7 +55,7 @@ class _JobV2(_BaseJob):
         submitted = self._job_api.submit_job(
             name=self.name,
             command=command,
-            cluster_id=self._cluster,
+            cluster_id=cluster,
             teamspace_id=self._teamspace.id,
             studio_id=studio_id,
             image=image,

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()