From 0d5f0023d0b2988515e2dd842032dd46a5642047 Mon Sep 17 00:00:00 2001
From: Verdi March <marcverd@amazon.com>
Date: Tue, 27 Feb 2024 18:33:28 +0800
Subject: [PATCH] Enable sacct on login nodes

---
 .../base-config/setup_mariadb_accounting.sh                 | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh b/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
index eea7b590..030966fe 100644
--- a/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
+++ b/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
@@ -68,12 +68,14 @@ create_slurmdbd_config() {
 # Append the accounting settings to accounting.conf, this file is empty by default and included into
 # slurm.conf. This is required for Slurm to enable accounting.
 add_accounting_to_slurm_config() {
-  cat >> $SLURM_ACCOUNTING_CONFIG_FILE << EOL
+    # `hostname -i` gave us "hostname: Name or service not known". So let's parse slurm.conf.
+    DBD_HOST=$(awk -F'[=(]' '/^SlurmctldHost=/ { print $NF }' /opt/slurm/etc/slurm.conf | tr -d ')')
+    cat >> $SLURM_ACCOUNTING_CONFIG_FILE << EOL
 # ACCOUNTING
 JobAcctGatherType=jobacct_gather/linux
 JobAcctGatherFrequency=30
 AccountingStorageType=accounting_storage/slurmdbd
-AccountingStorageHost=localhost
+AccountingStorageHost=$DBD_HOST
 AccountingStoragePort=6819
 EOL
 }