TemplatesModels
Build
Train
Inference
GPUsSolutionsDocsPricing
Request demoStart freeOpen app

Studio

The Lightning SDK provides programmatic control for the Studio. Use this to interact with the Studio programmatically.

Within a Studio, use the SDK to programatically start Jobs, sweeps, and plugins programatically.

Outside a Studio, use the SDK to integrate Studios into ML pipelines, CI/CD, production workflows, and more.

Install the SDK

The SDK is automatically installed inside the Studios.

To install the SDK outside of a Studio (on your local, or a CI/CD system) run this command:

pip install lightning-sdk

Once the SDK is installed, you'll now have full programmatic control over all your Studios on the Lightning AI platform.

Basic objects

The SDK follows principles you know from interacting with a Studio through your browser. It is implemented in object-oriented Python and has User, Organization, Teamspace and Studio classes.

from lightning_sdk import Studio, Teamspace, User, Organization

# allows interaction with the Studio named "my-studio" in the Teamspace "train-model" of user "zeus"
s = Studio(name="my-studio", teamspace="train-model", user="zeus")

# allows interaction with Teamspace "train-model" of organization "olymp"
t = Teamspace("train-model", org="olymp")

# allows interaction with user "zeus"
u = User("zeus")

# allows interaction with organization "olymp"
o = Organization("olymp")

If you are accessing a studio in an *organizational *Teamspace rather than a user (personal) Teamspace, be sure to specify the org.

s = Studio(name="my-studio", teamspace="train-model", org="my-company-org")

Get started

Exploring the SDK in a Studio is the easiest as you are already authenticated, and all the relevant environment variables are exported.

To get started on your laptop, you need to export two environment variables: LIGHTNING_USER_ID and LIGHTNING_API_KEY, both of which you can find in the "Global Settings" if you click on your profile icon on the top right. In the appearing Menu, select "Keys" on the left-hand side, and you will see the values for your unique set of keys under "Programmatic access" labeled as "Login via CLI". Export those variables in your terminal and you are good to go.

Keys are in your global settings

Class construction

You need to initialize the corresponding classes to use specific parts of the SDK. For example, when interacting with a Studio, you need to initialize the Studio class. Every class has several optional arguments and some environment variables that will be picked up if those arguments are not specified.

Example initializing a Studio

#! pip install lightning-sdk
# export the following variables for authentication (use your own ID and API key)
# LIGHTNING_USER_ID=000000-000000-000000-00000-000000000
# LIGHTNING_API_KEY=111111-111111-111111-11111-111111111
from lightning_sdk import Studio

# this picks up name, teamspace and owner from environment variables.
# running this within a Studio allows you to manage the Studio you're in.
s = Studio()

# this allows remote control of a Studio from outside
# replace name, teamspace and user/org with your own information
remote_control = Studio(name="my-studio", teamspace="my-teamspace", user="my-user")

To initialize a Studio, you could do Studio(name="my-studio", teamspace="my-teamspace", user="my-user") or you could just do Studio() if you have the correct environment variables exported. Inside a Studio these always point to the current Studio so that Studio() will always give you programmatic control over your current Studio. Same goes for Teamspace(), User() and Organization(): They always point to the current scope.

You can also only specify some of these variables and, provided you exported the related environment variables, the other ones will be picked out automatically. So Studio(name="some-other-studio") would point to another Studio in the same Teamspace within the same User/Organization.

Studio ID LIGHTNING_CLOUD_SPACE_ID
Teamspace Name LIGHTNING_TEAMSPACE
User Name LIGHTNING_USERNAME
Organization Name LIGHTNING_ORG

For the sake of readability, in this documentation, we will use the empty constructors, but you could also pass the respective arguments instead.

Start, stop and delete a Studio

After construction with studio = Studio(), we can start, stop and delete it with simple calls to studio.start(), studio.stop() and studio.delete(). That means that after initializing your Studio, it is not automatically starting. You have to explicitly start it with studio.start(). While auto-shutdown works on Studios started through the SDK as well, these Studios will be kept alive as long as you have a Python object referencing. Note that if the desired machine is unavailable, studio.start() will return an error indicating the unavailability.

While studio.stop() gently shuts down your Studio until you decide to spin it up again, studio.delete() is a destructive action that also deletes the Studio and all your files from that Studio, so you won't be able to use it again.

Example stopping a running Studio (caution: when running this from within a Studio, you will have to restart it to continue exploring):

#! pip install lightning-sdk
from lightning_sdk import Studio

# run from inside a Studio: control the studio you are running on
s = Studio()
s.stop()

# create a new Studio in same Teamspace
s = Studio(name="start-stop-delete-demo", create_ok=True)

# before we call start, the Studio exists but isn't running
s.start()

# now let's stop the Studio
s.stop()

# and finally, delete it to leave no traces.
s.delete()

Switch machines

The SDK has a built-in enum for different machine types. Switching machines with this enum is as easy as running studio.switch_machine(Machine.T4) or whatever other type of machine you want to use. In contrast to the browser-based interface, the SDK actually blocks until the new machine is ready, as it would be unsafe to continue using the machine otherwise.

Example switching machines using the SDK:

#! pip install lightning-sdk
from lightning_sdk import Studio, Machine

# run from inside a Studio: control the studio you are running on
s = Studio()

# the Studio is already running since we're running from within the Studio - no need to start it
s.switch_machine(Machine.T4)

Another option is that, unlike through the browser, you can start Studios on a specific machine via the SDK. Just pass the machine type to the start method like so: studio.start(Machine.T4).

#! pip install lightning-sdk
from lightning_sdk import Studio, Machine

# run from inside a Studio: create a Studio in the same Teamspace
s = Studio(name="start-on-machine-demo")

# start the Studio directly on an T4 instead of first going to a CPU-4
s.start(machine.T4)

# delete the Studio to leave no traces
s.delete()

Note: this only works if your Studio wasn't already running.

Run Commands

To run commands, the SDK provides you with two options:

  1. Using studio.run("my fancy command"), the Studio will execute the command and return the output. If your command doesn't succeed (e.g. returns a nonzero exit code) this method will raise an error.
  2. The studio.run_with_exit_code("my fancy command") will also run the command but instead of raising and error on failure, it will always return the exit code alongside the output for you to handle it manually.

Example comparing the two mechanisms to run things:

#! pip install lightning-sdk
from lightning_sdk import Studio

# run from within a Studio: control the Studio you're in
s = Studio()

# run a command and don't care for specific exit-codes.
output = s.run("echo hello")
assert output == "hello"

# run a command that fails
try:
    output = s.run("echo failed; exit 5")
    # we know this won't happen since we explicitly exited with 5
    print(output)
except RuntimeError:
    # now we can do error handling
    print("an error occurred")

# get output and exit_code separately
output, exit_code = s.run_with_exit_code("echo hello; exit 5")
assert output == hello

# allows for error-handling depending on specific exit-codes
if exit_code == 0:
    # we know this won't happen since we explicitly exited with 5
    print("command ran successfully")
elif exit_code == 5:
    # this could be a specific error case to catch
    print("expected manual exit code")
else:
    # we only expect 0 or 5 as valid exit-codes here
    print("an error occurred", output)

Enable/disable autosleep

To directly enable autosleep from the SDK:

from lightning_sdk import Studio

s = Studio()

# Disable auto-sleep
s.auto_sleep = False

# Re-enable auto-sleep
s.auto_sleep = True

Exposing ports

To expose port 8000:

from lightning_sdk import Studio

s = Studio()

s.add_ports(8000)

To get the url of the exposed port 8000

from lightning_sdk import Studio

s = Studio()

# get url directly
url = s.add_ports(8000)[0].urls[0]

# get port
url = s.get_port_url(8000)

API Reference

Studio

from lightning_sdk import Studio

studio = Studio(name='my-studio-name', teamspace='my-teamspace', org='my-org')
studio.start()
studio.stop()
duplicated_studio = studio.duplicate()

A Studio object has the following functions:

  • __init__: creation of Studios
    • name (Optional[str]): the name of the studio
    • teamspace (Optional[Union[Teamspace, str]]): the teamspace the studio is contained by
    • org (Optional[Union[str, Organization]]): the organization owning the teamspace
    • user (Optional[Union[str, User]]): the user owning the teamspace, Note: Since a teamspace can either be owned by an org or by a user directly, only one of the arguments can be provided.
  • start: start a Studio
    • machine (Machine): the machine to start a Studio on, default: CPU-4
    • interruptible (bool): starts the Studio on a interruptible instance. Note that this means the Studio can be terminated at any point in time.
  • stop: stop a running Studio
  • delete: delete an existing Studio
  • duplicate: duplicate an existing Studio
    • Returns:
      • Studio: the new duplicated Studio
  • switch_machine: switch to a given machine type
    • machine (Machine): the new machine type to switch to
    • interruptible (bool): switches the Studio to a interruptible instance. Note that this means the Studio can be terminated at any point in time.
  • run_with_exit_code: run a command and get output and exit code
    • *commands (str): the commands to run on the Studio in sequence
    • Returns:
      • str: the command's output
      • int: the command's exit code
  • run: run a command and get the output. Raises an error for nonzero exit codes.
    • *commands (str): the commands to run on the Studio in sequence
    • Returns:
      • str: the command's output
  • upload_file: upload a given file to a remote path on the Studio
    • file_path (str): path to the file to upload
    • remote_path (Optional[str]): where the path should be stored on the Studio. Defaults to just the filename in the home directory.
    • progress_bar (bool): whether to display a progress bar for this file or not. Defaults to True
  • download_file: download a file from the Studio to a given target path
    • remote_path (str): the path inside the Studio to download
    • file_path (str): the local path to store the file contents in. Defaults to the path relative to the home directory inside the Studio starting from the current working directory.
  • download_folder: download a folder from the Studio to a given target path
    • remote_path (str): the path inside the Studio to download
    • target_path (Optional[str]): the local path to store the directory contents in. Defaults to the path relative to the home directory inside the Studio starting from the current working directory.
  • install_plugin: install a given plugin by name
    • plugin_name (str): the plugin to install
  • uninstall_plugin: uninstalls a given plugin by name
    • plugin_name (str): the plugin to install
  • run_plugin: runs a given plugin with the given arguments
    • plugin_name (str): the plugin to run
    • *args (Any): the positional arguments passed to the plugins run method
    • **kwargs (Any): the keyword arguments passed to the plugins run method``

A Studio object also has the following properties:

  • name (str): the Studio's name
  • status (Status): the Status of the Studio. Can be one of { NotCreated | Pending | Running | Stopping | Stopped | Failed }
  • teamspace (Teamspace): the Teamspace the Studio belongs to
  • owner (Union[User, Organization]): the owner of the Teamspace the Studio belongs to
  • machine (Optional[Machine]): the machine type the Studio is running on or None if it's not running
  • available_plugins (Mapping[str, str]): a dictionary with names of available plugins as keys and their descriptions as values
  • installed_plugins (Mapping[str, Plugin]): a dictionary with names of installed plugins as keys and the corresponding Plugin objects as values

Teamspace

from lightning_sdk import Teamspace

teamspace = Teamspace('my-teamspace', user='my-user')

all_studios = teamspace.studios
print(teamspace.name)

A Teamspace object has the following functions:

  • __init__: creating a Teamspace object
    • name (Optional[str]): the name of the Teamspace
    • org (Optional[Union[str, Organization]]): the organization owning the teamspace
    • user (Optional[Union[str, User]]): the user owning the teamspace, Note: Since a teamspace can either be owned by an org or by a user directly, only one of the arguments can be provided.

A Teamspace object also has the following properties:

  • name (str): the Teamspace's name
  • id (str): the Teamspace's ID
  • owner (Union[User, Organization]): the Teamspace's owner
  • studios (List[Studio]): a list of all Studios in a Teamspace the authenticated user can access
  • clusters (List[str]): a list of all cluster names in that Teamspace

Organization

from lightning_sdk import Organization

org = Organization('my-org')

print(org.name)

all_teamspaces = org.teamspaces

An Organization object has the following functions:

  • __init__: creating an Organization object
    • name (Optional[str]): the name of the Organization

An Organization object also has the following properties:

  • name (str): the Organization's name
  • id (str): the Organization's ID
  • teamspaces (List[Teamspace]): A list of Teamspaces owned by this Organization

User

from lightning_sdk import User

user = User('my-user')

print(user.name)
all_teamspaces = user.teamspaces

A User object has the following functions:

  • __init__: creating a User object
    • name (Optional[str]): the name of the User

A User object also has the following properties:

  • name (str): the User's name
  • id (str): the User's ID
  • teamspaces (List[Teamspace]): A list of Teamspaces owned by this User

Machine & CloudProvider

lightning-sdk offers abstractions for machines and cloud providers. These can generally be used as follows:

from lightning_sdk import Machine, CloudProvider

studio = Studio(name="my-new-studio", create_ok=True, cloud_provider=CloudProvider.AWS)
studio.start(Machine.L4)

Some machines (or their variants) are only supported on certain cloud providers. If a machine has multiple variants (for example the A100 has 40GB and 80GB VRAM variants), the Machine class has attributes for all of them (e.g. Machine.A100_40GB and Machine.A100_80GB). With multiple variants, every cloud provider usually provides only one of them (if at all). For better cloud agnostic machine specification, there always exists a case that's variant independent (e.g. Machine.A100) that would match for all variants depending on the given cloud provider:

from lightning_sdk import Machine, CloudProvider

studio = Studio(name="my-new-aws-studio", create_ok=True, cloud_provider=CloudProvider.AWS)
studio.start(Machine.A100_X_8) # starts on Machine.A100_40GB_X_8 since AWS only provides the 40GB variant of A100s

studio = Studio(name="my-new-aws-studio", create_ok=True, cloud_provider=CloudProvider.GCP)
studio.start(Machine.A100_X_8) # starts on Machine.A100_80GB_X_8 since GCP only provides the 80GB variant of A100s

Some cloud providers (mainly GCP) provide certain instance types only for a previously determined amount of time. This amount of time can be specified in seconds with the max_runtime argument:

from lightning_sdk import Machine, CloudProvider

studio = Studio(name="my-new-aws-studio", create_ok=True, cloud_provider=CloudProvider.GCP)
studio.start(Machine.H200, max_runtime=3600) # runs for 1 hour

If not specified, the default for this argument is 3 hours where necessary and the argument will be ignored if it's set but not required.

Was this page helpful?

What went wrong?

↑↓ Navigate↵ Open⌘↵ New tabesc Close