benchmark sample run

Author: Doris26

dags/map_reproducibility/internal_runs/sample_a3ultra_maxtext_single_run.py

@@ -0,0 +1,54 @@
+# 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.
+
+"""Sample job to run Aotc reproducibility benchmarks."""
+import sys
+import os
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.abspath(os.path.join(script_dir, "..", "..", ".."))
+
+print(f"Script directory: {script_dir}")
+print(f"Project root: {project_root}")
+
+if project_root not in sys.path:
+  sys.path.insert(0, project_root)
+
+import datetime
+from dags.map_reproducibility.utils.constants import Image
+from dags.map_reproducibility.internal_runs.dag_configs import DAG_CONFIGS_ULTRA
+from dags.map_reproducibility.utils.sample_workload_utils import run_internal_sample_aotc_workload
+
+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}"
+RELEASE_IMAGE = f"{Image.MAXTEXT_JAX_STABLE_RELEASE}:2025-04-17"
+SAMPLE_RUN_BUCKET_NAME = "yujunzou-dev-supercomputer-testing"
+
+
+# Setup configuratio
+relative_config_yaml_path = (
+    "recipes/a3ultra/a3ultra_llama3.1-8b_8gpus_bf16_maxtext.yaml"
+)
+config_name = relative_config_yaml_path.replace(".yaml", "")
+timeout = DAG_CONFIGS_ULTRA[relative_config_yaml_path]["timeout_minutes"]
+base_recipe_repo_root = f"{project_root}/../internal-gpu-recipes"
+
+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,
+)

dags/map_reproducibility/utils/common_utils.py

@@ -22,8 +22,8 @@
 import string
 import time
 import subprocess
+import getpass
 
-from google.cloud import storage
 from airflow.decorators import task
 from airflow.hooks.subprocess import SubprocessHook
 from xlml.utils import metric
@@ -115,11 +115,14 @@ def get_internal_pre_workload_cmds(job_name):
   return prepare_workload_cmds
 
 
-def get_internal_pre_workload_job_name(model_id, framework):
+def get_internal_pre_workload_job_name(model_id, framework, is_sample_run=False):
   helm_model_id = model_id.replace(".", "-")
   random_id = "".join(random.choices(string.ascii_lowercase, k=4))
   now = int(time.time())
   job_name = f"coreml-{helm_model_id}-{now}-{framework}-{random_id}"
+  if is_sample_run:
+    job_name = f"{getpass.getuser()}-{job_name}"
+  print(f"NAME: {job_name}")
   return job_name
 
 
@@ -207,16 +210,17 @@ def helm_apply_cmds_internal_run(
     kueue_name: str = "a3-ultra",
     additional_cmds: str = "",
     test_run=False,
+    bucket_name=BUCKET_NAME,
 ):
   gcs_cmd = ""
   if framework == "maxtext":
-    gcs_cmd += f" --set volumes.gcsMounts[0].bucketName={BUCKET_NAME} "
+    gcs_cmd += f" --set volumes.gcsMounts[0].bucketName={bucket_name} "
 
   if hypercomputer == "a3ultra":
     if framework != "maxtext":
       gcs_cmd += f" --set queue={kueue_name} "
   else:
-    gcs_cmd += f" --set workload.gcsBucketForDataCataPath={BUCKET_NAME} "
+    gcs_cmd += f" --set workload.gcsBucketForDataCataPath={bucket_name} "
 
   cluster_cmd = ""
   if framework == "nemo" and hypercomputer == "a3ultra":
@@ -230,8 +234,9 @@ def helm_apply_cmds_internal_run(
   if aotc:
     set_aotc = " --set-string workload.aotc=true "
 
-  if test_run:
-    helm_template_path = f"/home/airflow/gcs/dags/dags/map_reproducibility/helm-charts/{hypercomputer}/{framework}-training"
+  local_helm_template_path = f"/home/airflow/gcs/dags/dags/map_reproducibility/helm-charts/{hypercomputer}/{framework}-training"
+  if test_run and os.path.exists(local_helm_template_path):
+    helm_template_path = local_helm_template_path
   else:
     helm_template_path = f"{recipe_repo_root}/src/helm-charts/{hypercomputer}/{framework}-training"
 
@@ -321,8 +326,8 @@ def copy_bucket_cmds_nemo(recipe_repo_root, hypercomputer: str = "a3mega"):
   return copy_bucket_contents
 
 
-def copy_bucket_cmds_maxtext(tmpdir, recipe_repo_root):
-  gcs_location = f"gs://{BUCKET_NAME}/maxtext/"
+def copy_bucket_cmds_maxtext(tmpdir, bucket_name=BUCKET_NAME):
+  gcs_location = f"gs://{bucket_name}/maxtext/"
 
   cmds = (
       f"METRICS_FILE={tmpdir}/tflog/metrics",
@@ -704,6 +709,7 @@ def parse_internal_config_filename(filename, config=None):
     return config
 
 
+@staticmethod
 def parse_internal_config_content(yaml_path, config=None):
   """
   Parse the internal content of a config YAML file and update the existing config.

dags/map_reproducibility/utils/internal_aotc_workload.py

@@ -148,9 +148,7 @@ def run_internal_aotc_workload(
                     test_run=test_run,
                 )
                 + internal_wait_for_jobs_cmds(timeout=container_timeout)
-                + copy_bucket_cmds_maxtext(
-                    tmpdir, recipe_repo_root=recipe_repo_root
-                )
+                + copy_bucket_cmds_maxtext(tmpdir)
                 + cleanup_cmds()
             ),
         ],