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

lightning_sdk/job/work.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Optional, Protocol
+from typing import TYPE_CHECKING, Any, Optional, Protocol, Union
 
 from lightning_sdk.api.job_api import JobApiV1
 
@@ -51,7 +51,7 @@ class Work:
         return self._job._name_filter(self._guaranteed_work.name)
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         return self._job_api.get_machine_from_work(self._guaranteed_work)
 
     @property

lightning_sdk/lightning_cloud/openapi/models/v1_cluster_spec.py

@@ -42,7 +42,6 @@ class V1ClusterSpec(object):
     """
     swagger_types = {
         'auth_token': 'str',
-        'available_instance_types': 'list[str]',
         'aws_v1': 'V1AWSDirectV1',
         'cluster_type': 'V1ClusterType',
         'deletion_options': 'V1ClusterDeletionOptions',
@@ -63,7 +62,6 @@ class V1ClusterSpec(object):
 
     attribute_map = {
         'auth_token': 'authToken',
-        'available_instance_types': 'availableInstanceTypes',
         'aws_v1': 'awsV1',
         'cluster_type': 'clusterType',
         'deletion_options': 'deletionOptions',
@@ -82,10 +80,9 @@ class V1ClusterSpec(object):
         'vultr_v1': 'vultrV1'
     }
 
-    def __init__(self, auth_token: 'str' =None, available_instance_types: 'list[str]' =None, aws_v1: 'V1AWSDirectV1' =None, cluster_type: 'V1ClusterType' =None, deletion_options: 'V1ClusterDeletionOptions' =None, desired_state: 'V1ClusterState' =None, domain: 'str' =None, freeze_accelerators: 'bool' =None, google_cloud_v1: 'V1GoogleCloudDirectV1' =None, insurer_disabled: 'bool' =None, lambda_labs_v1: 'V1LambdaLabsDirectV1' =None, overprovisioning: 'list[V1InstanceOverprovisioningSpec]' =None, pause_automation: 'bool' =None, security_options: 'V1ClusterSecurityOptions' =None, slurm_v1: 'V1SlurmV1' =None, tagging_options: 'V1ClusterTaggingOptions' =None, user_id: 'str' =None, vultr_v1: 'V1VultrDirectV1' =None):  # noqa: E501
+    def __init__(self, auth_token: 'str' =None, aws_v1: 'V1AWSDirectV1' =None, cluster_type: 'V1ClusterType' =None, deletion_options: 'V1ClusterDeletionOptions' =None, desired_state: 'V1ClusterState' =None, domain: 'str' =None, freeze_accelerators: 'bool' =None, google_cloud_v1: 'V1GoogleCloudDirectV1' =None, insurer_disabled: 'bool' =None, lambda_labs_v1: 'V1LambdaLabsDirectV1' =None, overprovisioning: 'list[V1InstanceOverprovisioningSpec]' =None, pause_automation: 'bool' =None, security_options: 'V1ClusterSecurityOptions' =None, slurm_v1: 'V1SlurmV1' =None, tagging_options: 'V1ClusterTaggingOptions' =None, user_id: 'str' =None, vultr_v1: 'V1VultrDirectV1' =None):  # noqa: E501
         """V1ClusterSpec - a model defined in Swagger"""  # noqa: E501
         self._auth_token = None
-        self._available_instance_types = None
         self._aws_v1 = None
         self._cluster_type = None
         self._deletion_options = None
@@ -105,8 +102,6 @@ class V1ClusterSpec(object):
         self.discriminator = None
         if auth_token is not None:
             self.auth_token = auth_token
-        if available_instance_types is not None:
-            self.available_instance_types = available_instance_types
         if aws_v1 is not None:
             self.aws_v1 = aws_v1
         if cluster_type is not None:
@@ -161,29 +156,6 @@ class V1ClusterSpec(object):
 
         self._auth_token = auth_token
 
-    @property
-    def available_instance_types(self) -> 'list[str]':
-        """Gets the available_instance_types of this V1ClusterSpec.  # noqa: E501
-
-        available_instance_types is a list of instance types that are available for the cluster. This is just a soft filter to prevent users from using instances that we haven't prepared for. If the list is empty, no filtering is done.  # noqa: E501
-
-        :return: The available_instance_types of this V1ClusterSpec.  # noqa: E501
-        :rtype: list[str]
-        """
-        return self._available_instance_types
-
-    @available_instance_types.setter
-    def available_instance_types(self, available_instance_types: 'list[str]'):
-        """Sets the available_instance_types of this V1ClusterSpec.
-
-        available_instance_types is a list of instance types that are available for the cluster. This is just a soft filter to prevent users from using instances that we haven't prepared for. If the list is empty, no filtering is done.  # noqa: E501
-
-        :param available_instance_types: The available_instance_types of this V1ClusterSpec.  # noqa: E501
-        :type: list[str]
-        """
-
-        self._available_instance_types = available_instance_types
-
     @property
     def aws_v1(self) -> 'V1AWSDirectV1':
         """Gets the aws_v1 of this V1ClusterSpec.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_lambda_labs_direct_v1.py

@@ -41,19 +41,24 @@ class V1LambdaLabsDirectV1(object):
                             and the value is json key in definition.
     """
     swagger_types = {
-        'credentials_secret_id': 'str'
+        'credentials_secret_id': 'str',
+        'parent_cluster_id': 'str'
     }
 
     attribute_map = {
-        'credentials_secret_id': 'credentialsSecretId'
+        'credentials_secret_id': 'credentialsSecretId',
+        'parent_cluster_id': 'parentClusterId'
     }
 
-    def __init__(self, credentials_secret_id: 'str' =None):  # noqa: E501
+    def __init__(self, credentials_secret_id: 'str' =None, parent_cluster_id: 'str' =None):  # noqa: E501
         """V1LambdaLabsDirectV1 - a model defined in Swagger"""  # noqa: E501
         self._credentials_secret_id = None
+        self._parent_cluster_id = None
         self.discriminator = None
         if credentials_secret_id is not None:
             self.credentials_secret_id = credentials_secret_id
+        if parent_cluster_id is not None:
+            self.parent_cluster_id = parent_cluster_id
 
     @property
     def credentials_secret_id(self) -> 'str':
@@ -78,6 +83,29 @@ class V1LambdaLabsDirectV1(object):
 
         self._credentials_secret_id = credentials_secret_id
 
+    @property
+    def parent_cluster_id(self) -> 'str':
+        """Gets the parent_cluster_id of this V1LambdaLabsDirectV1.  # noqa: E501
+
+        Note: LambdaLabs is missing object store in their offering, we will need to use either GCP or AWS S3 for that.  # noqa: E501
+
+        :return: The parent_cluster_id of this V1LambdaLabsDirectV1.  # noqa: E501
+        :rtype: str
+        """
+        return self._parent_cluster_id
+
+    @parent_cluster_id.setter
+    def parent_cluster_id(self, parent_cluster_id: 'str'):
+        """Sets the parent_cluster_id of this V1LambdaLabsDirectV1.
+
+        Note: LambdaLabs is missing object store in their offering, we will need to use either GCP or AWS S3 for that.  # noqa: E501
+
+        :param parent_cluster_id: The parent_cluster_id of this V1LambdaLabsDirectV1.  # noqa: E501
+        :type: str
+        """
+
+        self._parent_cluster_id = parent_cluster_id
+
     def to_dict(self) -> dict:
         """Returns the model properties as a dict"""
         result = {}

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -120,6 +120,7 @@ class V1UserFeatures(object):
         'teamspace_storage_tab': 'bool',
         'trainium2': 'bool',
         'use_rclone_mounts_only': 'bool',
+        'vultr': 'bool',
         'writable_data_connections': 'bool'
     }
 
@@ -203,10 +204,11 @@ class V1UserFeatures(object):
         'teamspace_storage_tab': 'teamspaceStorageTab',
         'trainium2': 'trainium2',
         'use_rclone_mounts_only': 'useRcloneMountsOnly',
+        'vultr': 'vultr',
         '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, 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_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
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
         self._advanced_deployment_autoscaling = None
         self._affiliate_links = None
@@ -287,6 +289,7 @@ class V1UserFeatures(object):
         self._teamspace_storage_tab = None
         self._trainium2 = None
         self._use_rclone_mounts_only = None
+        self._vultr = None
         self._writable_data_connections = None
         self.discriminator = None
         if advanced_deployment_autoscaling is not None:
@@ -447,6 +450,8 @@ class V1UserFeatures(object):
             self.trainium2 = trainium2
         if use_rclone_mounts_only is not None:
             self.use_rclone_mounts_only = use_rclone_mounts_only
+        if vultr is not None:
+            self.vultr = vultr
         if writable_data_connections is not None:
             self.writable_data_connections = writable_data_connections
 
@@ -2109,6 +2114,27 @@ class V1UserFeatures(object):
 
         self._use_rclone_mounts_only = use_rclone_mounts_only
 
+    @property
+    def vultr(self) -> 'bool':
+        """Gets the vultr of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The vultr of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._vultr
+
+    @vultr.setter
+    def vultr(self, vultr: 'bool'):
+        """Sets the vultr of this V1UserFeatures.
+
+
+        :param vultr: The vultr of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._vultr = vultr
+
     @property
     def writable_data_connections(self) -> 'bool':
         """Gets the writable_data_connections of this V1UserFeatures.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_vultr_direct_v1.py

@@ -42,24 +42,29 @@ class V1VultrDirectV1(object):
     """
     swagger_types = {
         'credentials_secret_id': 'str',
+        'parent_cluster_id': 'str',
         'primary_region': 'str',
         'regions': 'list[str]'
     }
 
     attribute_map = {
         'credentials_secret_id': 'credentialsSecretId',
+        'parent_cluster_id': 'parentClusterId',
         'primary_region': 'primaryRegion',
         'regions': 'regions'
     }
 
-    def __init__(self, credentials_secret_id: 'str' =None, primary_region: 'str' =None, regions: 'list[str]' =None):  # noqa: E501
+    def __init__(self, credentials_secret_id: 'str' =None, parent_cluster_id: 'str' =None, primary_region: 'str' =None, regions: 'list[str]' =None):  # noqa: E501
         """V1VultrDirectV1 - a model defined in Swagger"""  # noqa: E501
         self._credentials_secret_id = None
+        self._parent_cluster_id = None
         self._primary_region = None
         self._regions = None
         self.discriminator = None
         if credentials_secret_id is not None:
             self.credentials_secret_id = credentials_secret_id
+        if parent_cluster_id is not None:
+            self.parent_cluster_id = parent_cluster_id
         if primary_region is not None:
             self.primary_region = primary_region
         if regions is not None:
@@ -88,6 +93,27 @@ class V1VultrDirectV1(object):
 
         self._credentials_secret_id = credentials_secret_id
 
+    @property
+    def parent_cluster_id(self) -> 'str':
+        """Gets the parent_cluster_id of this V1VultrDirectV1.  # noqa: E501
+
+
+        :return: The parent_cluster_id of this V1VultrDirectV1.  # noqa: E501
+        :rtype: str
+        """
+        return self._parent_cluster_id
+
+    @parent_cluster_id.setter
+    def parent_cluster_id(self, parent_cluster_id: 'str'):
+        """Sets the parent_cluster_id of this V1VultrDirectV1.
+
+
+        :param parent_cluster_id: The parent_cluster_id of this V1VultrDirectV1.  # noqa: E501
+        :type: str
+        """
+
+        self._parent_cluster_id = parent_cluster_id
+
     @property
     def primary_region(self) -> 'str':
         """Gets the primary_region of this V1VultrDirectV1.  # noqa: E501

lightning_sdk/lit_container.py

@@ -55,3 +55,43 @@ class LitContainer:
             raise ValueError("Could not resolve teamspace") from e
         project_id = teamspace.id
         return self._api.delete_container(project_id, container)
+
+    def upload_container(
+        self, container: str, teamspace: str, org: Optional[str] = None, user: Optional[str] = None, tag: str = "latest"
+    ) -> None:
+        """Upload a container to the docker registry.
+
+        Args:
+            container: The name of the container to upload.
+            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
+
+        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

@@ -23,7 +23,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 +54,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,
@@ -199,7 +199,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,
@@ -250,9 +250,14 @@ class _BaseMMT(_BaseJob):
     def machines(self) -> Tuple[MMTMachine, ...]:
         """Returns the sub-jobs for each individual instance."""
 
+    @property
+    def num_machines(self) -> int:
+        """Returns the number of machines assigned to this multi-machine job."""
+        return len(self.machines)
+
     @property
     @abstractmethod
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
 
     @abstractmethod
@@ -303,7 +308,18 @@ class _BaseMMT(_BaseJob):
 
     def dict(
         self
-    ) -> Dict[str, Union[str, "Studio", "Status", "Machine", None, List[Dict[str, Union[str, "Status", "Machine"]]]]]:
+    ) -> Dict[
+        str,
+        Union[
+            str,
+            float,
+            "Studio",
+            "Status",
+            "Machine",
+            None,
+            List[Dict[str, Union[str, "Status", "Machine"]]],
+        ],
+    ]:
         """Dict representation of this job."""
         studio = self.studio
 
@@ -319,6 +335,7 @@ class _BaseMMT(_BaseJob):
                 {"name": d["name"], "status": d["status"], "machine": d["machine"]}
                 for d in (x.dict() for x in self.machines)
             ],
+            "total_cost": self.total_cost,
         }
 
     @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,
@@ -167,6 +167,7 @@ class MMT(_BaseMMT):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
             cluster=cluster,  # deprecated in favor of cloud_account
         )
         # required for typing with "MMT"
@@ -180,7 +181,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,
@@ -239,6 +240,7 @@ class MMT(_BaseMMT):
             cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            entrypoint=entrypoint,
         )
         return self
 
@@ -264,7 +266,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
 

lightning_sdk/mmt/v1.py

@@ -4,11 +4,11 @@ from lightning_sdk.api.mmt_api import MMTApiV1
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.job.v1 import _internal_status_to_external_status
 from lightning_sdk.job.work import Work
+from lightning_sdk.status import Status
 
 if TYPE_CHECKING:
     from lightning_sdk.machine import Machine
     from lightning_sdk.organization import Organization
-    from lightning_sdk.status import Status
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
     from lightning_sdk.user import User
@@ -43,7 +43,7 @@ class _MMTV1(_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,
@@ -133,6 +133,8 @@ class _MMTV1(_BaseMMT):
 
     def stop(self) -> None:
         """Stops the job."""
+        if self.status in (Status.Stopped, Status.Completed, Status.Failed):
+            return
         self._job_api.stop_job(self._guaranteed_job.id, self.teamspace.id)
 
     def delete(self) -> None:
@@ -164,7 +166,7 @@ class _MMTV1(_BaseMMT):
         return f"/teamspace/jobs/{self.name}/snapshot"
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
         return self.machines[0].machine
 

lightning_sdk/mmt/v2.py

@@ -1,13 +1,13 @@
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 
 from lightning_sdk.api.mmt_api import MMTApiV2
 from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.status import Status
 
 if TYPE_CHECKING:
     from lightning_sdk.job.job import Job
     from lightning_sdk.machine import Machine
     from lightning_sdk.organization import Organization
-    from lightning_sdk.status import Status
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
     from lightning_sdk.user import User
@@ -42,7 +42,7 @@ class _MMTV2(_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,
@@ -137,6 +137,8 @@ class _MMTV2(_BaseMMT):
 
     def stop(self) -> None:
         """Stops the job."""
+        if self.status in (Status.Stopped, Status.Completed, Status.Failed):
+            return
         self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
 
     def delete(self) -> None:
@@ -149,12 +151,6 @@ class _MMTV2(_BaseMMT):
             teamspace_id=self._teamspace.id,
         )
 
-    @property
-    def _latest_job(self) -> Any:
-        """Guarantees to fetch the latest version of a job before returning it."""
-        self._update_internal_job()
-        return self._job
-
     @property
     def status(self) -> "Status":
         """The current status of the job."""
@@ -173,7 +169,7 @@ class _MMTV2(_BaseMMT):
         raise NotImplementedError
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
         return self._job_api._get_job_machine_from_spec(self._guaranteed_job.spec)
 
@@ -220,3 +216,8 @@ class _MMTV2(_BaseMMT):
     def command(self) -> str:
         """The command the job is running."""
         return self._job_api.get_command(self._guaranteed_job)
+
+    @property
+    def num_machines(self) -> int:
+        """Returns the number of machines assigned to this multi-machine job."""
+        return self._job_api.get_num_machines(self._guaranteed_job)

{lightning_sdk-0.1.48.dist-info → lightning_sdk-0.1.50.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.1.48
+Version: 0.1.50
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License