Recipe Modularization, DAG Integration, and BigQuery Integration

[DannyLiCom]

Description

When integrating Recipe Modularization, DAG Integration, and BigQuery, the DAG will pass in a command that includes multiple command-line flags.

Tests

Users can use commands similar to the following to run tests

python3 -m benchmarks.recipes.pw_mcjax_benchmark_recipe \
--user='user' \
--cluster_name='pw-scale-test-v5e-32' \
--project='cloud-tpu-multipod-dev' \
--zone='us-south1-a' \
--device_type='v5litepod-32' \
--benchmark_steps=20 \
--num_slices_list='2' \
--server_image='gcr.io/tpu-prod-env-one-vm/lidanny/unsanitized_server:latest' \
--proxy_image='gcr.io/tpu-prod-env-one-vm/lidanny/unsanitized_proxy_server:latest' \
--runner='gcr.io/tpu-prod-env-one-vm/lidanny_latest:latest' \
--selected_model_framework='pathways' \
--selected_model_names='llama3_1_8b_8192_v5e_256' \
--priority='medium' \
--max_restarts=0 \
--bq_enable=True \
--bq_db_project='cloud-tpu-multipod-dev' \
--bq_db_dataset="chzheng_test_100steps"

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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[DannyLiCom]

Since the original recipe could only accept the --delete flag, this section had to be modified to allow it to accept multiple flags.

[DannyLiCom]

This logic is implemented to determine whether multiple flags are being used.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[aireenmei]

Thanks for the work! I run the PR through AI and got these comments: https://paste.googleplex.com/6085130735190016. I looked at some points and they make sense to me, so I left inline comments. But there are more points, could you take a look?

[github-actions]

🤖 Hi @aireenmei, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This pull request introduces significant refactoring to the benchmarking recipes, modularizing argument parsing and improving configuration management. The changes make the recipes more flexible and easier to use from the command line, especially for DAG integration.

🔍 General Feedback

• The introduction of parser_utils.py is a good step towards centralizing command-line argument handling.
• The use of a UserConfig class helps in managing configurations cleanly.
• The addition of check_and_create_bucket is a useful utility, though its error handling could be improved in the calling code.

[github-actions]

🤖 I'm sorry @aireenmei, but I was unable to process your request. Please see the logs for more details.

[DannyLiCom]

Hi @aireenmei I noticed a lot of these errors in Pylint.
E0611: No name 'multimodal' in module 'MaxText' (no-name-in-module)
And they were all from code that was just merged
https://screenshot.googleplex.com/8ddsSFkx9RuLo6u

[DannyLiCom]

@aireenmei Also, these failures don't seem to be related to my code at all.
https://screenshot.googleplex.com/4BRE3k8xXhoft4s

[DannyLiCom]

Hi @aireenmei I noticed a lot of these errors in Pylint. E0611: No name 'multimodal' in module 'MaxText' (no-name-in-module) And they were all from code that was just merged https://screenshot.googleplex.com/8ddsSFkx9RuLo6u

Resolved

[DannyLiCom]

@aireenmei Also, these failures don't seem to be related to my code at all. https://screenshot.googleplex.com/4BRE3k8xXhoft4s

Resolved

[DannyLiCom]

Hi @aireenmei Could you please take another look? If everything looks good, could you help approve it?

[aireenmei]

The PR makes change in multiple recipes, but in the description only benchmarks.recipes.pw_mcjax_benchmark_recipe is tested. Could you make sure other recipes are also tested?

[DannyLiCom]

The PR makes change in multiple recipes, but in the description only benchmarks.recipes.pw_mcjax_benchmark_recipe is tested. Could you make sure other recipes are also tested?

Hi @aireenmei Exactly. We are focusing on benchmarks.recipes.pw_mcjax_benchmark_recipe. Our current large-scale tests rely almost exclusively on this recipe. Once this is merged, we will standardize all other recipes to follow this model. In other words, the other recipes do not yet support multiple command-line flags. My changes to those recipes are simply because I modified some utils, which required minor adjustments to keep everything consistent. Currently, pw_headless_mode.py on the main branch may still encounter issues.

[DannyLiCom]

Another reason is that integrating every recipe following the pw_mcjax_benchmark_recipe pattern all at once would make the PR too difficult to review. Therefore, we will submit one PR per recipe integration to keep things manageable. This is the first and most important recipe to be integrated; it will serve as the template for all future recipes.

[aireenmei]

Thanks for revising and the clarification!