Add cluster-manager=local support. (#83)

jsleight · web-flow · commit 44cbfb2f2597 · 2022-08-03T10:09:11.000-04:00
Add cluster-manager=local support.
diff --git a/service_configuration_lib/spark_config.py b/service_configuration_lib/spark_config.py
@@ -71,11 +71,26 @@
     'spark.kubernetes.executor.label.paasta.yelp.com/cluster',
 }
 K8S_AUTH_FOLDER = '/etc/pki/spark'
+K8S_BASE_VOLUMES: List[Dict[str, str]] = [
+    {'containerPath': K8S_AUTH_FOLDER, 'hostPath': K8S_AUTH_FOLDER, 'mode': 'RO'},
+    {'containerPath': '/etc/passwd', 'hostPath': '/etc/passwd', 'mode': 'RO'},
+    {'containerPath': '/etc/group', 'hostPath': '/etc/group', 'mode': 'RO'},
+]
 
 log = logging.Logger(__name__)
 log.setLevel(logging.INFO)
 
 
+SUPPORTED_CLUSTER_MANAGERS = ['mesos', 'kubernetes', 'local']
+
+
+class UnsupportedClusterManagerException(Exception):
+
+    def __init__(self, manager: str):
+        msg = f'Unsupported cluster manager {manager}, must be one of {SUPPORTED_CLUSTER_MANAGERS}'
+        super().__init__(msg)
+
+
 def _load_aws_credentials_from_yaml(yaml_file_path) -> Tuple[str, str, Optional[str]]:
     with open(yaml_file_path, 'r') as yaml_file:
         try:
@@ -179,9 +194,7 @@ def _get_k8s_docker_volumes_conf(
     env = {}
     mounted_volumes = set()
     k8s_volumes = volumes or []
-    k8s_volumes.append({'containerPath': K8S_AUTH_FOLDER, 'hostPath': K8S_AUTH_FOLDER, 'mode': 'RO'})
-    k8s_volumes.append({'containerPath': '/etc/passwd', 'hostPath': '/etc/passwd', 'mode': 'RO'})
-    k8s_volumes.append({'containerPath': '/etc/group', 'hostPath': '/etc/group', 'mode': 'RO'})
+    k8s_volumes.extend(K8S_BASE_VOLUMES)
     _get_k8s_volume = functools.partial(_get_k8s_volume_hostpath_dict, count=itertools.count())
 
     for volume in k8s_volumes:
@@ -454,6 +467,13 @@ def _adjust_spark_requested_resources(
             'spark.scheduler.maxRegisteredResourcesWaitingTime',
             str(waiting_time) + 'min',
         )
+    elif cluster_manager == 'local':
+        executor_instances = int(
+            user_spark_opts.setdefault('spark.executor.instances', str(DEFAULT_EXECUTOR_INSTANCES)),
+        )
+        max_cores = executor_instances * executor_cores
+    else:
+        raise UnsupportedClusterManagerException(cluster_manager)
 
     if max_cores < executor_cores:
         raise ValueError(f'Total number of cores {max_cores} is less than per-executor cores {executor_cores}')
@@ -649,6 +669,27 @@ def _get_k8s_spark_env(
     return spark_env
 
 
+def _get_local_spark_env(
+    paasta_cluster: str,
+    paasta_service: str,
+    paasta_instance: str,
+    volumes: Optional[List[Mapping[str, str]]],
+    num_threads: int = 4,
+) -> Dict[str, str]:
+    return {
+        'spark.master': f'local[{num_threads}]',
+        'spark.executorEnv.PAASTA_SERVICE': paasta_service,
+        'spark.executorEnv.PAASTA_INSTANCE': paasta_instance,
+        'spark.executorEnv.PAASTA_CLUSTER': paasta_cluster,
+        'spark.executorEnv.PAASTA_INSTANCE_TYPE': 'spark',
+        'spark.executorEnv.SPARK_EXECUTOR_DIRS': '/tmp',
+        # Adding k8s docker volume params is a bit of a hack.  PaasSTA
+        # looks at the spark config to find the volumes it needs to mount
+        # and so we are borrowing the spark k8s configs.
+        **_get_k8s_docker_volumes_conf(volumes),
+    }
+
+
 def _get_k8s_resource_name_limit_size_with_hash(name: str, limit: int = 63, suffix: int = 4) -> str:
     """ Returns `name` unchanged if it's length does not exceed the `limit`.
         Otherwise, returns truncated `name` with it's hash of size `suffix`
@@ -721,7 +762,7 @@ def get_spark_conf(
 ) -> Dict[str, str]:
     """Build spark config dict to run with spark on paasta
 
-    :param cluster_manager: which manager to use, value supported: [`mesos`, `kubernetes`]
+    :param cluster_manager: which manager to use, must be in SUPPORTED_CLUSTER_MANAGERS
     :param spark_app_base_name: the base name to create spark app, we will append port
         and time to make the app name unique for easier to separate the output. Note that
         this is noop if `spark_opts_from_env` have `spark.app.name` configured.
@@ -812,8 +853,15 @@ def get_spark_conf(
             extra_volumes,
             paasta_pool,
         ))
+    elif cluster_manager == 'local':
+        spark_conf.update(_get_local_spark_env(
+            paasta_cluster,
+            paasta_service,
+            paasta_instance,
+            extra_volumes,
+        ))
     else:
-        raise ValueError('Unknown resource_manager, should be either [mesos,kubernetes]')
+        raise UnsupportedClusterManagerException(cluster_manager)
 
     # configure dynamic resource allocation configs
     spark_conf = get_dra_configs(spark_conf)
diff --git a/setup.py b/setup.py
@@ -17,7 +17,7 @@
 
 setup(
     name='service-configuration-lib',
-    version='2.10.8',
+    version='2.11.0',
     provides=['service_configuration_lib'],
     description='Start, stop, and inspect Yelp SOA services',
     url='https://github.com/Yelp/service_configuration_lib',
diff --git a/tests/spark_config_test.py b/tests/spark_config_test.py
@@ -149,10 +149,13 @@ class TestGetSparkConf:
     docker_image = 'docker-dev.yelp.com/test-image'
     executor_cores = '10'
     spark_app_base_name = 'test_app_base_name'
-    extra_volumes = [{'hostPath': '/tmp', 'containerPath': '/tmp', 'mode': 'RO'}]
     default_mesos_leader = 'mesos://some-url.yelp.com:5050'
     aws_provider_key = 'spark.hadoop.fs.s3a.aws.credentials.provider'
 
+    @pytest.fixture
+    def base_volumes(self):
+        return [{'hostPath': '/tmp', 'containerPath': '/tmp', 'mode': 'RO'}]
+
     @pytest.fixture
     def mock_paasta_volumes(self, monkeypatch, tmpdir):
         files = [f'paasta{i + 1}' for i in range(2)]
@@ -934,7 +937,7 @@ def verify(output):
         return verify
 
     @pytest.mark.parametrize('use_temp_provider', [True, False])
-    def test_get_spark_conf_aws_session(self, use_temp_provider):
+    def test_get_spark_conf_aws_session(self, use_temp_provider, base_volumes):
         other_spark_opts = {'spark.driver.memory': '2g', 'spark.executor.memoryOverhead': '1024'}
         not_allowed_opts = {'spark.executorEnv.PAASTA_SERVICE': 'random-service'}
         user_spark_opts = {
@@ -954,7 +957,7 @@ def test_get_spark_conf_aws_session(self, use_temp_provider):
             paasta_service=self.service,
             paasta_instance=self.instance,
             docker_img=self.docker_image,
-            extra_volumes=self.extra_volumes,
+            extra_volumes=base_volumes,
             aws_creds=aws_creds,
             auto_set_temporary_credentials_provider=use_temp_provider,
         )
@@ -968,6 +971,7 @@ def test_get_spark_conf_mesos(
         self,
         user_spark_opts,
         spark_opts_from_env,
+        base_volumes,
         ui_port,
         with_secret,
         mesos_leader,
@@ -1009,7 +1013,7 @@ def test_get_spark_conf_mesos(
             paasta_service=self.service,
             paasta_instance=self.instance,
             docker_img=self.docker_image,
-            extra_volumes=self.extra_volumes,
+            extra_volumes=base_volumes,
             aws_creds=aws_creds,
             extra_docker_params=extra_docker_params,
             with_secret=with_secret,
@@ -1038,7 +1042,7 @@ def test_get_spark_conf_mesos(
         )
         assert len(set(output.keys()) - verified_keys) == 0
         mock_get_mesos_docker_volumes_conf.mocker.assert_called_once_with(
-            mock.ANY, self.extra_volumes, True,
+            mock.ANY, base_volumes, True,
         )
         mock_adjust_spark_requested_resources_mesos.mocker.assert_called_once_with(
             mock.ANY, 'mesos', self.pool,
@@ -1053,8 +1057,15 @@ def test_get_spark_conf_mesos(
         (warning_msg,), _ = mock_log.warning.call_args
         assert next(iter(not_allowed_opts.keys())) in warning_msg
 
+    def _get_k8s_base_volumes(self):
+        """Helper needed to allow tests to pass in github CI checks."""
+        return [
+            volume for volume in spark_config.K8S_BASE_VOLUMES
+            if os.path.exists(volume['containerPath'])
+        ]
+
     @pytest.fixture
-    def assert_kubernetes_conf(self):
+    def assert_kubernetes_conf(self, base_volumes):
         expected_output = {
             'spark.master': f'k8s://https://k8s.{self.cluster}.paasta:6443',
             'spark.executorEnv.PAASTA_SERVICE': self.service,
@@ -1079,8 +1090,10 @@ def assert_kubernetes_conf(self):
             'spark.kubernetes.executor.label.paasta.yelp.com/cluster': self.cluster,
             'spark.kubernetes.node.selector.yelp.com/pool': self.pool,
             'spark.kubernetes.executor.label.yelp.com/pool': self.pool,
+            'spark.logConf': 'true',
+            'spark.ui.showConsoleProgress': 'true',
         }
-        for i, volume in enumerate(self.extra_volumes):
+        for i, volume in enumerate(base_volumes + self._get_k8s_base_volumes()):
             expected_output[f'spark.kubernetes.executor.volumes.hostPath.{i}.mount.path'] = volume['containerPath']
             expected_output[f'spark.kubernetes.executor.volumes.hostPath.{i}.mount.readOnly'] = str(
                 volume['mode'] == 'RO',
@@ -1093,11 +1106,12 @@ def verify(output):
             return list(expected_output.keys())
         return verify
 
-    def tes_leaderst_get_spark_conf_kubernetes(
+    def test_leaders_get_spark_conf_kubernetes(
         self,
         user_spark_opts,
         spark_opts_from_env,
         ui_port,
+        base_volumes,
         mock_append_event_log_conf,
         mock_append_aws_credentials_conf,
         mock_append_sql_shuffle_partitions_conf,
@@ -1125,7 +1139,7 @@ def tes_leaderst_get_spark_conf_kubernetes(
             paasta_service=self.service,
             paasta_instance=self.instance,
             docker_img=self.docker_image,
-            extra_volumes=self.extra_volumes,
+            extra_volumes=base_volumes,
             aws_creds=aws_creds,
             spark_opts_from_env=spark_opts_from_env,
         )
@@ -1140,7 +1154,7 @@ def tes_leaderst_get_spark_conf_kubernetes(
             list(mock_append_aws_credentials_conf.return_value.keys()) +
             list(mock_append_sql_shuffle_partitions_conf.return_value.keys()),
         )
-        assert len(set(output.keys()) - verified_keys) == 0
+        assert set(output.keys()) == verified_keys
         mock_adjust_spark_requested_resources_kubernetes.mocker.assert_called_once_with(
             mock.ANY, 'kubernetes', self.pool,
         )
@@ -1152,6 +1166,82 @@ def tes_leaderst_get_spark_conf_kubernetes(
             mock.ANY,
         )
 
+    @pytest.fixture
+    def assert_local_conf(self, base_volumes):
+        expected_output = {
+            'spark.master': 'local[4]',
+            'spark.executorEnv.PAASTA_SERVICE': self.service,
+            'spark.executorEnv.PAASTA_INSTANCE': self.instance,
+            'spark.executorEnv.PAASTA_CLUSTER': self.cluster,
+            'spark.executorEnv.PAASTA_INSTANCE_TYPE': 'spark',
+            'spark.executorEnv.SPARK_EXECUTOR_DIRS': '/tmp',
+            'spark.logConf': 'true',
+            'spark.ui.showConsoleProgress': 'true',
+        }
+        for i, volume in enumerate(base_volumes + self._get_k8s_base_volumes()):
+            expected_output[f'spark.kubernetes.executor.volumes.hostPath.{i}.mount.path'] = volume['containerPath']
+            expected_output[f'spark.kubernetes.executor.volumes.hostPath.{i}.mount.readOnly'] = str(
+                volume['mode'] == 'RO',
+            ).lower()
+            expected_output[f'spark.kubernetes.executor.volumes.hostPath.{i}.options.path'] = volume['hostPath']
+
+        def verify(output):
+            for key, value in expected_output.items():
+                assert output[key] == value
+            return list(expected_output.keys())
+        return verify
+
+    def test_local_spark(
+        self,
+        user_spark_opts,
+        spark_opts_from_env,
+        ui_port,
+        base_volumes,
+        mock_append_event_log_conf,
+        mock_append_aws_credentials_conf,
+        mock_append_sql_shuffle_partitions_conf,
+        mock_adjust_spark_requested_resources_kubernetes,
+        mock_time,
+        assert_ui_port,
+        assert_app_name,
+        assert_local_conf,
+        mock_log,
+    ):
+        aws_creds = (None, None, None)
+        output = spark_config.get_spark_conf(
+            cluster_manager='local',
+            spark_app_base_name=self.spark_app_base_name,
+            user_spark_opts=user_spark_opts or {},
+            paasta_cluster=self.cluster,
+            paasta_pool=self.pool,
+            paasta_service=self.service,
+            paasta_instance=self.instance,
+            docker_img=self.docker_image,
+            extra_volumes=base_volumes,
+            aws_creds=aws_creds,
+            spark_opts_from_env=spark_opts_from_env,
+        )
+        verified_keys = set(
+            assert_ui_port(output) +
+            assert_app_name(output) +
+            assert_local_conf(output) +
+            list(mock_append_event_log_conf.return_value.keys()) +
+            list(mock_adjust_spark_requested_resources_kubernetes.return_value.keys()) +
+            list(mock_append_aws_credentials_conf.return_value.keys()) +
+            list(mock_append_sql_shuffle_partitions_conf.return_value.keys()),
+        )
+        assert set(output.keys()) == verified_keys
+        mock_append_event_log_conf.mocker.assert_called_once_with(
+            mock.ANY, *aws_creds,
+        )
+        mock_adjust_spark_requested_resources_kubernetes.mocker.assert_called_once_with(
+            mock.ANY, 'local', self.pool,
+        )
+        mock_append_aws_credentials_conf.mocker.assert_called_once_with(mock.ANY, *aws_creds)
+        mock_append_sql_shuffle_partitions_conf.mocker.assert_called_once_with(
+            mock.ANY,
+        )
+
     @pytest.mark.parametrize('reason', ['mesos_leader', 'mesos_secret'])
     def test_get_spark_conf_mesos_error(self, reason, monkeypatch, mock_request_mesos_leader):
         if reason == 'mesos_leader':
