Lee-W commented on code in PR #75369: URL: https://github.com/apache/airflow/pull/75369#discussion_r4519788070
########## providers/common/ai/docs/self_hosted_models.rst: ########## @@ -257,6 +257,94 @@ model doesn't have to guess a path. That matters because ``HttpHook`` raises on any response outside the 2xx/3xx range: if the model guessed wrong, the task would fail outright rather than the agent recovering gracefully. +Waiting for an Ollama model to load +------------------------------------- + +Ollama loads a model into memory on the first request, which can take long +enough to time out that request. To load the model ahead of time, call +:meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.ensure_model_ready` +before the task builds its agent. It sends Ollama a load request for this +connection's model and waits until Ollama answers, which it does once the model +is available (right away if it is already loaded): + +.. code-block:: python + + from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook + + hook = PydanticAIHook(llm_conn_id="pydanticai_default") + hook.ensure_model_ready(timeout=300) + agent = hook.create_agent(instructions="Answer briefly.") + +Nothing calls it for you; it is opt-in. It has these limits: + +- Only Ollama is built in. For any other model prefix pass ``backend_factory=`` + (see below); without it, ``NotImplementedError`` is raised. +- Only the hook's own connection is checked. Models behind ``fallback_conn_ids`` + are not loaded. +- The call blocks and holds the worker slot for the whole wait. +- The built-in Ollama backend does not send the connection's password, so an Ollama + server behind authentication is not supported. + +``@task.llm`` and ``AgentOperator`` build their own hook from ``llm_conn_id`` and +their ``model_id``, so call ``ensure_model_ready`` from an upstream task with Review Comment: Added a paragraph: with an Ollama on each worker (`localhost`) or several replicas behind a load balancer, the warm-up request can reach a different server than the LLM call and then has no effect on it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
