dabla commented on PR #76317:
URL: https://github.com/apache/airflow/pull/76317#issuecomment-6607860746
Thanks for this, nice to see a Hugging Face provider coming in!
One suggestion for the hook: add native async variants of the metadata
calls, so they can be used from async tasks and triggers without blocking the
event loop. Deferrable operators aren't needed for this, async methods on the
hook are enough.
`HfApi` itself is sync only, but `huggingface_hub` already ships what's
needed: `huggingface_hub.utils.get_async_session()` returns the configured
`httpx.AsyncClient` that `AsyncInferenceClient` uses, and the Hub endpoints
behind the current methods are plain REST calls:
| Hook method | `HfApi` call | Endpoint |
|---|---|---|
| `get_dataset_info` | `dataset_info` | `GET
/api/datasets/{repo_id}[/revision/{revision}]` |
| `find_commit_by_marker` | `list_repo_refs` | `GET
/api/datasets/{repo_id}/refs?include_prs=1` |
| `find_commit_by_marker` | `list_repo_commits` | `GET
/api/datasets/{repo_id}/commits/{revision}` (paginated) |
Something along these lines:
```python
from urllib.parse import quote
from huggingface_hub import DatasetInfo, constants
from huggingface_hub.utils import build_hf_headers, get_async_session,
hf_raise_for_status
from airflow.providers.common.compat.connection import get_async_connection
async def aget_dataset_info(self, repo_id: str, revision: str | None = None)
-> DatasetInfo:
connection = await get_async_connection(self.huggingface_conn_id,
hook=self)
url = f"{constants.ENDPOINT}/api/datasets/{repo_id}"
if revision:
url += f"/revision/{quote(revision, safe='')}"
async with get_async_session() as client:
response = await client.get(url,
headers=build_hf_headers(token=connection.password))
hf_raise_for_status(response)
return DatasetInfo(**response.json())
```
Two things to watch for:
- The token has to come from `get_async_connection` (common.compat): calling
the sync `get_token()` / `get_connection()` inside a coroutine blocks the event
loop and can deadlock on the supervisor channel.
- Building `DatasetInfo` / `GitCommitInfo` from the JSON (as `HfApi` does
internally) keeps the async methods returning the same types as the sync ones,
so `_commit_is_complete` and the marker matching can be shared between both
paths.
`load_dataset` and `push_to_hub` can stay sync, since `datasets` has no
async API and those are file transfers anyway. When Inference Endpoints are
added later, `AsyncInferenceClient` would give the async side of that for free.
Happy to review if you pick this up, here or in a follow-up PR.
--
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]
