Plugins

skyward.plugins

Skyward plugins for third-party integrations.

All imports are lazy to avoid requiring optional dependencies at import time.

Usage: import skyward as sky

with sky.ComputePool(
    plugins=[sky.plugins.torch(backend="nccl")],
) as compute:
    ...

__all__ = ['AccelerateConfig', 'LaunchCommand', 'LaunchContext', 'Plugin', 'accelerate', 'torch', 'jax', 'keras', 'cuml', 'joblib', 'sklearn', 'mig', 'mps', 'around_app', 'around_client', 'around_process'] module-attribute

AccelerateConfig

Bases: TypedDict

Accelerate config in the same structure as the YAML produced by accelerate config.

Skyward injects topology fields automatically (num_machines, machine_rank, main_process_ip, main_process_port, compute_environment). distributed_type is inferred from the presence of fsdp or deepspeed when not set explicitly.

num_processes means total across all nodes (matching accelerate semantics). Defaults to the cluster node count (1 GPU per node).

distributed_type instance-attribute

mixed_precision instance-attribute

num_processes instance-attribute

num_cpu_threads_per_process instance-attribute

gpu_ids instance-attribute

same_network instance-attribute

debug instance-attribute

downcast_bf16 instance-attribute

rdzv_backend instance-attribute

fsdp instance-attribute

deepspeed instance-attribute

LaunchCommand dataclass

Structured worker launch command.

Parameters:

Name	Type	Description	Default
launcher	str	Wrapper tool that launches the worker process (e.g. "accelerate launch --num_machines 4"). Empty string means no wrapper — Python runs the entrypoint directly.	''
entrypoint	str	Worker script or -c invocation.	''
args	tuple[str, ...]	CLI arguments for the worker (--node-id, --port, etc.).	()

launcher = '' class-attribute instance-attribute

entrypoint = '' class-attribute instance-attribute

args = () class-attribute instance-attribute

render()

Render to a single shell command string.

__init__(launcher='', entrypoint='', args=())

LaunchContext dataclass

Context available to launcher transform hooks.

Parameters:

Name	Type	Description	Default
node_id	int	Index of this node in the cluster (0 = head).	required
num_nodes	int	Total number of nodes in the cluster.	required
head_addr	str	IP address of the head node.	required
casty_port	int	Port for the Casty actor system.	required
private_ip	str	Private IP of this node.	required
python_bin	str	Path to the Python binary inside the remote venv.	required
venv_dir	str	Path to the remote venv directory.	required
concurrency	int	Number of workers per node.	required
executor	str	Execution backend ("thread" or "process").	required

node_id instance-attribute

num_nodes instance-attribute

head_addr instance-attribute

casty_port instance-attribute

private_ip instance-attribute

python_bin instance-attribute

venv_dir instance-attribute

concurrency instance-attribute

executor instance-attribute

__init__(node_id, num_nodes, head_addr, casty_port, private_ip, python_bin, venv_dir, concurrency, executor)

Plugin dataclass

Declarative third-party integration bundle.

Bundles environment setup, bootstrap ops, worker lifecycle hooks, client-side hooks, and per-task wrapping into a single composable unit.

Parameters:

Name	Type	Description	Default
name	str	Plugin identifier.	required
transform	ImageTransform[Any] | None	(Image, Cluster[S]) -> Image transformer. Receives current Image and cluster metadata, returns modified copy.	None
bootstrap	BootstrapFactory[Any] | None	Factory that receives Cluster[S] and returns extra shell ops appended after Image-driven bootstrap phases.	None
decorate	TaskDecorator | None	Classic Python decorator: (fn) -> fn. Wraps each @sky.function function at execution time on the remote worker.	None
around_app	AppLifecycle | None	Worker lifecycle context manager: InstanceInfo -> ContextManager[None]. Entered once in the main worker process.	None
around_process	ProcessLifecycle | None	Subprocess lifecycle context manager: InstanceInfo -> ContextManager[None]. Entered once per subprocess when executor="process". Lazy -- enters on the first task execution in each subprocess, after env vars are propagated.	None
around_client	ClientLifecycle[Any] | None	Client lifecycle context manager: (Pool, Cluster[S]) -> ContextManager[None].	None
launcher	LauncherTransform | None	(LaunchCommand, LaunchContext) -> LaunchCommand transformer. Wraps the worker launch command with a process launcher like accelerate or torchrun.	None

name instance-attribute

transform = None class-attribute instance-attribute

bootstrap = None class-attribute instance-attribute

decorate = None class-attribute instance-attribute

around_app = None class-attribute instance-attribute

around_process = None class-attribute instance-attribute

around_client = None class-attribute instance-attribute

launcher = None class-attribute instance-attribute

create(name) staticmethod

Create an empty plugin with just a name.

Use the with_* builder methods to attach hooks.

Parameters:

Name	Type	Description	Default
name	str	Plugin identifier.	required

Returns:

Type	Description
Plugin	Empty plugin ready for hook attachment.

Examples:

>>> plugin = (
...     Plugin.create("my-plugin")
...     .with_image_transform(add_deps)
...     .with_decorator(wrap_fn)
... )

with_image_transform(transform)

Attach an image transform hook.

Parameters:

Name	Type	Description	Default
transform	ImageTransform[S]	(Image, Cluster[S]) -> Image that modifies the image before bootstrap.	required

Returns:

Type	Description
Plugin	New plugin instance with the transform attached.

with_bootstrap(factory)

Attach a bootstrap factory hook.

Parameters:

Name	Type	Description	Default
factory	BootstrapFactory[S]	Cluster[S] -> tuple[Op, ...] returning extra shell ops appended after image-driven bootstrap phases.	required

Returns:

Type	Description
Plugin	New plugin instance with the bootstrap factory attached.

with_decorator(decorate)

Attach a per-task decorator hook.

Parameters:

Name	Type	Description	Default
decorate	TaskDecorator[P, R]	Classic Python decorator (fn) -> fn applied to each @sky.function at execution time on the remote worker.	required

Returns:

Type	Description
Plugin	New plugin instance with the decorator attached.

with_around_app(around)

Attach a worker lifecycle hook.

Parameters:

Name	Type	Description	Default
around	AppLifecycle	InstanceInfo -> ContextManager[None] entered once in the main worker process.	required

Returns:

Type	Description
Plugin	New plugin instance with the lifecycle hook attached.

with_around_process(around)

Attach a subprocess lifecycle hook.

Parameters:

Name	Type	Description	Default
around	ProcessLifecycle	InstanceInfo -> ContextManager[None] entered once per subprocess when executor="process".	required

Returns:

Type	Description
Plugin	New plugin instance with the subprocess hook attached.

with_around_client(around)

Attach a client-side lifecycle hook.

Parameters:

Name	Type	Description	Default
around	ClientLifecycle[S]	(Pool, Cluster[S]) -> ContextManager[None] entered on the client at pool __enter__.	required

Returns:

Type	Description
Plugin	New plugin instance with the client hook attached.

with_launcher(launcher)

Attach a worker launch command transform.

Parameters:

Name	Type	Description	Default
launcher	LauncherTransform	(LaunchCommand, LaunchContext) -> LaunchCommand that wraps the worker startup with a process launcher (e.g. accelerate launch, torchrun).	required

Returns:

Type	Description
Plugin	New plugin instance with the launcher hook attached.

__init__(name, transform=None, bootstrap=None, decorate=None, around_app=None, around_process=None, around_client=None, launcher=None)

around_app(name, around)

Create a plugin with only a worker lifecycle hook.

Shortcut for Plugin.create(name).with_around_app(around).

Parameters:

Name	Type	Description	Default
name	str	Plugin identifier.	required
around	AppLifecycle	Worker lifecycle context manager.	required

Returns:

Type	Description
Plugin	Plugin with the around_app hook set.

around_client(name, around)

Create a plugin with only a client-side lifecycle hook.

Shortcut for Plugin.create(name).with_around_client(around).

Parameters:

Name	Type	Description	Default
name	str	Plugin identifier.	required
around	ClientLifecycle[S]	Client-side lifecycle context manager.	required

Returns:

Type	Description
Plugin	Plugin with the around_client hook set.

around_process(name, around)

Create a plugin with only a subprocess lifecycle hook.

Shortcut for Plugin.create(name).with_around_process(around).

Parameters:

Name	Type	Description	Default
name	str	Plugin identifier.	required
around	ProcessLifecycle	Subprocess lifecycle context manager.	required

Returns:

Type	Description
Plugin	Plugin with the around_process hook set.

__getattr__(name)