lightning-sdk 0.1.50__py3-none-any.whl → 0.1.51__py3-none-any.whl

lightning_sdk/__init__.py

@@ -29,5 +29,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.50"
+__version__ = "0.1.51"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/api/job_api.py

@@ -7,10 +7,9 @@ from lightning_sdk.api.utils import (
     _create_app,
     _machine_to_compute_name,
     remove_datetime_prefix,
+    resolve_path_mappings,
 )
-from lightning_sdk.api.utils import (
-    _get_cloud_url as _cloud_url,
-)
+from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
 from lightning_sdk.lightning_cloud.openapi import (
     AppinstancesIdBody,
@@ -214,9 +213,10 @@ class JobApiV2:
         env: Optional[Dict[str, str]],
         image_credentials: Optional[str],
         cloud_account_auth: bool,
-        artifacts_local: Optional[str],
-        artifacts_remote: Optional[str],
         entrypoint: str,
+        path_mappings: Optional[Dict[str, str]],
+        artifacts_local: Optional[str],  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
     ) -> V1Job:
         env_vars = []
         if env is not None:
@@ -227,6 +227,12 @@ class JobApiV2:
 
         run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if studio_id is not None else ""
 
+        path_mappings_list = resolve_path_mappings(
+            mappings=path_mappings or {},
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
@@ -239,8 +245,7 @@ class JobApiV2:
             spot=interruptible,
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
-            artifacts_source=artifacts_local or "",
-            artifacts_destination=artifacts_remote or "",
+            path_mappings=path_mappings_list,
         )
         body = ProjectIdJobsBody(name=name, spec=spec)
 

lightning_sdk/api/lit_container_api.py

@@ -1,23 +1,36 @@
 from typing import Generator, List
 
+import docker
+
 from lightning_sdk.api.utils import _get_registry_url
+from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
 from lightning_sdk.lightning_cloud.openapi.models import V1DeleteLitRepositoryResponse
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.teamspace import Teamspace
 
 
+class LCRAuthFailedError(Exception):
+    def __init__(self) -> None:
+        super().__init__("Failed to authenticate with Lightning Container Registry")
+
+
 class LitContainerApi:
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=3)
 
-        import docker
-
         try:
             self._docker_client = docker.from_env()
             self._docker_client.ping()
         except docker.errors.DockerException as e:
             raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
 
+    def authenticate(self) -> bool:
+        authed_user = self._client.auth_service_get_user()
+        username = authed_user.username
+        api_key = authed_user.api_key
+        resp = self._docker_client.login(username, password=api_key, registry=_get_registry_url())
+        return resp["Status"] == "Login Succeeded"
+
     def list_containers(self, project_id: str) -> List:
         project = self._client.lit_registry_service_get_lit_project_registry(project_id)
         return project.repositories
@@ -29,8 +42,6 @@ class LitContainerApi:
             raise ValueError(f"Could not delete container {container} from project {project_id}") from ex
 
     def upload_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
-        import docker
-
         try:
             self._docker_client.images.get(container)
         except docker.errors.ImageNotFound:
@@ -41,11 +52,17 @@ class LitContainerApi:
         tagged = self._docker_client.api.tag(container, repository, tag)
         if not tagged:
             raise ValueError(f"Could not tag container {container} with {repository}:{tag}")
-        return self._docker_client.api.push(repository, stream=True, decode=True)
+        lines = self._docker_client.api.push(repository, stream=True, decode=True)
+        for line in lines:
+            if "errorDetail" in line and "authorization failed" in line["error"]:
+                raise LCRAuthFailedError()
+            yield line
+        yield {
+            "finish": True,
+            "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container}",
+        }
 
     def download_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
-        import docker
-
         registry_url = _get_registry_url()
         repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
         try:

lightning_sdk/api/mmt_api.py

@@ -7,10 +7,9 @@ from lightning_sdk.api.utils import (
     _COMPUTE_NAME_TO_MACHINE,
     _create_app,
     _machine_to_compute_name,
+    resolve_path_mappings,
 )
-from lightning_sdk.api.utils import (
-    _get_cloud_url as _cloud_url,
-)
+from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
 from lightning_sdk.lightning_cloud.openapi import (
     Externalv1LightningappInstance,
@@ -85,9 +84,10 @@ class MMTApiV2:
         env: Optional[Dict[str, str]],
         image_credentials: Optional[str],
         cloud_account_auth: bool,
-        artifacts_local: Optional[str],
-        artifacts_remote: Optional[str],
         entrypoint: str,
+        path_mappings: Optional[Dict[str, str]],
+        artifacts_local: Optional[str],  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
     ) -> V1MultiMachineJob:
         env_vars = []
         if env is not None:
@@ -98,6 +98,12 @@ class MMTApiV2:
 
         run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if studio_id is not None else ""
 
+        path_mappings_list = resolve_path_mappings(
+            mappings=path_mappings or {},
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
@@ -110,8 +116,7 @@ class MMTApiV2:
             spot=interruptible,
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
-            artifacts_source=artifacts_local or "",
-            artifacts_destination=artifacts_remote or "",
+            path_mappings=path_mappings_list,
         )
         body = ProjectIdMultimachinejobsBody(
             name=name, spec=spec, cluster_id=cloud_account or "", machines=num_machines

lightning_sdk/api/utils.py

@@ -24,6 +24,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     UploadsUploadIdBody,
     V1CompletedPart,
     V1CompleteUpload,
+    V1PathMapping,
     V1PresignedUrl,
     V1SignedUrl,
     V1UploadProjectArtifactPartsResponse,
@@ -614,3 +615,54 @@ def remove_datetime_prefix(text: str) -> str:
     # lines looks something like
     # '[2025-01-08T14:15:03.797142418Z] ⚡  ~ echo Hello\n[2025-01-08T14:15:03.803077717Z] Hello\n'
     return re.sub(r"^\[.*?\] ", "", text, flags=re.MULTILINE)
+
+
+def resolve_path_mappings(
+    mappings: Dict[str, str],
+    artifacts_local: Optional[str],
+    artifacts_remote: Optional[str],
+) -> List[V1PathMapping]:
+    path_mappings_list = []
+    for k, v in mappings.items():
+        splitted = str(v).rsplit(":", 1)
+        connection_name: str
+        connection_path: str
+        if len(splitted) == 1:
+            connection_name = splitted[0]
+            connection_path = ""
+        else:
+            connection_name, connection_path = splitted
+
+        path_mappings_list.append(
+            V1PathMapping(
+                connection_name=connection_name,
+                connection_path=connection_path,
+                container_path=k,
+            )
+        )
+
+    if artifacts_remote:
+        splitted = str(artifacts_remote).rsplit(":", 2)
+        if len(splitted) not in (2, 3):
+            raise RuntimeError(
+                f"Artifacts remote need to be of format efs:connection_name[:path] but got {artifacts_remote}"
+            )
+        else:
+            if not artifacts_local:
+                raise RuntimeError("If Artifacts remote is specified, artifacts local should be specified as well")
+
+            if len(splitted) == 2:
+                _, connection_name = splitted
+                connection_path = ""
+            else:
+                _, connection_name, connection_path = splitted
+
+            path_mappings_list.append(
+                V1PathMapping(
+                    connection_name=connection_name,
+                    connection_path=connection_path,
+                    container_path=artifacts_local,
+                )
+            )
+
+    return path_mappings_list

lightning_sdk/cli/run.py

@@ -43,20 +43,28 @@ class _Run:
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
                 Required if the registry is part of a cloud provider (e.g. ECR).
-            artifacts_local: The path of inside the docker container, you want to persist images from.
+            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.
+            path_mappings: Maps path inside of containers to paths inside data-connections.
+                Should be a comma separated list of form:
+                <MAPPING_1>,<MAPPING_2>,...
+                where each mapping is of the form
+                <CONTAINER_PATH_1>:<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1> and
+                omitting the path inside the connection defaults to the connections root.
+            artifacts_local: Deprecated in favor of path_mappings.
+                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.
+            artifacts_remote: Deprecated in favor of path_mappings.
+                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.
-            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
-                just runs the provided command in a standard shell.
-                To use the pre-defined entrypoint of the provided image, set this to an empty string.
-                Only applicable when submitting docker jobs.
         """
         # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
         # might need to switch to explicit cli definition
@@ -87,20 +95,28 @@ class _Run:
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
                 Required if the registry is part of a cloud provider (e.g. ECR).
-            artifacts_local: The path of inside the docker container, you want to persist images from.
+            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.
+            path_mappings: Maps path inside of containers to paths inside data-connections.
+                Should be a comma separated list of form:
+                <MAPPING_1>,<MAPPING_2>,...
+                where each mapping is of the form
+                <CONTAINER_PATH_1>:<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1> and
+                omitting the path inside the connection defaults to the connections root.
+            artifacts_local: Deprecated in favor of path_mappings.
+                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.
+            artifacts_remote: Deprecated in favor of path_mappings.
+                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.
-            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
-                just runs the provided command in a standard shell.
-                To use the pre-defined entrypoint of the provided image, set this to an empty string.
-                Only applicable when submitting docker jobs.
         """
         # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
         # might need to switch to explicit cli definition
@@ -124,14 +140,15 @@ class _Run:
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: str = "",
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
     ) -> None:
         if not name:
             from datetime import datetime
 
-            timestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+            timestr = datetime.now().strftime("%b-%d-%H_%M")
             name = f"job-{timestr}"
 
         if machine is None:
@@ -149,6 +166,17 @@ class _Run:
             cloud_account = resolved_teamspace.default_cloud_account
         machine_enum = Machine(machine.upper())
 
+        path_mappings_dict = {}
+        for mapping in path_mappings.split(","):
+            splits = str(mapping).split(":", 1)
+            if len(splits) != 2:
+                raise RuntimeError(
+                    "Mapping needs to be of form <CONTAINER_PATH>:<CONNECTION_NAME>[:<PATH_WITHIN_CONNECTION>], "
+                    f"but got {mapping}"
+                )
+
+            path_mappings_dict[splits[0]] = splits[1]
+
         Job.run(
             name=name,
             machine=machine_enum,
@@ -163,9 +191,10 @@ class _Run:
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings_dict,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )
 
     # TODO: sadly, fire displays both Optional[type] and Union[type, None] as Optional[Optional]
@@ -186,14 +215,15 @@ class _Run:
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: str = "",
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
     ) -> None:
         if name is None:
             from datetime import datetime
 
-            timestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+            timestr = datetime.now().strftime("%b-%d-%H_%M")
             name = f"mmt-{timestr}"
 
         if machine is None:
@@ -212,6 +242,17 @@ class _Run:
         if image is None:
             raise RuntimeError("Image needs to be specified to run a multi-machine job")
 
+        path_mappings_dict = {}
+        for mapping in path_mappings.split(","):
+            splits = str(mapping).split(":", 1)
+            if len(splits) != 2:
+                raise RuntimeError(
+                    "Mapping needs to be of form <CONTAINER_PATH>:<CONNECTION_NAME>[:<PATH_WITHIN_CONNECTION>], "
+                    f"but got {mapping}"
+                )
+
+            path_mappings_dict[splits[0]] = splits[1]
+
         MMT.run(
             name=name,
             num_machines=num_machines,
@@ -227,7 +268,8 @@ class _Run:
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings_dict,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )

lightning_sdk/cli/serve.py

@@ -4,6 +4,7 @@ import warnings
 from pathlib import Path
 from typing import Optional, Union
 
+import docker
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.prompt import Confirm
@@ -89,11 +90,6 @@ class _LitServe:
         tag: str = "litserve-model",
         non_interactive: bool = False,
     ) -> None:
-        try:
-            import docker
-        except ImportError:
-            raise ImportError("docker-py is not installed. Please install it with `pip install docker`") from None
-
         try:
             client = docker.from_env()
             client.ping()

lightning_sdk/cli/upload.py

@@ -2,14 +2,15 @@ import concurrent.futures
 import json
 import os
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, Generator, List, Optional
 
+import rich
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from simple_term_menu import TerminalMenu
 from tqdm import tqdm
 
-from lightning_sdk.api.lit_container_api import LitContainerApi
+from lightning_sdk.api.lit_container_api import LCRAuthFailedError, LitContainerApi
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.exceptions import StudioCliError
 from lightning_sdk.cli.studios_menu import _StudiosMenu
@@ -162,21 +163,38 @@ class _Uploads(_StudiosMenu, _TeamspacesMenu):
             transient=False,
         ) as progress:
             push_task = progress.add_task("Pushing Docker image", total=None)
-            resp = api.upload_container(container, teamspace, tag)
-            for line in resp:
-                if "status" in line:
-                    console.print(line["status"], style="bright_black")
-                    progress.update(push_task, description="Pushing Docker image")
-                elif "aux" in line:
-                    console.print(line["aux"], style="bright_black")
-                elif "error" in line:
-                    progress.stop()
-                    console.print(f"\n[red]{line}[/red]")
-                    return
-                else:
-                    console.print(line, style="bright_black")
+            try:
+                lines = api.upload_container(container, teamspace, tag)
+                self._print_docker_push(lines, console, progress, push_task)
+            except LCRAuthFailedError:
+                console.print("Authenticating with Lightning Container Registry...")
+                if not api.authenticate():
+                    raise StudioCliError("Failed to authenticate with Lightning Container Registry") from None
+                console.print("Authenticated with Lightning Container Registry", style="green")
+                lines = api.upload_container(container, teamspace, tag)
+                self._print_docker_push(lines, console, progress, push_task)
             progress.update(push_task, description="[green]Container pushed![/green]")
 
+    @staticmethod
+    def _print_docker_push(
+        lines: Generator, console: Console, progress: Progress, push_task: rich.progress.TaskID
+    ) -> None:
+        for line in lines:
+            if "status" in line:
+                console.print(line["status"], style="bright_black")
+                progress.update(push_task, description="Pushing Docker image")
+            elif "aux" in line:
+                console.print(line["aux"], style="bright_black")
+            elif "error" in line:
+                progress.stop()
+                console.print(f"\n[red]{line}[/red]")
+                return
+            elif "finish" in line:
+                console.print(f"Container available at [i]{line['url']}[/i]")
+                return
+            else:
+                console.print(line, style="bright_black")
+
     def _start_parallel_upload(
         self, executor: concurrent.futures.ThreadPoolExecutor, studio: Studio, upload_state: Dict[str, str]
     ) -> List[concurrent.futures.Future]:

lightning_sdk/helpers.py

@@ -43,7 +43,7 @@ def _check_version_and_prompt_upgrade(curr_version: str) -> None:
         warnings.warn(
             f"A newer version of {__package_name__} is available ({new_version}). "
             f"Please consider upgrading with `pip install -U {__package_name__}`. "
-            "Not all functionalities of the platform can be guaranteed to work with the current version.",
+            "Not all platform functionality can be guaranteed to work with the current version.",
             UserWarning,
         )
     return

lightning_sdk/job/base.py

@@ -1,8 +1,9 @@
+import warnings
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union
 
 from lightning_sdk.api.utils import _get_cloud_url
-from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, in_studio
 
 if TYPE_CHECKING:
     from lightning_sdk.machine import Machine
@@ -83,6 +84,7 @@ class _BaseJob(ABC):
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
         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.
@@ -121,6 +123,7 @@ class _BaseJob(ABC):
                 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.lightning_cloud.openapi.rest import ApiException
         from lightning_sdk.studio import Studio
 
         cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
@@ -175,6 +178,12 @@ class _BaseJob(ABC):
                 raise RuntimeError(
                     "image and studio are mutually exclusive as both define the environment to run the job in"
                 )
+            if cloud_account is None and in_studio():
+                try:
+                    resolve_studio = Studio(teamspace=teamspace, user=user, org=org)
+                    cloud_account = resolve_studio.cloud_account
+                except (ValueError, ApiException):
+                    warnings.warn("Could not infer cloud account from studio. Using teamspace default.")
 
             # they either need to specified both or none of them
             if bool(artifacts_local) != bool(artifacts_remote):
@@ -200,6 +209,7 @@ class _BaseJob(ABC):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             entrypoint=entrypoint,
+            path_mappings=path_mappings,
         )
 
     @abstractmethod
@@ -217,6 +227,7 @@ class _BaseJob(ABC):
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
     ) -> "_BaseJob":
         """Submit a new job to the Lightning AI platform.
 

lightning_sdk/job/job.py

@@ -101,9 +101,10 @@ class Job(_BaseJob):
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        artifacts_local: Optional[str] = None,
-        artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        artifacts_local: Optional[str] = None,  # deprecated in terms of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in terms of path_mappings
         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.
@@ -127,20 +128,19 @@ class Job(_BaseJob):
             This should be the name of the respective credentials secret created on the Lightning AI platform.
         cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
             Required if the registry is part of a cloud 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.
         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.
+        path_mappings: Dictionary of path mappings. The keys are the path inside the container whereas the value
+            represents the data-connection name and the path inside that connection.
+            Should be of form
+                {
+                    "<CONTAINER_PATH_1>": "<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1>",
+                    "<CONTAINER_PATH_2>": "<CONNECTION_NAME_2>"
+                }
+            If the path inside the connection is omitted it's assumed to be the root path of that connection.
+            Only applicable when submitting docker jobs.
         """
         ret_val = super().run(
             name=name,
@@ -159,6 +159,7 @@ class Job(_BaseJob):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             entrypoint=entrypoint,
+            path_mappings=path_mappings,
             cluster=cluster,
         )
         # required for typing with "Job"
@@ -178,9 +179,10 @@ class Job(_BaseJob):
         cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        artifacts_local: Optional[str] = None,
-        artifacts_remote: Optional[str] = None,
         entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        artifacts_local: Optional[str] = None,  # deprecated in terms of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in terms of path_mappings
     ) -> "Job":
         """Submit a new job to the Lightning AI platform.
 
@@ -199,19 +201,18 @@ class Job(_BaseJob):
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
                 Required if the registry is part of a cloud 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.
             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.
+            path_mappings: Dictionary of path mappings. The keys are the path inside the container whereas the value
+                represents the data-connection name and the path inside that connection.
+                Should be of form
+                    {
+                        "<CONTAINER_PATH_1>": "<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1>",
+                        "<CONTAINER_PATH_2>": "<CONNECTION_NAME_2>"
+                    }
+                If the path inside the connection is omitted it's assumed to be the root path of that connection.
+                Only applicable when submitting docker jobs.
         """
         self._job = self._internal_job._submit(
             machine=machine,
@@ -223,9 +224,10 @@ class Job(_BaseJob):
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
         )
         return self