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

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -58,9 +58,9 @@ class V1UserFeatures(object):
         'custom_app_domain': 'bool',
         'custom_instance_types': 'bool',
         'default_one_cluster': 'bool',
+        'deployment_alerts': 'bool',
         'deployment_customize_api': 'bool',
         'deployment_data_path': 'bool',
-        'deployment_gallery': 'bool',
         'deployment_persistent_disk': 'bool',
         'deployment_version_visibility': 'bool',
         'docs_agent': 'bool',
@@ -70,9 +70,9 @@ class V1UserFeatures(object):
         'enable_storage_limits': 'bool',
         'featured_studios_admin': 'bool',
         'filesystem_optimisation': 'bool',
-        'gcp': 'bool',
         'inference_job_deployment_plugin': 'bool',
         'instant_capacity_reservation': 'bool',
+        'job_artifacts_v2': 'bool',
         'jobs_init': 'bool',
         'jobs_v2': 'bool',
         'landing_studios': 'bool',
@@ -82,7 +82,6 @@ class V1UserFeatures(object):
         'mmt_fault_tolerance': 'bool',
         'mmt_strategy_selector': 'bool',
         'mmt_v2': 'bool',
-        'model_store': 'bool',
         'multiple_deployment_versions': 'bool',
         'multiple_studio_versions': 'bool',
         'org_level_member_permissions': 'bool',
@@ -108,6 +107,7 @@ class V1UserFeatures(object):
         'restart_ide_on_hang': 'bool',
         'restartable_jobs': 'bool',
         'runnable_public_studio_page': 'bool',
+        'security_docs': 'bool',
         'show_dev_admin': 'bool',
         'slurm': 'bool',
         'slurm_machine_selector': 'bool',
@@ -142,9 +142,9 @@ class V1UserFeatures(object):
         'custom_app_domain': 'customAppDomain',
         'custom_instance_types': 'customInstanceTypes',
         'default_one_cluster': 'defaultOneCluster',
+        'deployment_alerts': 'deploymentAlerts',
         'deployment_customize_api': 'deploymentCustomizeApi',
         'deployment_data_path': 'deploymentDataPath',
-        'deployment_gallery': 'deploymentGallery',
         'deployment_persistent_disk': 'deploymentPersistentDisk',
         'deployment_version_visibility': 'deploymentVersionVisibility',
         'docs_agent': 'docsAgent',
@@ -154,9 +154,9 @@ class V1UserFeatures(object):
         'enable_storage_limits': 'enableStorageLimits',
         'featured_studios_admin': 'featuredStudiosAdmin',
         'filesystem_optimisation': 'filesystemOptimisation',
-        'gcp': 'gcp',
         'inference_job_deployment_plugin': 'inferenceJobDeploymentPlugin',
         'instant_capacity_reservation': 'instantCapacityReservation',
+        'job_artifacts_v2': 'jobArtifactsV2',
         'jobs_init': 'jobsInit',
         'jobs_v2': 'jobsV2',
         'landing_studios': 'landingStudios',
@@ -166,7 +166,6 @@ class V1UserFeatures(object):
         'mmt_fault_tolerance': 'mmtFaultTolerance',
         'mmt_strategy_selector': 'mmtStrategySelector',
         'mmt_v2': 'mmtV2',
-        'model_store': 'modelStore',
         'multiple_deployment_versions': 'multipleDeploymentVersions',
         'multiple_studio_versions': 'multipleStudioVersions',
         'org_level_member_permissions': 'orgLevelMemberPermissions',
@@ -192,6 +191,7 @@ class V1UserFeatures(object):
         'restart_ide_on_hang': 'restartIdeOnHang',
         'restartable_jobs': 'restartableJobs',
         'runnable_public_studio_page': 'runnablePublicStudioPage',
+        'security_docs': 'securityDocs',
         'show_dev_admin': 'showDevAdmin',
         'slurm': 'slurm',
         'slurm_machine_selector': 'slurmMachineSelector',
@@ -208,7 +208,7 @@ class V1UserFeatures(object):
         'writable_data_connections': 'writableDataConnections'
     }
 
-    def __init__(self, advanced_deployment_autoscaling: 'bool' =None, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, default_one_cluster: 'bool' =None, deployment_customize_api: 'bool' =None, deployment_data_path: 'bool' =None, deployment_gallery: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_version_visibility: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_crypto_crackdown: 'bool' =None, enable_efs: 'bool' =None, enable_storage_limits: 'bool' =None, featured_studios_admin: 'bool' =None, filesystem_optimisation: 'bool' =None, gcp: 'bool' =None, inference_job_deployment_plugin: 'bool' =None, instant_capacity_reservation: 'bool' =None, jobs_init: 'bool' =None, jobs_v2: 'bool' =None, landing_studios: 'bool' =None, lightning_registry: 'bool' =None, lit_logger: 'bool' =None, lit_logger_storage_v2: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, mmt_v2: 'bool' =None, model_store: 'bool' =None, multiple_deployment_versions: 'bool' =None, multiple_studio_versions: 'bool' =None, org_level_member_permissions: 'bool' =None, pipelines: 'bool' =None, plugin_biz_chat: 'bool' =None, plugin_distributed: 'bool' =None, plugin_fiftyone: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_lightning_apps: 'bool' =None, plugin_lightning_apps_distributed: 'bool' =None, plugin_mage_ai: 'bool' =None, plugin_milvus: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_react: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, plugin_weviate: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, restart_ide_on_hang: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, show_dev_admin: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, snapshotter_service: 'bool' =None, snowflake_connection: 'bool' =None, spot_v2: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, writable_data_connections: 'bool' =None):  # noqa: E501
+    def __init__(self, advanced_deployment_autoscaling: 'bool' =None, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, custom_instance_types: 'bool' =None, default_one_cluster: 'bool' =None, deployment_alerts: 'bool' =None, deployment_customize_api: 'bool' =None, deployment_data_path: 'bool' =None, deployment_persistent_disk: 'bool' =None, deployment_version_visibility: 'bool' =None, docs_agent: 'bool' =None, drive_v2: 'bool' =None, enable_crypto_crackdown: 'bool' =None, enable_efs: 'bool' =None, enable_storage_limits: 'bool' =None, featured_studios_admin: 'bool' =None, filesystem_optimisation: 'bool' =None, inference_job_deployment_plugin: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, jobs_init: 'bool' =None, jobs_v2: 'bool' =None, landing_studios: 'bool' =None, lightning_registry: 'bool' =None, lit_logger: 'bool' =None, lit_logger_storage_v2: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, mmt_v2: 'bool' =None, multiple_deployment_versions: 'bool' =None, multiple_studio_versions: 'bool' =None, org_level_member_permissions: 'bool' =None, pipelines: 'bool' =None, plugin_biz_chat: 'bool' =None, plugin_distributed: 'bool' =None, plugin_fiftyone: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_lightning_apps: 'bool' =None, plugin_lightning_apps_distributed: 'bool' =None, plugin_mage_ai: 'bool' =None, plugin_milvus: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_react: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, plugin_weviate: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, project_selector: 'bool' =None, restart_ide_on_hang: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, slurm: 'bool' =None, slurm_machine_selector: 'bool' =None, snapshotter_service: 'bool' =None, snowflake_connection: 'bool' =None, spot_v2: 'bool' =None, studio_config: 'bool' =None, studio_on_stop: 'bool' =None, studio_version_visibility: 'bool' =None, teamspace_storage_tab: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, writable_data_connections: 'bool' =None):  # noqa: E501
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
         self._advanced_deployment_autoscaling = None
         self._affiliate_links = None
@@ -227,9 +227,9 @@ class V1UserFeatures(object):
         self._custom_app_domain = None
         self._custom_instance_types = None
         self._default_one_cluster = None
+        self._deployment_alerts = None
         self._deployment_customize_api = None
         self._deployment_data_path = None
-        self._deployment_gallery = None
         self._deployment_persistent_disk = None
         self._deployment_version_visibility = None
         self._docs_agent = None
@@ -239,9 +239,9 @@ class V1UserFeatures(object):
         self._enable_storage_limits = None
         self._featured_studios_admin = None
         self._filesystem_optimisation = None
-        self._gcp = None
         self._inference_job_deployment_plugin = None
         self._instant_capacity_reservation = None
+        self._job_artifacts_v2 = None
         self._jobs_init = None
         self._jobs_v2 = None
         self._landing_studios = None
@@ -251,7 +251,6 @@ class V1UserFeatures(object):
         self._mmt_fault_tolerance = None
         self._mmt_strategy_selector = None
         self._mmt_v2 = None
-        self._model_store = None
         self._multiple_deployment_versions = None
         self._multiple_studio_versions = None
         self._org_level_member_permissions = None
@@ -277,6 +276,7 @@ class V1UserFeatures(object):
         self._restart_ide_on_hang = None
         self._restartable_jobs = None
         self._runnable_public_studio_page = None
+        self._security_docs = None
         self._show_dev_admin = None
         self._slurm = None
         self._slurm_machine_selector = None
@@ -326,12 +326,12 @@ class V1UserFeatures(object):
             self.custom_instance_types = custom_instance_types
         if default_one_cluster is not None:
             self.default_one_cluster = default_one_cluster
+        if deployment_alerts is not None:
+            self.deployment_alerts = deployment_alerts
         if deployment_customize_api is not None:
             self.deployment_customize_api = deployment_customize_api
         if deployment_data_path is not None:
             self.deployment_data_path = deployment_data_path
-        if deployment_gallery is not None:
-            self.deployment_gallery = deployment_gallery
         if deployment_persistent_disk is not None:
             self.deployment_persistent_disk = deployment_persistent_disk
         if deployment_version_visibility is not None:
@@ -350,12 +350,12 @@ class V1UserFeatures(object):
             self.featured_studios_admin = featured_studios_admin
         if filesystem_optimisation is not None:
             self.filesystem_optimisation = filesystem_optimisation
-        if gcp is not None:
-            self.gcp = gcp
         if inference_job_deployment_plugin is not None:
             self.inference_job_deployment_plugin = inference_job_deployment_plugin
         if instant_capacity_reservation is not None:
             self.instant_capacity_reservation = instant_capacity_reservation
+        if job_artifacts_v2 is not None:
+            self.job_artifacts_v2 = job_artifacts_v2
         if jobs_init is not None:
             self.jobs_init = jobs_init
         if jobs_v2 is not None:
@@ -374,8 +374,6 @@ class V1UserFeatures(object):
             self.mmt_strategy_selector = mmt_strategy_selector
         if mmt_v2 is not None:
             self.mmt_v2 = mmt_v2
-        if model_store is not None:
-            self.model_store = model_store
         if multiple_deployment_versions is not None:
             self.multiple_deployment_versions = multiple_deployment_versions
         if multiple_studio_versions is not None:
@@ -426,6 +424,8 @@ class V1UserFeatures(object):
             self.restartable_jobs = restartable_jobs
         if runnable_public_studio_page is not None:
             self.runnable_public_studio_page = runnable_public_studio_page
+        if security_docs is not None:
+            self.security_docs = security_docs
         if show_dev_admin is not None:
             self.show_dev_admin = show_dev_admin
         if slurm is not None:
@@ -812,6 +812,27 @@ class V1UserFeatures(object):
 
         self._default_one_cluster = default_one_cluster
 
+    @property
+    def deployment_alerts(self) -> 'bool':
+        """Gets the deployment_alerts of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The deployment_alerts of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._deployment_alerts
+
+    @deployment_alerts.setter
+    def deployment_alerts(self, deployment_alerts: 'bool'):
+        """Sets the deployment_alerts of this V1UserFeatures.
+
+
+        :param deployment_alerts: The deployment_alerts of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._deployment_alerts = deployment_alerts
+
     @property
     def deployment_customize_api(self) -> 'bool':
         """Gets the deployment_customize_api of this V1UserFeatures.  # noqa: E501
@@ -854,27 +875,6 @@ class V1UserFeatures(object):
 
         self._deployment_data_path = deployment_data_path
 
-    @property
-    def deployment_gallery(self) -> 'bool':
-        """Gets the deployment_gallery of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The deployment_gallery of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._deployment_gallery
-
-    @deployment_gallery.setter
-    def deployment_gallery(self, deployment_gallery: 'bool'):
-        """Sets the deployment_gallery of this V1UserFeatures.
-
-
-        :param deployment_gallery: The deployment_gallery of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._deployment_gallery = deployment_gallery
-
     @property
     def deployment_persistent_disk(self) -> 'bool':
         """Gets the deployment_persistent_disk of this V1UserFeatures.  # noqa: E501
@@ -1064,27 +1064,6 @@ class V1UserFeatures(object):
 
         self._filesystem_optimisation = filesystem_optimisation
 
-    @property
-    def gcp(self) -> 'bool':
-        """Gets the gcp of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The gcp of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._gcp
-
-    @gcp.setter
-    def gcp(self, gcp: 'bool'):
-        """Sets the gcp of this V1UserFeatures.
-
-
-        :param gcp: The gcp of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._gcp = gcp
-
     @property
     def inference_job_deployment_plugin(self) -> 'bool':
         """Gets the inference_job_deployment_plugin of this V1UserFeatures.  # noqa: E501
@@ -1127,6 +1106,27 @@ class V1UserFeatures(object):
 
         self._instant_capacity_reservation = instant_capacity_reservation
 
+    @property
+    def job_artifacts_v2(self) -> 'bool':
+        """Gets the job_artifacts_v2 of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The job_artifacts_v2 of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._job_artifacts_v2
+
+    @job_artifacts_v2.setter
+    def job_artifacts_v2(self, job_artifacts_v2: 'bool'):
+        """Sets the job_artifacts_v2 of this V1UserFeatures.
+
+
+        :param job_artifacts_v2: The job_artifacts_v2 of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._job_artifacts_v2 = job_artifacts_v2
+
     @property
     def jobs_init(self) -> 'bool':
         """Gets the jobs_init of this V1UserFeatures.  # noqa: E501
@@ -1316,27 +1316,6 @@ class V1UserFeatures(object):
 
         self._mmt_v2 = mmt_v2
 
-    @property
-    def model_store(self) -> 'bool':
-        """Gets the model_store of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The model_store of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._model_store
-
-    @model_store.setter
-    def model_store(self, model_store: 'bool'):
-        """Sets the model_store of this V1UserFeatures.
-
-
-        :param model_store: The model_store of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._model_store = model_store
-
     @property
     def multiple_deployment_versions(self) -> 'bool':
         """Gets the multiple_deployment_versions of this V1UserFeatures.  # noqa: E501
@@ -1862,6 +1841,27 @@ class V1UserFeatures(object):
 
         self._runnable_public_studio_page = runnable_public_studio_page
 
+    @property
+    def security_docs(self) -> 'bool':
+        """Gets the security_docs of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The security_docs of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._security_docs
+
+    @security_docs.setter
+    def security_docs(self, security_docs: 'bool'):
+        """Sets the security_docs of this V1UserFeatures.
+
+
+        :param security_docs: The security_docs of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._security_docs = security_docs
+
     @property
     def show_dev_admin(self) -> 'bool':
         """Gets the show_dev_admin of this V1UserFeatures.  # noqa: E501

lightning_sdk/lit_container.py

@@ -76,3 +76,22 @@ class LitContainer:
         resp = self._api.upload_container(container, teamspace, tag)
         for line in resp:
             print(line)
+
+    def download_container(
+        self, container: str, teamspace: str, org: Optional[str] = None, user: Optional[str] = None, tag: str = "latest"
+    ) -> None:
+        """Download a container from the docker registry.
+
+        Args:
+            container: The name of the container to download.
+            teamspace: The teamspace which contains the container.
+            org: The organization which contains the container.
+            user: The user which contains the container.
+            tag: The tag to use for the container.
+        """
+        try:
+            teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        except Exception as e:
+            raise ValueError(f"Could not resolve teamspace: {e}") from e
+
+        return self._api.download_container(container, teamspace, tag)

lightning_sdk/mmt/base.py

@@ -1,3 +1,4 @@
+import warnings
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Tuple, Union
 
@@ -11,7 +12,7 @@ if TYPE_CHECKING:
     from lightning_sdk.user import User
 
 from lightning_sdk.job.base import _BaseJob
-from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, in_studio
 
 
 class MMTMachine(Protocol):
@@ -23,7 +24,7 @@ class MMTMachine(Protocol):
         ...
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """The actual machine type this node is running on."""
         ...
 
@@ -54,7 +55,7 @@ class _BaseMMT(_BaseJob):
     def run(
         cls,
         name: str,
-        machine: "Machine",
+        machine: Union["Machine", str],
         num_machines: int,
         command: Optional[str] = None,
         studio: Union["Studio", str, None] = None,
@@ -67,9 +68,10 @@ class _BaseMMT(_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 favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseMMT":
         """Run async workloads using a docker image across multiple machines.
@@ -94,21 +96,21 @@ class _BaseMMT(_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.
         """
+        from lightning_sdk.lightning_cloud.openapi.rest import ApiException
         from lightning_sdk.studio import Studio
 
         cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
@@ -167,6 +169,13 @@ class _BaseMMT(_BaseJob):
                     "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):
                 raise ValueError("Artifact persistence requires both artifacts_local and artifacts_remote to be set")
@@ -189,9 +198,10 @@ class _BaseMMT(_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 inst
 
@@ -199,7 +209,7 @@ class _BaseMMT(_BaseJob):
     def _submit(
         self,
         num_machines: int,
-        machine: "Machine",
+        machine: Union["Machine", str],
         command: Optional[str] = None,
         studio: Optional["Studio"] = None,
         image: Optional[str] = None,
@@ -208,9 +218,10 @@ class _BaseMMT(_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 favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> None:
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -230,19 +241,18 @@ class _BaseMMT(_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.
         """
 
     @property
@@ -257,7 +267,7 @@ class _BaseMMT(_BaseJob):
 
     @property
     @abstractmethod
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
 
     @abstractmethod

lightning_sdk/mmt/mmt.py

@@ -97,7 +97,7 @@ class MMT(_BaseMMT):
         cls,
         name: str,
         num_machines: int,
-        machine: "Machine",
+        machine: Union["Machine", str],
         command: Optional[str] = None,
         studio: Union["Studio", str, None] = None,
         image: Union[str, None] = None,
@@ -109,9 +109,10 @@ class MMT(_BaseMMT):
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-        entrypoint: str = "sh -c",
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "MMT":
         """Run async workloads using a docker image across multiple machines.
@@ -136,19 +137,18 @@ class MMT(_BaseMMT):
                 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.
         """
         ret_val = super().run(
             name=name,
@@ -165,6 +165,8 @@ class MMT(_BaseMMT):
             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,
             cluster=cluster,  # deprecated in favor of cloud_account
@@ -180,7 +182,7 @@ class MMT(_BaseMMT):
     def _submit(
         self,
         num_machines: int,
-        machine: "Machine",
+        machine: Union["Machine", str],
         command: Optional[str] = None,
         studio: Optional["Studio"] = None,
         image: Optional[str] = None,
@@ -189,9 +191,10 @@ class MMT(_BaseMMT):
         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 favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> "MMT":
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -211,20 +214,19 @@ class MMT(_BaseMMT):
                 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.
         """
         self._job = self._internal_mmt._submit(
             num_machines=num_machines,
@@ -237,6 +239,8 @@ class MMT(_BaseMMT):
             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,
         )
@@ -264,7 +268,7 @@ class MMT(_BaseMMT):
         return self._internal_mmt.machines
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
         return self._internal_mmt.machine