use std::future::Future;
use std::sync::Arc;
use crate::client::HFClient;
use crate::error::{HFError, HFResult};
use crate::repository::{HFRepository, RepoType, RepoTypeDataset, RepoTypeKernel, RepoTypeModel, RepoTypeSpace};
pub(crate) struct RuntimeThread {
handle: tokio::runtime::Handle,
_shutdown: futures::channel::oneshot::Sender<()>,
}
impl RuntimeThread {
fn spawn() -> HFResult<Arc<Self>> {
let (handle_tx, handle_rx) = std::sync::mpsc::sync_channel(1);
let (shutdown_tx, shutdown_rx) = futures::channel::oneshot::channel::<()>();
std::thread::Builder::new()
.name("hf-hub-blocking-runtime".to_string())
.spawn(move || {
let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
Ok(runtime) => runtime,
Err(e) => {
let _ = handle_tx.send(Err(e));
return;
},
};
let _ = handle_tx.send(Ok(runtime.handle().clone()));
runtime.block_on(async {
let _ = shutdown_rx.await;
});
})
.map_err(|e| HFError::Other(format!("Failed to spawn tokio runtime thread: {e}")))?;
handle_rx
.recv()
.map_err(|_| HFError::Other("tokio runtime thread exited before initializing".to_string()))?
.map_err(|e| HFError::Other(format!("Failed to create tokio runtime: {e}")))
.map(|handle| {
Arc::new(Self {
handle,
_shutdown: shutdown_tx,
})
})
}
pub(crate) fn block_on<F>(&self, future: F) -> F::Output
where
F: Future + Send,
F::Output: Send,
{
std::thread::scope(|scope| {
scope
.spawn(|| self.handle.block_on(future))
.join()
.unwrap_or_else(|panic| std::panic::resume_unwind(panic))
})
}
}
#[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
#[derive(Clone)]
pub struct HFClientSync {
pub(crate) inner: HFClient,
pub(crate) runtime: Arc<RuntimeThread>,
}
#[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
pub struct HFRepositorySync<T: RepoType> {
pub(crate) inner: Arc<HFRepository<T>>,
pub(crate) runtime: Arc<RuntimeThread>,
}
impl<T: RepoType> Clone for HFRepositorySync<T> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
runtime: Arc::clone(&self.runtime),
}
}
}
#[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
#[derive(Clone)]
pub struct HFBucketSync {
pub(crate) inner: Arc<crate::buckets::HFBucket>,
pub(crate) runtime: Arc<RuntimeThread>,
}
impl std::fmt::Debug for HFClientSync {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HFClientSync").finish()
}
}
impl<T: RepoType> std::fmt::Debug for HFRepositorySync<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HFRepositorySync").field("inner", &self.inner).finish()
}
}
impl std::fmt::Debug for HFBucketSync {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("HFBucketSync").field("inner", &self.inner).finish()
}
}
impl HFClientSync {
pub fn new() -> HFResult<Self> {
Ok(Self {
inner: HFClient::new()?,
runtime: RuntimeThread::spawn()?,
})
}
pub fn from_inner(inner: HFClient) -> HFResult<Self> {
Ok(Self {
inner,
runtime: RuntimeThread::spawn()?,
})
}
pub fn repository<T: RepoType>(
&self,
repo_type: T,
owner: impl Into<String>,
name: impl Into<String>,
) -> HFRepositorySync<T> {
HFRepositorySync::new(self.clone(), repo_type, owner, name)
}
pub fn model(&self, owner: impl Into<String>, name: impl Into<String>) -> HFRepositorySync<RepoTypeModel> {
self.repository(RepoTypeModel, owner, name)
}
pub fn dataset(&self, owner: impl Into<String>, name: impl Into<String>) -> HFRepositorySync<RepoTypeDataset> {
self.repository(RepoTypeDataset, owner, name)
}
pub fn space(&self, owner: impl Into<String>, name: impl Into<String>) -> HFRepositorySync<RepoTypeSpace> {
self.repository(RepoTypeSpace, owner, name)
}
pub fn kernel(&self, owner: impl Into<String>, name: impl Into<String>) -> HFRepositorySync<RepoTypeKernel> {
self.repository(RepoTypeKernel, owner, name)
}
pub fn bucket(&self, owner: impl Into<String>, name: impl Into<String>) -> HFBucketSync {
HFBucketSync::new(self.clone(), owner, name)
}
}
impl<T: RepoType> HFRepositorySync<T> {
pub fn new(client: HFClientSync, repo_type: T, owner: impl Into<String>, name: impl Into<String>) -> Self {
Self {
inner: Arc::new(HFRepository::new(client.inner.clone(), repo_type, owner, name)),
runtime: client.runtime.clone(),
}
}
pub fn owner(&self) -> &str {
self.inner.owner()
}
pub fn name(&self) -> &str {
self.inner.name()
}
pub fn repo_path(&self) -> String {
self.inner.repo_path()
}
pub fn repo_type(&self) -> &T {
self.inner.repo_type()
}
}
impl HFBucketSync {
pub fn new(client: HFClientSync, owner: impl Into<String>, name: impl Into<String>) -> Self {
Self {
inner: Arc::new(crate::buckets::HFBucket::new(client.inner.clone(), owner, name)),
runtime: client.runtime.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hfapisync_creation() {
let client = HFClientSync::new();
assert!(client.is_ok());
}
#[test]
fn test_hfapisync_from_api() {
let async_client = HFClient::builder().build().unwrap();
let client = HFClientSync::from_inner(async_client);
assert!(client.is_ok());
}
#[test]
fn test_sync_repo_constructors() {
let client = HFClientSync::from_inner(HFClient::builder().build().unwrap()).unwrap();
let repo = client.model("openai-community", "gpt2");
let space = client.space("huggingface", "transformers-benchmarks");
assert_eq!(repo.owner(), "openai-community");
assert_eq!(repo.name(), "gpt2");
assert_eq!(repo.repo_type().singular(), "model");
assert_eq!(space.repo_type().singular(), "space");
}
fn unreachable_endpoint() -> String {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let port = listener.local_addr().unwrap().port();
drop(listener);
format!("http://127.0.0.1:{port}")
}
fn client_with_unreachable_endpoint() -> HFClientSync {
let inner = HFClient::builder()
.endpoint(unreachable_endpoint())
.retry_max_attempts(0)
.build()
.unwrap();
HFClientSync::from_inner(inner).unwrap()
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sync_call_inside_multi_thread_runtime_does_not_panic() {
let client = client_with_unreachable_endpoint();
let result = client.model("openai-community", "gpt2").info().send();
assert!(result.is_err());
}
#[tokio::test]
async fn sync_call_inside_current_thread_runtime_does_not_panic() {
let client = client_with_unreachable_endpoint();
let result = client.model("openai-community", "gpt2").info().send();
assert!(result.is_err());
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn construct_and_drop_inside_async_context_does_not_panic() {
let client = HFClientSync::from_inner(HFClient::builder().build().unwrap()).unwrap();
drop(client);
}
#[test]
fn clones_share_runtime_thread_and_survive_staggered_drops() {
let client = client_with_unreachable_endpoint();
let clone = client.clone();
let repo = client.model("openai-community", "gpt2");
assert!(Arc::ptr_eq(&client.runtime, &clone.runtime));
assert!(Arc::ptr_eq(&client.runtime, &repo.runtime));
drop(client);
drop(clone);
let result = repo.exists().send();
assert!(result.is_err());
}
#[test]
fn block_on_resumes_task_panics_on_the_caller() {
let runtime = RuntimeThread::spawn().unwrap();
let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
runtime.block_on(async {
panic!("boom");
})
}));
let panic = caught.unwrap_err();
assert_eq!(panic.downcast_ref::<&str>(), Some(&"boom"));
}
#[test]
fn block_on_accepts_borrowing_futures() {
let runtime = RuntimeThread::spawn().unwrap();
let data = String::from("borrowed");
let borrowed = &data;
let len = runtime.block_on(async move { borrowed.len() });
assert_eq!(len, data.len());
}
}