fix: creating blobs only once for publish trace retries (#14845)

This commit is contained in:
Douglas Yang
2025-12-10 22:40:06 -08:00
committed by GitHub
parent 32829b1638
commit 1a96e66493
+63 -23
View File
@@ -58,29 +58,32 @@ def verify_token_permissions(repo_owner, repo_name, token):
"""Verify that the token has necessary permissions for the repository""" """Verify that the token has necessary permissions for the repository"""
print("Verifying token permissions...") print("Verifying token permissions...")
# Check if we can access the repository checks = [
try: (
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}" f"https://api.github.com/repos/{repo_owner}/{repo_name}", # Check if we can access the repository
response = make_github_request(url, token) "Repository access verified",
repo_data = json.loads(response) ),
print(f"Repository access verified: {repo_data['full_name']}") (
except Exception as e: f"https://api.github.com/repos/{repo_owner}/{repo_name}/contents", # Check if we can read the repository contents
if is_rate_limit_error(e): "Repository contents access verified",
warnings.warn("GitHub API rate limit exceeded during token verification.") ),
return "rate_limited" ]
print(f"Failed to access repository: {e}")
return False
# Check if we can read the repository contents for url, success_message in checks:
try: try:
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/contents"
response = make_github_request(url, token) response = make_github_request(url, token)
print("Repository contents access verified") if success_message == "Repository access verified":
repo_data = json.loads(response)
print(f"{success_message}: {repo_data['full_name']}")
else:
print(success_message)
except Exception as e: except Exception as e:
if is_rate_limit_error(e): if is_rate_limit_error(e):
warnings.warn("GitHub API rate limit exceeded during token verification.") warnings.warn(
"GitHub API rate limit exceeded during token verification."
)
return "rate_limited" return "rate_limited"
print(f"Failed to access repository contents: {e}") print(f"Failed to verify permissions for {url}: {e}")
return False return False
return True return True
@@ -118,6 +121,10 @@ def create_blob(repo_owner, repo_name, content, token, max_retries=3):
response = make_github_request(url, token, method="POST", data=data) response = make_github_request(url, token, method="POST", data=data)
return json.loads(response)["sha"] return json.loads(response)["sha"]
except Exception as e: except Exception as e:
# Don't retry on rate limit errors - fail fast
if is_rate_limit_error(e):
raise
if attempt < max_retries - 1: if attempt < max_retries - 1:
wait_time = 2**attempt # Exponential backoff: 1s, 2s, 4s wait_time = 2**attempt # Exponential backoff: 1s, 2s, 4s
print( print(
@@ -128,10 +135,8 @@ def create_blob(repo_owner, repo_name, content, token, max_retries=3):
raise raise
def create_tree(repo_owner, repo_name, base_tree_sha, files, token, max_retries=3): def create_blobs(repo_owner, repo_name, files, token):
"""Create a new tree with files""" """Create blobs for all files and return tree items with blob SHAs"""
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/trees"
tree_items = [] tree_items = []
for i, (file_path, content) in enumerate(files): for i, (file_path, content) in enumerate(files):
# Create blob first to get SHA # Create blob first to get SHA
@@ -147,6 +152,12 @@ def create_tree(repo_owner, repo_name, base_tree_sha, files, token, max_retries=
# Progress indicator for large uploads # Progress indicator for large uploads
if (i + 1) % 10 == 0 or (i + 1) == len(files): if (i + 1) % 10 == 0 or (i + 1) == len(files):
print(f"Created {i + 1}/{len(files)} blobs...") print(f"Created {i + 1}/{len(files)} blobs...")
return tree_items
def create_tree(repo_owner, repo_name, base_tree_sha, tree_items, token, max_retries=3):
"""Create a new tree from pre-created blob SHAs"""
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/trees"
data = {"base_tree": base_tree_sha, "tree": tree_items} data = {"base_tree": base_tree_sha, "tree": tree_items}
@@ -155,6 +166,10 @@ def create_tree(repo_owner, repo_name, base_tree_sha, files, token, max_retries=
response = make_github_request(url, token, method="POST", data=data) response = make_github_request(url, token, method="POST", data=data)
return json.loads(response)["sha"] return json.loads(response)["sha"]
except Exception as e: except Exception as e:
# Don't retry on rate limit errors - fail fast
if is_rate_limit_error(e):
raise
if attempt < max_retries - 1: if attempt < max_retries - 1:
wait_time = 2**attempt wait_time = 2**attempt
print( print(
@@ -189,6 +204,10 @@ def create_commit(
return commit_sha return commit_sha
except Exception as e: except Exception as e:
# Don't retry on rate limit errors - fail fast
if is_rate_limit_error(e):
raise
if attempt < max_retries - 1: if attempt < max_retries - 1:
wait_time = 2**attempt wait_time = 2**attempt
print( print(
@@ -212,6 +231,10 @@ def update_branch_ref(repo_owner, repo_name, branch, commit_sha, token, max_retr
make_github_request(url, token, method="PATCH", data=data) make_github_request(url, token, method="PATCH", data=data)
return return
except HTTPError as e: except HTTPError as e:
# Don't retry on rate limit errors - fail fast
if is_rate_limit_error(e):
raise
# Check if this is an "Object does not exist" error # Check if this is an "Object does not exist" error
is_object_not_exist = False is_object_not_exist = False
if hasattr(e, "error_body"): if hasattr(e, "error_body"):
@@ -232,6 +255,10 @@ def update_branch_ref(repo_owner, repo_name, branch, commit_sha, token, max_retr
else: else:
raise raise
except Exception as e: except Exception as e:
# Don't retry on rate limit errors - fail fast
if is_rate_limit_error(e):
raise
if attempt < max_retries - 1: if attempt < max_retries - 1:
wait_time = 2**attempt wait_time = 2**attempt
print( print(
@@ -308,6 +335,19 @@ def publish_traces(traces_dir, run_id, run_number):
max_retries = 5 max_retries = 5
retry_delay = 5 # seconds retry_delay = 5 # seconds
# Create blobs once before retry loop to avoid re-uploading on failures
try:
tree_items = create_blobs(repo_owner, repo_name, files_to_upload, token)
except Exception as e:
# Check for rate limit errors during blob creation
if is_rate_limit_error(e):
warnings.warn(
"GitHub API rate limit exceeded during blob creation. Skipping trace upload."
)
return
print(f"Failed to create blobs: {e}")
raise
for attempt in range(max_retries): for attempt in range(max_retries):
try: try:
# Get current branch head # Get current branch head
@@ -318,9 +358,9 @@ def publish_traces(traces_dir, run_id, run_number):
tree_sha = get_tree_sha(repo_owner, repo_name, branch_sha, token) tree_sha = get_tree_sha(repo_owner, repo_name, branch_sha, token)
print(f"Current tree SHA: {tree_sha}") print(f"Current tree SHA: {tree_sha}")
# Create new tree with all files # Create new tree with pre-created blobs
new_tree_sha = create_tree( new_tree_sha = create_tree(
repo_owner, repo_name, tree_sha, files_to_upload, token repo_owner, repo_name, tree_sha, tree_items, token
) )
print(f"Created new tree: {new_tree_sha}") print(f"Created new tree: {new_tree_sha}")