Slurm prolog

.github/workflows/release.yaml

@@ -67,6 +67,7 @@ jobs:
         run: |
           aws s3 cp provisioner_scripts/create_users.py s3://${{ secrets.S3_BUCKET }}/scripts/create_users.py
           aws s3 cp provisioner_scripts/environment.sh s3://${{ secrets.S3_BUCKET }}/scripts/environment.sh
+          aws s3 cp provisioner_scripts/80_cloudwatch_agent_config_prolog.sh s3://${{ secrets.S3_BUCKET }}/scripts/80_cloudwatch_agent_config_prolog.sh
       - name: Release
         uses: softprops/action-gh-release@v2
         if: github.ref == 'refs/heads/main'

hpc_provisioner/src/hpc_provisioner/handlers.py

@@ -13,7 +13,7 @@
     PROJECT_TAG_KEY,
     VLAB_TAG_KEY,
 )
-from hpc_provisioner.utils import generate_public_key
+from hpc_provisioner.utils import generate_public_key, get_suffix
 
 from .logging_config import LOGGING_CONFIG
 from .pcluster_manager import (
@@ -130,7 +130,7 @@ def pcluster_create_request_handler(event, _context=None):
 
     logger.debug("calling create lambda async")
     boto3.client("lambda").invoke_async(
-        FunctionName="hpc-resource-provisioner-creator",
+        FunctionName=f"hpc-resource-provisioner-creator-{get_suffix()}",
         InvokeArgs=json.dumps(create_args),
     )
     logger.debug("called create lambda async")

hpc_provisioner/src/hpc_provisioner/pcluster_manager.py

@@ -208,6 +208,14 @@ def pcluster_create(
     pcluster_config["Scheduling"]["SlurmQueues"] = choose_tier(pcluster_config, options)
     if not options["include_lustre"]:
         pcluster_config["SharedStorage"].pop(1)
+    if benchmark:
+        pcluster_config["HeadNode"]["CustomActions"]["OnNodeConfigured"]["Sequence"].append(
+            {
+                "Script": "s3://sboinfrastructureassets-sandbox/scripts/80_cloudwatch_agent_config_prolog.sh",
+                "Args": [cluster_name],
+            }
+        )
+
     output_file_name = write_config(cluster_name, pcluster_config)
 
     try:

hpc_provisioner/src/hpc_provisioner/utils.py

@@ -30,6 +30,10 @@ def get_fsx_policy_arn() -> str:
     return _get_env_var("FSX_POLICY_ARN")
 
 
+def get_suffix() -> str:
+    return _get_env_var("SUFFIX")
+
+
 def generate_public_key(key_material):
     private_key = serialization.load_pem_private_key(key_material.encode(), password=None)
     public_key = private_key.public_key()

provisioner_scripts/80_cloudwatch_agent_config_prolog.sh

@@ -0,0 +1,80 @@
+#!/bin/bash
+
+export CLUSTER_NAME=${1}
+
+mkdir -p /opt/slurm/etc/scripts/prolog.d
+
+cat << _EOF_ >> /opt/slurm/etc/scripts/prolog.d/80_cloudwatch_agent_config_prolog.sh
+#!/bin/bash
+
+CWAGENT_CONFIG=/opt/slurm/CWAgent_config_$CLUSTER_NAME.json
+
+if [ ! -f \$CWAGENT_CONFIG ]; then
+	echo "\$CWAGENT_CONFIG not found"
+	exit 0
+fi
+
+sudo /opt/aws/amazon-cloudwatch-agent/bin/amazon-cloudwatch-agent-ctl -a fetch-config -m ec2 -c file:/\$CWAGENT_CONFIG -s
+_EOF_
+
+chmod 755 /opt/slurm/etc/scripts/prolog.d/80_cloudwatch_agent_config_prolog.sh
+
+cat << _EOF_ >> /opt/slurm/CWAgent_config_$CLUSTER_NAME.json
+
+{
+	"agent": {
+                "metrics_collection_interval": 60,
+                "run_as_user": "root"
+        },
+        "metrics": {
+		"namespace": "CustomMetrics_test",
+		"aggregation_dimensions": [
+       			["ClusterName", "InstanceId"]
+       		],
+                "append_dimensions": {
+                        "InstanceId": "\${aws:InstanceId}"
+                },
+	        "metrics_collected": {
+                	"disk": {
+       				"append_dimensions":{
+					"ClusterName": "$CLUSTER_NAME"
+				},
+                                "measurement": [
+                                        "used_percent", "used"
+                                ],
+                                "resources": [
+					"/sbo/data/scratch", "/sbo/home"
+                                ]
+                        },
+			"diskio": {
+                                "append_dimensions":{
+                                        "ClusterName": "$CLUSTER_NAME"
+                                },
+                                 "measurement": [
+                                        "reads", "read_bytes", "writes", "write_bytes"
+                                ],
+                                "resources": [
+                                        "*"
+                                ]
+                        },
+                        "mem": {
+				"append_dimensions":{
+					"ClusterName": "$CLUSTER_NAME"
+				},
+                                "measurement": [
+                                        "mem_used_percent"
+                                ]
+                        },
+			"cpu": {
+       				"append_dimensions":{
+					"ClusterName": "$CLUSTER_NAME"
+				},
+        			"measurement": [
+            				"cpu_usage_active"
+        			],
+        			"totalcpu": true
+    				}
+                }
+        }
+}
+_EOF_