AWS EMR中如何设置可被PySpark脚本访问的SSM环境变量

问题原因 #

bootstrap.sh在节点启动时执行，设置的环境变量仅存在于当前shell进程中。而Spark的Driver和Executor进程由YARN启动，不会自动继承这些shell环境变量，因此PySpark脚本无法直接获取。

解决方案 #

方案1：将环境变量写入Spark的spark-env.sh配置文件 #

Spark启动时会加载spark-env.sh中的环境变量，因此可以在bootstrap.sh中将解析后的密钥追加到该文件中。修改你的bootstrap.sh脚本如下：

SECRET_NAME="/myapp/dev/secrets"
SECRETS_JSON=$(aws secretsmanager get-secret-value --secret-id $SECRET_NAME --query SecretString --output text)

# Parse the secrets and set them as environment variables
SPARK_ENV_FILE="/usr/lib/spark/conf/spark-env.sh"
for key in $(echo "$SECRETS_JSON" | jq -r "keys[]"); do
  value=$(echo "$SECRETS_JSON" | jq -r ".[$key] // empty" | sed 's/"/\\"/g')
  echo "$value"
  if [ ! -z "$value" ]; then
    export "$key"="$value"
    # 追加到spark-env.sh，确保Spark进程能加载
    echo "export $key=\"$value\"" | sudo tee -a $SPARK_ENV_FILE
  fi
done

方案2：配置Spark的spark-defaults.conf传递环境变量 #

通过spark.driverEnv和spark.executorEnv前缀配置，让Driver和Executor进程都能获取到环境变量。修改bootstrap.sh脚本：

SECRET_NAME="/myapp/dev/secrets"
SECRETS_JSON=$(aws secretsmanager get-secret-value --secret-id $SECRET_NAME --query SecretString --output text)

# Parse the secrets and set them as environment variables
SPARK_DEFAULTS_FILE="/usr/lib/spark/conf/spark-defaults.conf"
for key in $(echo "$SECRETS_JSON" | jq -r "keys[]"); do
  value=$(echo "$SECRETS_JSON" | jq -r ".[$key] // empty" | sed 's/"/\\"/g')
  echo "$value"
  if [ ! -z "$value" ]; then
    export "$key"="$value"
    # 配置Driver和Executor的环境变量
    echo "spark.driverEnv.$key=$value" | sudo tee -a $SPARK_DEFAULTS_FILE
    echo "spark.executorEnv.$key=$value" | sudo tee -a $SPARK_DEFAULTS_FILE
  fi
done

方案3：在PySpark脚本中直接读取Secret Manager（备选） #

如果不想修改集群配置，也可以在PySpark脚本中使用boto3直接从Secret Manager获取密钥（需确保EMR实例角色拥有secretsmanager:GetSecretValue权限）：

import boto3
import json
import os

def get_secrets():
    secret_name = "/myapp/dev/secrets"
    client = boto3.client('secretsmanager', region_name='your-region')
    response = client.get_secret_value(SecretId=secret_name)
    secrets = json.loads(response['SecretString'])
    # 将密钥设置为当前进程的环境变量
    for key, value in secrets.items():
        os.environ[key] = value

# 调用函数加载密钥
get_secrets()

# 现在可以访问环境变量
postgres_url = os.environ.get("POSTGRES_URL")

注意事项 #

• 确保bootstrap.sh执行时拥有sudo权限，因为Spark配置文件通常属于root用户。
• 验证EMR实例角色是否具备访问Secret Manager的权限，否则aws secretsmanager命令会失败。
• 修改配置文件后无需手动重启Spark服务，EMR节点启动时会自动加载这些配置。

内容的提问来源于stack exchange，提问作者Vivek Mishra