lightning-sdk 0.1.57__py3-none-any.whl → 0.2.0__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/validate.py

@@ -45,6 +45,7 @@ class Validate(object):
         'check_is_public': 'bool',
         'cluster_ids': 'list[str]',
         'efs': 'V1EfsConfig',
+        'filestore': 'V1FilestoreDataConnection',
         'gcp': 'V1GcpDataConnection',
         'gcs_folder': 'V1GCSFolderDataConnection',
         's3_folder': 'V1S3FolderDataConnection'
@@ -55,17 +56,19 @@ class Validate(object):
         'check_is_public': 'checkIsPublic',
         'cluster_ids': 'clusterIds',
         'efs': 'efs',
+        'filestore': 'filestore',
         'gcp': 'gcp',
         'gcs_folder': 'gcsFolder',
         's3_folder': 's3Folder'
     }
 
-    def __init__(self, aws: 'V1AwsDataConnection' =None, check_is_public: 'bool' =None, cluster_ids: 'list[str]' =None, efs: 'V1EfsConfig' =None, gcp: 'V1GcpDataConnection' =None, gcs_folder: 'V1GCSFolderDataConnection' =None, s3_folder: 'V1S3FolderDataConnection' =None):  # noqa: E501
+    def __init__(self, aws: 'V1AwsDataConnection' =None, check_is_public: 'bool' =None, cluster_ids: 'list[str]' =None, efs: 'V1EfsConfig' =None, filestore: 'V1FilestoreDataConnection' =None, gcp: 'V1GcpDataConnection' =None, gcs_folder: 'V1GCSFolderDataConnection' =None, s3_folder: 'V1S3FolderDataConnection' =None):  # noqa: E501
         """Validate - a model defined in Swagger"""  # noqa: E501
         self._aws = None
         self._check_is_public = None
         self._cluster_ids = None
         self._efs = None
+        self._filestore = None
         self._gcp = None
         self._gcs_folder = None
         self._s3_folder = None
@@ -78,6 +81,8 @@ class Validate(object):
             self.cluster_ids = cluster_ids
         if efs is not None:
             self.efs = efs
+        if filestore is not None:
+            self.filestore = filestore
         if gcp is not None:
             self.gcp = gcp
         if gcs_folder is not None:
@@ -169,6 +174,27 @@ class Validate(object):
 
         self._efs = efs
 
+    @property
+    def filestore(self) -> 'V1FilestoreDataConnection':
+        """Gets the filestore of this Validate.  # noqa: E501
+
+
+        :return: The filestore of this Validate.  # noqa: E501
+        :rtype: V1FilestoreDataConnection
+        """
+        return self._filestore
+
+    @filestore.setter
+    def filestore(self, filestore: 'V1FilestoreDataConnection'):
+        """Sets the filestore of this Validate.
+
+
+        :param filestore: The filestore of this Validate.  # noqa: E501
+        :type: V1FilestoreDataConnection
+        """
+
+        self._filestore = filestore
+
     @property
     def gcp(self) -> 'V1GcpDataConnection':
         """Gets the gcp of this Validate.  # noqa: E501

lightning_sdk/lightning_cloud/rest_client.py

@@ -32,7 +32,8 @@ from lightning_sdk.lightning_cloud.openapi import (
     StorageServiceApi,
     DeploymentTemplatesServiceApi,
     ModelsStoreApi,
-    LitRegistryServiceApi
+    LitRegistryServiceApi,
+    PipelinesServiceApi,
 )
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.source_code.logs_socket_api import LightningLogsSocketAPI
@@ -92,7 +93,8 @@ class GridRestClient(
         StorageServiceApi,
         DeploymentTemplatesServiceApi,
         ModelsStoreApi,
-        LitRegistryServiceApi
+        LitRegistryServiceApi,
+        PipelinesServiceApi,
 ):
 
     def __init__(self, api_client: Optional[ApiClient] = None):

lightning_sdk/machine.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import ClassVar, Optional
+from typing import Any, ClassVar, Optional, Tuple
 
 
 @dataclass(frozen=True)
@@ -42,6 +42,30 @@ class Machine:
             return self.instance_type == other.instance_type
         return False
 
+    def is_cpu(self) -> bool:
+        """Whether the machine is a CPU."""
+        return (
+            self == Machine.CPU_SMALL
+            or self == Machine.CPU
+            or self == Machine.DATA_PREP
+            or self == Machine.DATA_PREP_MAX
+            or self == Machine.DATA_PREP_ULTRA
+        )
+
+    @classmethod
+    def from_str(cls, machine: str, *additional_machine_ids: Any) -> "Machine":
+        possible_values: Tuple["Machine", ...] = tuple(
+            [machine for machine in cls.__dict__.values() if isinstance(machine, cls)]
+        )
+        for m in possible_values:
+            for machine_id in [machine, *additional_machine_ids]:
+                if machine_id in (getattr(m, "name", None), getattr(m, "instance_type", None)):
+                    return m
+
+        if additional_machine_ids:
+            return cls(machine, *additional_machine_ids)
+        return cls(machine, machine)
+
 
 Machine.CPU_SMALL = Machine(name="CPU_SMALL", instance_type="m3.medium")
 Machine.CPU = Machine(name="CPU", instance_type="cpu-4")

lightning_sdk/models.py

@@ -7,7 +7,7 @@ from lightning_sdk.api import OrgApi, TeamspaceApi, UserApi
 from lightning_sdk.lightning_cloud.openapi.models import V1Membership, V1OwnerType
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _get_authed_user
+from lightning_sdk.utils.resolve import _get_authed_user, _resolve_teamspace
 
 if TYPE_CHECKING:
     from lightning_sdk.teamspace import Teamspace
@@ -76,20 +76,20 @@ def _parse_model_name_and_version(name: str) -> Tuple[str, str, str, str]:
     """Parse the name argument into its components."""
     try:
         org_name, teamspace_name, model_name = name.split("/")
-        parts = model_name.split(":")
-        if len(parts) == 1:
-            return org_name, teamspace_name, parts[0], "latest"
-        if len(parts) == 2:
-            return org_name, teamspace_name, parts[0], parts[1]
-        # The rest of the validation for name and version happens in the backend
-        raise ValueError(
-            "Model version is expected to be in the format `entity/modelname:version` separated by a"
-            f" single colon, but got: {name}"
-        )
     except ValueError as err:
         raise ValueError(
-            f"Model name must be in the format 'organization/teamspace/model' but you provided '{name}'."
+            f"Model name must be in the format `organization/teamspace/model_name` but you provided '{name}'."
         ) from err
+    parts = model_name.split(":")
+    if len(parts) == 1:
+        return org_name, teamspace_name, parts[0], "default"
+    if len(parts) == 2:
+        return org_name, teamspace_name, parts[0], parts[1]
+    # The rest of the validation for name and version happens in the backend
+    raise ValueError(
+        "Model version is expected to be in the format `organization/teamspace/model_name:version`"
+        f" separated by a single colon, but got: {name}"
+    )
 
 
 def download_model(
@@ -105,6 +105,9 @@ def download_model(
         download_dir: The directory where the Model should be downloaded.
         progress_bar: Whether to show a progress bar when downloading.
     """
+    if "/" not in name:  # do some magic if you run studio
+        teamspace = _resolve_teamspace(None, None, None)
+        name = f"{teamspace.owner.name}/{teamspace.name}/{name}"
     teamspace_owner_name, teamspace_name, model_name, version = _parse_model_name_and_version(name)
 
     download_dir = Path(download_dir)
@@ -144,6 +147,9 @@ def upload_model(
             If not provided, the default cloud account for the Teamspace will be used.
         progress_bar: Whether to show a progress bar for the upload.
     """
+    if "/" not in name:  # do some magic if you run studio
+        teamspace = _resolve_teamspace(None, None, None)
+        name = f"{teamspace.owner.name}/{teamspace.name}/{name}"
     org_name, teamspace_name, model_name, _ = _parse_model_name_and_version(name)
     teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
     return teamspace.upload_model(

lightning_sdk/pipeline/__init__.py

@@ -0,0 +1,4 @@
+from lightning_sdk.pipeline.pipeline import Pipeline
+from lightning_sdk.pipeline.types import MMT, Deployment, Job
+
+__all__ = ["Pipeline", "Job", "MMT", "Deployment"]

lightning_sdk/pipeline/pipeline.py

@@ -0,0 +1,109 @@
+from typing import List, Optional, Union
+
+from lightning_sdk.api import UserApi
+from lightning_sdk.api.pipeline_api import PipelineApi
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.organization import Organization
+from lightning_sdk.pipeline.types import MMT, Deployment, Job
+from lightning_sdk.pipeline.utils import prepare_steps
+from lightning_sdk.services.utilities import _get_cluster
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _resolve_org, _resolve_teamspace, _resolve_user
+
+
+class Pipeline:
+    def __init__(
+        self,
+        name: str,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cloud_account: Optional[str] = None,
+        shared_filesystem: Optional[bool] = None,
+    ) -> None:
+        """The Lightning Pipeline can be used to create complex DAG.
+
+        Arguments:
+            name: The desired name of the pipeline.
+            teamspace: The teamspace where the pipeline will be created.
+            org: The organization where the pipeline will be created.
+            user: The creator of the pipeline.
+            cloud_account: The cloud account to use for the entire pipeline.
+            shared_filesystem: Whether the pipeline should use a shared filesystem across all nodes.
+                Note: This forces the pipeline steps to be in the cloud_account and same region
+        """
+        self._auth = Auth()
+        self._user = None
+
+        try:
+            self._auth.authenticate()
+            if user is None:
+                self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
+        except ConnectionError as e:
+            raise e
+
+        self._name = name
+        self._org = _resolve_org(org)
+        self._user = _resolve_user(self._user or user)
+
+        self._teamspace = _resolve_teamspace(
+            teamspace=teamspace,
+            org=self._org,
+            user=self._user,
+        )
+
+        self._pipeline_api = PipelineApi()
+        self._cloud_account = _get_cluster(
+            client=self._pipeline_api._client, project_id=self._teamspace.id, cluster_id=cloud_account
+        )
+        self._shared_filesystem = shared_filesystem
+        self._is_created = False
+
+        pipeline = None
+
+        if name.startswith("pip_"):
+            pipeline = self._pipeline_api.get_pipeline_by_id(name, self._teamspace.id)
+
+        if pipeline:
+            self._name = pipeline.name
+            self._is_created = True
+            self._pipeline = pipeline
+
+    def run(self, steps: List[Union[Job, Deployment, MMT]]) -> None:
+        if len(steps) == 0:
+            raise ValueError("The provided steps is empty")
+
+        for step_idx, step in enumerate(steps):
+            if step.name in [None, ""]:
+                raise ValueError(f"The step {step_idx} requires a name")
+
+        steps = [
+            step.to_proto(self._teamspace, self._cloud_account.cluster_id or "", self._shared_filesystem)
+            for step in steps
+        ]
+
+        self._pipeline = self._pipeline_api.create_pipeline(
+            self._name,
+            self._teamspace.id,
+            prepare_steps(steps),
+            self._shared_filesystem or False,
+        )
+
+    def stop(self) -> None:
+        if self._pipeline is None:
+            return
+
+        self._pipeline_api.stop(self._pipeline)
+
+    def delete(self) -> None:
+        if self._pipeline is None:
+            return
+
+        self._pipeline_api.delete(self._teamspace.id, self._pipeline.id)
+
+    @property
+    def name(self) -> str:
+        if self._pipeline:
+            return self._pipeline.name
+        return None

lightning_sdk/pipeline/types.py

@@ -0,0 +1,268 @@
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from lightning_sdk.api.deployment_api import (
+    AutoScaleConfig,
+    AutoScalingMetric,
+    BasicAuth,
+    Env,
+    ExecHealthCheck,
+    HttpHealthCheck,
+    ReleaseStrategy,
+    Secret,
+    TokenAuth,
+    to_autoscaling,
+    to_endpoint,
+    to_spec,
+    to_strategy,
+)
+from lightning_sdk.job.v2 import JobApiV2
+from lightning_sdk.lightning_cloud.openapi.models import (
+    V1CreateDeploymentRequest,
+    V1PipelineStep,
+    V1PipelineStepType,
+)
+from lightning_sdk.mmt.v2 import MMTApiV2
+from lightning_sdk.studio import Studio
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import Machine
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.teamspace import Teamspace
+    from lightning_sdk.user import User
+
+
+from lightning_sdk.pipeline.utils import DEFAULT
+
+
+class Deployment:
+    # Note: This class is only temporary while pipeline is wip
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        machine: Optional["Machine"] = None,
+        image: Optional[str] = None,
+        autoscale: Optional["AutoScaleConfig"] = None,
+        ports: Optional[List[float]] = None,
+        release_strategy: Optional["ReleaseStrategy"] = None,
+        entrypoint: Optional[str] = None,
+        command: Optional[str] = None,
+        env: Union[List[Union["Secret", "Env"]], Dict[str, str], None] = None,
+        spot: Optional[bool] = None,
+        replicas: Optional[int] = None,
+        health_check: Optional[Union["HttpHealthCheck", "ExecHealthCheck"]] = None,
+        auth: Optional[Union["BasicAuth", "TokenAuth"]] = None,
+        cloud_account: Optional[str] = None,
+        custom_domain: Optional[str] = None,
+        quantity: Optional[int] = None,
+        wait_for: Union[str, List[str]] = DEFAULT,
+    ) -> None:
+        self.name = name
+        self.machine = machine
+        self.image = image
+        self.autoscale = autoscale or AutoScaleConfig(
+            min_replicas=0,
+            max_replicas=1,
+            target_metrics=[
+                AutoScalingMetric(
+                    name="CPU" if machine.is_cpu() else "GPU",
+                    target=80,
+                )
+            ],
+        )
+        self.ports = ports
+        self.release_strategy = release_strategy
+        self.entrypoint = entrypoint
+        self.command = command
+        self.env = env
+        self.spot = spot
+        self.replicas = replicas or 1
+        self.health_check = health_check
+        self.auth = auth
+        self.cloud_account = cloud_account or ""
+        self.custom_domain = custom_domain
+        self.quantity = quantity
+        self.wait_for = wait_for
+
+    def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+        _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
+        return V1PipelineStep(
+            name=self.name,
+            type=V1PipelineStepType.DEPLOYMENT,
+            wait_for=to_wait_for(self.wait_for),
+            deployment=V1CreateDeploymentRequest(
+                autoscaling=to_autoscaling(self.autoscale, self.replicas),
+                endpoint=to_endpoint(self.ports, self.auth, self.custom_domain),
+                name=self.name,
+                project_id=teamspace.id,
+                replicas=self.replicas,
+                spec=to_spec(
+                    cloud_account=self.cloud_account or cloud_account,
+                    command=self.command,
+                    entrypoint=self.entrypoint,
+                    env=self.env,
+                    image=self.image,
+                    spot=self.spot,
+                    machine=self.machine,
+                    health_check=self.health_check,
+                    quantity=self.quantity,
+                ),
+                strategy=to_strategy(self.release_strategy),
+            ),
+        )
+
+
+class Job:
+    # Note: This class is only temporary while pipeline is wip
+
+    def __init__(
+        self,
+        machine: Union["Machine", str],
+        name: Optional[str] = None,
+        command: Optional[str] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Union[str, None] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cloud_account: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        wait_for: Union[str, List[str]] = DEFAULT,
+    ) -> None:
+        self.name = name
+        self.machine = machine
+        self.command = command
+        self.studio = studio
+        self.image = image
+        self.teamspace = teamspace
+        self.org = org
+        self.user = user
+        self.cloud_account = cloud_account
+        self.env = env
+        self.interruptible = interruptible
+        self.image_credentials = image_credentials
+        self.cloud_account_auth = cloud_account_auth
+        self.entrypoint = entrypoint
+        self.path_mappings = path_mappings
+        self.wait_for = wait_for
+
+    def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+        _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
+        body = JobApiV2._create_job_body(
+            name=self.name,
+            command=self.command,
+            cloud_account=self.cloud_account or cloud_account,
+            studio_id=None,
+            image=self.image,
+            machine=self.machine,
+            interruptible=self.interruptible,
+            env=self.env,
+            image_credentials=self.image_credentials,
+            cloud_account_auth=self.cloud_account_auth,
+            entrypoint=self.entrypoint,
+            path_mappings=self.path_mappings,
+            artifacts_local=None,
+            artifacts_remote=None,
+        )
+
+        return V1PipelineStep(
+            name=self.name,
+            type=V1PipelineStepType.JOB,
+            wait_for=to_wait_for(self.wait_for),
+            job=body,
+        )
+
+
+class MMT:
+    # Note: This class is only temporary while pipeline is wip
+
+    def __init__(
+        self,
+        name: str,
+        machine: Union["Machine", str],
+        num_machines: Optional[int] = 2,
+        command: Optional[str] = None,
+        studio: Union["Studio", str, None] = None,
+        image: Optional[str] = None,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        cloud_account: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        wait_for: Union[str, List[str]] = DEFAULT,
+    ) -> None:
+        self.machine = machine
+        self.num_machines = num_machines
+        self.name = name
+        self.command = command
+        self.studio = studio
+        self.image = image
+        self.teamspace = teamspace
+        self.org = org
+        self.user = user
+        self.cloud_account = cloud_account
+        self.env = env
+        self.interruptible = interruptible
+        self.image_credentials = image_credentials
+        self.cloud_account_auth = cloud_account_auth
+        self.entrypoint = entrypoint
+        self.path_mappings = path_mappings
+        self.wait_for = wait_for
+
+    def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+        _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
+        body = MMTApiV2._create_mmt_body(
+            name=self.name,
+            num_machines=self.num_machines,
+            command=self.command,
+            cloud_account=self.cloud_account or cloud_account,
+            studio_id=self.studio.studio_id if isinstance(self.studio, Studio) else None,
+            image=self.image,
+            machine=self.machine,
+            interruptible=self.interruptible,
+            env=self.env,
+            image_credentials=self.image_credentials,
+            cloud_account_auth=self.cloud_account_auth,
+            entrypoint=self.entrypoint,
+            path_mappings=self.path_mappings,
+            artifacts_local=None,  # deprecated in favor of path_mappings
+            artifacts_remote=None,  # deprecated in favor of path_mappings
+        )
+
+        return V1PipelineStep(
+            name=self.name,
+            type=V1PipelineStepType.MMT,
+            wait_for=to_wait_for(self.wait_for),
+            mmt=body,
+        )
+
+
+def to_wait_for(wait_for: Optional[Union[str, List[str]]]) -> Optional[List[str]]:
+    if wait_for == DEFAULT:
+        return wait_for
+
+    if wait_for is None:
+        return []
+
+    return wait_for if isinstance(wait_for, list) else [wait_for]
+
+
+def _validate_cloud_account(pipeline_cloud_account: str, step_cloud_account: str, shared_filesystem: bool) -> None:
+    if not shared_filesystem:
+        return
+
+    if pipeline_cloud_account != "" and step_cloud_account != "" and pipeline_cloud_account != step_cloud_account:
+        raise ValueError(
+            "With shared filesystem enabled, all the pipeline steps wait_for to be on the same cluster."
+            f" Found {pipeline_cloud_account} and {step_cloud_account}"
+        )

lightning_sdk/pipeline/utils.py

@@ -0,0 +1,69 @@
+from typing import List
+
+from lightning_sdk.lightning_cloud.openapi.models import V1PipelineStep, V1PipelineStepType
+
+DEFAULT = "DEFAULT"
+
+
+def prepare_steps(steps: List["V1PipelineStep"]) -> List["V1PipelineStep"]:
+    """The prepare_steps function is responsible for creating dependencies between steps.
+
+    The dependencies are based on whether a step wait_for to be executed before another.
+    """
+    name_to_step = {}
+    name_to_idx = {}
+
+    for current_step_idx, current_step in enumerate(steps):
+        if current_step.name not in name_to_step:
+            name_to_step[current_step.name] = current_step
+            name_to_idx[current_step.name] = current_step_idx
+        else:
+            raise ValueError(f"A step with the name {current_step.name} already exists.")
+
+    if steps[0].wait_for != DEFAULT:
+        raise ValueError("The first step isn't allowed to receive `wait_for=...`.")
+
+    steps[0].wait_for = []
+
+    # This implements a linear dependency between the steps as the default behaviour
+    for current_step_idx, current_step in reversed(list(enumerate(steps))):
+        if current_step_idx == 0:
+            continue
+
+        if current_step.wait_for == DEFAULT:
+            prev_step_idx = current_step_idx - 1
+            wait_for = []
+            while prev_step_idx > -1:
+                prev_step = steps[prev_step_idx]
+                wait_for.insert(0, steps[prev_step_idx].name)
+                if prev_step.wait_for != []:
+                    break
+                prev_step_idx -= 1
+            current_step.wait_for = wait_for
+        else:
+            for name in current_step.wait_for:
+                if current_step.name == name:
+                    raise ValueError("You can only reference prior steps")
+
+                if name not in name_to_step:
+                    raise ValueError(f"The step {current_step_idx} doesn't have a valid wait_for. Found {name}")
+
+                if name_to_idx[name] >= name_to_idx[current_step.name]:
+                    raise ValueError("You can only reference prior steps")
+
+    print()
+    print("===== Generated Pipeline =====")
+    for step_idx, step in enumerate(steps):
+        step_type = ""
+        if step.type == V1PipelineStepType.DEPLOYMENT:
+            step_type = "Deployment"
+        elif step.type == V1PipelineStepType.JOB:
+            step_type = "Job"
+        else:
+            step_type = "MMT"
+        wait_for = "nothing" if len(step.wait_for) == 0 else step.wait_for
+        print(f"{step_idx} - {step_type}['{step.name}'] wait_for {wait_for}")
+    print("===== ================== =====")
+    print()
+
+    return steps

lightning_sdk/plugin.py

@@ -186,16 +186,15 @@ class MultiMachineTrainingPlugin(_Plugin):
 
         machine = _resolve_deprecated_cloud_compute(machine, cloud_compute)
 
-        with forced_v1(MMT) as v1mmt:
-            return v1mmt.run(
-                name=name,
-                num_machines=num_instances,
-                machine=machine,
-                command=command,
-                studio=self._studio,
-                teamspace=self._studio.teamspace,
-                interruptible=interruptible,
-            )
+        return MMT.run(
+            name=name,
+            num_machines=num_instances,
+            machine=machine,
+            command=command,
+            studio=self._studio,
+            teamspace=self._studio.teamspace,
+            interruptible=interruptible,
+        )
 
 
 class MultiMachineDataPrepPlugin(_Plugin):

lightning_sdk/services/utilities.py

@@ -6,7 +6,7 @@ import requests
 import urllib3
 
 from lightning_sdk.api.utils import _get_cloud_url
-from lightning_sdk.lightning_cloud.openapi import V1Membership
+from lightning_sdk.lightning_cloud.openapi import V1Membership, V1ProjectClusterBinding
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 _CHUNK_SIZE = 1024 * 1024
@@ -35,7 +35,7 @@ def _get_project(client: LightningClient, project_name: Optional[str] = None) ->
     raise ValueError("No valid projects found. Please reach out to lightning.ai team to create a project")
 
 
-def _get_cluster(client: LightningClient, project_id: str, cluster_id: Optional[str] = None) -> V1Membership:
+def _get_cluster(client: LightningClient, project_id: str, cluster_id: Optional[str] = None) -> V1ProjectClusterBinding:
     """Get a project membership for the user from the backend."""
     clusters = client.projects_service_list_project_cluster_bindings(project_id=project_id)
     if cluster_id: