Repository: codalab/codalab-worksheets

Issue:
Add the ability to specify a group for Worker Manager
Add the ability to specify a group for Worker Manager `--worker-group`. The new argument to specify a group for a worker is:

```
    parser.add_argument(
        '--group', default=None, help='Name of the group that can run jobs on this worker'
    )
```

Discussion:

Patch:
diff --git a/codalab/worker_manager/aws_batch_worker_manager.py b/codalab/worker_manager/aws_batch_worker_manager.py
index 104802de..ef5cb4d0 100644
--- a/codalab/worker_manager/aws_batch_worker_manager.py
+++ b/codalab/worker_manager/aws_batch_worker_manager.py
@@ -89,36 +89,7 @@ class AWSBatchWorkerManager(WorkerManager):
         )
         # This needs to be a unique directory since Batch jobs may share a host
         work_dir = os.path.join(work_dir_prefix, 'cl_worker_{}_work_dir'.format(worker_id))
-        worker_network_prefix = 'cl_worker_{}_network'.format(worker_id)
-        command = [
-            self.args.worker_executable,
-            '--server',
-            self.args.server,
-            '--verbose',
-            '--exit-when-idle',
-            '--idle-seconds',
-            str(self.args.worker_idle_seconds),
-            '--work-dir',
-            work_dir,
-            '--id',
-            worker_id,
-            '--network-prefix',
-            worker_network_prefix,
-        ]
-        if self.args.worker_tag:
-            command.extend(['--tag', self.args.worker_tag])
-        if self.args.worker_max_work_dir_size:
-            command.extend(['--max-work-dir-size', self.args.worker_max_work_dir_size])
-        if self.args.worker_delete_work_dir_on_exit:
-            command.extend(['--worker-delete-work-dir-on-exit'])
-        if self.args.worker_exit_after_num_runs and self.args.worker_exit_after_num_runs > 0:
-            command.extend(['--exit-after-num-runs', str(self.args.worker_exit_after_num_runs)])
-        if self.args.worker_exit_on_exception:
-            command.extend(['--exit-on-exception'])
-        if self.args.worker_pass_down_termination:
-            command.extend(['--pass-down-termination'])
-        if self.args.worker_tag_exclusive:
-            command.extend(['--tag-exclusive'])
+        command = self.build_command(worker_id, work_dir)
 
         # https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-batch-jobdefinition.html
         # Need to mount:
diff --git a/codalab/worker_manager/main.py b/codalab/worker_manager/main.py
index f3038ae0..8c07caf3 100644
--- a/codalab/worker_manager/main.py
+++ b/codalab/worker_manager/main.py
@@ -73,6 +73,11 @@ def main():
         action='store_true',
         help="If set, the CodaLab worker will only run bundles that match the worker\'s tag.",
     )
+    parser.add_argument(
+        '--worker-group',
+        type=str,
+        help="If set, the CodaLab worker will only run bundles for that group.",
+    )
     parser.add_argument(
         '--worker-executable', default="cl-worker", help="The CodaLab worker executable to run."
     )
diff --git a/codalab/worker_manager/slurm_batch_worker_manager.py b/codalab/worker_manager/slurm_batch_worker_manager.py
index 2d389461..a48a2b6e 100644
--- a/codalab/worker_manager/slurm_batch_worker_manager.py
+++ b/codalab/worker_manager/slurm_batch_worker_manager.py
@@ -236,8 +236,6 @@ class SlurmBatchWorkerManager(WorkerManager):
         :param worker_id: a string representing the worker id
         :return: the command to run on
         """
-        # This needs to be a unique directory since Batch jobs may share a host
-        worker_network_prefix = 'cl_worker_{}_network'.format(worker_id)
         # Codalab worker's work directory
         if self.args.worker_work_dir_prefix:
             work_dir_prefix = Path(self.args.worker_work_dir_prefix)
@@ -247,38 +245,12 @@ class SlurmBatchWorkerManager(WorkerManager):
         worker_work_dir = work_dir_prefix.joinpath(
             Path('{}-codalab-SlurmBatchWorkerManager-scratch'.format(self.username), worker_id)
         )
+        command = self.build_command(worker_id, str(worker_work_dir))
 
-        command = [
-            self.args.worker_executable,
-            '--server',
-            self.args.server,
-            '--verbose',
-            '--exit-when-idle',
-            '--idle-seconds',
-            str(self.args.worker_idle_seconds),
-            '--work-dir',
-            str(worker_work_dir),
-            '--id',
-            worker_id,
-            '--network-prefix',
-            worker_network_prefix,
-            # always set in the Slurm worker manager to ensure safe shutdown
-            '--pass-down-termination',
-        ]
-        if self.args.worker_tag:
-            command.extend(['--tag', self.args.worker_tag])
+        # --pass-down-termination should always be set for Slurm worker managers to ensure safe shutdown
+        command.append('--pass-down-termination')
         if self.args.password_file:
             command.extend(['--password-file', self.args.password_file])
-        if self.args.worker_exit_after_num_runs and self.args.worker_exit_after_num_runs > 0:
-            command.extend(['--exit-after-num-runs', str(self.args.worker_exit_after_num_runs)])
-        if self.args.worker_max_work_dir_size:
-            command.extend(['--max-work-dir-size', self.args.worker_max_work_dir_size])
-        if self.args.worker_delete_work_dir_on_exit:
-            command.extend(['--delete-work-dir-on-exit'])
-        if self.args.worker_exit_on_exception:
-            command.extend(['--exit-on-exception'])
-        if self.args.worker_tag_exclusive:
-            command.extend(['--tag-exclusive'])
 
         return command
 
diff --git a/codalab/worker_manager/worker_manager.py b/codalab/worker_manager/worker_manager.py
index fb7fb8a7..d7dc2c3e 100644
--- a/codalab/worker_manager/worker_manager.py
+++ b/codalab/worker_manager/worker_manager.py
@@ -82,6 +82,43 @@ class WorkerManager(object):
         """Start a new `WorkerJob`."""
         raise NotImplementedError
 
+    def build_command(self, worker_id, work_dir):
+        command = [
+            self.args.worker_executable,
+            '--server',
+            self.args.server,
+            '--verbose',
+            '--exit-when-idle',
+            '--idle-seconds',
+            str(self.args.worker_idle_seconds),
+            '--work-dir',
+            work_dir,
+            '--id',
+            worker_id,
+            '--network-prefix',
+            'cl_worker_{}_network'.format(worker_id),
+        ]
+
+        # Additional optional arguments
+        if self.args.worker_tag:
+            command.extend(['--tag', self.args.worker_tag])
+        if self.args.worker_group:
+            command.extend(['--group', self.args.worker_group])
+        if self.args.worker_exit_after_num_runs and self.args.worker_exit_after_num_runs > 0:
+            command.extend(['--exit-after-num-runs', str(self.args.worker_exit_after_num_runs)])
+        if self.args.worker_max_work_dir_size:
+            command.extend(['--max-work-dir-size', self.args.worker_max_work_dir_size])
+        if self.args.worker_delete_work_dir_on_exit:
+            command.extend(['--delete-work-dir-on-exit'])
+        if self.args.worker_exit_on_exception:
+            command.extend(['--exit-on-exception'])
+        if self.args.worker_tag_exclusive:
+            command.extend(['--tag-exclusive'])
+        if self.args.worker_pass_down_termination:
+            command.extend(['--pass-down-termination'])
+
+        return command
+
     def run_loop(self):
         while True:
             try:


Test patch:
diff --git a/tests/unit/worker_manager/__init__.py b/tests/unit/worker_manager/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/unit/worker_manager/slurm_batch_worker_manager_test.py b/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
new file mode 100644
index 00000000..841b2d95
--- /dev/null
+++ b/tests/unit/worker_manager/slurm_batch_worker_manager_test.py
@@ -0,0 +1,39 @@
+import unittest
+from types import SimpleNamespace
+
+from codalab.worker_manager.slurm_batch_worker_manager import SlurmBatchWorkerManager
+
+
+class SlurmBatchWorkerManagerTest(unittest.TestCase):
+    def test_base_command(self):
+        args = SimpleNamespace(
+            server="some_server",
+            user="some_user",
+            partition="some_partition",
+            worker_executable="cl-worker",
+            worker_idle_seconds="888",
+            worker_tag="some_tag",
+            worker_group="some_group",
+            worker_exit_after_num_runs=8,
+            worker_max_work_dir_size="88g",
+            worker_work_dir_prefix="/some/path",
+            worker_delete_work_dir_on_exit=False,
+            worker_exit_on_exception=False,
+            worker_tag_exclusive=False,
+            worker_pass_down_termination=False,
+            password_file=None,
+        )
+
+        worker_manager = SlurmBatchWorkerManager(args)
+        command = worker_manager.setup_codalab_worker("some_worker_id")
+
+        # --pass-down-termination should always be set for Slurm worker managers
+        self.assertTrue('--pass-down-termination' in command)
+
+        expected_command_str = (
+            "cl-worker --server some_server --verbose --exit-when-idle --idle-seconds 888 "
+            "--work-dir /some/path/some_user-codalab-SlurmBatchWorkerManager-scratch/some_worker_id "
+            "--id some_worker_id --network-prefix cl_worker_some_worker_id_network --tag some_tag "
+            "--group some_group --exit-after-num-runs 8 --max-work-dir-size 88g --pass-down-termination"
+        )
+        self.assertEqual(' '.join(command), expected_command_str)