Repository: codalab/codalab-worksheets

Issue:
[FEATURE REQUEST] workers that _only_ run bundles with a matching tag
**Is your feature request related to a problem? Please describe.**

Currently, tagged codalab workers will run bundles without a `request_queue` set. Would it be possible to have some flag that forces a worker to *only* run bundles with a particular tag?

Rationale: Sometimes, I have a lot of bundles running that are untagged. I've already developed these bundles, and I'm sure they'll work—I'm just letting them run (think 50+ staged jobs).

In parallel, though, I'd often like to set up the next round of experiments, which sometimes means running stuff on codalab and seeing if it works. However, the development cycle is super slow since it's hard to get a worker—they're all busy running the untagged jobs waiting in the queue (even if I give the worker a special tag). It'd be nice if, in this case, these dev jobs were able to instantly land somewhere (assuming there isn't another one before it).

**Describe the solution you'd like**

Not sure what a solution would look like here—maybe an extra flag on `cl-worker` that locks it to running bundles with a particular tag? If you use the flag with an untagged worker, it throws an error or something.

Discussion:
percyliang: Thanks for being proactive about this!  It would be nice if we can avoid making schema changes, so let's discuss in #2097.
nelson-liu: > One thing i was thinking was whether in the future we wanted a more powerful tag/matching system, where each worker / run had a set of tags, some of which were exclusive versus not. Not sure if there's a pressing real use case for this, but perhaps worth keeping in mind.

Not sure how the old slurm worker functioned, but a more full-featured tag matching system could be useful there—you could have a worker with tags for each of its slurm attributes (e.g., the gpu type), and bundles could request the combination of attributes they want (via `||` and `&&`).
nelson-liu: > Another thing: It'd be good to write somewhere, perhaps as a comment block or in the documentation, a description of how one typically uses --tag-exclusive.

I stuck this in the execution docs, unsure if it was the right place—feel free to let me know if it's not.

Patch:
diff --git a/alembic/versions/2020032904_add_tag_exclusive_workers_9d0f1ffb18e9.py b/alembic/versions/2020032904_add_tag_exclusive_workers_9d0f1ffb18e9.py
new file mode 100644
index 00000000..711a4364
--- /dev/null
+++ b/alembic/versions/2020032904_add_tag_exclusive_workers_9d0f1ffb18e9.py
@@ -0,0 +1,22 @@
+"""Add tag_exclusive workers
+
+Revision ID: 9d0f1ffb18e9
+Revises: 75d4288ae265
+Create Date: 2020-03-29 04:57:59.567422
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = '9d0f1ffb18e9'
+down_revision = '75d4288ae265'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    op.add_column('worker', sa.Column('tag_exclusive', sa.Boolean(), nullable=False, server_default=sa.false()))
+
+
+def downgrade():
+    op.drop_column('worker', 'tag_exclusive')
diff --git a/codalab/lib/bundle_cli.py b/codalab/lib/bundle_cli.py
index e4ee06f9..8db89033 100644
--- a/codalab/lib/bundle_cli.py
+++ b/codalab/lib/bundle_cli.py
@@ -1063,6 +1063,7 @@ class BundleCLI(object):
             'tag',
             'runs',
             'shared_file_system',
+            'tag_exclusive',
         ]
 
         data = []
@@ -1081,6 +1082,7 @@ class BundleCLI(object):
                     'tag': worker['tag'],
                     'runs': ",".join([uuid[0:8] for uuid in worker['run_uuids']]),
                     'shared_file_system': worker['shared_file_system'],
+                    'tag_exclusive': worker['tag_exclusive'],
                 }
             )
 
diff --git a/codalab/model/tables.py b/codalab/model/tables.py
index f117ae30..c0cbd29c 100644
--- a/codalab/model/tables.py
+++ b/codalab/model/tables.py
@@ -316,6 +316,9 @@ worker = Table(
     Column(
         'shared_file_system', Boolean, nullable=False
     ),  # Whether the worker and the server have a shared filesystem.
+    Column(
+        'tag_exclusive', Boolean, nullable=False
+    ),  # Whether worker runs bundles if and only if they match tags.
 )
 
 # Store information about all sockets currently allocated to each worker.
diff --git a/codalab/model/worker_model.py b/codalab/model/worker_model.py
index bf40ec51..a4408b67 100644
--- a/codalab/model/worker_model.py
+++ b/codalab/model/worker_model.py
@@ -47,6 +47,7 @@ class WorkerModel(object):
         free_disk_bytes,
         dependencies,
         shared_file_system,
+        tag_exclusive,
     ):
         """
         Adds the worker to the database, if not yet there. Returns the socket ID
@@ -61,6 +62,7 @@ class WorkerModel(object):
                 'free_disk_bytes': free_disk_bytes,
                 'checkin_time': datetime.datetime.utcnow(),
                 'shared_file_system': shared_file_system,
+                'tag_exclusive': tag_exclusive,
             }
             existing_row = conn.execute(
                 cl_worker.select().where(
@@ -185,6 +187,7 @@ class WorkerModel(object):
                 'dependencies': row.dependencies
                 and self._deserialize_dependencies(row.dependencies),
                 'shared_file_system': row.shared_file_system,
+                'tag_exclusive': row.tag_exclusive,
             }
             for row in worker_rows
         }
diff --git a/codalab/rest/workers.py b/codalab/rest/workers.py
index 98729b73..03e1ee23 100644
--- a/codalab/rest/workers.py
+++ b/codalab/rest/workers.py
@@ -34,6 +34,7 @@ def checkin(worker_id):
         request.json.get("free_disk_bytes"),
         request.json["dependencies"],
         request.json.get("shared_file_system", False),
+        request.json.get("tag_exclusive", False),
     )
 
     for run in request.json["runs"]:
diff --git a/codalab/server/bundle_manager.py b/codalab/server/bundle_manager.py
index bc81a861..0af04f11 100644
--- a/codalab/server/bundle_manager.py
+++ b/codalab/server/bundle_manager.py
@@ -388,6 +388,15 @@ class BundleManager(object):
         # Filter by tag.
         if bundle.metadata.request_queue:
             workers_list = self._get_matched_workers(bundle.metadata.request_queue, workers_list)
+        else:
+            # The bundle is untagged, so we want to keep workers that are not
+            # tag-exclusive or don't have a tag defined.
+            # (removing workers that are tag_exclusive and have tags defined).
+            workers_list = [
+                worker
+                for worker in workers_list
+                if not worker['tag_exclusive'] or not worker['tag']
+            ]
 
         # Filter by CPUs.
         workers_list = [
@@ -428,12 +437,14 @@ class BundleManager(object):
                 num_available_deps = len(needed_deps & deps)
 
             # Subject to the worker meeting the resource requirements of the bundle, we also want to:
-            # 1. prioritize workers with fewer GPUs (including zero).
-            # 2. prioritize workers that have more bundle dependencies.
-            # 3. prioritize workers with fewer CPUs.
-            # 4. prioritize workers with fewer running jobs.
-            # 5. break ties randomly by a random seed.
+            # 1. prioritize workers that are tag-exclusive.
+            # 2. prioritize workers with fewer GPUs (including zero).
+            # 3. prioritize workers that have more bundle dependencies.
+            # 4. prioritize workers with fewer CPUs.
+            # 5. prioritize workers with fewer running jobs.
+            # 6. break ties randomly by a random seed.
             return (
+                not worker['tag_exclusive'],
                 worker['gpus'],
                 -num_available_deps,
                 worker['cpus'],
diff --git a/codalab/worker/main.py b/codalab/worker/main.py
index 9b4add82..98b3bb4f 100644
--- a/codalab/worker/main.py
+++ b/codalab/worker/main.py
@@ -120,6 +120,11 @@ def parse_args():
         action='store_true',
         help='To be used when the server and the worker share the bundle store on their filesystems.',
     )
+    parser.add_argument(
+        '--tag-exclusive',
+        action='store_true',
+        help='To be used when the worker should only run bundles that match the worker\'s tag.',
+    )
     return parser.parse_args()
 
 
@@ -206,6 +211,7 @@ def main():
         args.idle_seconds,
         bundle_service,
         args.shared_file_system,
+        args.tag_exclusive,
         docker_runtime=docker_runtime,
         docker_network_prefix=args.network_prefix,
     )
diff --git a/codalab/worker/worker.py b/codalab/worker/worker.py
index d9f9271f..fe778929 100644
--- a/codalab/worker/worker.py
+++ b/codalab/worker/worker.py
@@ -58,6 +58,7 @@ class Worker:
         idle_seconds,  # type: int
         bundle_service,  # type: BundleServiceClient
         shared_file_system,  # type: bool
+        tag_exclusive,  # type: bool
         docker_runtime=docker_utils.DEFAULT_RUNTIME,  # type: str
         docker_network_prefix='codalab_worker_network',  # type: str
     ):
@@ -73,6 +74,7 @@ class Worker:
 
         self.id = worker_id
         self.tag = tag
+        self.tag_exclusive = tag_exclusive
 
         self.work_dir = work_dir
         self.local_bundles_dir = local_bundles_dir
@@ -239,6 +241,7 @@ class Worker:
             'hostname': socket.gethostname(),
             'runs': [run.as_dict for run in self.all_runs],
             'shared_file_system': self.shared_file_system,
+            'tag_exclusive': self.tag_exclusive,
         }
         try:
             response = self.bundle_service.checkin(self.id, request)
diff --git a/docs/Execution.md b/docs/Execution.md
index 03e23cfd..ca6d4354 100644
--- a/docs/Execution.md
+++ b/docs/Execution.md
@@ -135,3 +135,17 @@ Check the bundle's `stdout`, and you should see something similar to before:
     ...
 
 And that's all.  Congrats!
+
+## Frequently asked questions
+
+### Can I reserve a worker for a particular tag?
+
+Yes, you can use the `--tag-exclusive` flag of `cl worker` (along with a
+`--tag`) to prevent your worker from running untagged bundles, since
+tag-exclusive workers only run bundles that match their tag.
+
+To be concrete, a worker started with `cl worker --tag debug --tag-exclusive`
+will _only_ run bundles with `--request-queue debug`. This is especially
+useful for development, since it lets you set off a machine for fast-turnaround
+debugging of CodaLab bundles, even if you have many untagged jobs queued---just
+make sure to set `--request-queue debug`.


Test patch:
diff --git a/test_cli.py b/test_cli.py
index a469cbde..92bab5a9 100644
--- a/test_cli.py
+++ b/test_cli.py
@@ -1689,13 +1689,14 @@ def test(ctx):
             'tag',
             'runs',
             'shared_file_system',
+            'tag_exclusive',
         ],
         header,
     )
 
     # Check number of not null values. First 7 columns should be not null. Column "tag" and "runs" could be empty.
     worker_info = lines[2].split()
-    check_equals(True, len(worker_info) >= 7)
+    check_equals(True, len(worker_info) >= 8)
 
 
 @TestModule.register('rest1')
diff --git a/tests/server/bundle_manager_test.py b/tests/server/bundle_manager_test.py
index 4b894c3a..c753cd1e 100644
--- a/tests/server/bundle_manager_test.py
+++ b/tests/server/bundle_manager_test.py
@@ -48,6 +48,7 @@ class BundleManagerTest(unittest.TestCase):
                 'run_uuids': [1, 2],
                 'dependencies': [(1, ''), (2, '')],
                 'shared_file_system': False,
+                'tag_exclusive': False,
             },
             {
                 'worker_id': 1,
@@ -58,6 +59,7 @@ class BundleManagerTest(unittest.TestCase):
                 'run_uuids': [1, 2, 3],
                 'dependencies': [(1, ''), (2, ''), (3, ''), (4, '')],
                 'shared_file_system': False,
+                'tag_exclusive': False,
             },
             {
                 'worker_id': 2,
@@ -68,6 +70,7 @@ class BundleManagerTest(unittest.TestCase):
                 'run_uuids': [],
                 'dependencies': [(1, ''), (2, ''), (3, ''), (4, '')],
                 'shared_file_system': False,
+                'tag_exclusive': False,
             },
             {
                 'worker_id': 3,
@@ -78,6 +81,7 @@ class BundleManagerTest(unittest.TestCase):
                 'run_uuids': [1],
                 'dependencies': [(1, '')],
                 'shared_file_system': False,
+                'tag_exclusive': False,
             },
             {
                 'worker_id': 4,
@@ -88,6 +92,18 @@ class BundleManagerTest(unittest.TestCase):
                 'run_uuids': [],
                 'dependencies': [],
                 'shared_file_system': False,
+                'tag_exclusive': False,
+            },
+            {
+                'worker_id': 5,
+                'cpus': 6,
+                'gpus': 0,
+                'memory_bytes': 2 * 1000,
+                'tag': 'worker_X',
+                'run_uuids': [],
+                'dependencies': [],
+                'shared_file_system': False,
+                'tag_exclusive': True,
             },
         ]
         return workers_list
@@ -103,6 +119,26 @@ class BundleManagerTest(unittest.TestCase):
         self.assertEqual(sorted_workers_list[1]['worker_id'], 1)
         self.assertEqual(sorted_workers_list[-1]['worker_id'], 0)
 
+    def test__filter_and_sort_workers_tag_exclusive(self):
+        # Only non-tag_exclusive workers should appear in the returned sorted worker list.
+        sorted_workers_list = self.bundle_manager._filter_and_sort_workers(
+            self.workers_list, self.bundle, self.bundle_resources
+        )
+        self.assertEqual(len(sorted_workers_list), 5)
+        for worker in sorted_workers_list:
+            self.assertEqual(worker['tag_exclusive'], False)
+
+    def test__filter_and_sort_workers_tag_exclusive_priority(self):
+        # All other things being equal, tag_exclusive workers
+        # should appear in the top from the returned sorted workers list.
+        self.bundle.metadata.request_queue = "tag=worker_X"
+        sorted_workers_list = self.bundle_manager._filter_and_sort_workers(
+            self.workers_list, self.bundle, self.bundle_resources
+        )
+        self.assertEqual(len(sorted_workers_list), 2)
+        self.assertEqual(sorted_workers_list[0]['worker_id'], 5)
+        self.assertEqual(sorted_workers_list[1]['worker_id'], 4)
+
     def test__filter_and_sort_workers_cpus(self):
         # CPU workers should appear in the top from the returning sorted worker list
         self.bundle_resources.cpus = 1
@@ -117,8 +153,9 @@ class BundleManagerTest(unittest.TestCase):
         matched_workers = BundleManager._get_matched_workers(
             self.bundle.metadata.request_queue, self.workers_list
         )
-        self.assertEqual(len(matched_workers), 1)
+        self.assertEqual(len(matched_workers), 2)
         self.assertEqual(matched_workers[0]['worker_id'], 4)
+        self.assertEqual(matched_workers[1]['worker_id'], 5)
 
     def test__get_matched_workers_with_bad_formatted_tag(self):
         self.bundle.metadata.request_queue = "tag="
@@ -132,8 +169,9 @@ class BundleManagerTest(unittest.TestCase):
         matched_workers = BundleManager._get_matched_workers(
             self.bundle.metadata.request_queue, self.workers_list
         )
-        self.assertEqual(len(matched_workers), 1)
+        self.assertEqual(len(matched_workers), 2)
         self.assertEqual(matched_workers[0]['worker_id'], 4)
+        self.assertEqual(matched_workers[1]['worker_id'], 5)
 
     def test__get_matched_workers_not_exist_tag(self):
         self.bundle.metadata.request_queue = "worker_Y"