Recipe Modularization, DAG Integration, and BigQuery Integration by DannyLiCom · Pull Request #2715 · AI-Hypercomputer/maxtext

DannyLiCom · 2025-11-19T05:38:21Z

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 · 2025-11-28T03:02:24Z

benchmarks/recipes/args_helper.py



-def handle_cmd_args(cluster_config: XpkClusterConfig, *actions: str, **kwargs) -> bool:
+def handle_cmd_args(cluster_config: XpkClusterConfig, is_delete: bool, user: str, **kwargs) -> bool:


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

DannyLiCom · 2025-11-28T03:03:07Z

benchmarks/recipes/pw_mcjax_benchmark_recipe.py

+  parser_utils.add_arguments(parser)
+  args = parser.parse_args()
+
+  if len(sys.argv) > 2:


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

codecov · 2025-12-29T02:31:44Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

github-actions · 2026-01-28T16:11:54Z

This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.

github-actions · 2026-02-04T16:19:02Z

This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it.

src/MaxText/layers/attention_op.py

src/MaxText/layers/moe.py

tools/dev/code_style.sh

SujeethJinesh · 2026-02-05T18:30:42Z

LGTM, but I would just verify that the imports work properly at the top of the files. Since the MaxText modularization refactor, we may need to check that since a lot of these changes were from before that.

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?

benchmarks/recipes/parser_utils.py

tools/dev/code_style.sh

src/MaxText/layers/attention_op.py

src/MaxText/layers/moe.py

benchmarks/recipes/pw_suspend_resume.py

benchmarks/recipes/parser_utils.py

github-actions · 2026-02-05T22:02:14Z

🤖 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.

benchmarks/recipes/parser_utils.py

benchmarks/recipes/pw_mcjax_benchmark_recipe.py

benchmarks/recipes/pw_utils.py

github-actions · 2026-02-05T22:04:33Z

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

DannyLiCom · 2026-02-06T01:11:15Z

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 · 2026-02-06T01:13:45Z

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

DannyLiCom · 2026-02-07T05:40:01Z

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 · 2026-02-07T05:40:09Z

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

Resolved

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch from a961644 to f2a46a6 Compare November 19, 2025 05:46

DannyLiCom marked this pull request as ready for review November 19, 2025 06:05

DannyLiCom requested review from NuojCheng, Obliviour, RissyRan, SujeethJinesh, bvandermoon, gobbleturk, khatwanimohit, mitalisi, richjames0, shauryagup, shralex and vipannalla as code owners November 19, 2025 06:06

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch from f2a46a6 to 3840670 Compare November 23, 2025 16:39

DannyLiCom requested a review from notabee as a code owner November 23, 2025 16:39

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch 2 times, most recently from 8adde2d to b08ea95 Compare November 28, 2025 02:17

DannyLiCom commented Nov 28, 2025

View reviewed changes

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch 3 times, most recently from 0303bcf to 78b8073 Compare December 17, 2025 09:02

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch 3 times, most recently from ce0fc18 to c37edc0 Compare December 29, 2025 02:19

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch from c37edc0 to 57a9040 Compare December 29, 2025 02:43

DannyLiCom requested review from SurbhiJainUSC, gagika and hengtaoguo as code owners December 29, 2025 02:43

DannyLiCom requested review from NicoGrande and suexu1025 as code owners December 29, 2025 02:43

github-actions bot added the stale Automatically applied to stale PRs. label Jan 28, 2026

github-actions bot closed this Feb 4, 2026

ycchenzheng reopened this Feb 4, 2026

ycchenzheng requested a review from jesselu-google as a code owner February 4, 2026 22:20

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch from 57a9040 to 991aaf6 Compare February 5, 2026 06:32

DannyLiCom requested review from michelle-yooh, parambole and shuningjin as code owners February 5, 2026 07:16

DannyLiCom commented Feb 5, 2026

View reviewed changes

src/MaxText/layers/attention_op.py Show resolved Hide resolved

src/MaxText/layers/moe.py Show resolved Hide resolved

tools/dev/code_style.sh Outdated Show resolved Hide resolved

github-actions bot removed the stale Automatically applied to stale PRs. label Feb 5, 2026

SujeethJinesh approved these changes Feb 5, 2026

View reviewed changes

aireenmei reviewed Feb 5, 2026

View reviewed changes

aireenmei added the gemini-review label Feb 5, 2026

github-actions bot reviewed Feb 5, 2026

View reviewed changes

benchmarks/recipes/parser_utils.py Show resolved Hide resolved

benchmarks/recipes/pw_mcjax_benchmark_recipe.py Outdated Show resolved Hide resolved

benchmarks/recipes/pw_utils.py Show resolved Hide resolved

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch from 9bb39cd to a57612a Compare February 6, 2026 01:04

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch from a57612a to 2e66a21 Compare February 6, 2026 01:20

ycchenzheng requested a review from aireenmei February 6, 2026 16:01

Recipe Modularization, DAG Integration, and BigQuery Integration

cd65cbd

DannyLiCom force-pushed the lidanny/pw-DAG-integrate2 branch from 2e66a21 to cd65cbd Compare February 7, 2026 01:52



		def handle_cmd_args(cluster_config: XpkClusterConfig, actions: str, *kwargs) -> bool:
		def handle_cmd_args(cluster_config: XpkClusterConfig, is_delete: bool, user: str, **kwargs) -> bool:

Conversation

DannyLiCom commented Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

DannyLiCom Nov 28, 2025

Choose a reason for hiding this comment

Uh oh!

DannyLiCom Nov 28, 2025

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Dec 29, 2025

Codecov Report

Uh oh!

github-actions bot commented Jan 28, 2026

Uh oh!

github-actions bot commented Feb 4, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SujeethJinesh commented Feb 5, 2026

Uh oh!

aireenmei left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Feb 5, 2026

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Feb 5, 2026

Uh oh!

DannyLiCom commented Feb 6, 2026

Uh oh!

DannyLiCom commented Feb 6, 2026

Uh oh!

DannyLiCom commented Feb 7, 2026

Uh oh!

DannyLiCom commented Feb 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

DannyLiCom commented Nov 19, 2025 •

edited

Loading