Create torchx.specs.overlays module (#1171)

Summary:

Create `torchx.specs.overlays` module to standardize how scheduler request overlays work.

Currently overlays are JSON (`dict[str, object]`) that follows the same structure as the scheduler's request type.

Overlays are stored as metadata in two places:

1. `AppDef.metadata`: job level overlays
2. `Role.metadata`: gang-level overlays

Overlays follow the rules below:

1. value types must be the same as the scheduler's request type at the same attribute.
2. dicts get upserted
3. nested dicts get recursively overlayed
4. lists get appended, no deduplication attempts are made
5. lists are NOT recursed, that is, lists for which elements are containers are not recursively overlayed, they are simply treated as primitives.
6. primitive attributes are replaced

Differential Revision: D88545380

Author: kiukchung

docs/source/specs.rst

@@ -46,7 +46,6 @@ Macros
 .. autoclass:: macros
    :members:
 
-
 Run Configs
 --------------
 .. autoclass:: runopts
@@ -78,6 +77,11 @@ Mounts
 .. autoclass:: DeviceMount
    :members:
 
+Overlays
+------------
+
+.. automodule:: torchx.specs.overlays
+   :members:
 
 Component Linter
 -----------------

torchx/specs/overlays.py

@@ -0,0 +1,105 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+"""
+Overlays are JSON structs applied to :py:class:`~torchx.specs.AppDef` and :py:class:`~torchx.specs.Role`
+to specify attributes of the scheduler's submit-job request that are not currently representable
+as attributes of :py:class:`~torchx.specs.AppDef` and :py:class:`~torchx.specs.Role`.
+
+For end-uses, here are a few use-cases of overlays:
+
+1. A new version of the scheduler has concepts/features that have not yet been added to TorchX.
+2. A bespoke internal scheduler has custom features that do not generalize hence not in TorchX.
+3. Re-using a pre-built ``AppDef`` but need to make a small change to the resulting scheduler request.
+
+And for scheduler authors:
+
+1. Scheduler setting needs to be applied to a ``Role``, which makes it hard to add as ``runopts``
+   since ``runopts`` apply at the ``AppDef`` level.
+2. Scheduler setting cannot be represented naturally as the types supported by ``runopts``.
+3. Exposing the setting as a ``runopts`` obfuscates things.
+
+See :py:func:`~torchx.specs.overlays.apply_overlay` for rules on how overlays are applied.
+"""
+
+from typing import Any
+
+Json = dict[str, Any]
+
+
+def apply_overlay(base: Json, overlay: Json) -> None:
+    """Applies ``overlay`` on ``base``.
+
+    .. note:: this function mutates the ``base``!
+
+    Overlays follow these rules:
+
+    1. Dicts, upsert key, value in base with the ones in overlay.
+    2. Nested dicts, overlay recursively.
+    3. Lists, append the overlay values to the base values.
+    3. Nested lists DO NOT append recursively.
+    4. Primitives (bool, str, int, float), replace base with the value in overlay.
+
+    .. doctest::
+        from torchx.specs.overlays import apply_overlay
+
+        base = {
+            "scheduler": {"policy": "default"},
+            "resources": {"limits": {"cpu": "500m"}},
+            "tolerations": [{"key": "gpu"}],
+            "nodeSelectorTerms": [
+                [{"matchExpressions": []}]
+            ],
+            "maxPods": 110,
+        }
+        overlay = {
+            "scheduler": {"policy": "binpacking"},
+            "resources": {"limits": {"memory": "1Gi"}},
+            "tolerations": [{"key": "spot"}],
+            "nodeSelectorTerms": [
+                [{"matchExpressions": [{"key": "disk"}]}]
+            ],
+            "maxPods": 250,
+        }
+
+        apply_overlay(base, overlay)
+
+        assert {
+            "scheduler": {"policy": "binpacking"},
+            "resources": {"limits": {"cpu": "500m", "memory": "1Gi"}},
+            "tolerations": [{"key": "gpu"}, {"key": "spot"}],
+            "nodeSelectorTerms": [
+                [{"matchExpressions": []}],
+                [{"matchExpressions": [{"key": "disk"}]}],
+            ],
+            "maxPods": 250,
+        } == base
+
+    """
+
+    def assert_type_equal(key: str, o1: object, o2: object) -> None:
+        o1_type = type(o1)
+        o2_type = type(o2)
+        assert (
+            o1_type == o2_type
+        ), f"Type mismatch for attr: `{key}`. {o1_type.__qualname__} != {o2_type.__qualname__}"
+
+    for key, overlay_value in overlay.items():
+        if key in base:
+            base_value = base[key]
+
+            assert_type_equal(key, base_value, overlay_value)
+
+            if isinstance(base_value, dict) and isinstance(overlay_value, dict):
+                apply_overlay(base_value, overlay_value)
+            elif isinstance(base_value, list) and isinstance(overlay_value, list):
+                base_value.extend(overlay_value)
+            else:
+                base[key] = overlay_value
+        else:
+            base[key] = overlay_value

torchx/specs/test/overlays_test.py

@@ -0,0 +1,67 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+import unittest
+
+from torchx.specs.overlays import apply_overlay
+
+
+class OverlaysTest(unittest.TestCase):
+    def test_apply_overlay_empty_base(self) -> None:
+        base = {}
+        overlay = {
+            "a0": {"b": "c"},
+            "a1": ["b", "c"],
+            "a2": "b",
+        }
+        apply_overlay(base, overlay)
+        self.assertDictEqual(base, overlay)
+
+    def test_apply_overlay_dict_attr(self) -> None:
+        base = {"a0": {"d": "e"}}
+        overlay = {"a0": {"b": "c"}}
+
+        apply_overlay(base, overlay)
+        self.assertDictEqual(
+            base,
+            {
+                "a0": {
+                    "b": "c",
+                    "d": "e",
+                },
+            },
+        )
+
+        base = {"a0": {"b": "d"}}
+        apply_overlay(base, overlay)
+        self.assertDictEqual(base, {"a0": {"b": "c"}})
+
+    def test_apply_overlay_list_attr(self) -> None:
+        base = {"a0": []}
+        overlay = {"a0": ["b", "c"]}
+
+        apply_overlay(base, overlay)
+        self.assertDictEqual(base, {"a0": ["b", "c"]})
+
+        base = {"a0": ["1", "b", "2"]}
+        apply_overlay(base, overlay)
+        # lists simply append - they do not dedup
+        self.assertDictEqual(base, {"a0": ["1", "b", "2", "b", "c"]})
+
+        base = {"a0": ["1", ["2", "3"]]}
+        overlay = {"a0": ["b", ["c", "d"]]}
+        apply_overlay(base, overlay)
+        # lists simply append - they do recusively apply
+        self.assertDictEqual(base, {"a0": ["1", ["2", "3"], "b", ["c", "d"]]})
+
+    def test_overlay_type_mismatch(self) -> None:
+
+        with self.assertRaises(AssertionError):
+            apply_overlay({"a": [1, 2]}, {"a": "b"})
+
+        with self.assertRaises(AssertionError):
+            apply_overlay({"a": {"b": 1}}, {"a": {"b": "c"}})