API ReferenceOther APIsAituneDynamo

aitune.dynamo.worker

View as Markdown

Dynamo worker — serve AITune-tuned models as Dynamo endpoints.

Module Contents

Classes

NameDescription
DynamoWorkerBase class for AITune Dynamo workers (power-user API).
DynamoWorkerConfigConfiguration for the high-level :func:dynamo_worker entrypoint.
_HighLevelDynamoWorkerInternal worker built from a DynamoWorkerConfig. Not part of public API.

Functions

NameDescription
_get_wiringReturn (ModelInput, ModelType, request_class) for type_.
_import_dynamoImport dynamo runtime dependencies, raising a clear error if absent.
_pack_responseConvert a user function’s return value to the Dynamo wire-format dict.
_run_dynamo_workerLow-level runtime loop. Call setup() once, serve until SIGTERM/SIGINT.
dynamo_workerServe a tuned model as a Dynamo worker endpoint.

Data

_VALID_TYPES

logger

API

class aitune.dynamo.worker.DynamoWorker()

Base class for AITune Dynamo workers (power-user API).

Subclass, override :meth:setup and :meth:serve, then call :meth:run.

component
str = 'backend'
endpoint_name
str = 'generate'
namespace
str = 'aitune'
aitune.dynamo.worker.DynamoWorker.on_ready(
    runtime: typing.Any,
    endpoint: typing.Any
) -> None
async

Called after the Dynamo endpoint is registered.

Override to call register_model or perform post-startup work.

Parameters:

runtime
Any

The :class:DistributedRuntime instance.

endpoint
Any

The registered Dynamo :class:Endpoint object.

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

Start serving. Blocks until SIGTERM/SIGINT.

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. Async generator — yield response chunks.

Parameters:

request
Any

Incoming request payload.

Raises:

  • NotImplementedError: Must be overridden by subclasses.
aitune.dynamo.worker.DynamoWorker.setup() -> None

Initialize the model. Called once before serving starts.

Raises:

  • NotImplementedError: Must be overridden by subclasses.
class aitune.dynamo.worker.DynamoWorkerConfig(
    type: typing.Literal['image', 'video', 'embedding'],
    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
Literal['image', 'video', 'embedding']

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

model_path
str

HuggingFace model ID or local path passed to register_model.

mapping
Callable | NoneDefaults to None

Optional adapter fn(DynamoRequest) -> dict. The dict is unpacked as **kwargs when calling the user function. If None and the user passed a plain callable (not nn.Module), the raw Dynamo request object is passed as the sole positional argument.

namespace
strDefaults to 'aitune'

Dynamo service namespace. Default: "aitune".

component
strDefaults to 'backend'

Component name within the namespace. Default: "backend".

endpoint
strDefaults to 'generate'

Endpoint name within the component. Default: "generate".

enable_nats
boolDefaults to False

Enable NATS JetStream for KV cache events. Default: False.

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
Literal['image', 'video', 'embedding']
class aitune.dynamo.worker._HighLevelDynamoWorker(
    model_or_fn: torch.nn.Module | collections.abc.Callable,
    config: aitune.dynamo.worker.DynamoWorkerConfig
)

Bases: DynamoWorker

Internal worker built from a DynamoWorkerConfig. Not part of public API.

component
= config.component
endpoint_name
= config.endpoint
namespace
= config.namespace
aitune.dynamo.worker._HighLevelDynamoWorker._version_specific_register_model_kwargs() -> dict

From version 1.3.0 dynamo introduces a new WorkerType enum and made it required.

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

Register this worker with the Dynamo frontend.

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

Deserialize request, run user function in executor, pack and yield response.

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

No-op: the model is already initialized before dynamo_worker() is called.

aitune.dynamo.worker._get_wiring(
    type_: str
) -> tuple

Return (ModelInput, ModelType, request_class) for type_.

All imports are deferred because dynamo is optional.

Parameters:

type_
str

Modality type string ("image", "video", "embedding").

Returns: tuple

Tuple of (ModelInput enum value, ModelType enum value, request Pydantic class).

aitune.dynamo.worker._import_dynamo() -> tuple

Import dynamo runtime dependencies, raising a clear error if absent.

Returns: tuple

Tuple of (DistributedRuntime class, dynamo_worker decorator, uvloop module).

Raises:

  • ImportError: When ai-dynamo-runtime or uvloop is not installed.
aitune.dynamo.worker._pack_response(
    result: typing.Any,
    config: aitune.dynamo.worker.DynamoWorkerConfig
) -> dict

Convert a user function’s return value to the Dynamo wire-format dict.

Parameters:

result
Any

Return value from the user’s inference function.

config
DynamoWorkerConfig

Worker configuration (used for model name and type).

Returns: dict

A dict ready to yield to the Dynamo runtime.

Raises:

  • TypeError: When result cannot be converted for config.type.
aitune.dynamo.worker._run_dynamo_worker(
    setup: collections.abc.Callable[[], None],
    serve: collections.abc.Callable[[Any], collections.abc.AsyncGenerator[typing.Any, None]],
    namespace: str = 'aitune',
    component: str = 'backend',
    endpoint: str = 'generate',
    enable_nats: bool = False,
    on_ready: collections.abc.Callable[[Any, Any], collections.abc.Coroutine] | None = None
) -> None

Low-level runtime loop. Call setup() once, serve until SIGTERM/SIGINT.

Parameters:

setup
Callable[[], None]

Zero-argument initializer called before serving starts.

serve
Callable[[Any], AsyncGenerator[Any, None]]

Async generator serve(request) -> AsyncIterable[response].

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.

on_ready
Callable[[Any, Any], Coroutine] | NoneDefaults to None

Optional async callable invoked after endpoint registration.

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

Serve a tuned model as a Dynamo worker endpoint.

The minimal path to serving after ait.tune() or ait.load():

  1. Build a :class:DynamoWorkerConfig for your modality.
  2. Call this function. It blocks until SIGTERM/SIGINT.

Startup validation happens before any Dynamo runtime is started. Modality-specific request deserialization, register_model, sync-to-async wrapping, and response packing are handled automatically.

Parameters:

model_or_fn
nn.Module | Callable

A torch.nn.Module (requires config.mapping) or any callable. If callable and config.mapping is None, the raw Dynamo request object is passed as the sole argument.

config
DynamoWorkerConfig

Worker configuration including modality type, model path, and optional request adapter.

Raises:

  • ValueError: If config.type is not a supported P0 modality, or if a torch.nn.Module is passed without a mapping.
  • ImportError: If ai-dynamo-runtime is not installed.
aitune.dynamo.worker._VALID_TYPES: frozenset[str] = frozenset({'image', 'video', 'embedding'})
 
aitune.dynamo.worker.logger = getLogger(__name__)