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

lightning_sdk/studio.py

@@ -1,25 +1,41 @@
+import glob
 import os
+import threading
 import warnings
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
+from tqdm.auto import tqdm
+
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.api.studio_api import StudioApi
-from lightning_sdk.api.utils import _machine_to_compute_name
+from lightning_sdk.base_studio import BaseStudio
 from lightning_sdk.constants import _LIGHTNING_DEBUG
-from lightning_sdk.machine import Machine
+from lightning_sdk.lightning_cloud.openapi import V1ClusterType
+from lightning_sdk.machine import DEFAULT_MACHINE, CloudProvider, Machine
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
 from lightning_sdk.status import Status
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, _setup_logger
+from lightning_sdk.utils.logging import TrackCallsMeta
+from lightning_sdk.utils.names import random_unique_name
+from lightning_sdk.utils.resolve import (
+    _get_org_id,
+    _resolve_deprecated_cluster,
+    _resolve_deprecated_provider,
+    _resolve_teamspace,
+    _setup_logger,
+)
 
 if TYPE_CHECKING:
+    from lightning_sdk.job import Job
+    from lightning_sdk.mmt import MMT
     from lightning_sdk.plugin import Plugin
 
 _logger = _setup_logger(__name__)
 
 
-class Studio:
+class Studio(metaclass=TrackCallsMeta):
     """A single Lightning AI Studio.
 
     Allows to fully control a studio, including retrieving the status, running commands
@@ -32,7 +48,14 @@ class Studio:
         user: the name of the user owning the :param`teamspace` in case it is owned directly by a user instead of an org
         cloud_account: the name of the cloud account, the studio should be created on.
             Doesn't matter when the studio already exists.
+        cloud_account_provider: The provider to select the cloud-account from.
+            If set, must be in agreement with the provider from the cloud_account (if specified).
+            If not specified, falls back to the teamspace default cloud account.
         create_ok: whether the studio will be created if it does not yet exist. Defaults to True
+        provider: the provider of the machine, the studio should be created on.
+        studio_type: Type of studio to create. Only effective during initial creation;
+            ignored for existing studios.
+
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
         only one of the arguments can be provided.
@@ -40,7 +63,11 @@ class Studio:
     """
 
     # skips init of studio, only set when using this as a shell for names, ids etc.
-    _skip_init = False
+    _skip_init = threading.local()
+    _skip_setup = threading.local()
+
+    # whether to show progress bars during operations
+    show_progress = False
 
     def __init__(
         self,
@@ -49,35 +76,123 @@ class Studio:
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union[CloudProvider, str]] = None,
         create_ok: bool = True,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
+        source: Optional[str] = None,
+        disable_secrets: bool = False,
+        provider: Optional[Union[CloudProvider, str]] = None,  # deprecated in favor of cloud_provider
+        studio_type: Optional[str] = None,  # for base studio templates
     ) -> None:
         self._studio_api = StudioApi()
+        self._cloud_account_api = CloudAccountApi()
+
+        self._prevent_refetch = False
+        self._teamspace = None
 
-        self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
-        self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
-        self._setup_done = False
+        # don't resolve anything if we're skipping init
+        if not getattr(self._skip_init, "value", False):
+            _teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+            if _teamspace is None:
+                raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
+
+            self._teamspace = _teamspace
+
+        self._setup_done = getattr(self._skip_setup, "value", False)
+        self._disable_secrets = disable_secrets
 
         self._plugins = {}
+        self._studio = None
+
+        # Check to see if we're inside a studio
+        current_studio = None
+        studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
+        if studio_id is not None and self._teamspace is not None:
+            # We're inside a studio, get it by ID
+            current_studio = self._studio_api.get_studio_by_id(studio_id=studio_id, teamspace_id=self._teamspace.id)
+
+        if cloud_account or not cloud_provider:
+            cloud_account = _resolve_deprecated_cluster(
+                cloud_account, cluster, current_studio.cluster_id if current_studio else None
+            )
+            cloud_provider = _resolve_deprecated_provider(cloud_provider, provider)
+        else:
+            cloud_provider = _resolve_deprecated_provider(cloud_provider, provider)
 
-        if name is None:
-            studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
-            if studio_id is None:
-                raise ValueError("Cannot autodetect Studio. Either use the SDK from within a Studio or pass a name!")
-            self._studio = self._studio_api.get_studio_by_id(studio_id=studio_id, teamspace_id=self._teamspace.id)
+        cls_name = self._cls_name
+
+        # if we're skipping init, we don't need to resolve the cloud account as then we're not creating a studio
+        if self._teamspace is not None:
+            _cloud_account = self._cloud_account_api.resolve_cloud_account(
+                self._teamspace.id,
+                cloud_account=cloud_account,
+                cloud_provider=cloud_provider,
+                default_cloud_account=self._teamspace.default_cloud_account,
+            )
+
+        self._studio_type = None
+        if studio_type:
+            self._base_studio = BaseStudio(teamspace=self._teamspace)
+            self._available_base_studios = self._base_studio.list()
+            for bst in self._available_base_studios:
+                if (
+                    bst.id == studio_type
+                    or bst.name == studio_type
+                    or bst.name.lower().replace(" ", "-") == studio_type
+                ):
+                    self._studio_type = bst.id
+
+            if not self._studio_type:
+                raise ValueError(
+                    f"Could not find studio type with ID or name '{studio_type}'. "
+                    f"Available studio types: "
+                    f"{[bst.name.lower().replace(' ', '-') for bst in self._available_base_studios]}"
+                )
         else:
+            if current_studio:
+                self._studio_type = current_studio.environment_template_id
+
+        # Resolve studio name if not provided: explicit → env (LIGHTNING_CLOUD_SPACE_ID) → config defaults
+        if name is None and not getattr(self._skip_init, "value", False):
+            if current_studio:
+                name = current_studio.name
+            else:
+                # Try config defaults
+                from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+                config = Config()
+                name = config.get_value(DefaultConfigKeys.studio)
+                if name is None and not create_ok:
+                    raise ValueError(
+                        f"Cannot autodetect {cls_name}. Either use the SDK from within a {cls_name} or pass a name!"
+                    )
+
+        if self._studio is None and not getattr(self._skip_init, "value", False):
+            # If we have a name (explicit or from config), get studio by name
             try:
+                if name is None:
+                    # if we don't have a name, raise an error to get
+                    # to the exception path and optionally create a studio
+                    raise ValueError(
+                        f"Cannot autodetect {cls_name}. Either use the SDK from within a {cls_name} or pass a name!"
+                    )
                 self._studio = self._studio_api.get_studio(name, self._teamspace.id)
             except ValueError as e:
                 if create_ok:
+                    name = name or random_unique_name()
                     self._studio = self._studio_api.create_studio(
-                        name, self._teamspace.id, cloud_account=self._cloud_account
+                        name,
+                        self._teamspace.id,
+                        cloud_account=_cloud_account,
+                        source=source,
+                        disable_secrets=self._disable_secrets,
+                        cloud_space_environment_template_id=self._studio_type,
                     )
                 else:
-                    raise ValueError(f"Studio {name} does not exist.") from e
+                    raise e
 
         if (
-            not self._skip_init
+            not getattr(self._skip_init, "value", False)
             and _internal_status_to_external_status(
                 self._studio_api._get_studio_instance_status_from_object(self._studio)
             )
@@ -133,7 +248,20 @@ class Studio:
         """Returns the current machine type the Studio is running on."""
         if self.status != Status.Running:
             return None
-        return self._studio_api.get_machine(self._studio.id, self._teamspace.id)
+        return self._studio_api.get_machine(
+            self._studio.id,
+            self._teamspace.id,
+            self.cloud_account,
+            _get_org_id(self._teamspace),
+        )
+
+    @property
+    def public_ip(self) -> Optional[str]:
+        """Returns the public IP address of the machine the Studio is running on."""
+        return self._studio_api.get_public_ip(
+            self._studio.id,
+            self._teamspace.id,
+        )
 
     @property
     def interruptible(self) -> bool:
@@ -146,29 +274,104 @@ class Studio:
     @property
     def cluster(self) -> str:
         """Returns the cluster the Studio is running on."""
-        warnings.warn("Studio.cluster is deprecated. Use Studio.cloud_account instead", DeprecationWarning)
+        warnings.warn(
+            f"{self._cls_name}.cluster is deprecated. Use {self._cls_name}.cloud_account instead", DeprecationWarning
+        )
         return self.cloud_account
 
     @property
     def cloud_account(self) -> str:
         return self._studio.cluster_id
 
-    def start(self, machine: Union[Machine, str] = Machine.CPU, interruptible: bool = False) -> None:
-        """Starts a Studio on the specified machine type (default: CPU-4)."""
+    def start(
+        self,
+        machine: Optional[Union[Machine, str]] = None,
+        interruptible: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
+    ) -> None:
+        """Starts a Studio on the specified machine type (default: CPU-4).
+
+        Args:
+            machine: the machine type to start the studio on. Defaults to CPU-4
+            interruptible: whether to use interruptible machines
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
+
+        """
+        # Check to see if we're inside a studio and if its running
+        current_studio_machine = None
+        studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
+        if studio_id is not None:
+            # We're inside a studio, get the machine if it is running
+            current_studio = self._studio_api.get_studio_by_id(studio_id=studio_id, teamspace_id=self._teamspace.id)
+            current_status = self._studio_api._get_studio_instance_status_from_object(current_studio)
+
+            if current_status and _internal_status_to_external_status(current_status) == Status.Running:
+                current_studio_machine = self._studio_api.get_machine(
+                    current_studio.id,
+                    self._teamspace.id,
+                    current_studio.cluster_id,
+                    _get_org_id(self._teamspace),
+                )
+
         status = self.status
+
+        if interruptible is None:
+            interruptible_override = os.environ.get("LIGHTNING_INTERRUPTIBLE_OVERRIDE", None)
+            if interruptible_override is not None:
+                interruptible = interruptible_override.lower() == "true"
+            else:
+                interruptible = self.teamspace.start_studios_on_interruptible
+
+        new_machine = DEFAULT_MACHINE
+        if machine is not None:
+            new_machine = machine
+        elif current_studio_machine is not None:
+            new_machine = current_studio_machine
+
+        if not isinstance(new_machine, Machine):
+            new_machine = Machine.from_str(new_machine)
+
         if status == Status.Running:
-            curr_machine = _machine_to_compute_name(self.machine) if self.machine is not None else None
-            if curr_machine != _machine_to_compute_name(machine):
+            if new_machine != self.machine:
                 raise RuntimeError(
-                    f"Requested to start studio on {machine}, but studio is already running on {self.machine}."
+                    f"Requested to start {self._cls_name} on {new_machine}, "
+                    "but {self._cls_name} is already running on {self.machine}."
                     " Consider switching instead!"
                 )
-            _logger.info(f"Studio {self.name} is already running")
+            _logger.info(f"{self._cls_name} {self.name} is already running")
             return
 
         if status != Status.Stopped:
-            raise RuntimeError(f"Cannot start a studio that is not stopped. Studio {self.name} is {status}.")
-        self._studio_api.start_studio(self._studio.id, self._teamspace.id, machine, interruptible=interruptible)
+            raise RuntimeError(
+                f"Cannot start a {self._cls_name} that is not stopped. {self._cls_name} {self.name} is {status}."
+            )
+
+        # Show progress bar during startup
+        if self.show_progress:
+            from lightning_sdk.utils.progress import StudioProgressTracker
+
+            with StudioProgressTracker("start", show_progress=True) as progress:
+                # Start the studio without blocking
+                self._studio_api.start_studio_async(
+                    self._studio.id,
+                    self._teamspace.id,
+                    new_machine,
+                    interruptible=interruptible,
+                    max_runtime=max_runtime,
+                )
+
+                # Track progress through completion
+                progress.track_startup_phases(
+                    lambda: self._studio_api.get_studio_status(self._studio.id, self._teamspace.id)
+                )
+        else:
+            # Use the blocking version if no progress is needed
+            self._studio_api.start_studio(
+                self._studio.id, self._teamspace.id, new_machine, interruptible=interruptible, max_runtime=max_runtime
+            )
 
         self._setup()
 
@@ -183,17 +386,56 @@ class Studio:
         """Deletes the current Studio."""
         self._studio_api.delete_studio(self._studio.id, self._teamspace.id)
 
-    def duplicate(self) -> "Studio":
-        """Duplicates the existing Studio to the same teamspace."""
-        kwargs = self._studio_api.duplicate_studio(self._studio.id, self._teamspace.id, self._teamspace.id)
+    def duplicate(
+        self,
+        target_teamspace: Optional[Union["Teamspace", str]] = None,
+        machine: Machine = Machine.CPU,
+        name: Optional[str] = None,
+    ) -> "Studio":
+        """Duplicates the existing Studio.
+
+        Args:
+            target_teamspace: the teamspace to duplicate the studio to.
+                Must have the same owner as the source teamspace.
+                If not provided, defaults to current teamspace.
+            machine: the machine to start the duplicated studio on.
+                Defaults to CPU
+        """
+        if target_teamspace is None:
+            target_teamspace_id = self._teamspace.id
+        else:
+            target_teamspace = _resolve_teamspace(
+                target_teamspace,
+                org=self._teamspace.owner if isinstance(self._teamspace.owner, Organization) else None,
+                user=self._teamspace.owner if isinstance(self._teamspace.owner, User) else None,
+            )
+
+            if target_teamspace is None:
+                raise ValueError(
+                    f"Could not resolve target teamspace {target_teamspace} "
+                    f"with owner {self.teamspace.owner} for duplication!"
+                )
+
+            target_teamspace_id = target_teamspace.id
+
+        kwargs = self._studio_api.duplicate_studio(
+            studio_id=self._studio.id,
+            teamspace_id=self._teamspace.id,
+            target_teamspace_id=target_teamspace_id,
+            machine=machine,
+            new_name=name,
+        )
         return Studio(**kwargs)
 
-    def switch_machine(self, machine: Union[Machine, str], interruptible: bool = False) -> None:
+    def switch_machine(
+        self, machine: Union[Machine, str], interruptible: bool = False, cloud_provider: Optional[CloudProvider] = None
+    ) -> None:
         """Switches machine to the provided machine type/.
 
         Args:
             machine: the new machine type to switch to
             interruptible: determines whether to switch to an interruptible instance
+            cloud_provider: the cloud provider to switch to, has no effect if the Studio is not on Lightning Cloud
 
         Note:
             this call is blocking until the new machine is provisioned
@@ -202,12 +444,82 @@ class Studio:
         status = self.status
         if status != Status.Running:
             raise RuntimeError(
-                f"Cannot switch machine on a studio that is not running. Studio {self.name} is {status}."
+                f"Cannot switch machine on a {self._cls_name} that is not running. "
+                "{self._cls_name} {self.name} is {status}."
             )
-        self._studio_api.switch_studio_machine(
-            self._studio.id, self._teamspace.id, machine, interruptible=interruptible
+
+        current_cloud = self._cloud_account_api.get_cloud_account_non_org(
+            self._teamspace.id,
+            self._studio.cluster_id,
         )
 
+        cloud_account = ""
+        if cloud_provider is not None and current_cloud.spec.cluster_type == V1ClusterType.GLOBAL:
+            cloud_account = self._cloud_account_api.resolve_cloud_account(
+                self._teamspace.id,
+                cloud_account=None,
+                cloud_provider=cloud_provider,
+                default_cloud_account=None,
+            )
+
+        if self.show_progress:
+            from lightning_sdk.utils.progress import StudioProgressTracker
+
+            with StudioProgressTracker("switch", show_progress=True) as progress:
+                # Update progress before starting the switch
+                progress.update_progress(5, "Initiating machine switch...")
+
+                # Start the switch operation with progress tracking
+                self._studio_api.switch_studio_machine_with_progress(
+                    self._studio.id,
+                    self._teamspace.id,
+                    machine,
+                    interruptible=interruptible,
+                    progress=progress,
+                    cloud_account=cloud_account,
+                )
+        else:
+            self._studio_api.switch_studio_machine(
+                self._studio.id, self._teamspace.id, machine, interruptible=interruptible, cloud_account=cloud_account
+            )
+
+        if self._studio and cloud_account:
+            # TODO: get this from the API
+            self._studio.cluster_id = cloud_account
+
+    def run_and_detach(self, *commands: str, timeout: float = 10, check_interval: float = 1) -> str:
+        """Runs given commands on the Studio and returns immediately.
+
+        The command will continue to run in the background.
+
+        Args:
+            timeout: wait for this many seconds for the command to finish.
+            check_interval: check the status of the command every this many seconds.
+        """
+        if check_interval > timeout:
+            raise ValueError("check_interval must be less than timeout")
+
+        if _LIGHTNING_DEBUG:
+            print(f"Running {commands=}")
+        status = self.status
+        if status != Status.Running:
+            raise RuntimeError(
+                f"Cannot run a command in a {self._cls_name} that is not running. "
+                "{self._cls_name} {self.name} is {status}."
+            )
+
+        iter_output = self._studio_api.run_studio_commands_and_yield(
+            self._studio.id, self._teamspace.id, *commands, timeout=timeout, check_interval=check_interval
+        )
+
+        output = ""
+        code = None
+        for line, exit_code in iter_output:
+            print(line)
+            output += line
+            code = exit_code
+        return output, code
+
     def run_with_exit_code(self, *commands: str) -> Tuple[str, int]:
         """Runs given commands on the Studio while returning output and exit code.
 
@@ -220,7 +532,10 @@ class Studio:
 
         status = self.status
         if status != Status.Running:
-            raise RuntimeError(f"Cannot run a command in a studio that is not running. Studio {self.name} is {status}.")
+            raise RuntimeError(
+                f"Cannot run a command in a {self._cls_name} that is not running. "
+                "{self._cls_name} {self.name} is {status}."
+            )
         output, exit_code = self._studio_api.run_studio_commands(self._studio.id, self._teamspace.id, *commands)
         output = output.strip()
 
@@ -255,6 +570,34 @@ class Studio:
             progress_bar=progress_bar,
         )
 
+    def upload_folder(self, folder_path: str, remote_path: Optional[str] = None, progress_bar: bool = True) -> None:
+        """Uploads a given folder to a remote path on the Studio."""
+        if folder_path is None:
+            raise ValueError("Cannot upload a folder that is None.")
+        folder_path = os.path.normpath(folder_path)
+        if os.path.isfile(folder_path):
+            raise NotADirectoryError(f"Cannot upload a file as a folder. '{folder_path}' is a file.")
+        if not os.path.exists(folder_path):
+            raise NotADirectoryError(f"Cannot upload a folder that does not exist. '{folder_path}' is not a directory.")
+        all_files = []
+        for fp in glob.glob(os.path.join(folder_path, "**"), recursive=True):
+            if not os.path.isfile(fp):
+                continue
+            rel_path = os.path.relpath(fp, folder_path)
+            remote_file = os.path.join(remote_path, rel_path) if remote_path else rel_path
+            all_files.append((fp, remote_file))
+
+        if progress_bar:
+            progress_bar = tqdm(total=len(all_files), desc="Uploading files", unit="file")
+        for local_file, remote_path in sorted(all_files, key=lambda p: p[1]):
+            if progress_bar:
+                progress_bar.set_description(f"Uploading {local_file}")
+            self.upload_file(local_file, remote_path=remote_path, progress_bar=False)
+            if progress_bar:
+                progress_bar.update(1)
+        if progress_bar:
+            progress_bar.close()
+
     def download_file(self, remote_path: str, file_path: Optional[str] = None) -> None:
         """Downloads a file from the Studio to a given target path."""
         if file_path is None:
@@ -281,6 +624,91 @@ class Studio:
             cloud_account=self._studio.cluster_id,
         )
 
+    def run_job(
+        self,
+        name: str,
+        machine: Union["Machine", str],
+        command: str,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        reuse_snapshot: bool = True,
+    ) -> "Job":
+        """Run async workloads using the compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            command: The command to run inside your job.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
+        """
+        from lightning_sdk.job import Job
+
+        return Job.run(
+            name=name,
+            machine=machine,
+            command=command,
+            studio=self,
+            image=None,
+            teamspace=self.teamspace,
+            cloud_account=self.cloud_account,
+            env=env,
+            interruptible=interruptible,
+            reuse_snapshot=reuse_snapshot,
+        )
+
+    def run_mmt(
+        self,
+        name: str,
+        num_machines: int,
+        machine: Union["Machine", str],
+        command: str,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+    ) -> "MMT":
+        """Run async workloads using the compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            num_machines: The number of machines to run on.
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            command: The command to run inside your job.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+        """
+        from lightning_sdk.mmt import MMT
+
+        return MMT.run(
+            name=name,
+            num_machines=num_machines,
+            machine=machine,
+            command=command,
+            studio=self,
+            image=None,
+            teamspace=self.teamspace,
+            cloud_account=self.cloud_account,
+            env=env,
+            interruptible=interruptible,
+        )
+
+    def create_assistant(self, name: str, port: int) -> None:
+        assistant = self._studio_api.create_assistant(
+            studio_id=self._studio.id, teamspace_id=self._teamspace.id, port=port, assistant_name=name
+        )
+        assistant_info = f"Created assisant with name: {assistant.name}, ID: {assistant.id}"
+        self._assistant_id = assistant.id
+        _logger.info(assistant_info)
+
+    def rename(self, new_name: str) -> None:
+        """Renames the current Studio to the provided new name."""
+        if new_name == self._studio.name:
+            return
+
+        self._studio_api._update_cloudspace(self._studio, self._teamspace.id, "display_name", new_name)
+        self._update_studio_reference()
+
     @property
     def auto_sleep(self) -> bool:
         """Returns if a Studio has auto-sleep enabled."""
@@ -289,8 +717,8 @@ class Studio:
     @auto_sleep.setter
     def auto_sleep(self, value: bool) -> None:
         if not value and self.machine == Machine.CPU:
-            warnings.warn("Disabling auto-sleep will convert the Studio from free to paid!")
-        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, enabled=value, studio=self._studio)
+            warnings.warn(f"Disabling auto-sleep will convert the {self._cls_name} from free to paid!")
+        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, enabled=value)
         self._update_studio_reference()
 
     @property
@@ -300,10 +728,8 @@ class Studio:
 
     @auto_sleep_time.setter
     def auto_sleep_time(self, value: int) -> None:
-        warnings.warn("Setting auto-sleep time will convert the Studio from free to paid!")
-        self._studio_api.update_autoshutdown(
-            self._studio.id, self._teamspace.id, idle_shutdown_seconds=value, studio=self._studio
-        )
+        warnings.warn(f"Setting auto-sleep time will convert the {self._cls_name} from free to paid!")
+        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, idle_shutdown_seconds=value)
         self._update_studio_reference()
 
     @property
@@ -326,6 +752,22 @@ class Studio:
         warnings.warn("auto_shutdown_time is deprecated. Use auto_sleep_time instead", DeprecationWarning)
         self.auto_sleep_time = value
 
+    @property
+    def env(self) -> Dict[str, str]:
+        self._update_studio_reference()
+        return self._studio_api.get_env(self._studio)
+
+    def set_env(self, new_env: Dict[str, str], partial: bool = True) -> None:
+        """Set the environment variables for the Studio.
+
+        Args:
+            new_env: The new environment variables to set.
+            partial: Whether to only set the environment variables that are provided.
+                If False, existing environment variables that are not in new_env will be removed.
+                If True, existing environment variables that are not in new_env will be kept.
+        """
+        self._studio_api.set_env(self._studio, self._teamspace.id, new_env, partial=partial)
+
     @property
     def available_plugins(self) -> Mapping[str, str]:
         """All available plugins to install in the current Studio."""
@@ -370,6 +812,7 @@ class Studio:
     def _add_plugin(self, plugin_name: str) -> None:
         """Adds the just installed plugin to the internal list of plugins."""
         from lightning_sdk.plugin import (
+            CustomPortPlugin,
             InferenceServerPlugin,
             JobsPlugin,
             MultiMachineTrainingPlugin,
@@ -383,6 +826,7 @@ class Studio:
             "jobs": JobsPlugin,
             "multi-machine-training": MultiMachineTrainingPlugin,
             "inference-server": InferenceServerPlugin,
+            "custom-port": CustomPortPlugin,
         }.get(plugin_name, Plugin)
 
         description = self._list_installed_plugins()[plugin_name]
@@ -415,6 +859,36 @@ class Studio:
     def _update_studio_reference(self) -> None:
         self._studio = self._studio_api.get_studio_by_id(studio_id=self._studio.id, teamspace_id=self._teamspace.id)
 
+    @property
+    def _cls_name(self) -> str:
+        return self.__class__.__qualname__
+
+
+class VM(Studio):
+    """A single Lightning AI VM.
+
+    Allows to fully control a vm, including retrieving the status, running commands
+    and switching machine types.
+
+    Args:
+        name: the name of the vm
+        teamspace: the name of the teamspace the vm is contained by
+        org: the name of the organization owning the :param`teamspace` in case it is owned by an org
+        user: the name of the user owning the :param`teamspace` in case it is owned directly by a user instead of an org
+        cloud_account: the name of the cloud account, the vm should be created on.
+            Doesn't matter when the vm already exists.
+        cloud_account_provider: The provider to select the cloud-account from.
+            If set, must be in agreement with the provider from the cloud_account (if specified).
+            If not specified, falls backto the teamspace default cloud account.
+        create_ok: whether the vm will be created if it does not yet exist. Defaults to True
+        provider: the provider of the machine, the vm should be created on.
+
+    Note:
+        Since a teamspace can either be owned by an org or by a user directly,
+        only one of the arguments can be provided.
+
+    """
+
 
 def _internal_status_to_external_status(internal_status: str) -> Status:
     """Converts internal status strings from HTTP requests to external enums."""