in dags/map_reproducibility/utils/common_utils.py [0:0]
def get_scheduled_time(hardware: str, model: str, framework: str):
"""
Returns a cron expression for the DAG schedule based on
the given hardware, model, and framework.
Each model runs on Thursday on a unique time so
that we have free nodes for each.
The alloted time for these tests is 6 pm - 10 pm PST on Thursday.
6 PM pst - 0 2 * * 5
10 PM pst - 0 6 * * 5
Args:
hardware: The hardware type (e.g., "a3ultra", "a3mega").
model: The model ID (e.g., "mixtral-8x7b", "llama-3.1-70b").
framework: The framework (e.g., "nemo", "maxtext").
Returns:
A cron expression string (e.g., "0 12 * * 4") or None
if no schedule is defined
for the given combination.
"""
schedule_map = {
"a3ultra": {
"mixtral-8x7b": {
"nemo": "0 3 * * 5",
"maxtext": "0 2 * * 5", # 6 PM PST on Thursday
},
"llama3-1-70b": {
"nemo": "0 4 * * 5",
"maxtext": "0 4 * * 5",
},
"llama3-1-405b": {
"nemo": "0 5 * * 5",
"maxtext": "0 5 * * 5",
},
},
"a3mega": {
"mixtral-8x7b": {
"nemo": "0 4 * * 5",
"maxtext": "0 3 * * 5",
},
"llama3-70b": {
"nemo": "0 2 * * 5",
"maxtext": "0 5 * * 5",
},
"llama3-1-70b": {
"nemo": "0 2 * * 5",
"maxtext": "0 4 * * 5",
},
"gpt3-175b": {
"nemo": "0 4 * * 5",
},
},
"a4": {
"mixtral-8x7b": {
"nemo": "0 2 * * 5",
},
"llama3-1-70b": {
"nemo": "0 3 * * 5",
"maxtext": "0 3 * * 5",
},
"llama3-1-405b": {
"nemo": "0 4 * * 5",
"maxtext": "0 4 * * 5",
},
},
}
if hardware in schedule_map:
if model in schedule_map[hardware]:
if framework in schedule_map[hardware][model]:
return schedule_map[hardware][model][framework]
return None # Return None if no schedule is found for the given combination