Monitor cumulative goodput if enabled.
(config)
| 36 | |
| 37 | @contextlib.contextmanager |
| 38 | def maybe_monitor_goodput(config): |
| 39 | """Monitor cumulative goodput if enabled.""" |
| 40 | if not config.monitor_goodput or jax.process_index() != 0: |
| 41 | yield |
| 42 | return |
| 43 | goodput_monitor = None |
| 44 | try: |
| 45 | if config.report_performance_metric_for_gcp_monitoring: |
| 46 | config.enable_gcp_step_deviation_metrics = False |
| 47 | |
| 48 | gcp_options = monitoring.GCPOptions( |
| 49 | enable_gcp_goodput_metrics=config.enable_gcp_goodput_metrics, |
| 50 | enable_gcp_step_deviation_metrics=config.enable_gcp_step_deviation_metrics, |
| 51 | ) |
| 52 | goodput_monitor = monitoring.GoodputMonitor( |
| 53 | job_name=config.run_name, |
| 54 | logger_name=f"goodput_{config.run_name}", |
| 55 | tensorboard_dir=config.tensorboard_dir, |
| 56 | upload_interval=config.goodput_upload_interval_seconds, |
| 57 | monitoring_enabled=True, |
| 58 | pathway_enabled=config.enable_pathways_goodput, |
| 59 | include_badput_breakdown=True, |
| 60 | include_step_deviation=config.monitor_step_time_deviation, |
| 61 | step_deviation_interval_seconds=config.step_deviation_interval_seconds, |
| 62 | gcp_options=gcp_options, |
| 63 | ) |
| 64 | goodput_monitor.start_goodput_uploader() |
| 65 | max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!") |
| 66 | yield |
| 67 | finally: |
| 68 | if goodput_monitor: |
| 69 | goodput_monitor.stop_goodput_uploader() |
| 70 | max_logging.log("Flushed final metrics and safe exited from Goodput monitoring.") |
| 71 | |
| 72 | |
| 73 | @contextlib.contextmanager |