Skip to content

API

The Python API allows for running tasks, services, and managing runs programmatically.

dstack.api

dstack.api.Client

High-level API client for interacting with dstack server

Attributes:

Name Type Description
repos RepoCollection

operations with repos

backends BackendCollection

operations with backends

runs RunCollection

operations with runs

client APIClient

low-level server API client

project str

project name

__init__(api_client, project_name, ssh_identity_file=None)

Parameters:

Name Type Description Default
api_client APIClient

low-level server API client

required
project_name str

project name used for runs

required
ssh_identity_file Optional[PathLike]

SSH keypair to access instances

None

from_config(project_name=None, server_url=None, user_token=None, ssh_identity_file=None) staticmethod

Creates a Client using global config ~/.dstack/config.yml

Parameters:

Name Type Description Default
project_name Optional[str]

name of the project, required if server_url and user_token are specified

None
server_url Optional[str]

dstack server url, e.g. http://localhost:3000/

None
user_token Optional[str]

dstack user token

None
ssh_identity_file Optional[PathLike]

SSH keypair to access instances

None

Returns:

Type Description
Client

dstack Client

dstack.api.Task

Attributes:

Name Type Description
commands List[str]

The bash commands to run

ports List[PortMapping]

Port numbers/mapping to expose

env Dict[str, str]

The mapping or the list of environment variables

image Optional[str]

The name of the Docker image to run

python Optional[str]

The major version of Python

entrypoint Optional[str]

The Docker entrypoint

registry_auth Optional[RegistryAuth]

Credentials for pulling a private container

home_dir str

The absolute path to the home directory inside the container. Defaults to /root.

dstack.api.Service

Attributes:

Name Type Description
commands List[str]

The bash commands to run

port PortMapping

The port, that application listens to or the mapping

env Dict[str, str]

The mapping or the list of environment variables

image Optional[str]

The name of the Docker image to run

python Optional[str]

The major version of Python

entrypoint Optional[str]

The Docker entrypoint

registry_auth Optional[RegistryAuth]

Credentials for pulling a private container

home_dir str

The absolute path to the home directory inside the container. Defaults to /root.

dstack.api.FineTuningTask

This task configuration loads a given model from the Hugging Face hub and fine-tunes it on the provided dataset (also from Hugging Face hub), utilizing the SFT and QLoRA techniques. The final model is pushed to Hugging Face hub.

Parameters:

Name Type Description Default
model_name str

The model that you want to train from the Hugging Face hub. E.g. gpt2, gpt2-xl, bert, etc.

required
dataset_name str

The instruction dataset to use.

required
new_model_name Optional[str]

The name to use for pushing the fine-tuned model to the Hugging Face Hub. If unset, it defaults to the name of the run.

None
env Dict[str, str]

The list of environment variables, e.g. "HUGGING_FACE_HUB_TOKEN", "WANDB_API_KEY", "WANDB_PROJECT", etc.

None
report_to Optional[str]

Supported integrations include "wandb" and "tensorboard".

None
per_device_train_batch_size int

Batch size per GPU for training.

4
per_device_eval_batch_size int

Batch size per GPU for evaluation.

4
gradient_accumulation_steps int

Number of update steps to accumulate the gradients for.

1
learning_rate float

Initial learning rate (AdamW optimizer).

0.0002
max_grad_norm float

Maximum gradient normal (gradient clipping).

0.3
weight_decay float

Weight decay to apply to all layers except bias/LayerNorm weights.

0.001
lora_alpha int

Alpha parameter for LoRA scaling.

16
lora_dropout float

Dropout probability for LoRA layers.

0.1
lora_r int

LoRA attention dimension.

64
max_seq_length Optional[int]

Maximum sequence length to use.

None
use_4bit bool

Activate 4bit precision base model loading.

True
use_nested_quant bool

Activate nested quantization for 4bit base models.

True
bnb_4bit_compute_dtype str

Compute dtype for 4bit base models.

'float16'
bnb_4bit_quant_type str

Quantization type fp4 or nf4.

'nf4'
num_train_epochs float

The number of training epochs for the reward model.

1
fp16 bool

Enables fp16 training.

False
bf16 bool

Enables bf16 training.

False
packing bool

Use packing dataset creating.

False
gradient_checkpointing bool

Enables gradient checkpointing.

True
optim str

The optimizer to use.

'paged_adamw_32bit'
lr_scheduler_type str

Learning rate schedule. Constant a bit better than cosine, and has advantage for analysis

'constant'
max_steps int

How many optimizer update steps to take

-1
warmup_ratio float

Fraction of steps to do a warmup for

0.03
group_by_length bool

Group sequences into batches with same length. Saves memory and speeds up training considerably.

True
save_steps int

Save checkpoint every X updates steps.

0
logging_steps int

Log every X updates steps.

25

dstack.api.CompletionService

This service configuration loads a given model from the Hugging Face hub and deploys it as a public endpoint.

Parameters:

Name Type Description Default
model_name str

The model that you want to deploy from the Hugging Face hub. E.g. gpt2, gpt2-xl, bert, etc.

required
env Dict[str, str]

The list of environment variables, e.g. "HUGGING_FACE_HUB_TOKEN", "WANDB_API_KEY", "WANDB_PROJECT", etc.

None
quantize Optional[str]

Whether you want the model to be quantized. Supported values: "awq", "eetq", "gptq", and "bitsandbytes".

None
dtype Optional[str]

The dtype to be forced upon the model. This option cannot be used with quantize.

None

dstack.api.CompletionTask

This task configuration loads a specified model from the Hugging Face hub and runs it as a private endpoint while forwarding the port to the local machine.

Parameters:

Name Type Description Default
model_name str

The model that you want to deploy from the Hugging Face hub. E.g. gpt2, gpt2-xl, bert, etc.

required
env Dict[str, str]

The list of environment variables, e.g. "HUGGING_FACE_HUB_TOKEN", "WANDB_API_KEY", "WANDB_PROJECT", etc.

None
quantize Optional[str]

Whether you want the model to be quantized. Supported values: "awq", "eetq", "gptq", and "bitsandbytes".

None
dtype Optional[str]

The dtype to be forced upon the model. This option cannot be used with quantize.

None
local_port int

The local port to forward the traffic to.

80

dstack.api.Run

Attributes:

Name Type Description
name str

run name

ports Optional[Dict[int, int]]

ports mapping, if run is attached

backend Optional[BackendType]

backend type

status JobStatus

run status

hostname str

instance hostname

attach(ssh_identity_file=None)

Establish an SSH tunnel to the instance and update SSH config

Parameters:

Name Type Description Default
ssh_identity_file Optional[PathLike]

SSH keypair to access instances

None

Raises:

Type Description
PortUsedError

If ports are in use or the run is attached by another process.

detach()

Stop the SSH tunnel to the instance and update SSH config

logs(start_time=None, diagnose=False)

Iterate through run's log messages

Parameters:

Name Type Description Default
start_time Optional[datetime]

minimal log timestamp

None
diagnose bool

return runner logs if True

False

Yields:

Type Description
Iterable[bytes]

log messages

refresh()

Get up-to-date run info

stop(abort=False)

Terminate the instance and detach

Parameters:

Name Type Description Default
abort bool

gracefully stop the run if False

False

dstack.api.Client.runs

Operations with runs

exec_plan(run_plan, repo, reserve_ports=True)

Execute run plan

Parameters:

Name Type Description Default
run_plan RunPlan

result of get_plan call

required
repo Repo

repo to use for the run

required
reserve_ports bool

reserve local ports before submit

True

Returns:

Type Description
SubmittedRun

submitted run

get(run_name)

Get run by run name

Parameters:

Name Type Description Default
run_name str

run name

required

Returns:

Type Description
Optional[Run]

run or None if not found

get_plan(configuration, repo, configuration_path=None, backends=None, resources=None, spot_policy=None, retry_policy=None, max_duration=None, max_price=None, working_dir=None, run_name=None)

Get run plan. Same arguments as submit

Returns:

Type Description
RunPlan

run plan

list(all=False)

List runs

Parameters:

Name Type Description Default
all bool

show all runs (active and finished) if True

False

Returns:

Type Description
List[Run]

list of runs

submit(configuration, configuration_path=None, repo=None, backends=None, resources=None, spot_policy=None, retry_policy=None, max_duration=None, max_price=None, working_dir=None, run_name=None, reserve_ports=True)

Submit a run

Parameters:

Name Type Description Default
configuration AnyRunConfiguration

run configuration

required
configuration_path Optional[str]

run configuration path, relative to repo_dir

None
repo Optional[Repo]

repo to use for the run

None
backends Optional[List[BackendType]]

list of allowed backend for provisioning

None
resources Optional[ProfileResources]

minimal resources for provisioning

None
spot_policy Optional[SpotPolicy]

spot policy for provisioning

None
retry_policy Optional[ProfileRetryPolicy]

retry policy for interrupted jobs

None
max_duration Optional[Union[int, str]]

max instance running duration in seconds

None
max_price Optional[float]

max instance price in dollars per hour for provisioning

None
working_dir Optional[str]

working directory relative to repo_dir

None
run_name Optional[str]

desired run_name. Must be unique in the project

None
reserve_ports bool

reserve local ports before submit

True

Returns:

Type Description
SubmittedRun

submitted run

dstack.api.Client.repos

Operations with repos

init(repo, git_identity_file=None, oauth_token=None)

Upload credentials and initializes the repo in the project

Parameters:

Name Type Description Default
repo Repo

repo to initialize

required
git_identity_file Optional[PathLike]

SSH private key to access the remote repo

None
oauth_token Optional[str]

GitHub OAuth token to access the remote repo

None

is_initialized(repo)

Checks if the remote repo is initialized in the project

Parameters:

Name Type Description Default
repo Repo

repo to check

required

Returns:

Type Description
bool

repo is initialized

load(repo_dir, local=False, init=False, git_identity_file=None, oauth_token=None)

Loads the repo from the local directory using global config

Parameters:

Name Type Description Default
repo_dir PathLike

repo root directory

required
local bool

do not try to load RemoteRepo first

False
init bool

initialize the repo if it's not initialized

False
git_identity_file Optional[PathLike]

path to an SSH private key to access the remote repo

None
oauth_token Optional[str]

GitHub OAuth token to access the remote repo

None

Raises:

Type Description
ConfigurationError

if the repo is not initialized and init is False

Returns:

Name Type Description
repo Union[RemoteRepo, LocalRepo]

initialized repo

dstack.api.Client.backends

Operations with backends

list()

List available backends in the project

Returns:

Type Description
List[Backend]

backends