Skip to content

Commit

Permalink
Handle case when Github releases rate limit is reached. (#204)
Browse files Browse the repository at this point in the history
* Handle case when Github releases rate limit is reached.

* Attach GITHUB_TOKEN to releases request if available.

* Fix test.

* Fix test.

* Ruff fix.
  • Loading branch information
robinjhuang authored Nov 7, 2024
1 parent 93ab01b commit 6525338
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 4 deletions.
37 changes: 34 additions & 3 deletions comfy_cli/command/install.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,13 +188,17 @@ def execute(
clone_comfyui(url=url, repo_dir=repo_dir)

if version != "nightly":
checkout_stable_comfyui(version=version, repo_dir=repo_dir)
try:
checkout_stable_comfyui(version=version, repo_dir=repo_dir)
except GitHubRateLimitError as e:
rprint(f"[bold red]Error checking out ComfyUI version: {e}[/bold red]")
sys.exit(1)

elif not check_comfy_repo(repo_dir)[0]:
rprint(
f"[bold red]'{repo_dir}' already exists. But it is an invalid ComfyUI repository. Remove it and retry.[/bold red]"
)
exit(-1)
sys.exit(-1)

# checkout specified commit
if commit is not None:
Expand Down Expand Up @@ -281,12 +285,39 @@ def validate_version(version: str) -> Optional[str]:
) from exc


class GitHubRateLimitError(Exception):
"""Raised when GitHub API rate limit is exceeded"""


def fetch_github_releases(repo_owner: str, repo_name: str) -> List[Dict[str, str]]:
"""
Fetch the list of releases from the GitHub API.
Handles rate limiting by logging the wait time.
"""
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/releases"
response = requests.get(url)

headers = {}
if github_token := os.getenv("GITHUB_TOKEN"):
headers["Authorization"] = f"Bearer {github_token}"

response = requests.get(url, headers=headers, timeout=5)

# Handle rate limiting
if response.status_code in (403, 429):
# Check rate limit headers
remaining = int(response.headers.get("x-ratelimit-remaining", 0))
if remaining == 0:
reset_time = int(response.headers.get("x-ratelimit-reset", 0))
message = f"Primary rate limit from Github exceeded! Please retry after: {reset_time})"
raise GitHubRateLimitError(message)

if "retry-after" in response.headers:
wait_seconds = int(response.headers["retry-after"])
message = f"Rate limit from Github exceeded! Please wait {wait_seconds} seconds before retrying."
rprint(f"[yellow]{message}[/yellow]")
raise GitHubRateLimitError(message)

response.raise_for_status()
return response.json()


Expand Down
2 changes: 1 addition & 1 deletion tests/comfy_cli/test_install.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ def test_fetch_releases_success(mock_get):
assert len(releases) == 2
assert releases[0]["tag_name"] == "v1.0.0"
assert releases[1]["tag_name"] == "v1.1.0"
mock_get.assert_called_once_with("https://api.github.com/repos/owner/repo/releases")
mock_get.assert_called_once_with("https://api.github.com/repos/owner/repo/releases", headers={}, timeout=5)


@patch("requests.get")
Expand Down

0 comments on commit 6525338

Please sign in to comment.