add GCP workload observability feature

[jcyang43]

Description

Add option to enable GCP workload monitoring for MaxText workloads.

• GCP workload monitoring sends performance metrics (heartbeat & training step times) to cloud monarch for monitoring such that if a metric hits its pre-defined threshold, oncalls will be notified to see if any actions are needed. This is ideal for critical workloads sensitive to infrastructure changes.
• Each metric can be configured to be on or off based on configs. Examples are included in MaxText/configs/base.yml
• Documentation can be found at getting_started/GCP_Workload_Monitoring.md

Tests

Tested on trillium TPU and confirmed metrics sent to cloud monarch successfully if configs are enabled. No metrics will be sent to cloud monarch if configs are set to False.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gobbleturk]

Changes like this need to be handled with care - jax is by default lazy e.g. functions return instantly, they do not block on finishing computation. In this case the results before and after this PR are roughly the same as confirmed by this diff https://diff.googleplex.com/#key=TLvLsJisjPpi (LHS main, RHS this PR)

The real blocking function is when an array is either checkpointed or printed - in this case this is done by write_metrics which is below (not record_scalar_metrics or p_train_step). This is fine since write_metrics is done after last_step_completion, so the step_time_delta will have to wait a real train_step worth of time

[gobbleturk]

Awesome, thank you for the documentation!