File size: 5,519 Bytes
0c221a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
import requests
from environs import Env
# Initialize environment
env = Env()
env.read_env() # Looks for .env files in current directory
# Set your token (in a real environment, get this from env vars)
HF_TOKEN = env("HF_TOKEN")
ORG_NAME = env("ORG_NAME")
BASE_URL = env("BASE_URL")
# HTTP Headers for API requests
headers = {
"Authorization": f"Bearer {HF_TOKEN}",
"Content-Type": "application/json"
}
# First thing's first: create the repo (or use existing one)
def create_repository(org_name, repo_name):
"""Create repository or use existing one"""
# First check if it already exists
repo_url = f"{BASE_URL}/models/{org_name}/{repo_name}"
check_response = requests.get(repo_url, headers=headers)
if check_response.status_code == 200:
print(f"Good news! Repository {org_name}/{repo_name} already exists!")
return True
# Try to create it if it doesn't exist
create_url = f"{BASE_URL}/repos/create"
print(f"Houston, stand by: Creating repo with URL: {create_url}")
data = {
"type": "model",
"name": repo_name,
"organization": org_name,
"private": True
}
response = requests.post(create_url, headers=headers, json=data)
if response.status_code == 200:
print(f"Woohoo!!! Created repository: {org_name}/{repo_name}")
return True
else:
# Check if the error is because the repo already exists
error_data = response.json()
if "error" in error_data and "already created" in error_data["error"]:
print(f"Repository {org_name}/{repo_name} already exists! Moving on.")
return True
else:
print(f"Son of a Lambda Cold Start! Failed to create repository: {response.text}")
return False
def create_resource_group_for_repo(repo_name, creator_username):
"""
The magical function that creates a resource group for a new repository.
"""
# Step 0: Make the repo first or we're going nowhere fast
if not create_repository(ORG_NAME, repo_name):
print("Hey, yo! Slow your roll! Can't make a resource group for a repo that doesn't exist!")
return False
# Step 1: Create our glorious resource group
rg_name = f"rg-{repo_name}" # Simpler name, less chance for API to throw a fit
rg_url = f"{BASE_URL}/organizations/{ORG_NAME}/resource-groups"
print(f"Conjuring a resource group with URL: {rg_url}")
# Let's be direct - just create it with minimal data
data = {
"name": rg_name,
"description": f"Automagically created for {repo_name} by the Wizard of Automation"
}
# Throw the request into the void and pray for a 200 or 201
response = requests.post(rg_url, headers=headers, json=data)
# Print full response for debugging (you know it will happen)
print(f"Resource group creation response: {response.status_code} - {response.text}")
success = False
if response.status_code == 201 or response.status_code == 200:
try:
response_data = response.json()
if "id" in response_data:
success = True
rg_id = response_data.get("id")
print(f"Resource group successfully summoned! ID: {rg_id}")
else:
print("Response has unexpected format - missing 'id' field")
except Exception as e:
print(f"Error parsing response: {e}")
if not success:
print(f"Son of a Nutcracker! Resource group creation failed: {response.text}")
return False
# We're in business! Grab that ID
rg_id = response.json().get("id")
print(f"I stole, I mean grabbed that Resource Group ID because I need it: {rg_id}")
# Step 2: Connect the repo to our shiny new resource group
repo_url = f"{BASE_URL}/models/{ORG_NAME}/{repo_name}/resource-group"
repo_data = {
"resourceGroupId": rg_id
}
repo_response = requests.post(repo_url, headers=headers, json=repo_data)
if repo_response.status_code != 200:
print(f"Son of an API! Failed to connect repo: {repo_response.text}")
return False
print(f"Repository {repo_name} has joined the resource group party!")
# Step 3: Make the creator an admin - they brought this repo into the world after all
member_url = f"{BASE_URL}/organizations/{ORG_NAME}/resource-groups/{rg_id}/users"
member_data = {
"users": [
{
"user": creator_username,
"role": "admin"
}
]
}
member_response = requests.post(member_url, headers=headers, json=member_data)
if member_response.status_code != 200:
print(f"Son of a broken websocket! Creator couldn't be added: {member_response.text}")
# We'll continue anyway - the group and repo connection are working
else:
print(f"Creator {creator_username} is now the admin overlord of this resource group!")
# Victory dance
return True
# Demo time: This simulates what would happen when a new repo is created
def main():
repo_name = "testing-magic"
creator = "joshhayles"
print(f"A WIZARD created a new Repo: {ORG_NAME}/{repo_name}")
result = create_resource_group_for_repo(repo_name, creator)
if result:
print("Ohh, snap! Successfully set up resource group for the new repository!")
else:
print("Damn. Failed to set up resource group")
if __name__ == "__main__":
main() |