Skip to content

datarobot_genai.dragent.registry_warmup

registry_warmup

Startup prefetch for central agent card registry lookups.

Collects registry blocks from authenticated_a2a_client function groups in the loaded NAT config and batch-fetches all agent cards before the server accepts traffic.

RegistryLookupIds

Bases: NamedTuple

Registry lookup IDs collected from a workflow config, by ID kind.

Source code in datarobot_genai/dragent/registry_warmup.py 
59
60
61
62
63
64
65
66
67
68
class RegistryLookupIds(NamedTuple):
    """Registry lookup IDs collected from a workflow config, by ID kind."""

    deployment_ids: list[str]
    external_ids: list[str]
    workload_ids: list[str]

    def is_empty(self) -> bool:
        """Return whether no registry lookups were configured at all."""
        return not (self.deployment_ids or self.external_ids or self.workload_ids)

is_empty 

is_empty() -> bool

Return whether no registry lookups were configured at all.

Source code in datarobot_genai/dragent/registry_warmup.py 
66
67
68
def is_empty(self) -> bool:
    """Return whether no registry lookups were configured at all."""
    return not (self.deployment_ids or self.external_ids or self.workload_ids)

is_registry_warm 

is_registry_warm() -> bool

Return whether the latest startup prefetch completed successfully.

When there are no registry-backed function groups, this returns True (nothing to warm).

Source code in datarobot_genai/dragent/registry_warmup.py 
45
46
47
48
49
50
51
def is_registry_warm() -> bool:
    """Return whether the latest startup prefetch completed successfully.

    When there are no registry-backed function groups, this returns ``True``
    (nothing to warm).
    """
    return _WarmState.warm

reset_registry_warm_state 

reset_registry_warm_state() -> None

Reset warm state (for tests).

Source code in datarobot_genai/dragent/registry_warmup.py 
54
55
56
def reset_registry_warm_state() -> None:
    """Reset warm state (for tests)."""
    _WarmState.warm = False

collect_registry_lookup_ids 

collect_registry_lookup_ids(config: Config) -> RegistryLookupIds

Return deduplicated registry lookup IDs per kind from config.

Only authenticated_a2a_client function groups with a registry block contribute IDs. Order is preserved (first-seen).

Source code in datarobot_genai/dragent/registry_warmup.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def collect_registry_lookup_ids(config: Config) -> RegistryLookupIds:
    """Return deduplicated registry lookup IDs per kind from *config*.

    Only ``authenticated_a2a_client`` function groups with a ``registry`` block
    contribute IDs. Order is preserved (first-seen).
    """
    collected = RegistryLookupIds(deployment_ids=[], external_ids=[], workload_ids=[])
    seen: dict[str, set[str]] = {"deployment": set(), "external": set(), "workload": set()}

    function_groups = getattr(config, "function_groups", None) or {}
    for fg_config in function_groups.values():
        if not isinstance(fg_config, AuthenticatedA2AClientConfig):
            continue
        if fg_config.registry is None:
            continue
        for kind, lookup_id, ids in (
            ("deployment", fg_config.registry.deployment_id, collected.deployment_ids),
            ("external", fg_config.registry.external_id, collected.external_ids),
            ("workload", fg_config.registry.workload_id, collected.workload_ids),
        ):
            if lookup_id and lookup_id not in seen[kind]:
                seen[kind].add(lookup_id)
                ids.append(lookup_id)

    return collected

register_registry_lookup_ids 

register_registry_lookup_ids(registry: AgentCardRegistry, collected: RegistryLookupIds) -> None

Register all lookup IDs from collected on registry.

Config validation registers IDs on the first registry singleton. L2 MemorySpace provisioning resets that singleton at lifespan startup, so warmup and the background refresh loop must re-register from the workflow config before calling :meth:AgentCardRegistry.prefetch or :meth:AgentCardRegistry.refresh_all_registered.

Source code in datarobot_genai/dragent/registry_warmup.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def register_registry_lookup_ids(
    registry: AgentCardRegistry,
    collected: RegistryLookupIds,
) -> None:
    """Register all lookup IDs from *collected* on *registry*.

    Config validation registers IDs on the first registry singleton. L2
    MemorySpace provisioning resets that singleton at lifespan startup, so
    warmup and the background refresh loop must re-register from the workflow
    config before calling :meth:`AgentCardRegistry.prefetch` or
    :meth:`AgentCardRegistry.refresh_all_registered`.
    """
    for deployment_id in collected.deployment_ids:
        registry.register(deployment_id=deployment_id)
    for external_id in collected.external_ids:
        registry.register(external_id=external_id)
    for workload_id in collected.workload_ids:
        registry.register(workload_id=workload_id)

warmup_registry_from_config async 

warmup_registry_from_config(config: Config) -> None

Batch-prefetch agent cards for all registry-backed A2A clients in config.

No-op when no registry lookups are configured. On failure, logs an error and leaves :func:is_registry_warm as False.

Source code in datarobot_genai/dragent/registry_warmup.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
async def warmup_registry_from_config(config: Config) -> None:
    """Batch-prefetch agent cards for all registry-backed A2A clients in *config*.

    No-op when no registry lookups are configured. On failure, logs an error and
    leaves :func:`is_registry_warm` as ``False``.
    """
    await ensure_registry_l2_cache_provisioned_async(phase="lifespan-warmup")

    collected = collect_registry_lookup_ids(config)
    if collected.is_empty():
        logger.debug("No registry-backed A2A function groups; skipping agent card prefetch.")
        _WarmState.warm = True
        return

    _WarmState.warm = False
    logger.info(
        "Warming agent card cache for registry-backed A2A clients "
        "(deployment_ids=%s, external_ids=%s, workload_ids=%s)",
        collected.deployment_ids,
        collected.external_ids,
        collected.workload_ids,
    )

    try:
        registry = await get_default_registry()
        register_registry_lookup_ids(registry, collected)
        await registry.prefetch(
            deployment_ids=collected.deployment_ids or None,
            external_ids=collected.external_ids or None,
            workload_ids=collected.workload_ids or None,
        )
    except Exception:
        logger.exception(
            "Agent card registry prefetch failed; registry-backed A2A tools may "
            "degrade until the central registry is reachable."
        )
        return

    _WarmState.warm = True
    logger.info(
        "Agent card registry prefetch complete "
        "(%d deployment ID(s), %d external ID(s), %d workload ID(s)).",
        len(collected.deployment_ids),
        len(collected.external_ids),
        len(collected.workload_ids),
    )