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

lightning_sdk/__init__.py

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

lightning_sdk/ai_hub.py

@@ -47,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,
@@ -129,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_xxxx")
+            deployment = hub.run("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)
+            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,
@@ -169,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/deployment_api.py

@@ -548,8 +548,6 @@ def to_spec(
         spot=spot,
         instance_name=_MACHINE_TO_COMPUTE_NAME[machine],
         readiness_probe=to_health_check(health_check),
-        skip_data_connections_setup=True,
-        skip_filesystem_setup=True,
     )
 
 

lightning_sdk/api/job_api.py

@@ -141,13 +141,17 @@ 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,
+        artifacts_local: Optional[str],
+        artifacts_remote: Optional[str],
     ) -> V1Job:
         env_vars = []
         if env is not None:
@@ -160,13 +164,17 @@ 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 "",
+            artifacts_source=artifacts_local or "",
+            artifacts_destination=artifacts_remote 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,
@@ -256,12 +260,14 @@ class TeamspaceApi:
         name: str,
         version: str,
         download_dir: Path,
-        teamspace_id: str,
+        teamspace_name: str,
+        teamspace_owner_name: str,
         progress_bar: bool = True,
     ) -> List[str]:
         return _download_model_files(
             client=self._client,
-            teamspace_id=teamspace_id,
+            teamspace_name=teamspace_name,
+            teamspace_owner_name=teamspace_owner_name,
             name=name,
             version=version,
             download_dir=download_dir,

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,9 +493,28 @@ 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,
+    teamspace_name: str,
+    teamspace_owner_name: str,
     name: str,
     version: str,
     download_dir: Path,
@@ -502,7 +522,9 @@ def _download_model_files(
     num_workers: int = 20,
 ) -> List[str]:
     api = ModelsStoreApi(client.api_client)
-    response = api.models_store_get_model_files(project_id=teamspace_id, name=name, version=version)
+    response = api.models_store_get_model_files(
+        project_name=teamspace_name, project_owner_name=teamspace_owner_name, name=name, version=version
+    )
 
     pbar = None
     if progress_bar:
@@ -522,7 +544,7 @@ def _download_model_files(
             client=client,
             model_id=response.model_id,
             version=response.version,
-            teamspace_id=teamspace_id,
+            teamspace_id=response.project_id,
             remote_path=filepath,
             file_path=str(local_file),
             num_workers=num_workers,

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/download.py

@@ -4,8 +4,8 @@ from pathlib import Path
 from typing import Optional
 
 from lightning_sdk.cli.exceptions import StudioCliError
-from lightning_sdk.cli.models import _get_teamspace, _parse_model_name
 from lightning_sdk.cli.studios_menu import _StudiosMenu
+from lightning_sdk.models import download_model
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
 
@@ -21,10 +21,8 @@ class _Downloads(_StudiosMenu):
             This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
           download_dir: The directory where the Model should be downloaded.
         """
-        org_name, teamspace_name, model_name = _parse_model_name(name)
-        teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
-        teamspace.download_model(
-            name=model_name,
+        download_model(
+            name=name,
             download_dir=download_dir,
             progress_bar=True,
         )

lightning_sdk/cli/run.py

@@ -36,7 +36,23 @@ class _Run:
                 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).
+            artifacts_local: The path of inside the docker container, you want to persist images from.
+                CAUTION: When setting this to "/", it will effectively erase your container.
+                Only supported for jobs with a docker image compute environment.
+            artifacts_remote: The remote storage to persist your artifacts to.
+                Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
+                PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
+                E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
+                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.
         """
+        # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
+        # might need to switch to explicit cli definition
         self.job.__func__.__doc__ = docstr
 
     # TODO: sadly, fire displays both Optional[type] and Union[type, None] as Optional[Optional]
@@ -55,6 +71,10 @@ class _Run:
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
     ) -> None:
         machine_enum = Machine(machine.upper())
         Job.run(
@@ -69,4 +89,8 @@ class _Run:
             cluster=cluster,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
         )

lightning_sdk/cli/upload.py

@@ -9,8 +9,8 @@ from tqdm import tqdm
 
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.exceptions import StudioCliError
-from lightning_sdk.cli.models import _get_teamspace, _parse_model_name
 from lightning_sdk.cli.studios_menu import _StudiosMenu
+from lightning_sdk.models import upload_model
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
 
@@ -20,7 +20,7 @@ class _Uploads(_StudiosMenu):
 
     _studio_upload_status_path = "~/.lightning/studios/uploads"
 
-    def model(self, name: str, path: Optional[str] = None, cloud_account: Optional[str] = None) -> None:
+    def model(self, name: str, path: str = ".", cloud_account: Optional[str] = None) -> None:
         """Upload a Model.
 
         Args:
@@ -29,14 +29,7 @@ class _Uploads(_StudiosMenu):
           path: The path to the file or directory you want to upload. Defaults to the current directory.
           cloud_account: The name of the cloud account to store the Model in.
         """
-        org_name, teamspace_name, model_name = _parse_model_name(name)
-        teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
-        teamspace.upload_model(
-            path=path or ".",
-            name=model_name,
-            progress_bar=True,
-            cluster_id=cloud_account,
-        )
+        upload_model(name, path, cloud_account=cloud_account)
 
     def _resolve_studio(self, studio: Optional[str]) -> Studio:
         user = _get_authed_user()

lightning_sdk/job/base.py

@@ -50,6 +50,10 @@ class _BaseJob(ABC):
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
     ) -> "_BaseJob":
         from lightning_sdk.studio import Studio
 
@@ -80,12 +84,35 @@ class _BaseJob(ABC):
                     "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")
+
+            if artifacts_local is not None or artifacts_remote is not None:
+                raise ValueError(
+                    "Specifying artifacts persistence is supported for docker images only. "
+                    "Other jobs will automatically persist artifacts to the teamspace distributed filesystem."
+                )
+
         else:
             if studio is not None:
                 raise RuntimeError(
                     "image and studio are mutually exclusive as both define the environment to run the job in"
                 )
 
+            # they either need to specified both or none of them
+            if bool(artifacts_local) != bool(artifacts_remote):
+                raise ValueError("Artifact persistence requires both artifacts_local and artifacts_remote to be set")
+
+            if artifacts_remote and len(artifacts_remote.split(":")) != 3:
+                raise ValueError(
+                    "Artifact persistence requires exactly three arguments separated by colon of kind "
+                    f"<CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>, got {artifacts_local}"
+                )
+
         inst = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False)
         inst._submit(
             machine=machine,
@@ -95,6 +122,10 @@ class _BaseJob(ABC):
             image=image,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
         )
         return inst
 
@@ -108,6 +139,10 @@ class _BaseJob(ABC):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
     ) -> None:
         """Submits a job and updates the internal _job attribute as well as the _name attribute."""
 

lightning_sdk/job/job.py

@@ -58,6 +58,10 @@ class Job(_BaseJob):
         cluster: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
     ) -> "Job":
         ret_val = super().run(
             name=name,
@@ -71,6 +75,10 @@ class Job(_BaseJob):
             cluster=cluster,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
         )
         # required for typing with "Job"
         assert isinstance(ret_val, cls)
@@ -85,8 +93,12 @@ class Job(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
     ) -> None:
-        return self._internal_job._submit(
+        self._job = self._internal_job._submit(
             machine=machine,
             cluster=cluster,
             command=command,
@@ -94,7 +106,12 @@ class Job(_BaseJob):
             image=image,
             env=env,
             interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
         )
+        return self
 
     def stop(self) -> None:
         return self._internal_job.stop()

lightning_sdk/job/v1.py

@@ -54,6 +54,8 @@ class _JobV1(_BaseJob):
             cluster=cluster,
             env=None,
             interruptible=interruptible,
+            image_credentials=None,
+            cluster_auth=False,
         )
 
     def _submit(
@@ -65,13 +67,20 @@ class _JobV1(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
     ) -> None:
         if studio is None:
             raise ValueError("Studio is required for submitting jobs")
 
-        if image is not None:
+        if image is not None or image_credentials is not None or cluster_auth:
             raise ValueError("Image is not supported for submitting jobs")
 
+        if artifacts_local is not None or artifacts_remote is not None:
+            raise ValueError("Specifying how to persist artifacts is not yet supported with jobs")
+
         if env is not None:
             raise ValueError("Environment variables are not supported for submitting jobs")
 

lightning_sdk/job/v2.py

@@ -34,6 +34,10 @@ class _JobV2(_BaseJob):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cluster: Optional[str] = None,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
     ) -> None:
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -62,6 +66,10 @@ class _JobV2(_BaseJob):
             machine=machine,
             interruptible=interruptible,
             env=env,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
         )
         self._job = submitted
         self._name = submitted.name
@@ -108,10 +116,18 @@ class _JobV2(_BaseJob):
 
     @property
     def artifact_path(self) -> Optional[str]:
+        if self._guaranteed_job.spec.image != "":
+            if self._guaranteed_job.spec.artifacts_destination != "":
+                splits = self._guaranteed_job.spec.artifacts_destination.split(":")
+                return f"/teamspace/{splits[0]}_connections/{splits[1]}/{splits[2]}"
+            return None
+
         return f"/teamspace/jobs/{self._guaranteed_job.name}/artifacts"
 
     @property
     def snapshot_path(self) -> Optional[str]:
+        if self._guaranteed_job.spec.image != "":
+            return None
         return f"/teamspace/jobs/{self._guaranteed_job.name}/snapshot"
 
     @property