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,
},
}