Lee-W commented on code in PR #75369:
URL: https://github.com/apache/airflow/pull/75369#discussion_r4519708375
##########
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:
Change `OllamaBackend.load` to send one empty-prompt `/api/generate` with
the whole `timeout`. Removed `/api/ps` checks and the name comparison.
Checked against a local Ollama 0.40.9: a cold load blocked about 5 seconds
before returning and a loaded model returned in about 10 ms.
--
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]
