From e663c42b45d96166336cfb197765f65a4b6d73a5 Mon Sep 17 00:00:00 2001
From: ErinWeisbart <54687786+ErinWeisbart@users.noreply.github.com>
Date: Fri, 12 Jul 2024 15:06:35 -0700
Subject: [PATCH] support assumed roles and users

---
 worker/run-worker.sh | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/worker/run-worker.sh b/worker/run-worker.sh
index 9aef692..57e7a00 100644
--- a/worker/run-worker.sh
+++ b/worker/run-worker.sh
@@ -2,13 +2,16 @@
 echo "${BASH_VERSION}"
 echo "Region $AWS_REGION"
 echo "Queue $SQS_QUEUE_URL"
-if [[ -z "$AWS_BUCKET" ]]
-then
+if [[ -z "$SOURCE_BUCKET" ]]; then
   SOURCE_BUCKET=$AWS_BUCKET
 fi
 echo "Source Bucket $SOURCE_BUCKET"
 
+mkdir -p /home/ubuntu/bucket
+mkdir -p /home/ubuntu/local_output
+
 # 1. CONFIGURE AWS CLI
+echo "Configuring AWS CLI"
 aws configure set default.region $AWS_REGION
 MY_INSTANCE_ID=$(curl http://169.254.169.254/latest/meta-data/instance-id)
 echo "Instance ID $MY_INSTANCE_ID"
@@ -19,20 +22,24 @@ aws ec2 create-tags --resources $VOL_0_ID --tags Key=Name,Value=${APP_NAME}Worke
 VOL_1_ID=$(aws ec2 describe-instance-attribute --instance-id $MY_INSTANCE_ID --attribute blockDeviceMapping --output text --query BlockDeviceMappings[1].Ebs.[VolumeId])
 aws ec2 create-tags --resources $VOL_1_ID --tags Key=Name,Value=${APP_NAME}Worker
 
-
 # 2. MOUNT S3 
-echo $AWS_ACCESS_KEY_ID:$AWS_SECRET_ACCESS_KEY > /credentials.txt
-chmod 600 /credentials.txt
-mkdir -p /home/ubuntu/bucket
-mkdir -p /home/ubuntu/local_output
-stdbuf -o0 s3fs $AWS_BUCKET /home/ubuntu/bucket -o passwd_file=/credentials.txt 
-
+echo "Mounting S3 using S3FS"
+if [[ -z "$AWS_ACCESS_KEY_ID" ]]; then
+    echo "Using role credentials to mount S3"
+    s3fs $SOURCE_BUCKET /home/ubuntu/bucket -o iam_role
+else
+    echo "Using user credentials to mount S3"
+    echo $AWS_ACCESS_KEY_ID:$AWS_SECRET_ACCESS_KEY > /credentials.txt
+    chmod 600 /credentials.txt
+    s3fs $SOURCE_BUCKET /home/ubuntu/bucket -o passwd_file=/credentials.txt
+fi
 
 # 3. SET UP ALARMS
+echo "Setting up instance metric alarms"
 aws cloudwatch put-metric-alarm --alarm-name ${APP_NAME}_${MY_INSTANCE_ID} --alarm-actions arn:aws:swf:${AWS_REGION}:${OWNER_ID}:action/actions/AWS_EC2.InstanceId.Terminate/1.0 --statistic Maximum --period 60 --threshold 1 --comparison-operator LessThanThreshold --metric-name CPUUtilization --namespace AWS/EC2 --evaluation-periods 15 --dimensions "Name=InstanceId,Value=${MY_INSTANCE_ID}"
 
 # 4. RUN VM STAT MONITOR
-
+echo "Setting up instance monitor"
 python3.8 instance-monitor.py &
 
 # 5. UPDATE AND/OR INSTALL PLUGINS
@@ -49,6 +56,7 @@ if [[ ${INSTALL_REQUIREMENTS} == 'True' ]]; then
 fi 
 
 # 6. RUN CP WORKERS
+echo "Starting workers"
 for((k=0; k<$DOCKER_CORES; k++)); do
     python3.8 cp-worker.py |& tee $k.out &
     sleep $SECONDS_TO_START