lightning-sdk 0.1.36__py3-none-any.whl → 0.1.38__py3-none-any.whl

lightning_sdk/__init__.py

@@ -27,5 +27,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.36"
+__version__ = "0.1.38"
 _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.
@@ -38,17 +47,7 @@ class AIHub:
             including its name, description, creation and update timestamps,
             parameters, tags, job specifications, and autoscaling settings.
         """
-        template = self._api.api_info(api_id)
-
-        api_arguments = [
-            {
-                "name": param.name,
-                "short_description": param.short_description,
-                "required": param.required,
-                "default": param.input.default_value,
-            }
-            for param in template.parameter_spec.parameters
-        ]
+        template, api_arguments = self._api.api_info(api_id)
 
         return {
             "name": template.name,
@@ -70,8 +69,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 +95,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
@@ -120,33 +119,33 @@ class AIHub:
             raise ValueError("You need to pass a teamspace or an org for your deployment.")
         return teamspace
 
-    def deploy(
+    def run(
         self,
         api_id: str,
-        cluster: Optional[str] = None,
+        api_arguments: Optional[Dict[str, Any]] = None,
         name: Optional[str] = None,
+        cluster: Optional[str] = None,
         teamspace: Optional[Union[str, "Teamspace"]] = None,
         org: Optional[Union[str, "Organization"]] = None,
-        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
             hub = AIHub()
-            deployment = hub.deploy("temp_01jc37n6qpqkdptjpyep0z06hy")
+            deployment = hub.run("temp_xxxx")
 
             # Using API arguments
-            api_arugments = {"batch_size" 10, "batch_timeout": 0.001, "env_token": "lit_xxxx"}
-            deployment = hub.deploy("temp_01jc37n6qpqkdptjpyep0z06hy", api_arugments=api_arugments)
+            api_arugments = {"model": "unitary/toxic-bert", "batch_size": 10, "token": "lit_xxxx"}
+            deployment = hub.run("temp_xxxx", api_arugments=api_arugments)
 
         Args:
-            api_id: The ID of the API you want to deploy.
-            cluster: The cluster where you want to deploy the API, such as "lightning-public-prod". Defaults to None.
+            api_id: The ID of the AIHub template you want to run.
+            api_arguments: Additional API argument, such as model name, or batch size.
             name: Name for the deployed API. Defaults to None.
+            cluster: The cluster where you want to run the template, such as "lightning-public-prod". Defaults to None.
             teamspace: The team or group for deployment. Defaults to None.
             org: The organization for deployment. Defaults to None.
-            api_arguments: Additional API argument, such as model name, or batch size.
 
         Returns:
             A dictionary containing the name of the deployed API,
@@ -160,14 +159,22 @@ class AIHub:
         teamspace_id = teamspace.id
 
         api_arguments = api_arguments or {}
-        deployment = self._api.deploy_api(
+        deployment = self._api.run_api(
             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=":/()")
+        url = (
+            quote(
+                f"{LIGHTNING_CLOUD_URL}/{teamspace._org.name}/{teamspace.name}/jobs/{deployment.name}",
+                safe=":/()",
+            )
+            + "?app_id=deployment"
+        )
         print("Deployment available at:", url)
+
         return {
             "id": deployment.id,
             "name": deployment.name,
-            "base_url": deployment.status.urls[0],
+            "deployment_url": url,
+            "api_endpoint": deployment.status.urls[0],
             "interruptible": deployment.spec.spot,
         }

lightning_sdk/api/ai_hub_api.py

@@ -1,5 +1,5 @@
 import traceback
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import backoff
 
@@ -22,15 +22,40 @@ class AIHubApi:
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=3)
 
-    def api_info(self, api_id: str) -> "V1DeploymentTemplate":
+    def api_info(self, api_id: str) -> Tuple[V1DeploymentTemplate, List[Dict[str, str]]]:
         try:
-            return self._client.deployment_templates_service_get_deployment_template(api_id)
+            template = self._client.deployment_templates_service_get_deployment_template(api_id)
         except Exception as e:
             stack_trace = traceback.format_exc()
             if "record not found" in stack_trace:
                 raise ValueError(f"api_id={api_id} not found.") from None
             raise e
 
+        api_arguments = []
+        for param in template.parameter_spec.parameters:
+            default = None
+            if param.type == V1DeploymentTemplateParameterType.INPUT and param.input:
+                default = param.input.default_value
+            if param.type == V1DeploymentTemplateParameterType.SELECT and param.select:
+                default = param.select.options[0]
+            if param.type == V1DeploymentTemplateParameterType.CHECKBOX and param.checkbox:
+                default = (
+                    (param.checkbox.true_value or "True")
+                    if param.checkbox.is_checked
+                    else (param.checkbox.false_value or "False")
+                )
+
+            api_arguments.append(
+                {
+                    "name": param.name,
+                    "short_description": param.short_description,
+                    "required": param.required,
+                    "type": param.type,
+                    "default": default,
+                }
+            )
+        return template, api_arguments
+
     @backoff.on_predicate(backoff.expo, lambda x: not x, max_tries=5)
     def list_apis(self, search_query: str) -> List[V1DeploymentTemplateGalleryResponse]:
         kwargs = {"show_globally_visible": True}
@@ -84,7 +109,7 @@ class AIHubApi:
 
         return job
 
-    def deploy_api(
+    def run_api(
         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)

lightning_sdk/api/job_api.py

@@ -141,13 +141,15 @@ class JobApiV2:
         self,
         name: str,
         command: Optional[str],
-        cluster_id: str,
+        cluster_id: Optional[str],
         teamspace_id: str,
         studio_id: Optional[str],
         image: Optional[str],
         machine: Machine,
         interruptible: bool,
         env: Optional[Dict[str, str]],
+        image_credentials: Optional[str],
+        cluster_auth: bool,
     ) -> V1Job:
         env_vars = []
         if env is not None:
@@ -160,13 +162,15 @@ class JobApiV2:
 
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
-            cluster_id=cluster_id,
+            cluster_id=cluster_id or "",
             command=command or "",
             env=env_vars,
             image=image or "",
             instance_name=instance_name,
             run_id=run_id,
             spot=interruptible,
+            image_cluster_credentials=cluster_auth,
+            image_secret_ref=image_credentials or "",
         )
         body = ProjectIdJobsBody(name=name, spec=spec)
 

lightning_sdk/api/teamspace_api.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Dict, List, Optional
 
-from lightning_sdk.api.utils import _download_model_files, _DummyBody, _ModelFileUploader
+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 (
     ModelIdVersionsBody,
@@ -159,6 +159,16 @@ class TeamspaceApi:
 
         return self._client.assistants_service_create_assistant(body=body, project_id=teamspace_id)
 
+    # lazy property which is only created when needed
+    @property
+    def models(self) -> ModelsStoreApi:
+        if not self._models:
+            self._models = ModelsStoreApi(self._client.api_client)
+        return self._models
+
+    def get_model_version(self, name: str, version: str, teamspace_id: str) -> V1ModelVersionArchive:
+        return _get_model_version(client=self._client, name=name, version=version, teamspace_id=teamspace_id)
+
     def create_model(
         self,
         name: str,
@@ -167,17 +177,15 @@ class TeamspaceApi:
         teamspace_id: str,
         cluster_id: str,
     ) -> V1ModelVersionArchive:
-        if not self._models:
-            self._models = ModelsStoreApi(self._client.api_client)
         # ask if such model already exists by listing models with specific name
-        models = self._models.models_store_list_models(project_id=teamspace_id, name=name).models
+        models = self.models.models_store_list_models(project_id=teamspace_id, name=name).models
         if len(models) == 0:
-            return self._models.models_store_create_model(
+            return self.models.models_store_create_model(
                 body=ProjectIdModelsBody(cluster_id=cluster_id, metadata=metadata, name=name, private=private),
                 project_id=teamspace_id,
             )
         assert len(models) == 1, "Multiple models with the same name found"
-        return self._models.models_store_create_model_version(
+        return self.models.models_store_create_model_version(
             body=ModelIdVersionsBody(cluster_id=cluster_id),
             project_id=teamspace_id,
             model_id=models[0].id,
@@ -185,18 +193,16 @@ class TeamspaceApi:
 
     def delete_model(self, name: str, version: Optional[str], teamspace_id: str) -> None:
         """Delete a model or a version from the model store."""
-        if not self._models:
-            self._models = ModelsStoreApi(self._client.api_client)
-        models = self._models.models_store_list_models(project_id=teamspace_id, name=name).models
+        models = self.models.models_store_list_models(project_id=teamspace_id, name=name).models
         assert len(models) == 1, "Multiple models with the same name found"
         model_id = models[0].id
         # decide if delete only version of whole model
         if version:
             if version == "latest":
                 version = models[0].latest_version
-            self._models.models_store_delete_model_version(project_id=teamspace_id, model_id=model_id, version=version)
+            self.models.models_store_delete_model_version(project_id=teamspace_id, model_id=model_id, version=version)
         else:
-            self._models.models_store_delete_model(project_id=teamspace_id, model_id=model_id)
+            self.models.models_store_delete_model(project_id=teamspace_id, model_id=model_id)
 
     def upload_model_file(
         self,
@@ -242,9 +248,7 @@ class TeamspaceApi:
             )
 
     def complete_model_upload(self, model_id: str, version: str, teamspace_id: str) -> None:
-        if not self._models:
-            self._models = ModelsStoreApi(self._client.api_client)
-        self._models.models_store_complete_model_upload(
+        self.models.models_store_complete_model_upload(
             body=_DummyBody(),
             project_id=teamspace_id,
             model_id=model_id,

lightning_sdk/api/utils.py

@@ -28,6 +28,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1UploadProjectArtifactResponse,
     VersionUploadsBody,
 )
+from lightning_sdk.lightning_cloud.openapi.models.v1_model_version_archive import V1ModelVersionArchive
 
 try:
     from lightning_sdk.lightning_cloud.openapi import AppsIdBody1 as AppsIdBody
@@ -492,6 +493,24 @@ class _FileDownloader:
         os.rename(tmp_filename, self.local_path)
 
 
+def _get_model_version(client: LightningClient, teamspace_id: str, name: str, version: str) -> V1ModelVersionArchive:
+    api = ModelsStoreApi(client.api_client)
+    models = api.models_store_list_models(project_id=teamspace_id, name=name).models
+    if not models:
+        raise ValueError(f"Model `{name}` does not exist")
+    elif len(models) > 1:
+        raise ValueError("Multiple models with the same name found")
+    if version == "latest":
+        return models[0].latest_version
+    versions = api.models_store_list_model_versions(project_id=teamspace_id, model_id=models[0].id).versions
+    if not versions:
+        raise ValueError(f"Model `{name}` does not have any versions")
+    for ver in versions:
+        if ver.version == version:
+            return ver
+    raise ValueError(f"Model `{name}` does not have version `{version}`")
+
+
 def _download_model_files(
     client: LightningClient,
     teamspace_id: str,

lightning_sdk/cli/ai_hub.py

@@ -46,4 +46,4 @@ class _AIHub(_StudiosMenu):
           teamspace: Teamspace to deploy the API to. Defaults to user's default teamspace.
           org: Organization to deploy the API to. Defaults to user's default organization.
         """
-        return self._hub.deploy(api_id, cluster=cluster, name=name, teamspace=teamspace, org=org)
+        return self._hub.run(api_id, cluster=cluster, name=name, teamspace=teamspace, org=org)

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,80 @@
+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.
+            image_credentials: The credentials used to pull the image. Required if the image is private.
+                This should be the name of the respective credentials secret created on the Lightning AI platform.
+            cluster_auth: Whether to authenticate with the cluster to pull the image.
+                Required if the registry is part of a cluster provider (e.g. ECR).
+        """
+        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,
+        image_credentials: Optional[str] = None,
+        cluster_auth: 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,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+        )

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,27 @@ 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,
+        image_credentials: Optional[str] = None,
+        cluster_auth: 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 +74,39 @@ 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."
+                )
+
+            if image_credentials is not None:
+                raise ValueError("image_credentials is only supported when using a custom image")
+
+            if cluster_auth:
+                raise ValueError("cluster_auth is only supported when using a custom image")
+
+        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,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+        )
         return inst
 
     @abstractmethod
@@ -76,6 +118,9 @@ class _BaseJob(ABC):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
     ) -> None:
         """Submits a job and updates the internal _job attribute as well as the _name attribute."""
 
@@ -123,7 +168,3 @@ class _BaseJob(ABC):
     @property
     def teamspace(self) -> "Teamspace":
         return self._teamspace
-
-    @property
-    def cluster(self) -> Optional[str]:
-        return self._cluster