Merge branch 'develop' into develop

Author: hanwen-pcluste

CHANGELOG.md

@@ -9,7 +9,6 @@ CHANGELOG
 - Add support for adding and removing shared storages at cluster update by updating `SharedStorage` configuration.
 - Add new configuration parameter `DeletionPolicy` for EFS and FSx for Lustre shared storage
   to support storage retention on deletion.
-- Add support for AWS Trainium instances.
 - Add support for Slurm Accounting.
 - Add support for on-demand capacity reservations.
 - Add support for specifying the supported IMDS version in cluster and build image configurations via the `Imds/ImdsSettings` property.

tests/integration-tests/tests/trainium/test_trainium.py

@@ -29,13 +29,12 @@ def test_trainium(
     bucket = boto3.resource("s3", region_name=region).Bucket(bucket_name)
     bucket.upload_file(str(test_datadir / "neuron-installation.sh"), "neuron-installation.sh")
 
-    # FIXME remove suppress_validators after GA
     cluster_config = pcluster_config_reader(bucket_name=bucket_name)
-    cluster = clusters_factory(cluster_config, suppress_validators="type:InstanceTypeBaseAMICompatibleValidator")
+    cluster = clusters_factory(cluster_config)
     remote_command_executor = RemoteCommandExecutor(cluster)
     scheduler_commands = scheduler_commands_factory(remote_command_executor)
 
-    # TODO uncomment allreduce test
+    # TODO uncomment allreduce test when bug fix in collective library
     # _test_allreduce_single_node(test_datadir, remote_command_executor, scheduler_commands)
     _test_ccl_two_nodes(test_datadir, remote_command_executor, scheduler_commands)
 

tests/integration-tests/tests/trainium/test_trainium/test_trainium/neuron-allreduce.sh

@@ -24,7 +24,7 @@ import os
 import torch_xla.core.xla_model as xm
 import torch
 import torch.distributed as dist
-from torch_xla.neuron.distributed import xla_backend
+from torch_xla.distributed import xla_backend
 
 #os.environ["NEURON_RT_LOG_LEVEL"] = "INFO"
 #os.environ["NEURON_RT_LOG_LOCATION"] = "syslog"

tests/integration-tests/tests/trainium/test_trainium/test_trainium/neuron-installation.sh

@@ -33,9 +33,13 @@ TEMPORARY_ARTIFACTS_BUCKET_PATH=s3://aws-parallelcluster-beta/neuron/
 
 _ubuntu_installation() {
   # Configure Linux for Neuron repository updates
-  sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
+  sudo tee /etc/apt/sources.list.d/neuron-private.list > /dev/null <<EOF
 deb https://${REPO_USER}:${REPO_SECRET}@apt.${REPO_SUFFIX} focal main
 EOF
+
+  sudo apt-get update -y
+  sudo apt-get install -y aws-neuronx-runtime-lib=2.* aws-neuronx-collectives=2.*
+
   wget -qO - https://${REPO_USER}:${REPO_SECRET}@apt.${REPO_SUFFIX}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
 
   # Install packages from S3 --> FIXME they should be installed from configured repository
@@ -61,13 +65,14 @@ EOF
 }
 
 _rhel_installation() {
-  # Install dkms driver. This is not required, installation is performed at AMI creation time
-  sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<EOF
-[neuron]
+  sudo tee /etc/yum.repos.d/neuron-private.repo > /dev/null <<EOF
+[neuron-private]
 name=Neuron YUM Repository
 baseurl=https://${REPO_USER}:${REPO_SECRET}@yum.${REPO_SUFFIX}
 enabled=1
 EOF
+  sudo yum install -y aws-neuronx-runtime-lib-2.* aws-neuronx-collectives-2.*
+
   sudo rpm --import https://${REPO_USER}:${REPO_SECRET}@yum.${REPO_SUFFIX}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB
 
   # Install packages from S3 --> FIXME they should be installed from configured repository
@@ -79,7 +84,6 @@ EOF
 
 
 _dkms_ubuntu_installation() {
-  # Install dkms driver. This is not required, installation is performed at AMI creation time
   sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
 deb https://apt.repos.neuron.amazonaws.com focal main
 EOF
@@ -91,7 +95,6 @@ EOF
 
 
 _dkms_rhel_installation() {
-  # Install dkms driver. This is not required, installation is performed at AMI creation time
   sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<EOF
 [neuron]
 name=Neuron YUM Repository
@@ -111,12 +114,12 @@ function main() {
   local OS="$(grep "^ID=" /etc/os-release | cut -d"=" -f 2 | xargs)"
   case ${OS} in
     ubuntu)
-      _dkms_ubuntu_installation  # not needed, installed at AMI creation time
+      _dkms_ubuntu_installation
       _ubuntu_installation
       USER=ubuntu
       ;;
     amzn)
-      _dkms_rhel_installation  # not needed, installed at AMI creation time
+      _dkms_rhel_installation
       _rhel_installation
       USER=ec2-user
       ;;
@@ -131,9 +134,8 @@ function main() {
   pip3 install -U pip
   pip3 install pytest
 
-  # Install packages from beta repo --> FIXME they should be installed from official PyPI
-  python3 -m pip config set global.extra-index-url "https://${REPO_USER}:${REPO_SECRET}@pip.${REPO_SUFFIX}"
-  PIPS=$(aws secretsmanager get-secret-value --secret-id arn:aws:secretsmanager:us-east-1:447714826191:secret:TrainiumPreviewRepository --region us-east-1 --query 'SecretString' --output text | jq -r '.pips')
+  python3 -m pip config set global.extra-index-url "https://pip.repos.neuron.amazonaws.com"
+  PIPS='torch-neuronx==1.11.0.1.* neuronx-cc==2.* transformers'
   pip3 install ${PIPS}
 }