diff --git a/docs/source/specs.rst b/docs/source/specs.rst index 6fd1ea2c4..e3cba7c59 100644 --- a/docs/source/specs.rst +++ b/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 ----------------- diff --git a/torchx/specs/overlays.py b/torchx/specs/overlays.py new file mode 100644 index 000000000..d35e971a7 --- /dev/null +++ b/torchx/specs/overlays.py @@ -0,0 +1,106 @@ +# 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. + 4. Nested lists DO NOT append recursively. + 5. 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 diff --git a/torchx/specs/test/overlays_test.py b/torchx/specs/test/overlays_test.py new file mode 100644 index 000000000..7ebecbd80 --- /dev/null +++ b/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 NOT 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"}})