Fix generate_metrics_and_upload_to_big_query

[ycchenzheng]

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_enable to True
• Set bq_db_project and bq_db_dataset to the user's target project and the dataset in the target project
  Run the command python3 -m benchmarks.recipes.pw_mcjax_benchmark_recipe and after 20 steps, one big query record will be written to the run_summary table of the target dataset

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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.

[ycchenzheng]

This PR is another version of #2405 to avoid pinning flax.

[SujeethJinesh]

Thanks Zheng! I think if we do end up to going this route, it would involve a lot of additional maintenance burden on our part because we're going to own this.

I'm thinking about potentially a few things we can do:

1. (ideal) Extricate the things we need from the benchmark-db-writer and use the minimum dependencies for what we need
2. (will require approvals) Check in the exact difference of changes we need into aotc
3. (might not be agreeable to others) Make a branch off the main branch in the aotc repo and use that in the requirements.txt

I think if we can minimize what we need by going with option 1, that'd be the best. I'd like to see what we should be using to upload to bigquery and maybe we can look at what the other folks have used. We can chat about this today.