[deps] raydepsets import all config (2/3)

ci/raydepsets/BUILD.bazel

@@ -70,6 +70,7 @@ py_test(
     srcs = ["tests/test_workspace.py"],
     data = [
         "tests/test_data/test.depsets.yaml",
+        "tests/test_data/test2.depsets.yaml",
     ],
     tags = [
         "ci_unit",

ci/raydepsets/tests/test_cli.py

@@ -276,6 +276,7 @@ def test_check_if_subset_exists(self):
                 output="requirements_compiled_general.txt",
                 append_flags=[],
                 override_flags=[],
+                config_name="test.depsets.yaml",
             )
             with self.assertRaises(RuntimeError):
                 manager.check_subset_exists(
@@ -437,6 +438,7 @@ def test_build_graph_bad_operation(self):
                 operation="invalid_op",
                 requirements=["requirements_test.txt"],
                 output="requirements_compiled_invalid_op.txt",
+                config_name="test.depsets.yaml",
             )
             _overwrite_config_file(tmpdir, depset)
             with self.assertRaises(ValueError):
@@ -608,6 +610,7 @@ def test_copy_lock_files_to_temp_dir(self):
                 constraints=["requirement_constraints_test.txt"],
                 requirements=["requirements_test.txt"],
                 output="requirements_compiled_test.txt",
+                config_name="test.depsets.yaml",
             )
             _overwrite_config_file(tmpdir, depset)
             manager = _create_test_manager(tmpdir, check=True)
@@ -632,6 +635,7 @@ def test_diff_lock_files_out_of_date(self):
                 constraints=["requirement_constraints_test.txt"],
                 requirements=["requirements_test.txt"],
                 output="requirements_compiled_test.txt",
+                config_name="test.depsets.yaml",
             )
             _overwrite_config_file(tmpdir, depset)
             manager = _create_test_manager(tmpdir, check=True)
@@ -666,6 +670,7 @@ def test_diff_lock_files_up_to_date(self):
                 constraints=["requirement_constraints_test.txt"],
                 requirements=["requirements_test.txt"],
                 output="requirements_compiled_test.txt",
+                config_name="test.depsets.yaml",
             )
             _overwrite_config_file(tmpdir, depset)
             manager = _create_test_manager(tmpdir, check=True)

ci/raydepsets/tests/test_data/test2.depsets.yaml

@@ -0,0 +1,28 @@
+build_arg_sets:
+  py311_cpu:
+      CUDA_VERSION: cpu
+      PYTHON_VERSION: py311
+
+depsets:
+  - name: other_config_depset
+    operation: expand
+    depsets:
+      - expand_general_depset__${PYTHON_VERSION}_${CUDA_VERSION}
+      - expanded_depset__${PYTHON_VERSION}_${CUDA_VERSION}
+    output: requirements_compiled_other_config.txt
+    build_arg_sets:
+      - py311_cpu
+  - name: pre_hook_test_depset
+    operation: compile
+    requirements:
+      - requirements_test.txt
+    output: requirements_compiled_pre_hook.txt
+    pre_hooks:
+      - pre-hook-test.sh
+  - name: pre_hook_invalid_test_depset
+    operation: compile
+    requirements:
+      - requirements_test.txt
+    output: requirements_compiled_pre_hook_invalid.txt
+    pre_hooks:
+      - pre-hook-error-test.sh

ci/raydepsets/tests/test_workspace.py

@@ -18,15 +18,14 @@ def test_parse_build_arg_sets():
     with tempfile.TemporaryDirectory() as tmpdir:
         copy_data_to_tmpdir(tmpdir)
         workspace = Workspace(dir=tmpdir)
-        config = workspace.load_config(path=Path(tmpdir) / "test.depsets.yaml")
-        assert config.build_arg_sets["py311_cpu"].build_args == {
-            "CUDA_VERSION": "cpu",
-            "PYTHON_VERSION": "py311",
-        }
-        assert config.build_arg_sets["py311_cuda128"].build_args == {
-            "CUDA_VERSION": 128,
-            "PYTHON_VERSION": "py311",
-        }
+        config = workspace.load_config(config_path=Path(tmpdir) / "test.depsets.yaml")
+        assert "general_depset__py311_cpu" in [depset.name for depset in config.depsets]
+        assert "build_args_test_depset__py311_cpu" in [
+            depset.name for depset in config.depsets
+        ]
+        assert "expanded_depset__py311_cpu" in [
+            depset.name for depset in config.depsets
+        ]
 
 
 def test_substitute_build_args():
@@ -65,17 +64,78 @@ def test_invalid_build_arg_set():
             )
         with pytest.raises(KeyError):
             workspace = Workspace(dir=tmpdir)
-            workspace.load_config(path=Path(tmpdir) / "test.depsets.yaml")
+            workspace.load_config(config_path=Path(tmpdir) / "test.depsets.yaml")
 
 
 def test_parse_pre_hooks():
     with tempfile.TemporaryDirectory() as tmpdir:
         copy_data_to_tmpdir(tmpdir)
         workspace = Workspace(dir=tmpdir)
-        config = workspace.load_config(path=Path(tmpdir) / "test.depsets.yaml")
+        config = workspace.load_config(config_path=Path(tmpdir) / "test2.depsets.yaml")
         pre_hook_depset = get_depset_by_name(config.depsets, "pre_hook_test_depset")
         assert pre_hook_depset.pre_hooks == ["pre-hook-test.sh"]
 
 
+def test_load_first_config():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        copy_data_to_tmpdir(tmpdir)
+        workspace = Workspace(dir=tmpdir)
+        config = workspace.load_config(config_path=Path(tmpdir) / "test.depsets.yaml")
+        assert config.depsets is not None
+        assert len(config.depsets) == 8
+
+
+def test_load_second_config():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        copy_data_to_tmpdir(tmpdir)
+        workspace = Workspace(dir=tmpdir)
+        config = workspace.load_config(config_path=Path(tmpdir) / "test2.depsets.yaml")
+        assert config.depsets is not None
+        assert len(config.depsets) == 3
+
+
+# load all configs should always load all depsets
+def test_load_all_configs_first_config():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        copy_data_to_tmpdir(tmpdir)
+        workspace = Workspace(dir=tmpdir)
+        config = workspace.load_configs(config_path=Path(tmpdir) / "test.depsets.yaml")
+        assert config.depsets is not None
+        assert len(config.depsets) == 11
+
+
+# load all configs should always load all depsets
+def test_load_all_configs_second_config():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        copy_data_to_tmpdir(tmpdir)
+        workspace = Workspace(dir=tmpdir)
+        config = workspace.load_configs(config_path=Path(tmpdir) / "test2.depsets.yaml")
+        assert config.depsets is not None
+        assert len(config.depsets) == 11
+
+
+def test_merge_configs():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        copy_data_to_tmpdir(tmpdir)
+        workspace = Workspace(dir=tmpdir)
+        config = workspace.load_config(config_path=Path(tmpdir) / "test.depsets.yaml")
+        config2 = workspace.load_config(config_path=Path(tmpdir) / "test2.depsets.yaml")
+        merged_config = workspace.merge_configs([config, config2])
+        assert merged_config.depsets is not None
+        assert len(merged_config.depsets) == 11
+
+
+def test_get_configs_dir():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        copy_data_to_tmpdir(tmpdir)
+        workspace = Workspace(dir=tmpdir)
+        configs_dir = workspace.get_configs_dir(
+            configs_path=Path(tmpdir) / "test.depsets.yaml"
+        )
+        assert len(configs_dir) == 2
+        assert f"{tmpdir}/test.depsets.yaml" in configs_dir
+        assert f"{tmpdir}/test2.depsets.yaml" in configs_dir
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))

ci/raydepsets/workspace.py

@@ -16,6 +16,7 @@ class Depset:
     name: str
     operation: str
     output: str
+    config_name: str
     constraints: Optional[List[str]] = None
     override_flags: Optional[List[str]] = None
     append_flags: Optional[List[str]] = None
@@ -40,7 +41,7 @@ def _substitute_build_args(obj: Any, build_arg_set: BuildArgSet):
         return obj
 
 
-def _dict_to_depset(depset: dict) -> Depset:
+def _dict_to_depset(depset: dict, config_name: str) -> Depset:
     return Depset(
         name=depset.get("name"),
         requirements=depset.get("requirements", []),
@@ -53,16 +54,16 @@ def _dict_to_depset(depset: dict) -> Depset:
         append_flags=depset.get("append_flags", []),
         pre_hooks=depset.get("pre_hooks", []),
         packages=depset.get("packages", []),
+        config_name=config_name,
     )
 
 
 @dataclass
 class Config:
     depsets: List[Depset] = field(default_factory=list)
-    build_arg_sets: Dict[str, BuildArgSet] = field(default_factory=dict)
 
     @classmethod
-    def from_dict(cls, data: dict) -> "Config":
+    def from_dict(cls, data: dict, config_name: str) -> "Config":
         build_arg_sets = cls.parse_build_arg_sets(data.get("build_arg_sets", {}))
         raw_depsets = data.get("depsets", [])
         depsets = []
@@ -75,10 +76,10 @@ def from_dict(cls, data: dict) -> "Config":
                     if build_arg_set is None:
                         raise KeyError(f"Build arg set {build_arg_set_key} not found")
                     depset_yaml = _substitute_build_args(depset, build_arg_set)
-                    depsets.append(_dict_to_depset(depset_yaml))
+                    depsets.append(_dict_to_depset(depset_yaml, config_name))
             else:
-                depsets.append(_dict_to_depset(depset))
-        return Config(depsets=depsets, build_arg_sets=build_arg_sets)
+                depsets.append(_dict_to_depset(depset, config_name))
+        return Config(depsets=depsets)
 
     @staticmethod
     def parse_build_arg_sets(build_arg_sets: Dict[str, dict]) -> Dict[str, BuildArgSet]:
@@ -98,7 +99,29 @@ def __init__(self, dir: str = None):
         if self.dir is None:
             raise RuntimeError("BUILD_WORKSPACE_DIRECTORY is not set")
 
-    def load_config(self, path: str) -> Config:
-        with open(os.path.join(self.dir, path), "r") as f:
+    def load_configs(self, config_path: str = None) -> Config:
+        merged_configs = self.merge_configs(self.get_all_configs(config_path))
+        return merged_configs
+
+    def get_all_configs(self, config_path: str = None) -> List[Config]:
+        return [self.load_config(path) for path in self.get_configs_dir(config_path)]
+
+    def get_configs_dir(self, configs_path: str) -> List[str]:
+        configs_dir = os.path.dirname(os.path.join(self.dir, configs_path))
+        return [
+            os.path.join(self.dir, configs_dir, path)
+            for path in os.listdir(os.path.join(self.dir, configs_dir))
+            if path.endswith(".depsets.yaml")
+        ]
+
+    def load_config(self, config_path: str) -> Config:
+        with open(os.path.join(self.dir, config_path), "r") as f:
             data = yaml.safe_load(f)
-            return Config.from_dict(data)
+            config_name = os.path.basename(config_path)
+            config = Config.from_dict(data, config_name)
+        return config
+
+    def merge_configs(self, configs: List[Config]) -> Config:
+        return Config(
+            depsets=[depset for config in configs for depset in config.depsets]
+        )