def get_quota()

in community-content/vertex_model_garden/model_oss/notebook_util/common_util.py [0:0]

• 32 lines of code
• 10 McCabe index (conditional complexity)

def get_quota(project_id: str, region: str, resource_id: str) -> int:
  """Returns the quota for a resource in a region.

  Args:
    project_id: The project id.
    region: The region.
    resource_id: The resource id.

  Returns:
    The quota for the resource in the region. Returns -1 if can not figure out
    the quota.

  Raises:
    RuntimeError: If the command to get quota fails.
  """
  service_endpoint = "aiplatform.googleapis.com"

  command = (
      "gcloud alpha services quota list"
      f" --service={service_endpoint} --consumer=projects/{project_id}"
      f" --filter='{service_endpoint}/{resource_id}' --format=json"
  )
  process = subprocess.run(
      command, shell=True, capture_output=True, text=True, check=True
  )
  if process.returncode == 0:
    quota_data = json.loads(process.stdout)
  else:
    raise RuntimeError(f"Error fetching quota data: {process.stderr}")

  if not quota_data or "consumerQuotaLimits" not in quota_data[0]:
    return -1
  if (
      not quota_data[0]["consumerQuotaLimits"]
      or "quotaBuckets" not in quota_data[0]["consumerQuotaLimits"][0]
  ):
    return -1
  all_regions_data = quota_data[0]["consumerQuotaLimits"][0]["quotaBuckets"]
  for region_data in all_regions_data:
    if (
        region_data.get("dimensions")
        and region_data["dimensions"]["region"] == region
    ):
      if "effectiveLimit" in region_data:
        return int(region_data["effectiveLimit"])
      else:
        return 0
  return -1