[ci] raydepsets: refactor depsets - using depset map #55555
Conversation
elliot-barn
commented
Aug 12, 2025
elliot-barn
commented
- remove depsets from the config
- move build arg sets to the workspace (config is pure utility)
- replace get depset with get depset by id
- using a depset map to fetch and retrieve the depset and build arg combination (Using a tuple of both as the id)
Signed-off-by: elliot-barn <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the handling of build argument sets for dependency sets, which is a good improvement for modularity. The key idea of using a (depset_name, build_arg_set_name) tuple as a unique identifier is sound. My review focuses on a correctness issue in how depsets are retrieved, a few opportunities to improve code clarity and documentation, and a suggestion to optimize performance in the new depset expansion logic.
| build_arg_set = next( | ||
| ( | ||
| build_arg_set | ||
| for build_arg_set in build_arg_sets | ||
| if build_arg_set.name == build_arg_set_name | ||
| ), | ||
| None, | ||
| ) |
There was a problem hiding this comment.
This search for build_arg_set is inefficient. It iterates through the build_arg_sets list for every name in build_arg_set_matrix, leading to O(N*M) complexity. For better performance, you could convert the build_arg_sets list to a dictionary before this function is called, allowing for an O(1) lookup here.
This would involve:
- In
Workspace.build_depset_map, create adictfromself.build_arg_sets. - Pass this
dicttoConfig.to_depset_mapand then to this function. - Here, you could replace this block with a simple dictionary lookup, like
build_arg_set = build_arg_sets.get(build_arg_set_name).
There was a problem hiding this comment.
Will do this in another PR
There was a problem hiding this comment.
done here: #55557
Signed-off-by: elliot-barn <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Elliot Barnwell <[email protected]>
Signed-off-by: elliot-barn <[email protected]>
Signed-off-by: elliot-barn <[email protected]>
elliot-barn
changed the title
ci/raydepsets/cli.py
Outdated
| depset.name, operation="subset", depset=depset | ||
| self.build_graph.add_node(depset_id, operation="subset", depset=depset) | ||
| self.build_graph.add_edge( | ||
| (depset.source_depset, depset.build_arg_set_name), depset_id |
There was a problem hiding this comment.
what does this add edge mean now? I am a bit confused on a tuple.
ci/raydepsets/cli.py
Outdated
| try: | ||
| return self.depset_map[depset_id] | ||
| except KeyError: | ||
| raise KeyError( | ||
| f"Dependency set {depset_id[0]} with build arg set {depset_id[1]} not found" | ||
| ) |
There was a problem hiding this comment.
use get and check if it is None? don't use try catch.
| self.build_arg_sets = Config.parse_build_arg_sets( | ||
| data.get("build_arg_sets", []) | ||
| ) | ||
| return Config.to_depset_map(data, self.build_arg_sets) |
There was a problem hiding this comment.
de-indent so that the file can be closed?
| @staticmethod | ||
| def from_dict(data: dict) -> "Config": | ||
| build_arg_sets = Config.parse_build_arg_sets(data.get("build_arg_sets", [])) | ||
| def to_depset_map( |
There was a problem hiding this comment.
can we still keep the source config struct, and convert to map in another step function? having the source config around can be useful for debugging and error reporting.
like now this Config class is just empty and useless now.
Signed-off-by: elliot-barn <[email protected]>
