Add bastion_tier label in bastion code

[eshen1991]

Overview

When using tpu.googleapis.com/accelerator/tensorcore_utilization metric, there is no way to break down by VM tiers(i.e. reserved vs spot) as the information is not captured in the existing filtering labels. However the tier information can be injected with VM custom labels, which show up as user labels in the filtering and can be used for grouping by the tiers. The code changes in this PR injects a custom label, vmtier, for TPU VM provisioned via QRM. The custom label supports two values: reserved and spot.

Testing

Unit Testing

(axlearn) ➜  axlearn git:(vm-tier-label-2) pytest axlearn/cloud/gcp/jobs/tpu_runner_test.py -v | grep test_start
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start0 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start1 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start2 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start3 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start4 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start5 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start6 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start7 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerJobTest::test_start8 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerMainTest::test_start0 PASSED
axlearn/cloud/gcp/jobs/tpu_runner_test.py::TPURunnerMainTest::test_start1 PASSED

e2e Testing (manual)

1. run BASTION_TIER=1 axlearn gcp tpu start --name=$USER-test --tpu_type=v4-8 -- python3 -c "'import jax; print(jax.devices())'" to start axlearn job on TPU that is provisioned through queue resource.
2. Once the queued resource state is active and run gcloud alpha compute tpus tpu-vm describe $USER-test --zone us-central2-b to check the TPU VM info which should have lines as below:

labels:
  bastion_tier: spot 
schedulingConfig:
  spot: true

if BASTION_TIER=0, bastion_tier: reserved.

[markblee]

Suggested change

	labels = {}
	if reserved is None:
	reserved_tier = gcp_settings("reserved_tpu", default=False, required=False)
	else:
	reserved_tier = reserved
	if reserved_tier:
	labels["vmtier"] = "reserved"
	else:
	labels["vmtier"] = "spot"
	if reserved is None:
	reserved = gcp_settings("reserved_tpu", default=False, required=False)
	labels = {"bastion_tier": "reserved" if reserved else "spot"}

Not sure if vmtier label is arbitrary?

[eshen1991]

@markblee made the code changes as suggested. Custom label names are user defined. "bastion_tier" makes more sense as a label name in the context.

[ruomingp]

Will defer to @markblee for review.

[markblee]

Thanks -- for future PRs, feel free to "re-request" review when comments are addressed.

[Ethanlm]

LGTM

Do you have an example metric graph handy to show that this label information is used to distinguish between reserved vs spot usages?

[eshen1991]

LGTM

Do you have an example metric graph handy to show that this label information is used to distinguish between reserved vs spot usage

In the Metric Explore, you'd see the labels under "User metadata labels", which can be used in filter or groupBy for the query to create charts. Note that image was captured before I changed the label name from vmtier to bastion_tier

[eshen1991]

@markblee I somehow had the pre-commit turned off. I ran pre-commit and fixed the lint issues that broke the circleci pre-commit phase. Please review the latest commit, 1911f1b.