[Release-3.11] Add pyxis integration test (#6475)

* Add test for Pyxis and Enroot functionality after configuration.
* This test creates a cluster with the necessary custom actions to configure Pyxis and Enroot.
* It submits two consecutive containerized jobs and verifies that they run successfully.

tests/integration-tests/README.md

@@ -518,7 +518,7 @@ test cases then you can do it in the following way:
 
 ```python
 @pytest.mark.usefixtures("region", "os", "instance", "scheduler")
-@pytest.mark.parametrized("cluster_max_size", [5, 10])
+@pytest.mark.parametrize("cluster_max_size", [5, 10])
 def test_case_2(cluster_max_size):
 ```
 

tests/integration-tests/configs/develop.yaml

@@ -829,3 +829,10 @@ test-suites:
           instances: {{ common.INSTANCES_DEFAULT_X86 }}
           oss: ["ubuntu2004"]
           schedulers: ["slurm"]
+  pyxis:
+    test_pyxis.py::test_pyxis:
+      dimensions:
+        - regions: ["eu-west-1"]
+          instances: {{ common.INSTANCES_DEFAULT_X86 }}
+          oss: ["ubuntu2204"]
+          schedulers: ["slurm"]

tests/integration-tests/tests/pyxis/test_pyxis.py

@@ -0,0 +1,83 @@
+# Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file.
+# This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+import logging
+
+import boto3
+import pytest
+from assertpy import assert_that
+from remote_command_executor import RemoteCommandExecutor
+
+from tests.common.schedulers_common import SlurmCommands
+
+
+@pytest.mark.parametrize("scale_up_fleet", [False])
+@pytest.mark.usefixtures("region", "os", "instance", "scheduler")
+def test_pyxis(pcluster_config_reader, clusters_factory, test_datadir, s3_bucket_factory, region, scale_up_fleet):
+    """
+    Test Pyxis and Enroot functionality after configuration.
+
+
+    This test creates a cluster with the necessary custom actions to configure Pyxis and Enroot.
+    It submits two consecutive containerized jobs and verifies that they run successfully,
+    and the output contains the expected messages.
+    """
+    # Set max_queue_size based on scale_up_fleet
+    max_queue_size = 1000 if scale_up_fleet else 3
+
+    # Create an S3 bucket for custom action scripts
+    bucket_name = s3_bucket_factory()
+    bucket = boto3.resource("s3", region_name=region).Bucket(bucket_name)
+
+    # Pre-upload custom scripts that set up pyxis to S3
+    bucket.upload_file(str(test_datadir / "head_node_configure.sh"), "head_node_configure.sh")
+    bucket.upload_file(str(test_datadir / "compute_node_start.sh"), "compute_node_start.sh")
+
+    cluster_config = pcluster_config_reader(bucket_name=bucket_name, max_queue_size=max_queue_size)
+    cluster = clusters_factory(cluster_config)
+
+    remote_command_executor = RemoteCommandExecutor(cluster)
+    slurm_commands = SlurmCommands(remote_command_executor)
+
+    # Submit the first containerized job with dynamic 3 or 1000 nodes
+    logging.info("Submitting first containerized job")
+
+    result = slurm_commands.submit_command(
+        command="srun --container-image docker://ubuntu:22.04 hostname",
+        nodes=max_queue_size,
+    )
+    job_id = slurm_commands.assert_job_submitted(result.stdout)
+    slurm_commands.wait_job_completed(job_id, timeout=30 if scale_up_fleet else 12)
+    slurm_commands.assert_job_succeeded(job_id)
+
+    # Fetch the job output and check for the expected messages
+    logging.info("Checking output of the first job")
+    slurm_out_1 = remote_command_executor.run_remote_command("cat slurm-1.out").stdout
+
+    logging.info("Checking for expected messages in first job output")
+    assert_that(slurm_out_1).contains("pyxis: imported docker image: docker://ubuntu:22.04")
+
+    # Submit the second containerized job with fixed 3 nodes after the first one completes
+    logging.info("Submitting second containerized job")
+    result = slurm_commands.submit_command(
+        command="srun --container-image docker://ubuntu:22.04 hostname",
+        nodes=3,
+    )
+    job_id = slurm_commands.assert_job_submitted(result.stdout)
+    slurm_commands.wait_job_completed(job_id)
+    slurm_commands.assert_job_succeeded(job_id)
+
+    # Fetch the job output and check for the expected messages
+    logging.info("Checking output of the second job")
+    slurm_out_2 = remote_command_executor.run_remote_command("cat slurm-2.out").stdout
+
+    logging.info("Checking for expected messages in second job output")
+    assert_that(slurm_out_2).contains("pyxis: imported docker image: docker://ubuntu:22.04")

tests/integration-tests/tests/pyxis/test_pyxis/test_pyxis/compute_node_start.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+set -e
+
+echo "Executing $0"
+
+# Configure Enroot
+ENROOT_PERSISTENT_DIR="/var/enroot"
+ENROOT_VOLATILE_DIR="/run/enroot"
+
+sudo mkdir -p $ENROOT_PERSISTENT_DIR
+sudo chmod 1777 $ENROOT_PERSISTENT_DIR
+sudo mkdir -p $ENROOT_VOLATILE_DIR
+sudo chmod 1777 $ENROOT_VOLATILE_DIR
+sudo mv /opt/parallelcluster/examples/enroot/enroot.conf /etc/enroot/enroot.conf
+sudo chmod 0644 /etc/enroot/enroot.conf
+
+# Configure Pyxis
+PYXIS_RUNTIME_DIR="/run/pyxis"
+
+sudo mkdir -p $PYXIS_RUNTIME_DIR
+sudo chmod 1777 $PYXIS_RUNTIME_DIR

tests/integration-tests/tests/pyxis/test_pyxis/test_pyxis/head_node_configure.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+set -e
+
+echo "Executing $0"
+
+# Configure Enroot
+ENROOT_PERSISTENT_DIR="/var/enroot"
+ENROOT_VOLATILE_DIR="/run/enroot"
+
+sudo mkdir -p $ENROOT_PERSISTENT_DIR
+sudo chmod 1777 $ENROOT_PERSISTENT_DIR
+sudo mkdir -p $ENROOT_VOLATILE_DIR
+sudo chmod 1777 $ENROOT_VOLATILE_DIR
+sudo mv /opt/parallelcluster/examples/enroot/enroot.conf /etc/enroot/enroot.conf
+sudo chmod 0644 /etc/enroot/enroot.conf
+
+# Configure Pyxis
+PYXIS_RUNTIME_DIR="/run/pyxis"
+
+sudo mkdir -p $PYXIS_RUNTIME_DIR
+sudo chmod 1777 $PYXIS_RUNTIME_DIR
+
+sudo mkdir -p /opt/slurm/etc/plugstack.conf.d/
+sudo mv /opt/parallelcluster/examples/spank/plugstack.conf /opt/slurm/etc/
+sudo mv /opt/parallelcluster/examples/pyxis/pyxis.conf /opt/slurm/etc/plugstack.conf.d/
+sudo -i scontrol reconfigure

tests/integration-tests/tests/pyxis/test_pyxis/test_pyxis/pcluster.config.yaml

@@ -0,0 +1,33 @@
+Image:
+  Os: {{ os }}
+HeadNode:
+  InstanceType: {{ instance }}
+  Networking:
+    SubnetId: {{ public_subnet_id }}
+  Ssh:
+    KeyName: {{ key_name }}
+  CustomActions:
+    OnNodeConfigured:
+      Script: s3://{{ bucket_name }}/head_node_configure.sh
+  Iam:
+    S3Access:
+      - BucketName: {{ bucket_name }}
+Scheduling:
+  Scheduler: {{ scheduler }}
+  SlurmQueues:
+  - Name: queue-0
+    ComputeResources:
+      - Name: compute-resource-0
+        Instances:
+          - InstanceType: t3.small
+        MinCount: 0
+        MaxCount: {{ max_queue_size }}
+    Networking:
+      SubnetIds:
+        - {{ private_subnet_id }}
+    CustomActions:
+      OnNodeStart:
+        Script: s3://{{ bucket_name }}/compute_node_start.sh
+    Iam:
+      S3Access:
+        - BucketName: {{ bucket_name }}