kaxil commented on code in PR #75369:
URL: https://github.com/apache/airflow/pull/75369#discussion_r4467663489
##########
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
Review Comment:
pydantic-ai's default request timeout is 600s (`DEFAULT_HTTP_TIMEOUT`),
longer than the 300s default here, so with defaults the agent's first request
already waits out the load. Is the case this covers a short `timeout` in
`model_settings`, like the 30s in `retry_policies.rst`? If so, could this
section say that, so people know when they need it instead of raising the
timeout?
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -62,6 +64,9 @@
_GOOGLE_MODEL_PREFIX_RE = re.compile(r"google[\w-]*(?=:)")
+def build_backend(*, model_name: str, base_url: str | None) ->
ModelReadinessBackend: ...
Review Comment:
The MyPy providers check fails on this stub (`Missing return statement
[empty-body]`). A `raise NotImplementedError` body fixes it, or if
`backend_factory` gets a Protocol, autospec that and drop the stub.
##########
providers/common/ai/docs/connections/pydantic_ai.rst:
##########
@@ -61,7 +61,8 @@ API Key (Password field)
``GOOGLE_APPLICATION_CREDENTIALS``).
Host (optional)
- Base URL for the provider's API. Only needed for custom endpoints:
+ Base URL for the provider's API. Only needed for custom endpoints. For
Ollama it can
+ instead come from the ``OLLAMA_BASE_URL`` environment variable:
Review Comment:
Only pydantic-ai's `OllamaProvider` reads `OLLAMA_BASE_URL`, so this holds
for `ollama:` models. The Ollama example on this page uses `openai:llama3`,
which reads `OPENAI_BASE_URL` instead. Could this say `ollama:` models?
##########
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:
This assumes the upstream task and the LLM task reach the same Ollama
process. With Ollama on each worker (`localhost`, as in the connection example)
or several replicas behind a load balancer, the warm-up can land on a different
server than the LLM call. Worth a line saying so?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/ollama.py:
##########
@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import httpx2
+
+
+def build_ollama_native_root(base_url: str) -> str:
+ """Return the server root of an Ollama URL, dropping a trailing ``/v1``
(the OpenAI-compatible path)."""
+ return base_url.rstrip("/").removesuffix("/v1")
+
+
+class OllamaBackend:
+ """
+ Load control for a model served by Ollama, through its native
``/api/generate`` endpoint.
+
+ :param base_url: Ollama URL; a trailing ``/v1`` is dropped to reach the
native API.
+ :param model_name: Model name without the ``ollama:`` prefix.
+ :param transport: Optional httpx2 transport, mainly to stub the server in
tests.
+
+ Known limitation: the connection password is not sent, so no
``Authorization`` header is
+ added and servers behind authentication are not supported.
+ """
+
+ def __init__(
+ self,
+ base_url: str,
+ model_name: str,
+ *,
+ transport: httpx2.BaseTransport | None = None,
+ ) -> None:
+ self.root = build_ollama_native_root(base_url)
+ self.model_name = model_name
+ self._transport = transport
+
+ def load(self, *, timeout: float) -> None:
+ """
+ Load the model, waiting at most *timeout* seconds.
+
+ :raises TimeoutError: If the server cannot be reached or does not
finish loading within *timeout*.
+ """
+ # A /api/generate request with no prompt and stream=False returns once
the model is
+ # available. keep_alive only controls how long the model stays loaded
afterwards.
+ try:
+ with httpx2.Client(transport=self._transport, timeout=timeout) as
client:
+ response = client.post(
+ f"{self.root}/api/generate", json={"model":
self.model_name, "stream": False}
+ )
+ response.raise_for_status()
Review Comment:
`raise_for_status()` only puts the status line and URL in the message, so a
model that hasn't been pulled surfaces as `Client error '404 Not Found' for url
'.../api/generate'`, with no model name and without Ollama's `error` text.
That's probably the most common failure here. Could this catch
`HTTPStatusError` and re-raise it with the model name and the server's `error`
field?
--
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]
