Merge pull request #847 from mlcommons/dev

Dev -> main

Author: priyakasimbeg

.github/workflows/regression_tests.yml

@@ -13,6 +13,9 @@ do
     esac
 done
 
+# Artifact repostiory
+ARTIFACT_REPO="europe-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo"
+
 if [[ -z ${GIT_BRANCH+x} ]]
 then 
 GIT_BRANCH='main' # Set default argument
@@ -22,9 +25,9 @@ for FRAMEWORK in "jax" "pytorch" "both"
 do
     IMAGE_NAME="algoperf_${FRAMEWORK}_${GIT_BRANCH}"
     DOCKER_BUILD_COMMAND="docker build --no-cache -t $IMAGE_NAME . --build-arg framework=$FRAMEWORK --build-arg branch=$GIT_BRANCH"
-    DOCKER_TAG_COMMAND="docker tag $IMAGE_NAME us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/$IMAGE_NAME"
-    DOCKER_PUSH_COMMAND="docker push us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/$IMAGE_NAME"
-    DOCKER_PULL_COMMAND="docker pull us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/$IMAGE_NAME"
+    DOCKER_TAG_COMMAND="docker tag $IMAGE_NAME $ARTIFACT_REPO/$IMAGE_NAME"
+    DOCKER_PUSH_COMMAND="docker push $ARTIFACT_REPO/$IMAGE_NAME"
+    DOCKER_PULL_COMMAND="docker pull $ARTIFACT_REPO/$IMAGE_NAME"
 
     echo "On branch: ${GIT_BRANCH}"
     echo $DOCKER_BUILD_COMMAND

.github/workflows/regression_tests_variants.yml

@@ -40,8 +40,8 @@ write_files:
       ExecStartPre=mount --bind /var/lib/nvidia /var/lib/nvidia
       ExecStartPre=mount -o remountexec /var/lib/nvidia
       ExecStartPre=/usr/bin/docker-credential-gcr configure-docker --registries us-central1-docker.pkg.dev
-      ExecStartPre=/usr/bin/docker pull us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/base_image:latest
-      ExecStart=/usr/bin/docker run --rm --name=mlcommons --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 --volume /var/lib/nvidia/bin:/usr/local/nvidia/bin --device /dev/nvidia0:/dev/nvidia0  --device /dev/nvidia1:/dev/nvidia1 --device /dev/nvidia2:/dev/nvidia2 --device /dev/nvidia3:/dev/nvidia3 --device /dev/nvidia4:/dev/nvidia4 --device /dev/nvidia5:/dev/nvidia5 --device /dev/nvidia6:/dev/nvidia6 --device /dev/nvidia7:/dev/nvidia7 --device /dev/nvidia-uvm:/dev/nvidia-uvm --device /dev/nvidiactl:/dev/nvidiactl us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/base_image:latest -b true
+      ExecStartPre=/usr/bin/docker pull europe-west4-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo/base_image:latest
+      ExecStart=/usr/bin/docker run --rm --name=mlcommons --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 --volume /var/lib/nvidia/bin:/usr/local/nvidia/bin --device /dev/nvidia0:/dev/nvidia0  --device /dev/nvidia1:/dev/nvidia1 --device /dev/nvidia2:/dev/nvidia2 --device /dev/nvidia3:/dev/nvidia3 --device /dev/nvidia4:/dev/nvidia4 --device /dev/nvidia5:/dev/nvidia5 --device /dev/nvidia6:/dev/nvidia6 --device /dev/nvidia7:/dev/nvidia7 --device /dev/nvidia-uvm:/dev/nvidia-uvm --device /dev/nvidiactl:/dev/nvidiactl europe-west4-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo/base_image:latest -b true
       StandardOutput=journal+console
       StandardError=journal+console
 

docker/build_docker_images.sh

@@ -51,6 +51,8 @@ RSYNC_DATA="true"
 OVERWRITE="false"
 SAVE_CHECKPOINTS="true"
 TUNING_RULESET="external"
+ROOT_DATA_BUCKET="algoperf-data"
+LOGS_BUCKET="algoperf-runs"
 
 # Pass flag
 while [ "$1" != "" ]; do
@@ -136,6 +138,14 @@ while [ "$1" != "" ]; do
             shift
             ADDITIONAL_REQUIREMENTS_PATH=$1
             ;;
+        --data_bucket)
+            shift
+            ROOT_DATA_BUCKET=$1
+            ;;
+        --logs_bucket)
+            shift
+            LOGS_BUCKET=$1
+            ;;
         *) 
             usage 
             exit 1
@@ -179,11 +189,11 @@ VALID_WORKLOADS=("criteo1tb" "imagenet_resnet" "imagenet_resnet_silu" "imagenet_
 VALID_RULESETS=("self" "external")
 
 # Set data and experiment paths
-ROOT_DATA_BUCKET="gs://mlcommons-data"
 ROOT_DATA_DIR="${HOME_DIR}/data"
+ROOT_DATA_BUCKET="gs://${ROOT_DATA_BUCKET}"
 
-EXPERIMENT_BUCKET="gs://mlcommons-runs"
 EXPERIMENT_DIR="${HOME_DIR}/experiment_runs"
+EXPERIMENT_LOGS_BUCKET="gs://${LOGS_BUCKET}"
 
 if [[ -n ${DATASET+x} ]]; then 
     if [[ ! " ${VALID_DATASETS[@]} " =~ " $DATASET " ]]; then
@@ -313,8 +323,8 @@ if [[ ! -z ${SUBMISSION_PATH+x} ]]; then
     RETURN_CODE=$?
 
     if [[ $INTERNAL_CONTRIBUTOR_MODE == "true" ]]; then 
-        /google-cloud-sdk/bin/gsutil -m cp -r ${EXPERIMENT_DIR}/${EXPERIMENT_NAME}/${WORKLOAD}_${FRAMEWORK} ${EXPERIMENT_BUCKET}/${EXPERIMENT_NAME}/
-        /google-cloud-sdk/bin/gsutil -m cp ${LOG_FILE} ${EXPERIMENT_BUCKET}/${EXPERIMENT_NAME}/${WORKLOAD}_${FRAMEWORK}/
+        /google-cloud-sdk/bin/gsutil -m cp -r ${EXPERIMENT_DIR}/${EXPERIMENT_NAME}/${WORKLOAD}_${FRAMEWORK} ${EXPERIMENT_LOGS_BUCKET}/${EXPERIMENT_NAME}/
+        /google-cloud-sdk/bin/gsutil -m cp ${LOG_FILE} ${EXPERIMENT_LOGS_BUCKET}/${EXPERIMENT_NAME}/${WORKLOAD}_${FRAMEWORK}/
     fi
 
 fi

docker/scripts/cloud-init.cfg

@@ -88,7 +88,7 @@ gcloud auth configure-docker $ARTIFACT_REGISTRY_URL
 To pull the latest prebuilt image:
 
 ```bash
-docker pull us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/<image_name>
+docker pull europe-west4-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo/<image_name>
 ```
 
 The naming convention for `image_name` is `algoperf_<framework>_<branch>`.
@@ -102,7 +102,7 @@ Currently maintained images on the repository are:
 - `algoperf_both_dev`
 
 To reference the pulled image you will have to use the full `image_path`, e.g.
-`us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/algoperf_jax_main`.
+`europe-west4-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo/algoperf_jax_main`.
 
 ### Trigger Rebuild and Push of Maintained Images
 

docker/scripts/startup.sh

@@ -26,7 +26,7 @@
 
 flags.DEFINE_string(
     'docker_image_url',
-    'us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo/algoperf_jax_dev',
+    'europe-west4-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo/algoperf_jax_dev',
     'URL to docker image')
 flags.DEFINE_integer(
     'run_percentage',

docs/CONTRIBUTING.md

@@ -1,4 +1,5 @@
 from flax import jax_utils
+from flax.core import FrozenDict
 import jax
 import numpy as np
 import torch
@@ -16,6 +17,8 @@ def torch2jax(jax_workload,
   jax_params, model_state = jax_workload.init_model_fn(jax.random.PRNGKey(0),
                                                        **init_kwargs)
   pytorch_model, _ = pytorch_workload.init_model_fn([0], **init_kwargs)
+  if isinstance(jax_params, dict):
+    jax_params = FrozenDict(jax_params)
   jax_params = jax_utils.unreplicate(jax_params).unfreeze()
   if model_state is not None:
     model_state = jax_utils.unreplicate(model_state)

scoring/run_workloads.py

@@ -50,7 +50,7 @@ def test_workload(self, workload):
     pyt_logs = '/tmp/pyt_log.pkl'
     try:
       run(
-          f'XLA_PYTHON_CLIENT_ALLOCATOR=platform python3 -m tests.reference_algorithm_tests --workload={workload} --framework=jax --global_batch_size={GLOBAL_BATCH_SIZE} --log_file={jax_logs}'
+          f'XLA_PYTHON_CLIENT_ALLOCATOR=platform python -m tests.reference_algorithm_tests --workload={workload} --framework=jax --global_batch_size={GLOBAL_BATCH_SIZE} --log_file={jax_logs}'
           f' --submission_path=tests/modeldiffs/vanilla_sgd_jax.py --identical=True --tuning_search_space=None --num_train_steps={NUM_TRAIN_STEPS}',
           shell=True,
           stdout=DEVNULL,