Python Bootstrapper.set_yarn_spark_resource_configの例

コード例 #1

ファイル: test_bootstrapper.py プロジェクト: shunsunsun/sagemaker-spark-container

def test_set_yarn_spark_resource_config_fallback(
    patched_virtual_memory,
    patched_cpu_count,
    patched_yarn_config,
    patched_spark_config,
    default_bootstrapper: Bootstrapper,
) -> None:
    mocked_virtual_memory_total = PropertyMock(return_value=123 * 1024 * 1024)
    type(patched_virtual_memory.return_value
         ).total = mocked_virtual_memory_total
    patched_cpu_count.return_value = 456

    default_bootstrapper.load_processing_job_config = MagicMock(
        return_value=None)
    default_bootstrapper.load_instance_type_info = MagicMock(return_value=None)
    default_bootstrapper.get_yarn_spark_resource_config = MagicMock(
        return_value=(patched_yarn_config, patched_spark_config))

    default_bootstrapper.set_yarn_spark_resource_config()

    patched_virtual_memory.assert_called_once()
    mocked_virtual_memory_total.assert_called_once()
    patched_cpu_count.assert_called_once()

    default_bootstrapper.load_processing_job_config.assert_called_once()
    default_bootstrapper.load_instance_type_info.assert_called_once()
    default_bootstrapper.get_yarn_spark_resource_config.assert_called_once_with(
        1, 123, 456)
    patched_yarn_config.write_config.assert_called_once()
    patched_spark_config.write_config.assert_called_once()

コード例 #2

ファイル: test_bootstrapper.py プロジェクト: Satish615/sagemaker-spark-container

def test_set_yarn_spark_resource_config(
    patched_yarn_config, patched_spark_config, default_bootstrapper: Bootstrapper
) -> None:
    processing_job_config = {
        "ProcessingResources": {"ClusterConfig": {"InstanceType": "foo.xbar", "InstanceCount": 123}}
    }
    instance_type_info = {"foo.xbar": {"MemoryInfo": {"SizeInMiB": 456}, "VCpuInfo": {"DefaultVCpus": 789}}}
    default_bootstrapper.load_processing_job_config = MagicMock(return_value=processing_job_config)
    default_bootstrapper.load_instance_type_info = MagicMock(return_value=instance_type_info)
    default_bootstrapper.get_yarn_spark_resource_config = MagicMock(
        return_value=(patched_yarn_config, patched_spark_config)
    )

    default_bootstrapper.set_yarn_spark_resource_config()

    default_bootstrapper.load_processing_job_config.assert_called_once()
    default_bootstrapper.load_instance_type_info.assert_called_once()
    default_bootstrapper.get_yarn_spark_resource_config.assert_called_once_with(123, 456, 789)
    patched_yarn_config.write_config.assert_called_once()
    patched_spark_config.write_config.assert_called_once()