maxtext
maxtext copied to clipboard
Fix generate_metrics_and_upload_to_big_query
Description
Fix generate_metrics_and_upload_to_big_query flag and related dependencies.
FIXES: b/441984274, b/446097400
Tests
Project preparation:
- Create a dataset in BigQuery or use an existed one in the target project
benchmarks/recipes/user_configs.py change:
Adding the following flags to
# Define the required configuration here
USER_CONFIG = UserConfig(
- Set
bq_enabletoTrue - Set
bq_db_projectandbq_db_datasetto the user's target project and the dataset in the target project
Run the commandpython3 -m benchmarks.recipes.pw_mcjax_benchmark_recipeand after 20 steps, one big query record will be written to therun_summarytable of the target dataset
Checklist
Before submitting this PR, please make sure (put X in square brackets):
- [x] I have performed a self-review of my code. For an optional AI review, add the
gemini-reviewlabel. - [x] I have necessary comments in my code, particularly in hard-to-understand areas.
- [x] I have run end-to-end tests tests and provided workload links above if applicable.
- [x] I have made or will make corresponding changes to the doc if needed.