Skip to main content
 首页 » 编程设计

python之在本地和 Yarn 模式下运行 PySpark 的参数化 pytest fixture

2024年10月24日27leader

我认为这是一个相当常见的用例。我有为 PySpark 编写的函数,我想在本地和 Yarn 模式下进行测试。 pytest fixtures 对此非常有用——您可以简单地编写一个在两种模式下运行的参数化 fixture。

@pytest.fixture(scope="session", 
                params=[pytest.mark.spark_local('local'), 
                        pytest.mark.spark_yarn('yarn')]) 
def spark_context(request): 
    if request.param == 'local': 
        conf = (SparkConf() 
                .setMaster("local[2]") 
                .setAppName("pytest-pyspark-local-testing") 
                ) 
    elif request.param == 'yarn': 
        conf = (SparkConf() 
                .setMaster("yarn-client") 
                .setAppName("pytest-pyspark-yarn-testing") 
                .set("spark.executor.memory", "1g") 
                .set("spark.executor.instances", 2) 
                ) 
    request.addfinalizer(lambda: sc.stop()) 
 
    sc = SparkContext(conf=conf) 
    return sc 

问题是您不能同时启动多个 Spark 上下文。如果您尝试运行这个固定装置,您将得到一个异常。任何想法如何解决这一问题?也许范围?

我使用的另一种 hack 是标记参数化并单独运行它们:

py.test -m spark_local 
... 
py.test -m spark_yarn 
... 

有没有人遇到过这个问题或者想过如何通过一次调用 py.test 来完成这个问题?

请您参考如下方法:

原来有 an issue在本地模式和 yarn 模式下创建 SparkContext。以下代码似乎适用于在 yarn 和本地模式下运行测试。只需简单地切换顺序即可!

@pytest.fixture(scope="session", 
                params=[pytest.mark.spark_yarn('yarn'), 
                        pytest.mark.spark_local('local')]) 
def spark_context(request): 
    if request.param == 'local': 
        conf = (SparkConf() 
                .setMaster("local[2]") 
                .setAppName("pytest-pyspark-local-testing") 
                ) 
    elif request.param == 'yarn': 
        conf = (SparkConf() 
                .setMaster("yarn-client") 
                .setAppName("pytest-pyspark-yarn-testing") 
                .set("spark.executor.memory", "1g") 
                .set("spark.executor.instances", 2) 
                ) 
    request.addfinalizer(lambda: sc.stop()) 
 
    sc = SparkContext(conf=conf) 
    return sc