in dags/map_reproducibility/internal_runs/sample_a4_maxtext_single_run.py [0:0]
def main():
utc_date = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
NIGHTLY_IMAGE = f"{Image.MAXTEXT_JAX_STABLE_NIGHTLY}:{utc_date}"
RELEASE_IMAGE = f"{Image.MAXTEXT_JAX_STABLE_RELEASE}:{utc_date}"
# The dynamic date-based image is defined above, but can be overridden below
# with specific image versions and bucket names as required for testing
RELEASE_IMAGE = f"{Image.MAXTEXT_JAX_STABLE_RELEASE}:2025-04-17"
SAMPLE_RUN_BUCKET_NAME = "yujunzou-dev-supercomputer-testing"
# Setup configuration
relative_config_yaml_path = "recipes/a4/a4_llama3.1-8b_8gpus_fp8_maxtext.yaml"
timeout = DAG_CONFIGS_A4[relative_config_yaml_path]["timeout_minutes"]
run_internal_sample_aotc_workload(
relative_config_yaml_path=relative_config_yaml_path,
base_recipe_repo_root=base_recipe_repo_root,
timeout=timeout,
image_version=RELEASE_IMAGE,
sample_run_bucket_name=SAMPLE_RUN_BUCKET_NAME,
)