Skip to content

Commit

Permalink
meta_proxy_generator update
Browse files Browse the repository at this point in the history
  • Loading branch information
jstzwj committed Aug 6, 2024
1 parent d451f2e commit ea97ba0
Show file tree
Hide file tree
Showing 4 changed files with 44 additions and 33 deletions.
49 changes: 20 additions & 29 deletions olah/proxy/meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,37 +75,28 @@ async def meta_proxy_generator(
allow_cache: bool,
save_path: str,
):
try:
temp_file_path = None
async with httpx.AsyncClient(follow_redirects=True) as client:
with tempfile.NamedTemporaryFile(mode="wb", delete=True) as temp_file:
temp_file_path = temp_file.name
if not allow_cache:
write_temp_file = False
else:
write_temp_file = True
async with client.stream(
method="GET",
url=meta_url,
headers=headers,
timeout=WORKER_API_TIMEOUT,
) as response:
response_headers = response.headers
yield response_headers
async with httpx.AsyncClient(follow_redirects=True) as client:
content_chunks = []
async with client.stream(
method="GET",
url=meta_url,
headers=headers,
timeout=WORKER_API_TIMEOUT,
) as response:
response_headers = response.headers
yield response_headers

async for raw_chunk in response.aiter_raw():
if not raw_chunk:
continue
if write_temp_file:
temp_file.write(raw_chunk)
yield raw_chunk
if temp_file_path is not None:
temp_file.flush()
shutil.copyfile(temp_file_path, save_path)
finally:
if temp_file_path is not None and os.path.exists(temp_file_path):
os.remove(temp_file_path)
async for raw_chunk in response.aiter_raw():
if not raw_chunk:
continue
content_chunks.append(raw_chunk)
yield raw_chunk

content = bytearray()
for chunk in content_chunks:
content += chunk
with open(save_path, "wb") as f:
f.write(bytes(content))

async def meta_generator(
app: FastAPI,
Expand Down
22 changes: 20 additions & 2 deletions olah/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,25 @@

import git
import httpx
from pydantic import BaseSettings

BASE_SETTINGS = False
if not BASE_SETTINGS:
try:
from pydantic import BaseSettings
BASE_SETTINGS = True
except ImportError:
BASE_SETTINGS = False

if not BASE_SETTINGS:
try:
from pydantic_settings import BaseSettings
BASE_SETTINGS = True
except ImportError:
BASE_SETTINGS = False

if not BASE_SETTINGS:
raise Exception("Cannot import BaseSettings from pydantic or pydantic-settings")

from olah.configs import OlahConfig
from olah.errors import error_repo_not_found, error_page_not_found
from olah.mirror.repos import LocalMirrorRepo
Expand Down Expand Up @@ -61,7 +79,7 @@ async def check_connection(url: str) -> bool:
return False


@repeat_every(seconds=60)
@repeat_every(seconds=60*5)
async def check_hf_connection() -> None:
if app.app_settings.config.offline:
return
Expand Down
5 changes: 3 additions & 2 deletions olah/utils/repo_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -311,5 +311,6 @@ async def check_commit_hf(
if authorization is not None:
headers["authorization"] = authorization
async with httpx.AsyncClient() as client:
response = await client.get(url, headers=headers, timeout=WORKER_API_TIMEOUT)
return response.status_code in [200, 307]
response = await client.request(method="HEAD", url=url, headers=headers, timeout=WORKER_API_TIMEOUT)
status_code = response.status_code
return status_code in [200, 307]
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ fastapi-utils==0.7.0
GitPython==3.1.43
httpx==0.27.0
pydantic==2.8.2
pydantic-setting==2.2.1
toml==0.10.2
huggingface_hub==0.23.4
pytest==8.2.2
Expand Down

0 comments on commit ea97ba0

Please sign in to comment.