instance data IMDSv2 fix for spark public hostname

nchammas · Sep 20, 2023 · e3c394d · e3c394d
1 parent 74a6fad
commit e3c394d
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/flintrock/templates/spark/conf/spark-env.sh b/flintrock/templates/spark/conf/spark-env.sh
@@ -14,7 +14,20 @@ export HADOOP_CONF_DIR="$HOME/hadoop/conf"
 
 # TODO: Make this non-EC2-specific.
 # Bind Spark's web UIs to this machine's public EC2 hostname
-export SPARK_PUBLIC_DNS="$(curl --silent http://169.254.169.254/latest/meta-data/public-hostname)"
+spark_public_hostname="$(curl --silent http://169.254.169.254/latest/meta-data/public-hostname)" #IMDSv1 check
+if [[ -z "$spark_public_hostname" ]] || [[ "$spark_public_hostname" == *"DOCTYPE"*"html"*"head"*"body"* ]]
+then
+      TOKEN="$(curl --silent -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600")"
+      spark_public_hostname="$(curl --silent -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169.254/latest/meta-data/public-hostname)" #IMDSv2 check
+      if [[ -z "$spark_public_hostname" ]] || [[ "$spark_public_hostname" == *"DOCTYPE"*"html"*"head"*"body"* ]]
+      then
+          true #skip setting SPARK_PUBLIC_DNS
+      else
+          export SPARK_PUBLIC_DNS="$spark_public_hostname"
+      fi
+else
+      export SPARK_PUBLIC_DNS="$spark_public_hostname"
+fi
 
 # TODO: Set a high ulimit for large shuffles
 # Need to find a way to do this, since "sudo ulimit..." doesn't fly.