Update common util (#3868)

Co-authored-by: Rayan Dasoriya <dasoriya@google.com>
This commit is contained in:
Rayan Dasoriya
2025-03-11 12:33:53 +00:00
committed by GitHub
co-authored by Rayan Dasoriya
parent 69fab881a9
commit 040fc424fe
@@ -533,6 +533,7 @@ def get_resource_id(
"NVIDIA_TESLA_A100": "nvidia_a100_gpus",
"NVIDIA_A100_80GB": "nvidia_a100_80gb_gpus",
"NVIDIA_H100_80GB": "nvidia_h100_gpus",
"NVIDIA_H100_MEGA_80GB": "nvidia_h100_mega_gpus",
"NVIDIA_TESLA_T4": "nvidia_t4_gpus",
"TPU_V5e": "tpu_v5e",
"TPU_V3": "tpu_v3",
@@ -620,3 +621,16 @@ def check_quota(
f" {accelerator_count}. {quota_request_instruction}"
)
def get_deploy_source() -> str:
"""Gets deploy_source string based on running environment."""
vertex_product = os.environ.get("VERTEX_PRODUCT", "")
match vertex_product:
case "COLAB_ENTERPRISE":
return "notebook_colab_enterprise"
case "WORKBENCH_INSTANCE":
return "notebook_workbench"
case _:
# Legacy workbench, legacy colab, or other custom environments.
return "notebook_environment_unspecified"