如何在Step Functions中传递环境变量给PySpark程序？

你之前尝试的yarn-env配置仅作用于Yarn守护进程，无法被PySpark的Driver和Executor进程读取，所以os.getenv()拿不到变量。以下是几种可行的方案：

方案1：通过Spark配置直接注入环境变量 #

针对Glue Job或EMR上的Spark任务，在Step Functions调用任务时，通过spark.driverEnv和spark.executorEnv配置项传递环境变量，确保Driver和所有Executor都能读取到：

针对AWS Glue Job的配置示例 #

在Step Functions的StartJobRun任务参数中添加Spark配置：

{
  "StartJobRun": {
    "Type": "Task",
    "Resource": "arn:aws:states:::glue:startJobRun.sync",
    "Parameters": {
      "JobName": "YourPySparkJob",
      "Arguments": {
        "--conf": "spark.driverEnv.ENVIRONMENT=prod",
        "--conf": "spark.executorEnv.ENVIRONMENT=prod"
      }
    }
  }
}

针对EMR Spark任务的配置示例 #

在Step Functions的AddStep任务中，给spark-submit命令添加配置参数：

{
  "AddStep": {
    "Type": "Task",
    "Resource": "arn:aws:states:::elasticmapreduce:addStep.sync",
    "Parameters": {
      "JobFlowId": "your-emr-cluster-id",
      "Steps": [
        {
          "Name": "RunPySparkJob",
          "ActionOnFailure": "CONTINUE",
          "HadoopJarStep": {
            "Jar": "command-runner.jar",
            "Args": [
              "spark-submit",
              "--conf", "spark.driverEnv.ENVIRONMENT=prod",
              "--conf", "spark.executorEnv.ENVIRONMENT=prod",
              "s3://your-bucket/path/to/job.py"
            ]
          }
        }
      ]
    }
  }
}

配置完成后，PySpark代码中直接用os.getenv("ENVIRONMENT")即可获取变量值。

方案2：通过EMR Bootstrap Action设置全局环境变量 #

如果你的PySpark任务都运行在同一个EMR集群上，可以通过Bootstrap Action在集群初始化时全局设置环境变量，所有任务无需单独配置：

1. 编写一个Shell脚本（比如s3://your-bucket/bootstrap/env-setup.sh）：

#!/bin/bash
# 给Spark进程添加环境变量
echo "export ENVIRONMENT=prod" >> /etc/spark/conf/spark-env.sh
# 给Yarn相关进程添加（可选）
echo "export ENVIRONMENT=prod" >> /etc/hadoop/conf/hadoop-env.sh

2. 在Step Functions创建EMR集群的CreateCluster任务中，添加Bootstrap Action配置：

{
  "CreateCluster": {
    "Type": "Task",
    "Resource": "arn:aws:states:::elasticmapreduce:createCluster.sync",
    "Parameters": {
      "Name": "YourEMRCluster",
      "ReleaseLabel": "emr-6.10.0",
      "Applications": [{"Name": "Spark"}],
      "BootstrapActions": [
        {
          "Name": "SetGlobalEnvVars",
          "ScriptBootstrapAction": {
            "Path": "s3://your-bucket/bootstrap/env-setup.sh"
          }
        }
      ],
      "Instances": {
        "InstanceGroups": [
          {
            "Name": "Master",
            "InstanceRole": "MASTER",
            "InstanceType": "m5.xlarge",
            "InstanceCount": 1
          },
          {
            "Name": "Core",
            "InstanceRole": "CORE",
            "InstanceType": "m5.xlarge",
            "InstanceCount": 2
          }
        ]
      }
    }
  }
}

这种方式会在集群所有节点上配置环境变量，所有运行在该集群的PySpark任务都能通过os.getenv()读取到。

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