ParagEkbote commented on PR #76317:
URL: https://github.com/apache/airflow/pull/76317#issuecomment-6636177959
> 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.
Thank you for reaching out and for the support. Currently, I would prefer to
finalize the HF provider's public API and focus on the features stated in the
[mailing
list](https://lists.apache.org/thread.html/3sj8p6smt31nm712yyfjzy7bhh4lxrq3).
I've created an [issue](https://github.com/apache/airflow/issues/76470) for the
requested feature and will certainly ask for your feedback in a follow up PR.
Regards,
Parag Ekbote
--
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]
