aitune.dynamo.worker | NVIDIA AITune Documentation

Serve AITune-tuned models as NVIDIA Dynamo endpoints.

Module Contents

Classes

Name	Description
`DynamoWorker`	Base class for custom AITune Dynamo workers.
`DynamoWorkerConfig`	Configuration for the high-level :func:`dynamo_worker` entrypoint.
`_HighLevelDynamoWorker`	Adapt a callable and :class:`DynamoWorkerConfig` to the worker API.

Functions

Name	Description
`dynamo_worker`	Serve a tuned model through a local or collective Dynamo worker.

Data

logger

API

class aitune.dynamo.worker.DynamoWorker()

Base class for custom AITune Dynamo workers.

Override :meth:setup and :meth:serve, then call :meth:run. An initialized multi-rank process group starts one endpoint on rank zero while every rank executes each request.

_collective

bool = False

component

str = 'backend'

endpoint_name

str = 'generate'

namespace

str = 'aitune'

aitune.dynamo.worker.DynamoWorker.on_ready(
    runtime: typing.Any,
    endpoint: typing.Any
) -> None

async

Run optional registration work after the endpoint is available.

aitune.dynamo.worker.DynamoWorker.run(
    enable_nats: bool = False
) -> None

Start serving and block until shutdown.

Parameters:

enable_nats

boolDefaults to False

Enable NATS JetStream for KV cache events.

aitune.dynamo.worker.DynamoWorker.serve(
    request: typing.Any
) -> collections.abc.AsyncGenerator[typing.Any, None]

async

Handle one request and yield response chunks.

aitune.dynamo.worker.DynamoWorker.setup() -> None

Initialize the model once before serving starts.

aitune.dynamo.worker.DynamoWorker.warmup() -> None

Run optional model warmup after every rank completes setup.

class aitune.dynamo.worker.DynamoWorkerConfig(
    type: aitune.dynamo.protocol.Modality,
    model_path: str,
    mapping: collections.abc.Callable | None = None,
    namespace: str = 'aitune',
    component: str = 'backend',
    endpoint: str = 'generate',
    enable_nats: bool = False,
    model_name: str | None = None
)

Dataclass

Configuration for the high-level :func:dynamo_worker entrypoint.

Parameters:

type

Modality

Modality type. One of "image", "video", "audio", "embedding".

model_path

str

Hugging Face model ID or local path passed to register_model.

mapping

Callable | NoneDefaults to None

Optional adapter fn(DynamoRequest) -> dict. The result is unpacked as keyword arguments when calling the model. Without a mapping, a plain callable receives the deserialized request object.

namespace

strDefaults to 'aitune'

Dynamo service namespace.

component

strDefaults to 'backend'

Component name within the namespace.

endpoint

strDefaults to 'generate'

Endpoint name within the component.

enable_nats

boolDefaults to False

Enable NATS JetStream for KV cache events.

model_name

str | NoneDefaults to None

Name advertised to the Dynamo frontend. Defaults to model_path.

component

str = 'backend'

enable_nats

bool = False

endpoint

str = 'generate'

mapping

Callable | None = None

model_name

str | None = None

model_path

str

namespace

str = 'aitune'

type

Modality

class aitune.dynamo.worker._HighLevelDynamoWorker(
    model_or_fn: torch.nn.Module | collections.abc.Callable,
    config: aitune.dynamo.worker.DynamoWorkerConfig,
    setup: collections.abc.Callable[[], None] | None = None,
    warmup: collections.abc.Callable[[], None] | None = None
)

Bases: DynamoWorker

Adapt a callable and :class:DynamoWorkerConfig to the worker API.

component

= config.component

endpoint_name

= config.endpoint

namespace

= config.namespace

aitune.dynamo.worker._HighLevelDynamoWorker.on_ready(
    runtime: typing.Any,
    endpoint: typing.Any
) -> None

async

aitune.dynamo.worker._HighLevelDynamoWorker.serve(
    request: typing.Any
) -> collections.abc.AsyncGenerator[typing.Any, None]

async

Deserialize one request, execute the callable, and pack rank-zero output.

aitune.dynamo.worker._HighLevelDynamoWorker.setup() -> None

Resolve protocol types on every rank before serving.

aitune.dynamo.worker._HighLevelDynamoWorker.warmup() -> None

Run the configured warmup callback after setup succeeds on every rank.

aitune.dynamo.worker.dynamo_worker(
    model_or_fn: torch.nn.Module | collections.abc.Callable,
    config: aitune.dynamo.worker.DynamoWorkerConfig,
    setup: collections.abc.Callable[[], None] | None = None,
    warmup: collections.abc.Callable[[], None] | None = None
) -> None

Serve a tuned model through a local or collective Dynamo worker.

Parameters:

model_or_fn

nn.Module | Callable

A callable, or a torch.nn.Module with config.mapping.

config

DynamoWorkerConfig

Modality, endpoint, request mapping, and execution configuration.

setup

Callable[[], None] | NoneDefaults to None

Optional rank-local model initialization. In collective execution, every rank must finish this callback before warmup begins.

warmup

Callable[[], None] | NoneDefaults to None

Optional model warmup run after setup succeeds on every rank.

Raises:

ValueError: If the modality is unsupported or a module has no mapping.
ImportError: If the optional Dynamo dependencies are unavailable.

aitune.dynamo.worker.logger = getLogger(__name__)