Plugins¶
A plugin is a frozen, serializable value included in the compute specification. The daemon writes its kind and parameters with the compute and reconstructs it on the worker. Plugins must not contain live clients, sockets, closures, or other process-local handles.
import skyward as sky
with sky.Compute(
provider=sky.AWS(),
accelerator="A100",
plugins=[sky.plugins.Torch(backend="nccl")],
) as compute:
train(data) >> compute
The provider belongs on Compute; plugins belong in plugins=[...]. A plugin can transform the image, add bootstrap phases, prepare the worker, wrap each task, and install a client-side context.
Built-in plugins¶
The worker registers these plugin kinds under sky.plugins:
| Plugin | Purpose | Main fields |
|---|---|---|
Torch |
PyTorch distributed initialization | backend, cuda, version |
HuggingFace |
Hugging Face installation and token setup | token |
Joblib |
Joblib parallel backend | version |
Jax |
JAX distributed initialization | cuda |
Keras |
Keras backend setup | backend |
Cuml |
RAPIDS cuML installation and initialization | cuda |
Sklearn |
scikit-learn client integration | version |
Accelerate |
Hugging Face Accelerate environment and collectives | config, env |
Mig |
NVIDIA MIG setup | profile |
Mps |
Apple Metal performance settings | active_thread_percentage, pinned_memory_limit |
Collective plugins such as Torch, Jax, and Accelerate make the worker topology part of the job. A compute using one cannot be elastically resized while the collective is active.
Plugin lifecycle¶
The base plugin hooks have separate scopes:
image(image)changes packages or image settings before provisioning;bootstrap(image, concurrency)adds generated node bootstrap phases;setup(info)runs once for the worker lifetime;run(call, info)wraps each task;client(compute)runs once in the client process while the compute is ready.
skyward.plugins.Plugin
¶
Bases: Struct
A plugin, with every hook optional and none of them doing anything by default.
Attributes:
| Name | Type | Description |
|---|---|---|
kind |
str
|
Its name on the wire, and how it is found again on the node. |
collective |
bool
|
Whether the plugin makes the nodes depend on each other. A collective freezes the world when the last rank joins it, so a compute running one cannot be resized: taking a rank away does not shrink the job, it hangs it at the next all-reduce on a peer that is never going to answer. The reconciler reads this and refuses to scale such a compute at all. |
kind
class-attribute
¶
collective = False
class-attribute
¶
image(image)
¶
What the machine needs installed before the plugin can run at all.
The daemon calls this once, when the compute is provisioned. It is a transform rather than a list of packages because plugins compose: each is handed what the ones before it asked for.
bootstrap(image, concurrency)
¶
Extra shell phases, appended after the image's own bootstrap.
Runs on the daemon at script-generation time, like :meth:image — not on
the node — so it may only return the phases the script will run, never do
anything itself. concurrency is the worker's width, the one datum a
phase needs that the image does not carry: a plugin that partitions the
machine has to know how many ways.
setup(info)
¶
The worker's lifetime, on the node.
Entered once, before the worker takes a task, and left when it stops. This is where a process group is formed and where an environment variable that a library reads at import time is set — both being things that must exist before the first task, not around each one.
run(call, info)
¶
One task, on the node.
Plugins wrap in the order they were listed: the first is outermost, and therefore the one that sees the others' work.
client(compute)
¶
The plugin's say on the client, for as long as the pool is up.
Unlike the others, this hook does not travel: it runs on the instance the
user constructed, in the process that opened the with block, and never
on a node. It is entered once the compute is ready and left before it is
torn down — the place a plugin reaches back into the live pool, as joblib
does to point its parallel backend at it.
ref()
¶
skyward.plugins.Torch
¶
Bases: Plugin
Install torch, and form the process group before the first task runs it.
The rendezvous is rank zero, because torch insists on being told where it is — the compute has no head, and this is the convention that satisfies a library that believes it does.
It is formed in the process that runs the task, and on that process's first
task, not at worker start. init_process_group is a collective — every node
blocks in it until the last one arrives — and the node that arrives there must
be the one that will run the collective code afterwards. Under a subprocess
executor that is the child, not the worker; forming it in the worker would leave
the child holding a group it never joined. Doing it on the first task, once and
under a lock, is what lets the same plugin serve either executor.
Attributes:
| Name | Type | Description |
|---|---|---|
backend |
Literal['nccl', 'gloo']
|
|
cuda |
str
|
The CUDA build to install torch from, as a |
version |
str | None
|
Pin, if the code needs one. Otherwise whatever the index has. |
skyward.plugins.HuggingFace
¶
Bases: Plugin
Put huggingface_hub on the machine and the token in its environment.
The token is set in setup rather than baked into the image because the image
is a description of a machine and the token is a secret about a person. It is
still carried in the spec, and the spec is still written to the database in the
clear — which is a thing to fix, and is not fixed by putting it somewhere else
on the same machine.
Attributes:
| Name | Type | Description |
|---|---|---|
token |
str | None
|
Read from |
skyward.plugins.Joblib
¶
skyward.plugins.Jax
¶
Bases: Plugin
Install jax, and join the process cluster before the first task runs it.
The rendezvous is rank zero, because jax insists on being told where the coordinator is — the compute has no head, and this is the convention that satisfies a library that believes it does.
It is joined in the process that runs the task, and on that process's first
task, not at worker start. jax.distributed.initialize is a collective —
every node blocks in it until the last one arrives — and the node that arrives
there must be the one that will run the collective code afterwards. Under a
subprocess executor that is the child, not the worker; joining in the worker
would leave the child outside a cluster it never entered. Doing it on the first
task, once and under a lock, is what lets the same plugin serve either executor.
Attributes:
| Name | Type | Description |
|---|---|---|
cuda |
str
|
The CUDA build to install jax from, as a |
skyward.plugins.Keras
¶
Bases: Plugin
Install Keras and its backend, and set the backend before the first import.
Keras reads KERAS_BACKEND once, when it is imported, and is stuck with what
it finds — so the backend is set in setup, in the worker process and before
any task imports keras. The image cannot do it: its env reaches the bootstrap
shell, which exits, and never the worker the tasks run in. The backend value
doubles as its package name, which is why one field installs both.
Multi-node training on the jax backend is data-parallel: every node runs the
same graph over its own shard, and the only thing they must agree on is the
random state, which setup synchronizes. Forming the JAX process group is not
this plugin's job — pair it with the jax plugin, which is the collective and
the one that knows the rendezvous.
Attributes:
| Name | Type | Description |
|---|---|---|
backend |
Literal['jax', 'tensorflow', 'torch']
|
The framework Keras runs on, and the package installed to carry it. |
skyward.plugins.Cuml
¶
Bases: Plugin
Install cuML from NVIDIA's index, and put its sklearn accelerator on before the first task runs.
The accelerator is installed in the process that runs the task, on that
process's first task, not at worker start. cuml.accel.install() rewrites
scikit-learn's estimators in the interpreter that imports them, and under a
subprocess executor that is the child, not the worker; installing it in the
worker would leave the child running plain CPU sklearn. Doing it on the first
task, once and under a lock, is what lets the same plugin serve either
executor.
Attributes:
| Name | Type | Description |
|---|---|---|
cuda |
str
|
The CUDA suffix of the RAPIDS wheel to install, e.g. |
skyward.plugins.Sklearn
¶
skyward.plugins.Accelerate
¶
Bases: Plugin
Install accelerate, and set its distributed environment before the first task.
The rendezvous is rank zero, because accelerate insists on being told where
the main process is — the compute has no head, and this is the convention that
satisfies a library that believes it does.
The group is formed in the process that runs the task, and on that process's
first task, not at worker start. init_process_group is a collective — every
node blocks in it until the last one arrives — and the node that arrives there
must be the one that will run the collective code afterwards. Under a subprocess
executor that is the child, not the worker; forming it in the worker would leave
the child holding a group it never joined. Doing it on the first task, once and
under a lock, is what lets the same plugin serve either executor.
It composes with :class:~skyward.worker.plugins.torch.Torch: the group is only formed
when nobody has formed one yet, so a user may list both and pay for the group once.
Attributes:
| Name | Type | Description |
|---|---|---|
config |
dict[str, Any]
|
The accelerate settings, in the same shape as the YAML |
skyward.plugins.Mig
¶
Bases: Plugin
Partition the GPU with MIG and give each subprocess its own slice.
The GPU is put in MIG mode and cut into concurrency instances during
bootstrap; at run time each worker subprocess reads the slice UUIDs and pins the
one at its index through CUDA_VISIBLE_DEVICES. That indexing is the whole
contract, and it only holds under executor='process' with reuse=True:
there each concurrent slot is a distinct, long-lived child with a stable
info.worker, so slice k belongs to child k for the child's life. Under
the thread executor every task shares one process and one info.worker of
zero — they would all pin the same slice — and without reuse a child dies
after its task, so the pinning buys nothing.
The pin is set on the process's first task, once and under a lock, because the
process that must see CUDA_VISIBLE_DEVICES is the one that will import the
GPU library and run the task — the child, not the worker that spawned it.
Attributes:
| Name | Type | Description |
|---|---|---|
profile |
str
|
The MIG profile every slice is cut to, e.g. |
skyward.plugins.Mps
¶
Bases: Plugin
Bring up the MPS control daemon before the worker takes its first task.
Multi-Process Service lets several CUDA processes share one GPU context instead
of each taking the whole device, which is what a worker running tasks in parallel
needs the GPU to allow. The daemon is started in setup, in the worker process
and before any child is spawned, so every task inherits the pipe directory it
rendezvous on. The image cannot carry it: MPS ships with the CUDA driver, so
there is nothing to install, and env reaches only the bootstrap shell that has
since exited — the daemon and its variables have to be put up where the tasks run.
Starting the daemon is best-effort. On a machine without the control binary the call fails and is swallowed, because a worker that cannot share its GPU should still run the task on the whole one, not refuse to start.
Attributes:
| Name | Type | Description |
|---|---|---|
active_thread_percentage |
int | None
|
Ceiling on the share of GPU compute one client may use, 1-100. Left to MPS's default when unset. |
pinned_memory_limit |
str | None
|
Per-device pinned memory limit, e.g. |