diff --git a/components/src/dynamo/common/configuration/arg_group.py b/components/src/dynamo/common/configuration/arg_group.py index 3e984500fb9d..e1df406aa8c1 100644 --- a/components/src/dynamo/common/configuration/arg_group.py +++ b/components/src/dynamo/common/configuration/arg_group.py @@ -11,6 +11,24 @@ class ArgGroup(ABC): Base interface for configuration groups. Each ArgGroup represents a domain of configuration parameters with clear ownership. + + Examples: + >>> import argparse + >>> from dynamo.common.configuration.arg_group import ArgGroup + >>> from dynamo.common.configuration.utils import add_argument + >>> + >>> class MyArgGroup(ArgGroup): + ... def add_arguments(self, parser) -> None: + ... add_argument( + ... parser, + ... flag_name="--my-option", + ... env_var="DYN_MY_OPTION", + ... default="value", + ... help="A custom option.", + ... ) + >>> + >>> parser = argparse.ArgumentParser() + >>> MyArgGroup().add_arguments(parser) """ @abstractmethod diff --git a/components/src/dynamo/common/configuration/config_base.py b/components/src/dynamo/common/configuration/config_base.py index 8ae450a4cc20..24f8450fd569 100644 --- a/components/src/dynamo/common/configuration/config_base.py +++ b/components/src/dynamo/common/configuration/config_base.py @@ -5,7 +5,24 @@ class ConfigBase: - """Base configuration class that allows properties with and without defaults in arbitrary order.""" + """Base configuration class that allows properties with and without defaults in arbitrary order. + + Examples: + >>> import argparse + >>> from dynamo.common.configuration.config_base import ConfigBase + >>> + >>> class MyConfig(ConfigBase): + ... model_name: str + ... http_port: int = 8000 + >>> + >>> parser = argparse.ArgumentParser() + >>> parser.add_argument("--model-name", default="llama") + >>> parser.add_argument("--http-port", type=int, default=8000) + >>> args = parser.parse_args([]) + >>> config = MyConfig.from_cli_args(args) + >>> config.http_port + 8000 + """ @classmethod def from_cli_args(cls, args: argparse.Namespace) -> Self: diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index b2342bb82749..087d23d8a4d8 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -38,7 +38,25 @@ class KvRouterConfigBase(ConfigBase): - """Mixin carrying the 16 KvRouterConfig fields.""" + """Mixin carrying the 16 KvRouterConfig fields. + + Examples: + >>> import argparse + >>> from dynamo.common.configuration.groups.kv_router_args import ( + ... KvRouterArgGroup, + ... KvRouterConfigBase, + ... ) + >>> + >>> class MyRouterConfig(KvRouterConfigBase): + ... endpoint: str = "http://localhost:8080" + >>> + >>> parser = argparse.ArgumentParser() + >>> KvRouterArgGroup().add_arguments(parser) + >>> config = MyRouterConfig.from_cli_args(parser.parse_args([])) # doctest: +SKIP + >>> kwargs = config.kv_router_kwargs() # doctest: +SKIP + >>> "overlap_score_weight" in kwargs # doctest: +SKIP + True + """ overlap_score_weight: float router_temperature: float @@ -63,7 +81,21 @@ def kv_router_kwargs(self) -> dict: class KvRouterArgGroup(ArgGroup): - """CLI arguments for the 16 KvRouterConfig parameters.""" + """CLI arguments for the 16 KvRouterConfig parameters. + + Examples: + >>> import argparse + >>> from dynamo.common.configuration.groups.kv_router_args import KvRouterArgGroup + >>> + >>> parser = argparse.ArgumentParser() + >>> KvRouterArgGroup().add_arguments(parser) + >>> args = parser.parse_args([ + ... "--router-kv-overlap-score-weight", "0.8", + ... "--router-temperature", "0.5", + ... ]) + >>> args.overlap_score_weight + 0.8 + """ def add_arguments(self, parser) -> None: g = parser.add_argument_group("KV Router Options") diff --git a/components/src/dynamo/common/configuration/utils.py b/components/src/dynamo/common/configuration/utils.py index 5262ebc63062..e00bdf56cafa 100644 --- a/components/src/dynamo/common/configuration/utils.py +++ b/components/src/dynamo/common/configuration/utils.py @@ -68,6 +68,23 @@ def add_argument( """ Add a CLI argument with env var default, optional alias and dest, and help message construction. + Examples: + >>> import argparse + >>> from dynamo.common.configuration.utils import add_argument + >>> + >>> parser = argparse.ArgumentParser() + >>> add_argument( + ... parser, + ... flag_name="--http-port", + ... env_var="DYN_HTTP_PORT", + ... default=8000, + ... help="HTTP port for the engine.", + ... arg_type=int, + ... ) + >>> args = parser.parse_args([]) + >>> args.http_port + 8000 + Args: parser: ArgumentParser or argument group flag_name: Primary flag (must start with '--', e.g., "--foo") diff --git a/components/src/dynamo/common/constants.py b/components/src/dynamo/common/constants.py index 5b1aa0141656..cacf584b327c 100644 --- a/components/src/dynamo/common/constants.py +++ b/components/src/dynamo/common/constants.py @@ -7,7 +7,15 @@ class DisaggregationMode(Enum): - """Disaggregation mode for LLM workers.""" + """Disaggregation mode for LLM workers. + + Examples: + >>> from dynamo.common.constants import DisaggregationMode + >>> DisaggregationMode.PREFILL.value + 'prefill' + >>> DisaggregationMode("agg") == DisaggregationMode.AGGREGATED + True + """ AGGREGATED = "agg" PREFILL = "prefill" @@ -15,7 +23,15 @@ class DisaggregationMode(Enum): class EmbeddingTransferMode(Enum): - """Embedding transfer mode for LLM workers.""" + """Embedding transfer mode for LLM workers. + + Examples: + >>> from dynamo.common.constants import EmbeddingTransferMode + >>> EmbeddingTransferMode.NIXL_WRITE.value + 'nixl-write' + >>> EmbeddingTransferMode("local") == EmbeddingTransferMode.LOCAL + True + """ LOCAL = "local" NIXL_WRITE = "nixl-write" diff --git a/components/src/dynamo/common/lora/manager.py b/components/src/dynamo/common/lora/manager.py index eb3ef1b02d28..63912860955f 100644 --- a/components/src/dynamo/common/lora/manager.py +++ b/components/src/dynamo/common/lora/manager.py @@ -33,6 +33,16 @@ class LoRAManager: The manager uses the Rust-based LoRADownloader for S3 and local file sources, and allows registering custom Python sources for other protocols. + + Examples: + >>> import asyncio + >>> from dynamo.common.lora import LoRAManager + >>> manager = LoRAManager() + >>> result = asyncio.run( + ... manager.download_lora("s3://bucket/adapters/lora-v1") + ... ) # doctest: +SKIP + >>> manager.is_cached("s3://bucket/adapters/lora-v1") # doctest: +SKIP + True """ def __init__(self, cache_path: Optional[Path] = None): diff --git a/components/src/dynamo/common/storage.py b/components/src/dynamo/common/storage.py index 2cec28570606..06e83db64a45 100644 --- a/components/src/dynamo/common/storage.py +++ b/components/src/dynamo/common/storage.py @@ -46,6 +46,17 @@ def get_fs(fs_url: str) -> DirFileSystem: Args: fs_url: The URL of the filesystem to initialize. e.g. s3://bucket, gs://bucket, file:///local/path + Examples: + >>> from dynamo.common.storage import get_fs + >>> + >>> fs = get_fs("file:///data/media") # doctest: +SKIP + >>> protocol = fs.fs.protocol # doctest: +SKIP + >>> protocol if isinstance(protocol, str) else protocol[0] # doctest: +SKIP + 'file' + >>> fs = get_fs("s3://my-bucket") # doctest: +SKIP + >>> fs.path # doctest: +SKIP + 'my-bucket' + Returns: The initialized DirFileSystem wrapper for the filesystem. @@ -81,6 +92,15 @@ def get_media_url( ``{base_url}/{storage_path}``. When *None*, the URL is constructed from the filesystem's protocol and root path. + Examples: + >>> from dynamo.common.storage import get_fs, get_media_url # doctest: +SKIP + >>> + >>> fs = get_fs("file:///data/media") # doctest: +SKIP + >>> get_media_url(fs, "videos/req-123.mp4") # doctest: +SKIP + 'file:///data/media/videos/req-123.mp4' + >>> get_media_url(fs, "img.png", base_url="https://cdn.example.com/media") # doctest: +SKIP + 'https://cdn.example.com/media/img.png' + Returns: Public URL string for the uploaded file. """ @@ -110,6 +130,16 @@ async def upload_to_fs( data: Raw bytes to upload. base_url: Optional CDN / proxy base URL for URL rewriting. + Examples: + >>> import asyncio + >>> from dynamo.common.storage import get_fs, upload_to_fs + >>> + >>> fs = get_fs("s3://my-media-bucket") + >>> image_bytes = b"\\x89PNG..." + >>> url = asyncio.run( + ... upload_to_fs(fs, "images/req-123/output.png", image_bytes) + ... ) # doctest: +SKIP + Returns: Public URL string for the uploaded file. """ diff --git a/components/src/dynamo/frontend/frontend_args.py b/components/src/dynamo/frontend/frontend_args.py index 4e2414c35e70..cec6610bf693 100644 --- a/components/src/dynamo/frontend/frontend_args.py +++ b/components/src/dynamo/frontend/frontend_args.py @@ -40,7 +40,22 @@ def validate_model_path(value: str) -> str: class FrontendConfig(KvRouterConfigBase): - """Configuration for the Dynamo frontend.""" + """Configuration for the Dynamo frontend. + + Examples: + >>> import argparse + >>> from dynamo.frontend.frontend_args import FrontendArgGroup, FrontendConfig + >>> + >>> parser = argparse.ArgumentParser() + >>> FrontendArgGroup().add_arguments(parser) + >>> args = parser.parse_args([ + ... "--model-name", "Llama-3.2-1B-Instruct", + ... "--http-port", "8080", + ... ]) + >>> config = FrontendConfig.from_cli_args(args) # doctest: +SKIP + >>> config.model_name # doctest: +SKIP + 'Llama-3.2-1B-Instruct' + """ interactive: bool kv_cache_block_size: Optional[int] @@ -93,7 +108,23 @@ def _preprocess_for_encode_config(config: FrontendConfig) -> Dict[str, Any]: class FrontendArgGroup(ArgGroup): - """Frontend configuration parameters.""" + """Frontend configuration parameters. + + Examples: + >>> import argparse + >>> from dynamo.frontend.frontend_args import FrontendArgGroup + >>> + >>> parser = argparse.ArgumentParser() + >>> group = FrontendArgGroup() + >>> group.add_arguments(parser) + >>> args = parser.parse_args([ + ... "--model-name", "Llama-3.2-1B-Instruct", + ... "--router-mode", "kv", + ... "--http-port", "8080", + ... ]) + >>> args.model_name + 'Llama-3.2-1B-Instruct' + """ def add_arguments(self, parser) -> None: parser.add_argument( diff --git a/components/src/dynamo/mocker/main.py b/components/src/dynamo/mocker/main.py index 1b830ad36a6c..3791298aba36 100644 --- a/components/src/dynamo/mocker/main.py +++ b/components/src/dynamo/mocker/main.py @@ -70,6 +70,15 @@ async def worker(): Each mocker gets its own DistributedRuntime instance for true isolation, while still sharing the same event loop and tokio runtime. + + Examples: + Launch via CLI (the standard invocation pattern): + + >>> # python -m dynamo.mocker --model-path /data/models/Qwen3-0.6B + >>> + >>> import uvloop # doctest: +SKIP + >>> from dynamo.mocker.main import worker # doctest: +SKIP + >>> uvloop.run(worker()) # doctest: +SKIP """ args = parse_args() diff --git a/components/src/dynamo/planner/defaults.py b/components/src/dynamo/planner/defaults.py index 5705f922a87e..cbf981844995 100644 --- a/components/src/dynamo/planner/defaults.py +++ b/components/src/dynamo/planner/defaults.py @@ -49,6 +49,23 @@ class BasePlannerDefaults: class SLAPlannerDefaults(BasePlannerDefaults): + """SLA-based planner defaults for throughput and latency targets. + + Provides default values for SLA objectives (TTFT, ITL), load prediction, + and scaling parameters used by the planner. + + Examples: + >>> from dynamo.planner.defaults import SLAPlannerDefaults + >>> SLAPlannerDefaults.ttft + 500.0 + >>> SLAPlannerDefaults.itl + 50.0 + >>> SLAPlannerDefaults.mode + 'disagg' + >>> SLAPlannerDefaults.max_gpu_budget + 8 + """ + # Prometheus endpoint URL for pulling/querying metrics metric_pulling_prometheus_endpoint = os.environ.get( "PROMETHEUS_ENDPOINT", @@ -143,6 +160,16 @@ class MockerComponentName: class SubComponentType(str, Enum): + """Type of sub-component in a disaggregated deployment. + + Examples: + >>> from dynamo.planner.defaults import SubComponentType + >>> SubComponentType.PREFILL.value + 'prefill' + >>> SubComponentType("decode") == SubComponentType.DECODE + True + """ + PREFILL = "prefill" DECODE = "decode" diff --git a/components/src/dynamo/planner/global_planner_connector.py b/components/src/dynamo/planner/global_planner_connector.py index 1485befafc94..1c2bd0fd3f4a 100644 --- a/components/src/dynamo/planner/global_planner_connector.py +++ b/components/src/dynamo/planner/global_planner_connector.py @@ -28,6 +28,19 @@ class GlobalPlannerConnector(PlannerConnector): This connector wraps RemotePlannerClient and implements the PlannerConnector interface, allowing planner_core.py to treat global-planner environment mode consistently with kubernetes and virtual modes. + + Examples: + >>> from dynamo.planner.global_planner_connector import GlobalPlannerConnector # doctest: +SKIP + >>> + >>> connector = GlobalPlannerConnector( # doctest: +SKIP + ... runtime=runtime, + ... dynamo_namespace="test-ns", + ... global_planner_namespace="global-ns", + ... model_name="Llama-3.2-1B-Instruct", + ... ) + >>> await connector._async_init() # doctest: +SKIP + >>> connector.get_model_name() # doctest: +SKIP + 'Llama-3.2-1B-Instruct' """ def __init__( diff --git a/components/src/dynamo/planner/kubernetes_connector.py b/components/src/dynamo/planner/kubernetes_connector.py index 1b7d3f20be83..74ad0ebb3621 100644 --- a/components/src/dynamo/planner/kubernetes_connector.py +++ b/components/src/dynamo/planner/kubernetes_connector.py @@ -46,6 +46,26 @@ class TargetReplica(BaseModel): class KubernetesConnector(PlannerConnector): + """Connector that manages scaling via the Kubernetes DynamoGraphDeployment API. + + Implements the PlannerConnector interface for Kubernetes-based deployments, + providing add/remove component operations, deployment validation, and + model name discovery from DGD service specs. + + Examples: + >>> import asyncio + >>> from dynamo.planner.kubernetes_connector import KubernetesConnector + >>> from dynamo.planner.defaults import SubComponentType + >>> + >>> connector = KubernetesConnector( + ... dynamo_namespace="dynamo", + ... k8s_namespace="default", + ... parent_dgd_name="my-dgd", + ... ) + >>> asyncio.run(connector.validate_deployment()) # doctest: +SKIP + >>> model_name = connector.get_model_name() # doctest: +SKIP + """ + def __init__( self, dynamo_namespace: str, diff --git a/components/src/dynamo/planner/planner_connector.py b/components/src/dynamo/planner/planner_connector.py index 7b40266abd49..37b540976629 100644 --- a/components/src/dynamo/planner/planner_connector.py +++ b/components/src/dynamo/planner/planner_connector.py @@ -20,6 +20,20 @@ # TODO: add ability to scale component to X replicas class PlannerConnector(ABC): + """Abstract base class for planner connectors that manage scaling operations. + + Concrete implementations include KubernetesConnector, VirtualConnector, + and GlobalPlannerConnector. + + Examples: + >>> from dynamo.planner import KubernetesConnector, VirtualConnector + >>> from dynamo.planner.planner_connector import PlannerConnector + >>> issubclass(KubernetesConnector, PlannerConnector) + True + >>> issubclass(VirtualConnector, PlannerConnector) + True + """ + @abstractmethod async def add_component( self, sub_component_type: SubComponentType, blocking: bool = True diff --git a/components/src/dynamo/planner/remote_planner_client.py b/components/src/dynamo/planner/remote_planner_client.py index d52abad4b240..71b39294a81d 100644 --- a/components/src/dynamo/planner/remote_planner_client.py +++ b/components/src/dynamo/planner/remote_planner_client.py @@ -15,7 +15,17 @@ class RemotePlannerClient: - """Client for delegating scaling requests to centralized planner""" + """Client for delegating scaling requests to centralized planner. + + Examples: + >>> from dynamo.planner.remote_planner_client import RemotePlannerClient + >>> client = RemotePlannerClient( # doctest: +SKIP + ... runtime=runtime, + ... central_namespace="global-planner", + ... central_component="GlobalPlanner", + ... connection_timeout=30.0, + ... ) + """ def __init__( self, diff --git a/components/src/dynamo/planner/scale_protocol.py b/components/src/dynamo/planner/scale_protocol.py index e912f4a7907f..a30dc084ac35 100644 --- a/components/src/dynamo/planner/scale_protocol.py +++ b/components/src/dynamo/planner/scale_protocol.py @@ -20,7 +20,25 @@ class ScaleStatus(str, Enum): class ScaleRequest(BaseModel): - """Request to scale a deployment""" + """Request to scale a deployment. + + Examples: + >>> from dynamo.planner.scale_protocol import ScaleRequest + >>> from dynamo.planner.defaults import SubComponentType + >>> from dynamo.planner.kubernetes_connector import TargetReplica + >>> + >>> request = ScaleRequest( + ... caller_namespace="dynamo", + ... graph_deployment_name="my-dgd", + ... k8s_namespace="default", + ... target_replicas=[ + ... TargetReplica(sub_component_type=SubComponentType.PREFILL, desired_replicas=2), + ... TargetReplica(sub_component_type=SubComponentType.DECODE, desired_replicas=4), + ... ], + ... ) + >>> request.caller_namespace + 'dynamo' + """ # Caller identification caller_namespace: str diff --git a/components/src/dynamo/planner/virtual_connector.py b/components/src/dynamo/planner/virtual_connector.py index 90e213a01173..6441f9cfa06b 100644 --- a/components/src/dynamo/planner/virtual_connector.py +++ b/components/src/dynamo/planner/virtual_connector.py @@ -30,6 +30,18 @@ class VirtualConnector(PlannerConnector): This is a virtual connector for planner to output scaling decisions to non-native environments This virtual connector does not actually scale the deployment, instead, it communicates with the non-native environment through dynamo-runtime's VirtualConnectorCoordinator. The deployment environment needs to use VirtualConnectorClient (in the Rust/Python bindings) to read from the scaling decisions and update report scaling status. + + Examples: + >>> from dynamo.planner.virtual_connector import VirtualConnector # doctest: +SKIP + >>> from dynamo.planner.defaults import SubComponentType # doctest: +SKIP + >>> + >>> connector = VirtualConnector( # doctest: +SKIP + ... runtime=runtime, + ... dynamo_namespace="dynamo", + ... model_name="Llama-3.2-1B-Instruct", + ... ) + >>> await connector._async_init() # doctest: +SKIP + >>> await connector.add_component(SubComponentType.PREFILL) # doctest: +SKIP """ def __init__( diff --git a/components/src/dynamo/router/args.py b/components/src/dynamo/router/args.py index ccb22268a2b9..4cd213e4b89a 100644 --- a/components/src/dynamo/router/args.py +++ b/components/src/dynamo/router/args.py @@ -16,7 +16,19 @@ class DynamoRouterConfig(KvRouterConfigBase): - """Typed configuration for the standalone KV router (router-owned options only).""" + """Typed configuration for the standalone KV router (router-owned options only). + + Examples: + >>> from dynamo.router.args import parse_args, DynamoRouterConfig + >>> config = parse_args([ + ... "--endpoint", "dynamo.prefill.generate", + ... "--router-mode", "kv", + ... ]) + >>> config.endpoint + 'dynamo.prefill.generate' + >>> config.namespace + 'dynamo' + """ namespace: str endpoint: str @@ -39,7 +51,22 @@ def validate(self) -> None: class DynamoRouterArgGroup(ArgGroup): - """CLI argument group for standalone router options.""" + """CLI argument group for standalone router options. + + Examples: + >>> import argparse + >>> from dynamo.router.args import DynamoRouterArgGroup, DynamoRouterConfig + >>> + >>> parser = argparse.ArgumentParser() + >>> DynamoRouterArgGroup().add_arguments(parser) + >>> args = parser.parse_args([ + ... "--endpoint", "dynamo.prefill.generate", + ... "--router-mode", "kv", + ... ]) + >>> config = DynamoRouterConfig.from_cli_args(args) + >>> config.endpoint + 'dynamo.prefill.generate' + """ name = "dynamo-router" @@ -71,7 +98,15 @@ def add_arguments(self, parser) -> None: def build_kv_router_config(router_config: DynamoRouterConfig) -> KvRouterConfig: - """Build KvRouterConfig from DynamoRouterConfig.""" + """Build KvRouterConfig from DynamoRouterConfig. + + Examples: + >>> from dynamo.router.args import parse_args, build_kv_router_config + >>> router_config = parse_args([ + ... "--endpoint", "dynamo.prefill.generate", + ... ]) + >>> kv_config = build_kv_router_config(router_config) + """ return KvRouterConfig(**router_config.kv_router_kwargs()) diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 058009ae43c6..a7513998c87e 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -34,6 +34,11 @@ def get_reasoning_parser_names() -> list[str]: class JsonLike: """ Any PyObject which can be serialized to JSON + + Examples: + >>> request: JsonLike = {"prompt": "Hello", "max_tokens": 128} + >>> request_list: JsonLike = [1, 2, 3] + >>> request_str: JsonLike = "plain text input" """ ... @@ -43,6 +48,13 @@ RequestHandler = Callable[[JsonLike], AsyncGenerator[JsonLike, None]] class DistributedRuntime: """ The runtime object for dynamo applications + + Examples: + >>> import asyncio + >>> loop = asyncio.get_event_loop() + >>> runtime = DistributedRuntime(loop, "etcd", "nats") + >>> endpoint = runtime.endpoint("myns.backend.generate") + >>> client = await endpoint.client() """ def __new__( @@ -122,6 +134,12 @@ class DistributedRuntime: class Endpoint: """ An Endpoint is a single API endpoint + + Examples: + >>> async def handler(request): # doctest: +SKIP + ... yield {"text": "hello"} + >>> endpoint = runtime.endpoint("dynamo.backend.generate") # doctest: +SKIP + >>> await endpoint.serve_endpoint(handler) # doctest: +SKIP """ ... @@ -187,6 +205,13 @@ class Endpoint: class Client: """ A client capable of calling served instances of an endpoint + + Examples: + >>> client = await endpoint.client() + >>> async for chunk in await client.round_robin("hello world"): + ... print(chunk.get("data")) + >>> async for resp in await client.random({"prompt": "hi"}): + ... print(resp) """ ... @@ -280,14 +305,12 @@ def compute_block_hash_for_seq( Returns: List of block hashes (one per block) - Example: + Examples: >>> tokens = [1, 2, 3, 4] * 8 # 32 tokens = 1 block - >>> mm_info = { - ... "mm_objects": [{ - ... "mm_hash": 0xDEADBEEF, - ... }] - ... } - >>> hashes = compute_block_hash_for_seq(tokens, 32, [mm_info]) + >>> hashes = compute_block_hash_for_seq(tokens, 32) + >>> mm_info = {"mm_objects": [{"mm_hash": 0xDEADBEEF}]} + >>> hashes_mm = compute_block_hash_for_seq(tokens, 32, [mm_info]) + >>> hashes_lora = compute_block_hash_for_seq(tokens, 32, lora_name="my-adapter") """ ... @@ -296,6 +319,14 @@ class Context: """ Context wrapper around AsyncEngineContext for Python bindings. Provides tracing and cancellation capabilities for request handling. + + Examples: + >>> from dynamo._core import Context + >>> context = Context() + >>> context = Context(id="req-42") + >>> print(context.id()) + >>> context.stop_generating() + >>> assert context.is_stopped() """ def __init__(self, id: Optional[str] = None) -> None: @@ -382,6 +413,11 @@ class Context: class WorkerMetricsPublisher: """ A metrics publisher will provide metrics to the router for load monitoring. + + Examples: + >>> publisher = WorkerMetricsPublisher() + >>> await publisher.create_endpoint(endpoint) + >>> publisher.publish(dp_rank=0, active_decode_blocks=128) """ ... @@ -419,6 +455,12 @@ class WorkerMetricsPublisher: class ModelDeploymentCard: """ A model deployment card is a collection of model information + + Examples: + >>> from dynamo._core import ModelDeploymentCard # doctest: +SKIP + >>> card = ModelDeploymentCard() # doctest: +SKIP + >>> json_str = card.to_json_str() # doctest: +SKIP + >>> restored = ModelDeploymentCard.from_json_str(json_str) # doctest: +SKIP """ def to_json_str(self) -> str: @@ -445,6 +487,16 @@ class ModelDeploymentCard: class ModelRuntimeConfig: """ A model runtime configuration is a collection of runtime information + + Examples: + >>> from dynamo._core import ModelRuntimeConfig + >>> config = ModelRuntimeConfig() + >>> config.total_kv_blocks = 2048 + >>> config.max_num_seqs = 256 + >>> config.max_num_batched_tokens = 8192 + >>> config.enable_local_indexer = True + >>> config.tool_call_parser = "hermes" + >>> config.set_engine_specific("tp_size", 4) """ total_kv_blocks: int | None @@ -492,6 +544,13 @@ class OverlapScores: """ A collection of prefix matching scores of workers for a given token ids. 'scores' is a map of worker id to the score which is the number of matching blocks. + + Examples: + >>> from dynamo._core import RadixTree + >>> tree = RadixTree() + >>> scores = tree.find_matches([42, 43]) + >>> worker_scores = scores.scores # Dict[int, int] + >>> freqs = scores.frequencies # List[int] """ @property @@ -521,6 +580,17 @@ class RadixTree: Thread-safe: operations route to a dedicated background thread and long calls release the Python GIL. + + Examples: + >>> import json + >>> from dynamo._core import RadixTree + >>> tree = RadixTree() + >>> tree_ttl = RadixTree(expiration_duration_secs=120.0) + >>> event = {"event_id": 1, "data": {"stored": {"parent_hash": None, + ... "blocks": [{"block_hash": 42, "tokens_hash": 42}]}}} + >>> tree.apply_event(0, json.dumps(event).encode()) + >>> scores = tree.find_matches([42]) + >>> print(scores.scores) """ def __init__(self, expiration_duration_secs: Optional[float] = None) -> None: @@ -591,6 +661,13 @@ class RadixTree: class KvIndexer: """ A KV Indexer that tracks KV Events emitted by workers. Events include add_block and remove_block. + + Examples: + >>> from dynamo._core import KvIndexer + >>> indexer = KvIndexer(endpoint=endpoint, block_size=64) + >>> scores = indexer.find_matches_for_request(token_ids=[1, 2, 3]) + >>> print(scores.scores) + >>> print(indexer.block_size()) """ ... @@ -636,6 +713,15 @@ class ApproxKvIndexer: - Backend engines don't emit KV events - You want to reduce event processing overhead - Lower routing accuracy is acceptable + + Examples: + >>> from dynamo._core import ApproxKvIndexer + >>> indexer = ApproxKvIndexer( + ... endpoint=endpoint, kv_block_size=64, + ... router_ttl_secs=120.0, router_max_tree_size=1048576, + ... ) + >>> scores = indexer.find_matches_for_request(token_ids=[1, 2, 3]) + >>> await indexer.process_routing_decision_for_request([1, 2, 3], worker_id=0) """ ... @@ -704,6 +790,17 @@ class ApproxKvIndexer: class KvEventPublisher: """ A KV event publisher will publish KV events corresponding to the component. + + Examples: + >>> publisher = KvEventPublisher( # doctest: +SKIP + ... endpoint=endpoint, kv_block_size=64, dp_rank=0, + ... zmq_endpoint="tcp://127.0.0.1:5557", zmq_topic="", + ... ) + >>> publisher.publish_stored( # doctest: +SKIP + ... token_ids=[1, 2, 3, 4], num_block_tokens=[4], + ... block_hashes=[123456], + ... ) + >>> publisher.publish_removed(block_hashes=[123456]) # doctest: +SKIP """ ... @@ -783,6 +880,12 @@ class HttpService: """ A HTTP service for dynamo applications. It is a OpenAI compatible http ingress into the Dynamo Distributed Runtime. + + Examples: + >>> service = HttpService(port=8000) # doctest: +SKIP + >>> service.add_chat_completions_model("my-model", "checksum", engine) # doctest: +SKIP + >>> await service.run(runtime) # doctest: +SKIP + >>> service.shutdown() # doctest: +SKIP """ def __init__(self, port: Optional[int] = None) -> None: @@ -812,6 +915,14 @@ class HttpService: class PythonAsyncEngine: """ Bridge a Python async generator onto Dynamo's AsyncEngine interface. + + Examples: + >>> import asyncio + >>> from dynamo._core import PythonAsyncEngine + >>> async def my_generator(request): + ... yield {"text": "hello"} + >>> loop = asyncio.get_running_loop() + >>> engine = PythonAsyncEngine(my_generator, loop) """ def __init__(self, generator: Any, event_loop: Any) -> None: @@ -833,6 +944,12 @@ class KserveGrpcService: """ A gRPC service implementing the KServe protocol for dynamo applications. Provides model management for completions, chat completions, and tensor-based models. + + Examples: + >>> from dynamo._core import KserveGrpcService # doctest: +SKIP + >>> service = KserveGrpcService(port=8787, host="0.0.0.0") # doctest: +SKIP + >>> service.add_completions_model("model", "checksum", engine) # doctest: +SKIP + >>> await service.run(runtime) # doctest: +SKIP """ def __init__(self, port: Optional[int] = None, host: Optional[str] = None) -> None: @@ -964,14 +1081,28 @@ class KserveGrpcService: ... class ModelInput: - """What type of request this model needs: Text, Tokens or Tensor""" + """What type of request this model needs: Text, Tokens or Tensor + + Examples: + >>> from dynamo._core import ModelInput + >>> input_type = ModelInput.Tokens + >>> input_type = ModelInput.Text + >>> input_type = ModelInput.Tensor + """ Text: ModelInput Tokens: ModelInput Tensor: ModelInput class ModelType: - """What type of request this model needs: Chat, Completions, Embedding, Tensor, Images, Videos or Prefill""" + """What type of request this model needs: Chat, Completions, Embedding, TensorBased, Images, Audios, Videos, or Prefill + + Examples: + >>> from dynamo._core import ModelType + >>> model_type = ModelType.Chat | ModelType.Completions + >>> model_type = ModelType.TensorBased + >>> model_type = ModelType.Prefill + """ Chat: ModelType Completions: ModelType Embedding: ModelType @@ -989,7 +1120,14 @@ class ModelType: ... class RouterMode: - """Router mode for load balancing requests across workers""" + """Router mode for load balancing requests across workers + + Examples: + >>> from dynamo._core import RouterMode + >>> mode = RouterMode.RoundRobin + >>> mode = RouterMode.KV + >>> mode = RouterMode.Direct + """ RoundRobin: "RouterMode" Random: "RouterMode" KV: "RouterMode" @@ -997,7 +1135,19 @@ class RouterMode: ... class RouterConfig: - """How to route the request""" + """How to route the request + + Examples: + >>> from dynamo._core import RouterConfig, RouterMode, KvRouterConfig + >>> config = RouterConfig(mode=RouterMode.RoundRobin) + >>> kv_config = KvRouterConfig(overlap_score_weight=0.5) + >>> config = RouterConfig( + ... mode=RouterMode.KV, + ... config=kv_config, + ... active_decode_blocks_threshold=0.9, + ... enforce_disagg=True, + ... ) + """ router_mode: RouterMode kv_router_config: KvRouterConfig @@ -1024,7 +1174,17 @@ class RouterConfig: ... class KvRouterConfig: - """Values for KV router""" + """ + Values for KV router + + Examples: + >>> config = KvRouterConfig() + >>> config_custom = KvRouterConfig( + ... overlap_score_weight=0.5, + ... router_temperature=0.5, + ... router_track_active_blocks=True, + ... ) + """ def __init__( self, @@ -1107,6 +1267,20 @@ async def register_model( For TensorBased models (using ModelInput.Tensor), HuggingFace downloads are skipped and a minimal model card is registered directly. Use model_path as the display name for these models. + + Examples: + >>> from dynamo._core import ModelInput, ModelType, register_model + >>> await register_model( + ... ModelInput.Tokens, + ... ModelType.Chat | ModelType.Completions, + ... endpoint, + ... "Qwen/Qwen3-0.6B", + ... ) + >>> await register_model( + ... ModelInput.Tensor, ModelType.TensorBased, + ... endpoint, "echo", + ... runtime_config=runtime_config, + ... ) """ ... @@ -1159,7 +1333,10 @@ async def fetch_model(remote_name: str, ignore_weights: bool = False) -> str: """ Download a model from Hugging Face, returning its local path. If `ignore_weights` is True, only fetches tokenizer and config files. - Example: `model_path = await fetch_model("Qwen/Qwen3-0.6B")` + + Examples: + >>> model_path = await fetch_model("Qwen/Qwen3-0.6B") + >>> tokenizer_path = await fetch_model("meta-llama/Llama-3-8B", ignore_weights=True) """ ... @@ -1169,15 +1346,34 @@ register_llm = register_model unregister_llm = unregister_model class EngineConfig: - """Holds internal configuration for a Dynamo engine.""" + """Holds internal configuration for a Dynamo engine. + + Examples: + >>> args = EntrypointArgs(engine_type=EngineType.Dynamic, model_path="/models/llama") + >>> engine_config = await make_engine(runtime, args) + >>> await run_input(runtime, "http", engine_config) + """ ... async def make_engine(distributed_runtime: DistributedRuntime, args: EntrypointArgs) -> EngineConfig: - """Make an engine matching the args""" + """Make an engine matching the args + + Examples: + >>> from dynamo._core import EntrypointArgs, EngineType, make_engine + >>> args = EntrypointArgs(engine_type=EngineType.Dynamic, model_path="/models/llama") + >>> engine_config = await make_engine(runtime, args) + """ ... async def run_input(runtime: DistributedRuntime, input: str, engine_config: EngineConfig) -> None: - """Start an engine, connect it to an input, and run until stopped.""" + """Start an engine, connect it to an input, and run until stopped. + + Examples: + >>> from dynamo._core import run_input + >>> await run_input(runtime, "http", engine_config) + >>> await run_input(runtime, "grpc", engine_config) + >>> await run_input(runtime, "text", engine_config) + """ ... class Layer: @@ -1202,6 +1398,12 @@ class Layer: class Block: """ A KV cache block + + Examples: + >>> block = blocks[0] # get first block from a BlockList + >>> num_layers = len(block) + >>> layer = block[0] + >>> all_layers = block.to_list() """ ... @@ -1289,6 +1491,11 @@ class BlockList: class BlockManager: """ A KV cache block manager + + Examples: + >>> bm = BlockManager(worker_id=0, num_layer=32, page_size=16, inner_dim=128) + >>> blocks = await bm.allocate_device_blocks(4) + >>> layer = blocks[0][0] # first block, first layer """ def __init__( @@ -1401,6 +1608,14 @@ class KvbmRequest: class KvRouter: """ A KV-aware router that performs intelligent routing based on KV cache overlap. + + Examples: + >>> config = KvRouterConfig() + >>> kv_router = KvRouter(endpoint=endpoint, block_size=64, kv_router_config=config) + >>> stream = await kv_router.generate(token_ids=[1, 2, 3], model="my-model") + >>> async for chunk in stream: + ... print(chunk) + >>> worker_id, dp_rank, overlap = await kv_router.best_worker([1, 2, 3]) """ def __init__( @@ -1583,7 +1798,14 @@ class KvRouter: ... class EngineType: - """Engine type for Dynamo workers""" + """Engine type for Dynamo workers + + Examples: + >>> from dynamo._core import EngineType + >>> engine = EngineType.Dynamic + >>> engine = EngineType.Echo + >>> engine = EngineType.Mocker + """ Echo: "EngineType" Dynamic: "EngineType" Mocker: "EngineType" @@ -1593,6 +1815,21 @@ class EntrypointArgs: """ Settings to connect an input to a worker and run them. Use by `dynamo run`. + + Examples: + >>> from dynamo._core import EntrypointArgs, EngineType, RouterConfig, RouterMode # doctest: +SKIP + >>> args = EntrypointArgs( # doctest: +SKIP + ... engine_type=EngineType.Dynamic, + ... model_path="/models/llama-3-8b", + ... model_name="dyn://dynamo.backend.generate", + ... http_port=8000, + ... router_config=RouterConfig(mode=RouterMode.KV), + ... ) + >>> mocker_args = EntrypointArgs( # doctest: +SKIP + ... engine_type=EngineType.Mocker, + ... model_path="/models/llama-3-8b", + ... is_prefill=True, + ... ) """ def __init__( @@ -1650,6 +1887,12 @@ class PlannerDecision: Fields: num_prefill_workers, num_decode_workers, decision_id. -1 in any of those fields mean not set, usually because planner hasn't decided anything yet. Call VirtualConnectorClient.complete(event) when action is completed. + + Examples: + >>> client = VirtualConnectorClient(runtime, "my-namespace") + >>> decision = await client.get() + >>> print(decision.num_prefill_workers, decision.num_decode_workers) + >>> await client.complete(decision) """ num_prefill_workers: int num_decode_workers: int diff --git a/lib/bindings/python/src/dynamo/health_check.py b/lib/bindings/python/src/dynamo/health_check.py index 38076eb0adf0..07c9178239bb 100644 --- a/lib/bindings/python/src/dynamo/health_check.py +++ b/lib/bindings/python/src/dynamo/health_check.py @@ -31,6 +31,22 @@ def load_health_check_from_env( Args: env_var: Name of the environment variable to check (default: DYN_HEALTH_CHECK_PAYLOAD) + Examples: + >>> import os + >>> from dynamo.health_check import load_health_check_from_env + >>> + >>> _prev = os.environ.get("DYN_HEALTH_CHECK_PAYLOAD") + >>> os.environ["DYN_HEALTH_CHECK_PAYLOAD"] = ( + ... '{"prompt": "test", "max_tokens": 1}' + ... ) + >>> payload = load_health_check_from_env() + >>> payload + {'prompt': 'test', 'max_tokens': 1} + >>> if _prev is None: # doctest: +SKIP + ... os.environ.pop("DYN_HEALTH_CHECK_PAYLOAD", None) + ... else: + ... os.environ["DYN_HEALTH_CHECK_PAYLOAD"] = _prev + Returns: Dict containing the health check payload, or None if not set. """ @@ -68,6 +84,23 @@ class HealthCheckPayload: in their __init__ method. Environment variable DYN_HEALTH_CHECK_PAYLOAD can override the default. + + Examples: + >>> import os + >>> from dynamo.health_check import HealthCheckPayload + >>> + >>> os.environ.pop("DYN_HEALTH_CHECK_PAYLOAD", None) + >>> class MyBackendHealthCheck(HealthCheckPayload): + ... def __init__(self): + ... self.default_payload = { + ... "prompt": "health check", + ... "max_tokens": 1, + ... } + ... super().__init__() + >>> + >>> hc = MyBackendHealthCheck() + >>> hc.to_dict() + {'prompt': 'health check', 'max_tokens': 1} """ default_payload: Dict[str, Any] # Type hint for mypy - set by subclasses diff --git a/lib/bindings/python/src/dynamo/logits_processing/base.py b/lib/bindings/python/src/dynamo/logits_processing/base.py index c5ee5658972d..9b937e46bf18 100644 --- a/lib/bindings/python/src/dynamo/logits_processing/base.py +++ b/lib/bindings/python/src/dynamo/logits_processing/base.py @@ -20,6 +20,21 @@ class BaseLogitsProcessor(Protocol): All logits processors must implement this interface to be compatible with backend adapters (TRT-LLM, vLLM, SGLang). + + Examples: + >>> import torch + >>> from typing import Sequence + >>> from dynamo.logits_processing import BaseLogitsProcessor + >>> + >>> class TemperatureProcessor: + ... def __init__(self, temperature: float = 0.7): + ... self.temperature = temperature + ... def __call__( + ... self, input_ids: Sequence[int], logits: torch.Tensor + ... ) -> None: + ... logits.div_(self.temperature) + >>> + >>> assert isinstance(TemperatureProcessor(), BaseLogitsProcessor) """ def __call__( diff --git a/lib/bindings/python/src/dynamo/runtime/__init__.py b/lib/bindings/python/src/dynamo/runtime/__init__.py index 70beba65c896..7c498f3c4e9b 100644 --- a/lib/bindings/python/src/dynamo/runtime/__init__.py +++ b/lib/bindings/python/src/dynamo/runtime/__init__.py @@ -24,6 +24,16 @@ def dynamo_worker(enable_nats: bool = True): enable_nats: Whether to enable NATS for KV events. Defaults to True. If request_plane is "nats", NATS is always enabled. Pass False (via --no-kv-events flag) to disable NATS initialization. + + Examples: + >>> from dynamo.runtime import DistributedRuntime, dynamo_worker + >>> + >>> @dynamo_worker() + ... async def worker(runtime: DistributedRuntime): + ... endpoint = runtime.endpoint("dynamo.backend.generate") + ... await endpoint.serve_endpoint(handler.generate) + >>> + >>> asyncio.run(worker()) """ def decorator(func): @@ -61,6 +71,33 @@ async def wrapper(*args, **kwargs): def dynamo_endpoint( request_model: Union[Type[BaseModel], Type[Any]], response_model: Type[BaseModel] ) -> Callable: + """ + Decorator that parses incoming requests into Pydantic models on an async generator endpoint. + + Currently validates and converts the *request* payload (JSON string or dict) + into ``request_model``. ``response_model`` is accepted for forward + compatibility but response validation is not yet implemented. + + Args: + request_model: Pydantic model class (or Any) for incoming requests. + response_model: Pydantic model class reserved for future response validation. + + Examples: + >>> from pydantic import BaseModel + >>> from dynamo.runtime import dynamo_endpoint + >>> + >>> class Request(BaseModel): + ... data: str + >>> class Response(BaseModel): + ... char: str + >>> + >>> class RequestHandler: + ... @dynamo_endpoint(Request, Response) + ... async def generate(self, request): + ... for char in request.data: + ... yield char + """ + def decorator( func: Callable[..., AsyncGenerator[Any, None]], ) -> Callable[..., AsyncGenerator[Any, None]]: