Multi-machine jobs

The MMT class in the Lightning SDK allows you to run multi-machine jobs asynchronously inside a Studio. Use this to offload tasks like data processing, training runs, or inference jobs that do not require interactive execution but multiple nodes communicating with each other.

The SDK is automatically available inside Studios. To use the SDK outside a Studio, install it with:

pip install --upgrade lightning-sdk

Ensure you authenticate using environment variables or login via CLI.

export LIGHTNING_USER_ID=your-user-id
export LIGHTNING_API_KEY=your-api-key

Submit a MMT

Here is a full example of how to submit a multi-machine job.

# Install the Lightning SDK
# pip install lightning-sdk

# Export authentication variables
# export LIGHTNING_USER_ID=your-user-id
# export LIGHTNING_API_KEY=your-api-key

from lightning_sdk import Studio, Machine, MMT

# Step 1: Initialize a Studio
studio = Studio(name="batch-processing", teamspace="my-teamspace", user="my-user")
studio.start()

# Step 3: Run a Multi-machine job
job = MMT.run(command="python process_data.py", name="data-job", machine=Machine.CPU, studio=studio, num_machines=2)

# Step 4: Monitor Job Status
print(job.status)  # Running, Completed, Failed

# Step 5: Stop or Delete a Job
# job.stop()  # Gracefully stop a running job
# job.delete()  # Cancel and remove the job

# Step 6: Shut Down the Studio
studio.stop()

Initialize a Studio

Before running a job, first initialize and start a Studio:

from lightning_sdk import Studio

studio = Studio(name="batch-processing", teamspace="my-teamspace", user="my-user")
studio.start()

Submit the Job

With the Studio defining the compute environment, submit the multi-machine job:

job = MMT.run(command="python process_data.py", name="data-job", machine=Machine.CPU, studio=studio, num_machines=2)

List running multi-machine jobs

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

for job in teamspace.multi_machine_jobs:
    print(job.name, job.status)

Check multi-machine status

job = MMT("data-job", teamspace="my-teamspace", user="my-user")
print(job.status)  # Running, Completed, Failed

Stop or cancel a multi-machine job

If a job is no longer needed, you can stop or delete it:

job.stop()  # Gracefully stops a job
job.delete()  # Cancels and removes the job

Run multiple jobs in parallel

You can launch multiple jobs on different machine types:

MMT.run(command="python preprocess.py", name="preprocess", machine=Machine.CPU, studio=studio, num_machines=2)
MMT.run(command="python train.py", name="train-job", machine=Machine.A10G, studio=studio, num_machines=2)
MMT.run(command="python infer.py", name="inference", machine=Machine.L4, num_machines=2)

Automate Jobs

GitHub actions

You can schedule jobs using GitHub Actions to run at a fixed time:

name: Scheduled Job

on:
  schedule:
    - cron: '0 3 * * *'  # Runs at 3 AM UTC daily

jobs:
  run_batch:
    runs-on: ubuntu-22.04
    steps:
      - name: Set up Python
        uses: actions/setup-python@v2
        with:
          python-version: "3.10"

      - name: Install Lightning SDK
        run: pip install lightning-sdk

      - name: Start a Job
        run: |
          python -c " \
              from lightning_sdk import Studio, MMT, Machine; \
              studio = Studio(name='batch-pipeline', teamspace='my-team', user='ci-bot'); \
              studio.start(); \
              MMT.run(command='python daily_task.py', name='daily-task', studio=studio, machine=Machine.A10, num_machines=2); \
              studio.stop();"
        env:
          LIGHTNING_USER_ID: ${{ secrets.LIGHTNING_USER_ID }}
          LIGHTNING_API_KEY: ${{ secrets.LIGHTNING_API_KEY }}

Use CPU machines for lightweight preprocessing tasks.
Use GPU machines for deep learning inference and training.
Enable interruptible mode for cost-saving on long-running multi-machine jobs.
Chain multi-machine jobs using status checks to create multi-step pipelines.

API Reference

init()

The init() method fetches an existing multi-machine job and accepts the following arguments:

Parameter	Type	Description	Default
name	str	The name of the Job	-
teamspace	str	Teamspace	The teamspace the job is part of
org	str	Organization	The name of the organization owning the teamspace in case it is owned by an org
user	str	User	The name of the user owning the teamspace in case it is owned by a user instead of an org

Example:

job = MMT("data-job", teamspace="my-teamspace", user="my-user")

run()

The run() method submits a new multi-machine job and accepts the following arguments:

Parameter	Type	Description	Default
name	str	The name of the job	-
machine	Machine	Machine to run the Job on	-
num_machines	int	The number of machines to run the job on	-
command	str	The command to run inside the Studio	None, The command to run inside your job. Required is using a Studio as compute environment. Optional for docker jobs. In case it's not specified for docker jobs, these will run the default command specified by the image.
studio	Studio	str	None
image	str	None	The docker image to run the job with. Mutually exclusive with studio.
teamspace	Teamspace	str	None
org	str	Organization	None
user	str	User	None
cloud_account	str	None	The cloud acocunt to run the job on
env	Dict[str, str]	None	Environment variables to set inside the job.
interruptible	bool	Whether the job should run on interruptible instances. They are cheaper but can be preempted.	False
image_credentials	str	None	The credentials used to pull the image. Required if the image is private. This should be the name of the respective credentials secret created on the Lightning AI platform.
cloud_account_auth	bool	Whether to authenticate with the cloud account to pull the image. Required if the registry is part of a cloud provider (e.g. ECR).	False
entrypoint	str	The entrypoint of your docker container. To use the pre-defined entrypoint of the provided image, set this to an empty string. Only applicable when submitting docker jobs.	"sh -c", which just runs the provided command in a standard shell.
path_mappings	Dict[str, str]	None	Dictionary of path mappings. The keys are the path inside the container whereas the value represents the data-connection name and the path inside that connection. Should be of form{ "": ":", "": ""}If the path inside the connection is omitted it's assumed to be the root path of that connection.Only applicable when submitting docker jobs.

Example:

job = MMT.run(command="python process_data.py", name="data-job", machine=Machine.T4, num_machines=2, studio=studio)

stop()

Stops the respective job. Accepts no arguments.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
job.stop()

delete()

Deletes the respective multi-machine job. Note that this also deletes all respective artifacts and is non-reversible. Accepts no arguments.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
job.delete()

status

Property returning the current status of the multi-machine job.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.status) # -> Status.Pending, Status.Running, Status.Completed, Status.Stopped, Status.Failed

machine

Property returning the machine the multi-machine job is/was running on.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.machine) # -> Machine.A10G

machines

Property returning objects referring to the individual machines the job is running on

artifacts_uri

Always None on a multi-machine job: each machine has its own artifacts, so there is no single address. Read one machine's address from machine.artifacts_uri, or use list_artifacts() and download_artifacts() below.

Each machine resolves its own address exactly as a single-machine job does, giving three outcomes:

A Studio-backed machine writes to jobs/<job-name>-<rank>.
A machine running a container image launched with artifacts_destination writes to the connection it names, at <type>_connections/<name>/<path>.
A machine running a container image launched without one keeps no artifacts, and its artifacts_uri is None.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
for machine in job.machines:
    # "lit://my-org/my-teamspace/jobs/data-job-0"                  Studio-backed
    # "lit://my-org/my-teamspace/s3_connections/my-bucket/run-7"   container + artifacts_destination
    # None                                                         container, no artifacts_destination
    print(machine.artifacts_uri)

list_artifacts()

Lists what the job wrote to the teamspace drive, across every machine. Returns one entry per file or folder, each with a path, is_dir, size and last_modified. Every path is prefixed with the name of the machine that wrote it, so the listing and download_artifacts() agree. A machine that keeps no artifacts contributes no entries, so the list is empty for a container-image job launched without artifacts_destination.

Parameter	Type	Description	Default
path	str	Subfolder of each machine's artifacts to list.	"", all of them
recursive	bool	Whether to descend into subfolders.	False

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
for entry in job.list_artifacts(recursive=True):
    if not entry.is_dir:
        print(f"{entry.path} ({entry.size} bytes)") # -> "data-job-0/metrics.json (1024 bytes)"

download_artifacts()

Downloads what the job wrote to a local directory, giving each machine a folder named after it.

Raises RuntimeError at the first machine that keeps no artifacts — every machine of a container-image job launched without artifacts_destination. The machines are taken in rank order, so the error can surface after earlier ones have already downloaded.

Parameter	Type	Description	Default
target_dir	str \| Path	Local directory to download into. Created if it does not exist.	".", the working directory
path	str	Subfolder of each machine's artifacts to download.	"", all of them

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
job.download_artifacts("./artifacts") # -> ./artifacts/data-job-0, ./artifacts/data-job-1, ...

# Take one subfolder, so a long run does not have to come down whole
job.download_artifacts("./checkpoints", path="checkpoints")

name

Property returning the name of the multi-machine job.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.name) # -> "data-job"

teamspace

Property returning the teamspace the multi-machine job belongs to.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.teamspace) # -> Teamspace("my-teamspace", org="my-org")

studio

Property returning the studio used as compute environment if it wasn't launched with a docker image.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.studio) # -> Studio("my-studio", teamspace="my-teamspace", org="my-org") or None

image

Property returning the image used as compute environment if it wasn't launched with a Studio.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.image) # -> "ubuntu:latest" or None

command

Property returning the command ran within the multi-machine job.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.command) # -> "python train.py"

logs

Property returning the logs generated by running the multi-machine job. Will return an error if the job is still pending or running (check job.status)

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.logs) # -> "Epoch 1/2: ################"

link

Property returning the link to view the job in the UI.

Example:

job = MMT("data-job", teamspace="my-teamspace", org="my-org")
print(job.link) # -> "https://lightning.ai/my-org/my-teamspace/studios/my-studio/app?app_id=jobs&job_name=data-job"

Was this page helpful?

Multi-machine jobs