Provide option to specify cpu and memory limits in project config YAML.

Summary:
`bazel.cpus` - Maximum number of CPUs available to the container. Also, #jobs = 2x #cpus
`bazel.mem`  - Maximum memory available on a container, in MB. Container is expected to issue OOMs beyond that.

Test Plan: Unit

Reviewers: kylec

Reviewed By: kylec

Subscribers: changesbot, naphat

Differential Revision: https://tails.corp.dropbox.com/D223645

changes/buildsteps/default.py

@@ -14,7 +14,7 @@
 from changes.artifacts.xunit import XunitHandler
 from changes.buildsteps.base import BuildStep, LXCConfig
 from changes.config import db
-from changes.constants import Cause, Status
+from changes.constants import Cause, Status, DEFAULT_CPUS, DEFAULT_MEMORY_MB
 from changes.db.utils import get_or_create
 from changes.jobs.sync_job_step import sync_job_step
 from changes.models.command import CommandType, FutureCommand
@@ -68,7 +68,7 @@ class DefaultBuildStep(BuildStep):
     # - teardown_commands
     def __init__(self, commands=None, repo_path=None, path=None, env=None,
                  artifacts=DEFAULT_ARTIFACTS, release=DEFAULT_RELEASE,
-                 max_executors=10, cpus=4, memory=8 * 1024, clean=True,
+                 max_executors=10, cpus=DEFAULT_CPUS, memory=DEFAULT_MEMORY_MB, clean=True,
                  debug_config=None, test_stats_from=None, cluster=None,
                  other_repos=None, artifact_search_path=None,
                  **kwargs):

changes/config.py

@@ -416,6 +416,18 @@ def create_app(_read_config=True, **config):
     # name of default cluster to use for autogenerated jobs
     app.config['DEFAULT_CLUSTER'] = None
 
+    # Maximum number of cpus allowed for a bazel executor. Since we expose `bazel.cpus` to
+    # the user, this number needs to be bounded to avoid runaway resource allocation (by always
+    # allocating large chunks of resources, like 12-16 cores), and to avoid invalid configuration
+    # (like, requesting more cpus than available on a single slave, typically 32)
+    app.config['MAX_CPUS_PER_EXECUTOR'] = 16
+
+    # Minimum memory allowed per executor (in MB)
+    app.config['MIN_MEM_MB_PER_EXECUTOR'] = 1024
+
+    # Maximum memory allowed per executor (in MB)
+    app.config['MAX_MEM_MB_PER_EXECUTOR'] = 16384
+
     app.config.update(config)
     if _read_config:
         if os.environ.get('CHANGES_CONF'):

changes/constants.py

@@ -9,6 +9,12 @@
 # if not specified in changes.conf.py
 CUSTOM_CSS_FILE = "changes.less"
 
+# Default number of CPUs assigned to an executor
+DEFAULT_CPUS = 4
+
+# Default amount of RAM (in MBs) assigned to an executor
+DEFAULT_MEMORY_MB = 8 * 1024
+
 
 class Status(Enum):
     unknown = 0

changes/models/jobplan.py

@@ -198,16 +198,38 @@ def get_build_step_for_job(cls, job_id):
                 return jobplan, None
 
             bazel_exclude_tags = project_config['bazel.exclude-tags']
+            bazel_cpus = project_config['bazel.cpus']
+            if bazel_cpus < 1 or bazel_cpus > current_app.config['MAX_CPUS_PER_EXECUTOR']:
+                logging.error('Project config for project %s requests invalid number of CPUs: constraint 1 <= %d <= %d' % (
+                            job.project.slug,
+                            bazel_cpus,
+                            current_app.config['MAX_CPUS_PER_EXECUTOR']))
+                return jobplan, None
+
+            bazel_memory = project_config['bazel.mem']
+            if bazel_memory < current_app.config['MIN_MEM_MB_PER_EXECUTOR'] or \
+               bazel_memory > current_app.config['MAX_MEM_MB_PER_EXECUTOR']:
+                logging.error('Project config for project %s requests invalid memory requirements: constraint %d <= %d <= %d' % (
+                            job.project.slug,
+                            current_app.config['MIN_MEM_MB_PER_EXECUTOR'],
+                            bazel_memory,
+                            current_app.config['MAX_MEM_MB_PER_EXECUTOR']))
+                return jobplan, None
 
             implementation = LXCBuildStep(
                 cluster=current_app.config['DEFAULT_CLUSTER'],
                 commands=[
                     {'script': get_bazel_setup(), 'type': 'setup'},
                     {'script': sync_encap_pkgs(project_config), 'type': 'setup'},  # TODO(anupc): Make this optional
-                    {'script': collect_bazel_targets(project_config['bazel.targets'], bazel_exclude_tags), 'type': 'collect_tests'},
+                    {'script': collect_bazel_targets(
+                        bazel_targets=project_config['bazel.targets'],
+                        bazel_exclude_tags=bazel_exclude_tags,
+                        max_jobs=2 * bazel_cpus), 'type': 'collect_tests'},
                 ],
                 artifacts=['*.xml'],
                 artifact_search_path=current_app.config['BAZEL_TEST_OUTPUT_RELATIVE_PATH'],
+                cpus=bazel_cpus,
+                memory=bazel_memory,
             )
             return jobplan, implementation
 

changes/models/project.py

@@ -14,6 +14,7 @@
 from changes.db.types.guid import GUID
 from changes.db.types.enum import Enum
 from changes.utils.slugs import slugify
+from changes.constants import DEFAULT_CPUS, DEFAULT_MEMORY_MB
 
 
 class ProjectConfigError(Exception):
@@ -61,7 +62,9 @@ def get(cls, id):
 
     _default_config = {
         'build.file-blacklist': [],
-        'bazel.exclude-tags': []
+        'bazel.exclude-tags': ['manual'],  # Ignore tests with manual tag
+        'bazel.cpus': DEFAULT_CPUS,
+        'bazel.mem': DEFAULT_MEMORY_MB,
     }
 
     def get_config_path(self):

changes/utils/bazel_setup.py

@@ -48,7 +48,7 @@ def get_bazel_setup():
     )
 
 
-def collect_bazel_targets(bazel_targets, bazel_exclude_tags):
+def collect_bazel_targets(bazel_targets, bazel_exclude_tags, max_jobs):
     """Construct a command to query the Bazel dependency graph to expand bazel project
     config into a set of individual test targets.
 
@@ -58,8 +58,9 @@ def collect_bazel_targets(bazel_targets, bazel_exclude_tags):
       are additive, meaning a union of all targets matching any of the patterns will be
       returned.
     - exclude-tags: List of target tags. Targets matching any of these tags are not returned.
-      By default, this list is empty.  # TODO(anupc): Should we remove `manual` test targets?
+      By default, this list is empty.
     """
+    # type: (List[str], List[str], int) -> str
     package_dir = os.path.dirname(__file__)
     bazel_target_py = os.path.join(package_dir, "collect_bazel_targets.py")
 
@@ -90,7 +91,7 @@ def collect_bazel_targets(bazel_targets, bazel_exclude_tags):
             apt_spec=current_app.config['APT_SPEC'],
             bazel_apt_pkgs=' '.join(current_app.config['BAZEL_APT_PKGS']),
             bazel_targets=' + '.join(bazel_targets),
-            jsonify_script=jsonify_script.read(),
+            jsonify_script=jsonify_script.read() % dict(max_jobs=max_jobs),
             exclusion_subquery=exclusion_subquery,
         )
 

changes/utils/collect_bazel_targets.py

@@ -4,7 +4,7 @@
 
 targets = sys.stdin.read().splitlines()
 out = {
-    'cmd': '/usr/bin/bazel test {test_names}',
+    'cmd': '/usr/bin/bazel test --jobs=%(max_jobs)s {test_names}',
     'tests': targets,
 }
 json.dump(out, sys.stdout)

tests/changes/models/test_job.py

@@ -12,22 +12,7 @@
 
 
 class AutogeneratedJobTest(TestCase):
-    @mock.patch('changes.models.project.Project.get_config')
-    def test_autogenerated_commands(self, get_config):
-        get_config.return_value = {
-            'bazel.targets': [
-                '//aa/bb/cc/...',
-                '//aa/abc/...',
-            ],
-            'bazel.dependencies': {
-                'encap': [
-                    'package1',
-                    'pkg-2',
-                ]
-            },
-            'bazel.exclude-tags': [],  # Default, specified in models/project.py
-        }
-
+    def _create_job_and_jobplan(self):
         current_app.config['APT_SPEC'] = 'deb http://example.com/debian distribution component1'
         current_app.config['ENCAP_RSYNC_URL'] = 'rsync://example.com/encap/'
         current_app.config['BAZEL_APT_PKGS'] = ['bazel']
@@ -42,8 +27,27 @@ def test_autogenerated_commands(self, get_config):
         build = self.create_build(project)
         job = self.create_job(build)
         jobplan = self.create_job_plan(job, plan)
+        return job
+
+    @mock.patch('changes.models.project.Project.get_config')
+    def test_autogenerated_commands(self, get_config):
+        get_config.return_value = {
+            'bazel.targets': [
+                '//aa/bb/cc/...',
+                '//aa/abc/...',
+            ],
+            'bazel.dependencies': {
+                'encap': [
+                    'package1',
+                    'pkg-2',
+                ]
+            },
+            'bazel.exclude-tags': [],
+            'bazel.cpus': 4,           # Default
+            'bazel.mem': 8192,         # Default
+        }
 
-        _, implementation = JobPlan.get_build_step_for_job(job.id)
+        _, implementation = JobPlan.get_build_step_for_job(self._create_job_and_jobplan().id)
 
         bazel_setup_expected = """#!/bin/bash -eux
 # Clean up any existing apt sources
@@ -82,7 +86,7 @@ def test_autogenerated_commands(self, get_config):
 
 targets = sys.stdin.read().splitlines()
 out = {
-    'cmd': '/usr/bin/bazel test {test_names}',
+    'cmd': '/usr/bin/bazel test --jobs=8 {test_names}',
     'tests': targets,
 }
 json.dump(out, sys.stdout)
@@ -114,23 +118,14 @@ def test_autogenerated_commands_with_exclusions(self, get_config):
                 'flaky',
                 'another_tag',
             ],
+            'bazel.cpus': 2,
+            'bazel.mem': 1234,
         }
 
-        current_app.config['APT_SPEC'] = 'deb http://example.com/debian distribution component1'
-        current_app.config['BAZEL_APT_PKGS'] = ['bazel']
+        _, implementation = JobPlan.get_build_step_for_job(self._create_job_and_jobplan().id)
 
-        project = self.create_project()
-        plan = self.create_plan(project)
-        option = self.create_option(
-            item_id=plan.id,
-            name='bazel.autogenerate',
-            value='1',
-        )
-        build = self.create_build(project)
-        job = self.create_job(build)
-        jobplan = self.create_job_plan(job, plan)
-
-        _, implementation = JobPlan.get_build_step_for_job(job.id)
+        assert implementation.resources['cpus'] == 2
+        assert implementation.resources['mem'] == 1234
 
         collect_tests_expected = """#!/bin/bash -eu
 # Clean up any existing apt sources
@@ -150,7 +145,7 @@ def test_autogenerated_commands_with_exclusions(self, get_config):
 
 targets = sys.stdin.read().splitlines()
 out = {
-    'cmd': '/usr/bin/bazel test {test_names}',
+    'cmd': '/usr/bin/bazel test --jobs=4 {test_names}',
     'tests': targets,
 }
 json.dump(out, sys.stdout)
@@ -163,3 +158,66 @@ def test_autogenerated_commands_with_exclusions(self, get_config):
         assert implementation.commands[1].type == CommandType.setup
         assert implementation.commands[2].type == CommandType.collect_tests
         assert implementation.commands[2].script == collect_tests_expected
+
+    @mock.patch('changes.models.project.Project.get_config')
+    def test_invalid_cpus(self, get_config):
+        get_config.return_value = {
+            'bazel.targets': [
+                '//aa/bb/cc/...',
+                '//aa/abc/...',
+            ],
+            'bazel.exclude-tags': [],
+            'bazel.cpus': 0,           # 0 CPUs is not valid
+            'bazel.mem': 8192,
+        }
+
+        _, implementation = JobPlan.get_build_step_for_job(self._create_job_and_jobplan().id)
+        assert implementation is None
+
+        get_config.return_value = {
+            'bazel.targets': [
+                '//aa/bb/cc/...',
+                '//aa/abc/...',
+            ],
+            'bazel.exclude-tags': [],
+            'bazel.cpus': 2,           # Too many
+            'bazel.mem': 8192,
+        }
+
+        current_app.config['MAX_CPUS_PER_EXECUTOR'] = 1
+        _, implementation = JobPlan.get_build_step_for_job(self._create_job_and_jobplan().id)
+        assert implementation is None
+
+    @mock.patch('changes.models.project.Project.get_config')
+    def test_invalid_mems(self, get_config):
+        get_config.return_value = {
+            'bazel.targets': [
+                '//aa/bb/cc/...',
+                '//aa/abc/...',
+            ],
+            'bazel.exclude-tags': [],
+            'bazel.cpus': 1,
+            'bazel.mem': 1025,         # Too high
+        }
+
+        current_app.config['MIN_MEM_MB_PER_EXECUTOR'] = 1
+        current_app.config['MAX_MEM_MB_PER_EXECUTOR'] = 10
+        _, implementation = JobPlan.get_build_step_for_job(self._create_job_and_jobplan().id)
+
+        assert implementation is None
+
+        get_config.return_value = {
+            'bazel.targets': [
+                '//aa/bb/cc/...',
+                '//aa/abc/...',
+            ],
+            'bazel.exclude-tags': [],
+            'bazel.cpus': 1,
+            'bazel.mem': 1025,         # Too low
+        }
+
+        current_app.config['MIN_MEM_MB_PER_EXECUTOR'] = 2000
+        current_app.config['MAX_MEM_MB_PER_EXECUTOR'] = 3000
+        _, implementation = JobPlan.get_build_step_for_job(self._create_job_and_jobplan().id)
+
+        assert implementation is None