maxtext
maxtext copied to clipboard
Recipe Modularization, DAG Integration, and BigQuery Integration
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='lidanny' \
--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 \ # Waiting for BQ PR to be merged
# --bq_db_project='cloud-tpu-multipod-dev' \ # Waiting for BQ PR to be merged
# --bq_db_dataset="chzheng_test_100steps" # Waiting for BQ PR to be merged
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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.