Skip to content

Commit

Permalink
Update main.py
Browse files Browse the repository at this point in the history 
adjusted functions to either use v1 or v2 apis
  • Loading branch information
@jbrache
jbrache committed Mar 1, 2024
1 parent 9f2a5b3 commit 965b171
Showing 1 changed file with 22 additions and 102 deletions.
124 changes: 22 additions & 102 deletions demo/terraform/function/notebooks_auto_shutdown/main.py
View file
Expand Up @@ -23,6 +23,7 @@
PROJECT_ID = os.environ.get("PROJECT_ID", "the-foo-bar") # Update placeholder value when this environment variable is not set. 'Specified environment variable is not set.'
SHUTDOWN_SECONDS_METADATA_KEY = os.environ.get("SHUTDOWN_SECONDS_METADATA_KEY", "auto-shutdown-seconds") # Look for this key in notebooks metadata value, defaults to "auto-shutdown-seconds"
REGION_LIST = ["us-central1"] # will only stop instances in this list of regions
NOTEBOOKS_API_VERSION = os.environ.get("NOTEBOOKS_API_VERSION", "v1") # defaults to "v1" notebooks API which is User-Managed Notebooks (UmN). Workbench Instances (WbI) uses v2.

# https://cloud.google.com/python/docs/reference/cloudresourcemanager/latest/google.cloud.resourcemanager_v3.services.projects.ProjectsClient#google_cloud_resourcemanager_v3_services_projects_ProjectsClient_search_projects
def get_project_ids():
Expand All @@ -42,33 +43,15 @@ def get_project_ids():
return project_ids

# Gets supported locations for the notebooks service given a Project ID - V1 API
def get_umn_location_ids(project_id):
credentials, project = auth.default(scopes = ['https://www.googleapis.com/auth/cloud-platform'])
authed_session = AuthorizedSession(credentials)

# Get zone resources available to the project
# https://cloud.google.com/ai-platform/notebooks/docs/reference/rest/v1/projects.locations/list
response = authed_session.get(f"https://notebooks.googleapis.com/v1/projects/{project_id}/locations")

locations_json = response.json()
location_ids = []
try:
locations = locations_json["locations"]
for location in locations:
location_ids.append(location["locationId"])
except:
location_ids = []
return location_ids

# https://cloud.google.com/ai-platform/notebooks/docs/reference/rest/v1/projects.locations/list
# Gets supported locations for the notebooks service given a Project ID - V2 API
def get_wbi_location_ids(project_id):
# https://cloud.google.com/ai-platform/notebooks/docs/reference/rest/v2/projects.locations/list
def get_instance_location_ids(project_id, version = 'v1'):
credentials, project = auth.default(scopes = ['https://www.googleapis.com/auth/cloud-platform'])
authed_session = AuthorizedSession(credentials)

# Get zone resources available to the project
# https://cloud.google.com/ai-platform/notebooks/docs/reference/rest/v2/projects.locations/list
response = authed_session.get(f"https://notebooks.googleapis.com/v2/projects/{project_id}/locations")

response = authed_session.get(f"https://notebooks.googleapis.com/{version}/projects/{project_id}/locations")
locations_json = response.json()
location_ids = []
try:
Expand All @@ -79,95 +62,28 @@ def get_wbi_location_ids(project_id):
location_ids = []
return location_ids

# Stop User-Managed Notebooks (UmN) using REST API Commands - V1 API
# https://cloud.google.com/vertex-ai/docs/workbench/reference/rest
def stop_umn_rest(request):
credentials, project = auth.default(scopes = ['https://www.googleapis.com/auth/cloud-platform'])
authed_session = AuthorizedSession(credentials)
project_ids = get_project_ids()
return_response = {}

for project_id in project_ids:
now_utc = datetime.now(timezone.utc)
location_ids = get_umn_location_ids(project_id)

for location_id in location_ids:
# this example only inspects instances in a region, exclude this to run for all zones
skip = False
for region in REGION_LIST:
if (region not in location_id) or (region == location_id):
skip = True
if skip:
continue

# Get AI Platform notebook instances
response = authed_session.get(f"https://notebooks.googleapis.com/v1/projects/{project_id}/locations/{location_id}/instances")
instances_json = response.json()
# If the response is empty, continue
if not instances_json:
continue
instances = instances_json["instances"]

# Documentation
# https://cloud.google.com/ai-platform/notebooks/docs/reference/rest/v1/projects.locations.instances/stop
for instance in instances:
instance_name = instance["name"]
print("-------", "Instance name:", instance_name, "-------")
print(instance["metadata"])
auto_shutdown_seconds = 0

try:
auto_shutdown_seconds = int(instance["metadata"][SHUTDOWN_SECONDS_METADATA_KEY])
print(SHUTDOWN_SECONDS_METADATA_KEY, "metadata set to:", auto_shutdown_seconds)
except:
print(SHUTDOWN_SECONDS_METADATA_KEY, "metadata not set, skip...")
continue

if instance["state"] != "ACTIVE":
print("Instance isn't ACTIVE, skip...")
continue

update_time = instance["updateTime"]
update_time_iso = update_time.replace("Z", "+00:00")
print("timeNow:", now_utc.isoformat())
print("updatedTime:", update_time_iso)

t1 = isoparse(update_time_iso)
t2 = isoparse(now_utc.isoformat())
# get difference
delta = t2 - t1
delta_seconds = delta.total_seconds()
print("Difference in seconds between updatedTime and timeNow:", delta_seconds)

if (delta_seconds >= auto_shutdown_seconds) and delta_seconds >= 0:
print("Stopping server...")
response = authed_session.post(f"https://notebooks.googleapis.com/v1/{instance_name}:stop")
print(response.json())
return_response[instance_name] = response.json()
else:
print("Shutdown threshold not hit, skip...")

return return_response

# Delete User-Managed Notebooks (UmN) using REST API Commands - V1 API
# Delete Workbench Instances (WbI) using REST API Commands - V2 API
# https://cloud.google.com/vertex-ai/docs/workbench/reference/rest
def delete_wbi_rest(instance_name):
def delete_instance_rest(instance_name, version = 'v1'):
credentials, project = auth.default(scopes = ['https://www.googleapis.com/auth/cloud-platform'])
authed_session = AuthorizedSession(credentials)
response = authed_session.delete(f"https://notebooks.googleapis.com/v2/{instance_name}")
response = authed_session.delete(f"https://notebooks.googleapis.com/{version}/{instance_name}")
return response

# Stop User-Managed Notebooks (UmN) using REST API Commands - V1 API
# Stop Workbench Instances (WbI) using REST API Commands - V2 API
# https://cloud.google.com/vertex-ai/docs/workbench/reference/rest
def stop_wbi_rest(request):
def stop_instance_rest(request):
version = NOTEBOOKS_API_VERSION
credentials, project = auth.default(scopes = ['https://www.googleapis.com/auth/cloud-platform'])
authed_session = AuthorizedSession(credentials)
project_ids = get_project_ids()
return_response = {}

for project_id in project_ids:
now_utc = datetime.now(timezone.utc)
location_ids = get_wbi_location_ids(project_id)
location_ids = get_instance_location_ids(project_id, version)

for location_id in location_ids:
# this example only inspects instances in a region, exclude this to run for all zones
Expand All @@ -179,24 +95,29 @@ def stop_wbi_rest(request):
continue

# Get AI Platform notebook instances
response = authed_session.get(f"https://notebooks.googleapis.com/v2/projects/{project_id}/locations/{location_id}/instances")
response = authed_session.get(f"https://notebooks.googleapis.com/{version}/projects/{project_id}/locations/{location_id}/instances")
instances_json = response.json()
# If the response is empty, continue
if not instances_json:
continue
instances = instances_json["instances"]

# Documentation
# https://cloud.google.com/ai-platform/notebooks/docs/reference/rest/v1/projects.locations.instances/stop
# https://cloud.google.com/ai-platform/notebooks/docs/reference/rest/v2/projects.locations.instances/stop
for instance in instances:
instance_name = instance["name"]
print("-------", "Instance name:", instance_name, "-------")
print(json.dumps(instance, indent=4))
print(instance["gceSetup"]["metadata"])
# print(json.dumps(instance, indent=4))
auto_shutdown_seconds = 0

try:
auto_shutdown_seconds = int(instance["gceSetup"]["metadata"][SHUTDOWN_SECONDS_METADATA_KEY])
if version == 'v2':
print(instance["gceSetup"]["metadata"])
auto_shutdown_seconds = int(instance["gceSetup"]["metadata"][SHUTDOWN_SECONDS_METADATA_KEY])
else:
print(instance["metadata"])
auto_shutdown_seconds = int(instance["metadata"][SHUTDOWN_SECONDS_METADATA_KEY])
print(SHUTDOWN_SECONDS_METADATA_KEY, "metadata set to:", auto_shutdown_seconds)
except:
print(SHUTDOWN_SECONDS_METADATA_KEY, "metadata not set, skip...")
Expand All @@ -218,10 +139,9 @@ def stop_wbi_rest(request):
delta_seconds = delta.total_seconds()
print("Difference in seconds between updatedTime and timeNow:", delta_seconds)

delta_seconds = 1000000
if (delta_seconds >= auto_shutdown_seconds) and delta_seconds >= 0:
print("Stopping server...")
response = authed_session.post(f"https://notebooks.googleapis.com/v2/{instance_name}:stop")
response = authed_session.post(f"https://notebooks.googleapis.com/{version}/{instance_name}:stop")
print(response.json())
return_response[instance_name] = response.json()
else:
Expand Down

0 comments on commit 965b171

Please sign in to comment.