def main()

in dags/map_reproducibility/internal_runs/sample_a3mega_maxtext_single_run.py [0:0]

• 17 lines of code
• 3 McCabe index (conditional complexity)

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/a3mega/a3mega_llama3.1-8b_8gpus_fp8_maxtext.yaml"
  )
  timeout = DAG_CONFIGS_MEGA[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,
  )