Feature/gcp workload monitoring #1150
Conversation
jcyang43
requested review from
gobbleturk,
khatwanimohit,
bvandermoon,
vipannalla and
RissyRan
as code owners
jcyang43
force-pushed
the
feature/gcp-workload-monitoring
branch
from
87341b4 to
a251b3e
Compare
There was a problem hiding this comment.
This file looks independent of maxtext (other than the maxtext logging plubmed in). Is there any opportunity to move this file / APIs to one of the packages we are also depending on (google-cloud-monitoring, google-api-core
google-api-python-client)
There was a problem hiding this comment.
Hi Matthew, thanks for the suggestion. We're planning to make this improvement in a subsequent version where we'll build an API and python package for workload observability. We'll send a follow-up PR when the API is ready to use
| @@ -456,6 +456,10 @@ monitor_goodput: True | |||
| goodput_upload_interval_seconds: 60 | |||
| enable_pathways_goodput: False | |||
|
|
|||
| # GCP workload monitoring | |||
| report_heartbeat_metric_for_gcp_monitoring: True | |||
There was a problem hiding this comment.
Can we default these to false? Most runs of maxtext will not be prod runs that want heartbeating or alerting
There was a problem hiding this comment.
We would like to default to true as it would enable us to collect telemetry data. We don't necessarily need prod runs/alerting. The metrics can just be used to monitor performance of maxtext workloads.
There was a problem hiding this comment.
Keeping this on by default might lead to errors in internal testing that rely on google3 libraries and some (such as cloud logging and perhaps even monitoring_v3) are not available in google3 yet. MaxText does not support dependency injection in testing so GCP services such as monitoring and logging do not work as intended with tests yet.
I would suggest running some internal tests with the google3 repo as well to validate the above.
| This guide provides an overview on how to enable GCP workload monitoring for your MaxText workload. | ||
|
|
||
| ## Overview | ||
| To address Google Cloud's lack of visibility into user workload performance, users now have the option to enable GCP workload monitoring for all of their MaxText workloads. This workload performance monitoring feature is ideal for critical workloads sensitive to infrastructure changes. |
There was a problem hiding this comment.
Can this be worded more positively =)
We offer a monitoring and alerting feature that is well suited for critical workloads sensitive to...
requirements.txt
Outdated
| @@ -10,6 +10,9 @@ datasets | |||
| gcsfs | |||
| google-cloud-aiplatform==1.61.0 | |||
| google-cloud-storage | |||
| google-cloud-monitoring==2.20.0 | |||
There was a problem hiding this comment.
Do these need to be pinned? Can we use latest?
jcyang43
force-pushed
the
feature/gcp-workload-monitoring
branch
2 times, most recently
from
5e4817b to
caef832
Compare
jcyang43
force-pushed
the
feature/gcp-workload-monitoring
branch
from
caef832 to
25e43fb
Compare
| @@ -852,6 +860,12 @@ def train_loop(config, state=None): | |||
| example_batch = None | |||
| last_step_completion = datetime.datetime.now() | |||
|
|
|||
| prof = profiler.Profiler(config) | |||
| if gcp_workload_monitor and heartbeat_reporting_stop_event: | |||
| max_logging.log("Starting background thread for reporting heartbeat") | |||
There was a problem hiding this comment.
Can we expose an API to start this thread based on configs?
| step_time_delta = datetime.datetime.now() - last_step_completion | ||
| record_scalar_metrics(metrics, step_time_delta, per_device_tflops, learning_rate_schedule(step), per_device_tokens) | ||
| if gcp_workload_monitor and config.report_performance_metric_for_gcp_monitoring: | ||
| gcp_workload_monitor.report_performance(step_time_delta.total_seconds()) |
There was a problem hiding this comment.
This should ideally be in its own API.
| global_rank = jax.process_index() | ||
| while not stop_event.is_set(): | ||
| self.report_heartbeat(local_rank, str(global_rank)) | ||
| time.sleep(5) |
There was a problem hiding this comment.
Can we make this configurable?
Description
Add option to enable GCP workload monitoring for MaxText workloads.
MaxText/configs/base.ymlgetting_started/GCP_Workload_Monitoring.mdTests
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):