dags/map_reproducibility/internal_runs/dag_configs.py (197 lines of code) (raw):

# Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from dags.map_reproducibility.utils.constants import Schedule DAG_CONFIGS_MEGA = { "recipes/a3mega/a3mega_llama3.1-8b_8gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "nightly_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_llama3.1-8b_8gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "nightly_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_llama3.1-8b_16gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "nightly_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_llama3.1-8b_16gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "nightly_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_mixtral-8x7b_8gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "nightly_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_mixtral-8x7b_16gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "nightly_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_llama3.1-70b_256gpus_bf16_maxtext.yaml": { "timeout_minutes": 25, "backfill_group_nightly": 2, "backfill_group_release": 2, "nightly_schedule": Schedule.WEEKDAY_PDT_12_30AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_12_30AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_llama3.1-70b_256gpus_fp8_maxtext.yaml": { "timeout_minutes": 20, "backfill_group_nightly": 3, "backfill_group_release": 3, "nightly_schedule": Schedule.WEEKDAY_PDT_1AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_1AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_llama3.1-405b_512gpus_fp8_maxtext.yaml": { "timeout_minutes": 40, "backfill_group_nightly": 4, "backfill_group_release": 5, "nightly_schedule": Schedule.WEEKDAY_PDT_1_30AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_2AM_EXCEPT_THURSDAY, }, "recipes/a3mega/a3mega_llama3.1-405b_512gpus_bf16_maxtext.yaml": { "timeout_minutes": 50, "backfill_group_nightly": 6, "backfill_group_release": 7, "nightly_schedule": Schedule.WEEKDAY_PDT_2_30AM_EXCEPT_THURSDAY, "release_schedule": Schedule.WEEKDAY_PDT_3AM_EXCEPT_THURSDAY, }, } DAG_CONFIGS_ULTRA = { "recipes/a3ultra/a3ultra_llama3.1-8b_8gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_llama3.1-8b_8gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_llama3.1-8b_16gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_llama3.1-8b_16gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_mixtral-8x7b_8gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_mixtral-8x7b_16gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_llama3.1-70b_256gpus_bf16_maxtext.yaml": { "timeout_minutes": 20, "backfill_group_nightly": 2, "backfill_group_release": 2, "schedule": Schedule.WEEKDAY_PDT_12_30AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_llama3.1-70b_256gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 3, "backfill_group_release": 3, "schedule": Schedule.WEEKDAY_PDT_1AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_llama3.1-405b_256gpus_fp8_maxtext.yaml": { "timeout_minutes": 30, "backfill_group_nightly": 4, "backfill_group_release": 4, "schedule": Schedule.WEEKDAY_PDT_1_30AM_EXCEPT_THURSDAY, }, "recipes/a3ultra/a3ultra_llama3.1-405b_256gpus_bf16_maxtext.yaml": { "timeout_minutes": 40, "backfill_group_nightly": 5, "backfill_group_release": 5, "schedule": Schedule.WEEKDAY_PDT_2AM_EXCEPT_THURSDAY, }, } DAG_CONFIGS_A4 = { "recipes/a4/a4_llama3.1-8b_8gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_llama3.1-8b_8gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_llama3.1-8b_16gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_llama3.1-8b_16gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_mixtral-8x7b_8gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_mixtral-8x7b_16gpus_bf16_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 1, "backfill_group_release": 1, "schedule": Schedule.WEEKDAY_PDT_12AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_llama3.1-70b_256gpus_bf16_maxtext.yaml": { "timeout_minutes": 20, "backfill_group_nightly": 2, "backfill_group_release": 2, "schedule": Schedule.WEEKDAY_PDT_12_30AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_llama3.1-70b_256gpus_fp8_maxtext.yaml": { "timeout_minutes": 15, "backfill_group_nightly": 3, "backfill_group_release": 3, "schedule": Schedule.WEEKDAY_PDT_1AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_llama3.1-405b_256gpus_fp8_maxtext.yaml": { "timeout_minutes": 30, "backfill_group_nightly": 4, "backfill_group_release": 4, "schedule": Schedule.WEEKDAY_PDT_1_30AM_EXCEPT_THURSDAY, }, "recipes/a4/a4_llama3.1-405b_256gpus_bf16_maxtext.yaml": { "timeout_minutes": 40, "backfill_group_nightly": 5, "backfill_group_release": 5, "schedule": Schedule.WEEKDAY_PDT_2AM_EXCEPT_THURSDAY, }, }