kaxil commented on code in PR #75369:
URL: https://github.com/apache/airflow/pull/75369#discussion_r4450442200
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -562,6 +586,92 @@ def _provider_factory(pname: str) -> Any:
return infer_model(model_name)
+ def _build_missing_model_error(self) -> ValueError:
+ return ValueError(
+ f"No model specified for connection '{self.llm_conn_id}'. Set
model_id on the "
+ "hook or the Model field on the connection."
+ )
+
+ def _resolve_readiness_backend(self, model_name: str) ->
ModelReadinessBackend:
+ """Build the built-in readiness backend for *model_name*, which only
exists for Ollama."""
+ prefix, _, ollama_model = model_name.partition(":")
+ if prefix != "ollama":
+ raise NotImplementedError(
+ f"ensure_model_ready has no built-in readiness backend for
'{model_name}' on "
+ f"connection '{self.llm_conn_id}'; only Ollama is built in. "
+ "Pass backend= an object with is_ready() and load()."
+ )
+ conn, _ = self._get_conn_and_extra()
+ if not conn.host:
Review Comment:
`get_conn()` works for an Ollama connection with no host, since
pydantic-ai's `OllamaProvider` falls back to `OLLAMA_BASE_URL`, but this raises
`ValueError` for that same connection. Should it use `conn.host or
os.getenv("OLLAMA_BASE_URL")` and raise only when both are empty?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -45,6 +51,27 @@
P = ParamSpec("P")
R = TypeVar("R")
+
+class ModelReadinessBackend(Protocol):
Review Comment:
The hook resolves the model name and connection, but the backend never gets
them, so a custom backend has to carry its own URL and model name. That also
hits the `openai:<model>` against Ollama's `/v1` setup that
`self_hosted_models.rst` calls valid: it gets `NotImplementedError`, and the
only way around it is building `OllamaBackend` by hand, which isn't documented.
Would passing `model_name` and `base_url` to the backend (or taking a factory)
work better?
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -1830,6 +1832,305 @@ def test_reports_provider_class_build_failure(
assert "unexpected keyword argument 'api_key'" in message
+class TestPydanticAIHookEnsureModelReady:
+ _HOOK_MODULE = "airflow.providers.common.ai.hooks.pydantic_ai"
+
+ @staticmethod
+ def _build_hook(
+ *,
+ model_id: str | None = "ollama:llama3",
+ host: str | None = "http://ollama:11434/v1";,
+ **hook_kwargs,
+ ) -> tuple[PydanticAIHook, Connection]:
+ hook = PydanticAIHook(llm_conn_id="local_conn", model_id=model_id,
**hook_kwargs)
+ conn = Connection(conn_id="local_conn", conn_type="pydanticai",
host=host)
+ return hook, conn
+
+ @staticmethod
+ def _install_fake_clock(module: str):
Review Comment:
Now that the polling is gone, the `sleep` wiring here and
`mock_time.sleep.assert_not_called()` below check a call the method can't make.
Could these two helpers become one fixture that only patches `time.monotonic`?
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -1830,6 +1832,305 @@ def test_reports_provider_class_build_failure(
assert "unexpected keyword argument 'api_key'" in message
+class TestPydanticAIHookEnsureModelReady:
+ _HOOK_MODULE = "airflow.providers.common.ai.hooks.pydantic_ai"
+
+ @staticmethod
+ def _build_hook(
+ *,
+ model_id: str | None = "ollama:llama3",
Review Comment:
Every test sets `model_id`, so the setup the docs show (model in the
connection extra, `PydanticAIHook(llm_conn_id=...)` alone) isn't covered, and
neither is a bare name with the `model_provider` extra. Could one case use
`model_id=None` with the model in `extra` and assert what `OllamaBackend` is
built with?
##########
providers/common/ai/docs/self_hosted_models.rst:
##########
@@ -257,6 +257,59 @@ 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 asks Ollama which models are loaded,
+sends a load request if this connection's model is missing, waits for the load
+to finish, then checks once that Ollama lists the model:
+
+.. 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:
Review Comment:
How would `@task.llm` or `AgentOperator` users use this? They build their
own hook, so I guess it'd be an upstream task. But the load doesn't send
`keep_alive`, so Ollama drops the model again after 5 minutes (its default) if
the LLM task sits in the queue longer. Should the docs mention that?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -562,6 +586,92 @@ def _provider_factory(pname: str) -> Any:
return infer_model(model_name)
+ def _build_missing_model_error(self) -> ValueError:
+ return ValueError(
+ f"No model specified for connection '{self.llm_conn_id}'. Set
model_id on the "
+ "hook or the Model field on the connection."
+ )
+
+ def _resolve_readiness_backend(self, model_name: str) ->
ModelReadinessBackend:
+ """Build the built-in readiness backend for *model_name*, which only
exists for Ollama."""
+ prefix, _, ollama_model = model_name.partition(":")
+ if prefix != "ollama":
+ raise NotImplementedError(
+ f"ensure_model_ready has no built-in readiness backend for
'{model_name}' on "
+ f"connection '{self.llm_conn_id}'; only Ollama is built in. "
+ "Pass backend= an object with is_ready() and load()."
+ )
+ conn, _ = self._get_conn_and_extra()
+ if not conn.host:
+ raise ValueError(f"Connection '{self.llm_conn_id}' needs a host to
check model readiness.")
+ return OllamaBackend(conn.host, ollama_model)
+
+ def ensure_model_ready(
+ self, *, timeout: float = 300.0, backend: ModelReadinessBackend | None
= None
+ ) -> None:
+ """
+ Block until this connection's local model is loaded, asking the server
to load it if needed.
+
+ Asks the server whether the model is loaded; if not, sends a load
request that returns once
+ the model is loaded, then checks once more that the server lists it.
Any error from the
+ server is raised as is.
+
+ Ollama is built in; for any other server pass *backend*. Only this
hook's own connection is
+ checked -- connections in ``fallback_conn_ids`` are not, so a fallback
model is never loaded
+ ahead of need.
+
+ The built-in Ollama backend does not send the connection password, so
a server behind
+ authentication is not supported.
+
+ Blocks and holds the worker slot for the whole wait.
+
+ :param timeout: Seconds to wait in total for the model to become ready.
+ :param backend: Object with is_ready() and load() (see
ModelReadinessBackend) to check and
+ load the model; when given, only the model prefix and connection
host checks are skipped.
+ The model name is still resolved from the hook argument or the
connection, and a
+ missing model still raises ``ValueError``.
+ :raises ValueError: If *timeout* is not a positive finite number, no
model is configured,
+ or *backend* is not given and the connection has no host.
+ :raises NotImplementedError: If *backend* is not given and the model's
provider prefix has
+ no built-in backend (only ``ollama`` has one).
+ :raises TimeoutError: If the time runs out before the model is
confirmed ready, or a
+ request times out.
+ :raises RuntimeError: If the load request finishes but the server does
not list the model
+ as loaded.
+ :raises httpx2.TransportError: If a request cannot reach the server.
+ :raises httpx2.HTTPStatusError: If the server answers with an error
status.
+ """
+ if not math.isfinite(timeout) or timeout <= 0:
+ raise ValueError(f"timeout must be positive and finite, got
{timeout}.")
+ if not (configured_name := self._get_configured_model_name()):
+ raise self._build_missing_model_error()
+ model_name = self._qualify_model_name(configured_name)
+ if backend is None:
+ backend = self._resolve_readiness_backend(model_name)
+
+ started = time.monotonic()
+ deadline = started + timeout
+ if backend.is_ready(timeout=min(timeout, DEFAULT_REQUEST_TIMEOUT)):
Review Comment:
`/api/ps` doesn't look like a safe readiness signal while a load is in
flight. On Ollama v0.30.0 and earlier it lists a model as soon as the load
starts (the runner goes into `s.sched.loaded` before `WaitUntilRunning`), so a
mapped sibling task can see `True` here mid-load. On v0.40.13 `loadedModels()`
waits on each runner's `refMu`, which the scheduler holds until the load
finishes, so this call blocks instead. If any model on the server takes over
30s to load, this raises `TimeoutError`.
An empty-prompt `/api/generate` already waits for the load, and returns
right away if the model is loaded. Could we send just that with the remaining
budget and drop both `/api/ps` checks? That'd also get rid of the name compare,
which won't match `/api/ps` for names like `library/llama3`.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_ollama.py:
##########
@@ -0,0 +1,156 @@
+# 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 json
+
+import httpx2
+import pytest
+
+from airflow.providers.common.ai.hooks.ollama import (
+ OllamaBackend,
+ build_ollama_native_root,
+ normalize_ollama_model_name,
+)
+
+
+def _build_backend(
+ handler, base_url: str = "http://ollama:11434/v1";, model_name: str =
"llama3"
+) -> OllamaBackend:
+ return OllamaBackend(base_url, model_name,
transport=httpx2.MockTransport(handler))
+
+
+class TestBuildNativeRoot:
+ @pytest.mark.parametrize(
+ "base_url",
+ [
+ "http://h:11434/v1";,
+ "http://h:11434/v1/";,
+ "http://h:11434";,
+ "http://h:11434/";,
+ ],
+ )
+ def test_returns_server_root(self, base_url):
+ assert build_ollama_native_root(base_url) == "http://h:11434";
+
+
+class TestNormalizeModelName:
+ @pytest.mark.parametrize(
+ ("model_name", "expected"),
+ [
+ ("llama3", "llama3:latest"),
+ ("llama3:8b", "llama3:8b"),
+ ("hf.co/org/model:q4", "hf.co/org/model:q4"),
+ ("localhost:5000/library/llama3",
"localhost:5000/library/llama3:latest"),
+ ],
+ )
+ def test_adds_latest_only_when_tag_is_missing(self, model_name, expected):
+ assert normalize_ollama_model_name(model_name) == expected
+
+
+class TestOllamaBackendIsReady:
+ @pytest.mark.parametrize(
+ ("model_name", "loaded", "expected"),
+ [
+ ("llama3", ["llama3:latest"], True),
+ ("llama3:latest", ["llama3:latest"], True),
+ ("llama3:8b", ["llama3:8b", "other:latest"], True),
+ ("llama3:8b", ["llama3:70b"], False),
+ ("llama3", ["llama3:8b"], False),
+ ("llama3", [], False),
+ ],
+ )
+ def test_matches_model_name_against_loaded_models(self, model_name,
loaded, expected):
+ requests: list[httpx2.Request] = []
+
+ def handler(request: httpx2.Request) -> httpx2.Response:
+ requests.append(request)
+ return httpx2.Response(200, json={"models": [{"name": name} for
name in loaded]})
+
+ assert _build_backend(handler, model_name=model_name).is_ready() is
expected
+ assert [(r.method, r.url.path) for r in requests] == [("GET",
"/api/ps")]
+
+ def test_http_error_propagates(self):
+ backend = _build_backend(lambda request: httpx2.Response(503))
+
+ with pytest.raises(httpx2.HTTPStatusError):
+ backend.is_ready()
+
+ def test_connection_error_propagates(self):
+ def handler(request: httpx2.Request) -> httpx2.Response:
+ raise httpx2.ConnectError("refused", request=request)
+
+ with pytest.raises(httpx2.ConnectError):
+ _build_backend(handler).is_ready()
+
+ def test_timeout_is_raised_as_timeout_error_naming_model_and_seconds(self):
+ def handler(request: httpx2.Request) -> httpx2.Response:
+ raise httpx2.ReadTimeout("slow", request=request)
+
+ with pytest.raises(TimeoutError, match=r"'llama3'.*5\.0 seconds") as
exc_info:
+ _build_backend(handler).is_ready(timeout=5.0)
+
+ assert isinstance(exc_info.value.__cause__, httpx2.ReadTimeout)
+
+ def test_non_json_body_propagates_decode_error(self):
+ backend = _build_backend(lambda request: httpx2.Response(200,
content=b"<html>"))
+
+ with pytest.raises(json.JSONDecodeError):
+ backend.is_ready()
+
+
+class TestOllamaBackendLoad:
+ def test_posts_model_to_generate_without_keep_alive(self):
+ requests: list[httpx2.Request] = []
+
+ def handler(request: httpx2.Request) -> httpx2.Response:
+ requests.append(request)
+ return httpx2.Response(200, json={"response": "", "done": True,
"done_reason": "load"})
+
+ _build_backend(handler).load(timeout=12.0)
+
+ (request,) = requests
+ assert (request.method, request.url.path) == ("POST", "/api/generate")
+ body = json.loads(request.content)
+ assert body["model"] == "llama3"
Review Comment:
Nothing here checks the timeout the client actually gets: the timeout tests
raise `ReadTimeout` from the handler, and the seconds in the message come from
the argument. Could this assert that `request.extensions["timeout"]["read"]`
equals the timeout passed to `load()`, and `body == {"model": "llama3",
"stream": False}` while at it? `load()` quietly falling back to 30s is the
failure this feature exists to avoid.
--
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]
