wehub-resource-sync
596 lines
22 KiB
Python
596 lines
22 KiB
Python
import difflib
|
|
import os
|
|
import platform
|
|
import shlex
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import List, Optional
|
|
|
|
import click
|
|
import runfiles
|
|
from networkx import DiGraph, ancestors as networkx_ancestors, topological_sort
|
|
from pip_requirements_parser import RequirementsFile
|
|
|
|
from ci.raydepsets.workspace import Depset, Workspace
|
|
|
|
DEFAULT_UV_FLAGS = """
|
|
--no-header
|
|
--generate-hashes
|
|
--index-strategy unsafe-best-match
|
|
--no-strip-markers
|
|
--emit-index-url
|
|
--emit-find-links
|
|
--quiet
|
|
""".split()
|
|
|
|
|
|
@click.group(name="raydepsets")
|
|
def cli():
|
|
"""Manage Python dependency sets."""
|
|
|
|
|
|
@cli.command()
|
|
@click.argument("config_path", default="ci/raydepsets/configs/*.depsets.yaml")
|
|
@click.option(
|
|
"--workspace-dir",
|
|
default=None,
|
|
help="The path to the workspace directory. If not specified, $BUILD_WORKSPACE_DIRECTORY will be used.",
|
|
)
|
|
@click.option(
|
|
"--name",
|
|
default=None,
|
|
help="The name of the dependency set to load. If not specified, all dependency sets will be loaded.",
|
|
)
|
|
@click.option(
|
|
"--uv-cache-dir", default=None, help="The directory to cache uv dependencies"
|
|
)
|
|
@click.option(
|
|
"--check",
|
|
is_flag=True,
|
|
help="Check the the compiled dependencies are valid. Only compatible with generating all dependency sets.",
|
|
)
|
|
@click.option(
|
|
"--all-configs",
|
|
is_flag=True,
|
|
help="Build all configs",
|
|
)
|
|
def build(
|
|
config_path: str,
|
|
workspace_dir: Optional[str],
|
|
name: Optional[str],
|
|
uv_cache_dir: Optional[str],
|
|
check: Optional[bool],
|
|
all_configs: Optional[bool],
|
|
):
|
|
"""
|
|
Build dependency sets from a config file.
|
|
Args:
|
|
config_path: The path to the config file. If not specified, ci/raydepsets/configs/ray.depsets.yaml will be used.
|
|
"""
|
|
manager = DependencySetManager(
|
|
config_path=config_path,
|
|
workspace_dir=workspace_dir,
|
|
uv_cache_dir=uv_cache_dir,
|
|
check=check,
|
|
build_all_configs=all_configs,
|
|
)
|
|
manager.execute(name)
|
|
if check:
|
|
try:
|
|
manager.diff_lock_files()
|
|
except RuntimeError as e:
|
|
click.echo(e, err=True)
|
|
sys.exit(1)
|
|
finally:
|
|
manager.cleanup()
|
|
|
|
|
|
class DependencySetManager:
|
|
def __init__(
|
|
self,
|
|
config_path: str = None,
|
|
workspace_dir: Optional[str] = None,
|
|
uv_cache_dir: Optional[str] = None,
|
|
check: Optional[bool] = False,
|
|
build_all_configs: Optional[bool] = False,
|
|
):
|
|
"""Initialize the dependency set manager.
|
|
|
|
Args:
|
|
config_path: Path to the depsets config file.
|
|
workspace_dir: Path to the workspace directory.
|
|
uv_cache_dir: Directory to cache uv dependencies.
|
|
check: Whether to check if lock files are up to date.
|
|
build_all_configs: Whether to build all configs or just the specified one.
|
|
"""
|
|
self.workspace = Workspace(workspace_dir)
|
|
self.config = self.workspace.load_configs(config_path)
|
|
self.config_name = os.path.basename(config_path)
|
|
self.build_graph = DiGraph()
|
|
self._build(build_all_configs)
|
|
self._uv_binary = _uv_binary()
|
|
self._uv_cache_dir = uv_cache_dir
|
|
if check:
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.output_paths = self.get_output_paths()
|
|
self.copy_to_temp_dir()
|
|
|
|
def get_output_paths(self) -> List[Path]:
|
|
"""Get all output paths for depset nodes in topological order."""
|
|
output_paths = []
|
|
for node in topological_sort(self.build_graph):
|
|
if self.build_graph.nodes[node]["node_type"] == "depset":
|
|
output_paths.append(Path(self.build_graph.nodes[node]["depset"].output))
|
|
return output_paths
|
|
|
|
def copy_to_temp_dir(self):
|
|
"""Copy the lock files from source file paths to temp dir."""
|
|
for output_path in self.output_paths:
|
|
source_fp, target_fp = self.get_source_and_dest(output_path)
|
|
target_fp.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(
|
|
source_fp,
|
|
target_fp,
|
|
)
|
|
|
|
def get_diffs(self) -> List[str]:
|
|
"""Compare current lock files with previously saved copies and return unified diffs."""
|
|
diffs = []
|
|
for output_path in self.output_paths:
|
|
new_lock_file_fp, old_lock_file_fp = self.get_source_and_dest(output_path)
|
|
old_lock_file_contents = self.read_lock_file(old_lock_file_fp)
|
|
new_lock_file_contents = self.read_lock_file(new_lock_file_fp)
|
|
for diff in difflib.unified_diff(
|
|
old_lock_file_contents,
|
|
new_lock_file_contents,
|
|
fromfile=new_lock_file_fp.as_posix(),
|
|
tofile=old_lock_file_fp.as_posix(),
|
|
lineterm="",
|
|
):
|
|
diffs.append(diff)
|
|
return diffs
|
|
|
|
def diff_lock_files(self):
|
|
"""Check if lock files are up to date and raise an error if not."""
|
|
diffs = self.get_diffs()
|
|
if len(diffs) > 0:
|
|
raise RuntimeError(
|
|
f"Lock files are not up to date for config: {self.config_name}. Please update lock files and push the changes.\n"
|
|
+ "".join(diffs)
|
|
)
|
|
click.echo("Lock files are up to date.")
|
|
|
|
def get_source_and_dest(self, output_path: str) -> tuple[Path, Path]:
|
|
"""Get the source workspace path and temporary destination path for a lock file."""
|
|
return (self.get_path(output_path), (Path(self.temp_dir) / output_path))
|
|
|
|
def _build(self, build_all_configs: Optional[bool] = False):
|
|
"""Build the dependency graph from config depsets."""
|
|
# First pass: add all depset nodes so we can validate edges
|
|
depset_names = {depset.name for depset in self.config.depsets}
|
|
|
|
for depset in self.config.depsets:
|
|
if depset.operation == "compile":
|
|
self.build_graph.add_node(
|
|
depset.name,
|
|
operation="compile",
|
|
depset=depset,
|
|
node_type="depset",
|
|
config_name=depset.config_name,
|
|
)
|
|
elif depset.operation == "subset":
|
|
if depset.source_depset not in depset_names:
|
|
raise ValueError(
|
|
f"Depset '{depset.name}' references source_depset '{depset.source_depset}' which does not exist. "
|
|
f"Available depsets: {sorted(depset_names)}"
|
|
)
|
|
self.build_graph.add_node(
|
|
depset.name,
|
|
operation="subset",
|
|
depset=depset,
|
|
node_type="depset",
|
|
config_name=depset.config_name,
|
|
)
|
|
self.build_graph.add_edge(depset.source_depset, depset.name)
|
|
elif depset.operation == "expand":
|
|
for dep_name in depset.depsets:
|
|
if dep_name not in depset_names:
|
|
raise ValueError(
|
|
f"Depset '{depset.name}' references depset '{dep_name}' which does not exist. "
|
|
f"Available depsets: {sorted(depset_names)}"
|
|
)
|
|
self.build_graph.add_node(
|
|
depset.name,
|
|
operation="expand",
|
|
depset=depset,
|
|
node_type="depset",
|
|
config_name=depset.config_name,
|
|
)
|
|
for depset_name in depset.depsets:
|
|
self.build_graph.add_edge(depset_name, depset.name)
|
|
elif depset.operation == "relax":
|
|
if depset.source_depset not in depset_names:
|
|
raise ValueError(
|
|
f"Depset '{depset.name}' references source_depset '{depset.source_depset}' which does not exist. "
|
|
f"Available depsets: {sorted(depset_names)}"
|
|
)
|
|
self.build_graph.add_node(
|
|
depset.name,
|
|
operation="relax",
|
|
depset=depset,
|
|
node_type="depset",
|
|
config_name=depset.config_name,
|
|
)
|
|
self.build_graph.add_edge(depset.source_depset, depset.name)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid operation: {depset.operation} for depset {depset.name} in config {depset.config_name}"
|
|
)
|
|
if depset.pre_hooks:
|
|
for ind, hook in enumerate(depset.pre_hooks):
|
|
hook_name = f"{depset.name}_pre_hook_{ind+1}"
|
|
self.build_graph.add_node(
|
|
hook_name,
|
|
operation="pre_hook",
|
|
pre_hook=hook,
|
|
node_type="pre_hook",
|
|
config_name=depset.config_name,
|
|
)
|
|
self.build_graph.add_edge(hook_name, depset.name)
|
|
if not build_all_configs:
|
|
self.subgraph_config_nodes()
|
|
|
|
def subgraph_dependency_nodes(self, depset_name: str):
|
|
"""Reduce the build graph to only include the specified depset and its ancestors."""
|
|
dependency_nodes = networkx_ancestors(self.build_graph, depset_name)
|
|
nodes = dependency_nodes | {depset_name}
|
|
self.build_graph = self.build_graph.subgraph(nodes).copy()
|
|
|
|
def subgraph_config_nodes(self):
|
|
"""Reduce the build graph to nodes matching the current config and their ancestors."""
|
|
# Get all nodes that have the target config name
|
|
config_nodes = [
|
|
node
|
|
for node in self.build_graph.nodes
|
|
if self.build_graph.nodes[node]["config_name"] == self.config_name
|
|
]
|
|
# Get all ancestors of the target config nodes
|
|
ancestors_by_confg_node = {
|
|
n: networkx_ancestors(self.build_graph, n) for n in config_nodes
|
|
}
|
|
# Union all the ancestors of the target config nodes
|
|
config_nodes_ancestors = set().union(
|
|
*(ancestors_by_confg_node[n] for n in config_nodes)
|
|
)
|
|
nodes = set(config_nodes) | config_nodes_ancestors
|
|
self.build_graph = self.build_graph.subgraph(nodes).copy()
|
|
|
|
def execute(self, single_depset_name: Optional[str] = None):
|
|
"""Execute all depsets in topological order, optionally limited to a single depset."""
|
|
if single_depset_name:
|
|
# check if the depset exists
|
|
_get_depset(self.config.depsets, single_depset_name)
|
|
self.subgraph_dependency_nodes(single_depset_name)
|
|
for node in topological_sort(self.build_graph):
|
|
node_type = self.build_graph.nodes[node]["node_type"]
|
|
if node_type == "pre_hook":
|
|
pre_hook = self.build_graph.nodes[node]["pre_hook"]
|
|
self.execute_pre_hook(pre_hook)
|
|
elif node_type == "depset":
|
|
depset = self.build_graph.nodes[node]["depset"]
|
|
self.execute_depset(depset)
|
|
|
|
def exec_uv_cmd(
|
|
self, cmd: str, args: List[str], stdin: Optional[bytes] = None
|
|
) -> str:
|
|
"""Execute a uv pip command with the given arguments."""
|
|
cmd = [self._uv_binary, "pip", cmd, *args]
|
|
click.echo(f"Executing command: {' '.join(cmd)}")
|
|
status = subprocess.run(
|
|
cmd, cwd=self.workspace.dir, input=stdin, capture_output=True
|
|
)
|
|
if status.returncode != 0:
|
|
raise RuntimeError(
|
|
f"Failed to execute command: {' '.join(cmd)} with error: {status.stderr.decode('utf-8')}"
|
|
)
|
|
return status.stdout.decode("utf-8")
|
|
|
|
def execute_pre_hook(self, pre_hook: str):
|
|
"""Execute a pre-hook shell command."""
|
|
status = subprocess.run(
|
|
shlex.split(pre_hook),
|
|
cwd=self.workspace.dir,
|
|
capture_output=True,
|
|
)
|
|
if status.returncode != 0:
|
|
raise RuntimeError(
|
|
f"Failed to execute pre_hook {pre_hook} with error: {status.stderr.decode('utf-8')}",
|
|
)
|
|
click.echo(f"{status.stdout.decode('utf-8')}")
|
|
click.echo(f"Executed pre_hook {pre_hook} successfully")
|
|
|
|
def execute_depset(self, depset: Depset):
|
|
"""Execute a single depset based on its operation type (compile, subset, or expand)."""
|
|
if depset.operation == "compile":
|
|
self.compile(
|
|
constraints=depset.constraints,
|
|
requirements=depset.requirements,
|
|
name=depset.name,
|
|
output=depset.output,
|
|
append_flags=depset.append_flags,
|
|
override_flags=depset.override_flags,
|
|
packages=depset.packages,
|
|
include_setuptools=depset.include_setuptools,
|
|
)
|
|
elif depset.operation == "subset":
|
|
self.subset(
|
|
source_depset=depset.source_depset,
|
|
requirements=depset.requirements,
|
|
append_flags=depset.append_flags,
|
|
override_flags=depset.override_flags,
|
|
name=depset.name,
|
|
output=depset.output,
|
|
include_setuptools=depset.include_setuptools,
|
|
)
|
|
elif depset.operation == "expand":
|
|
self.expand(
|
|
depsets=depset.depsets,
|
|
requirements=depset.requirements,
|
|
constraints=depset.constraints,
|
|
append_flags=depset.append_flags,
|
|
override_flags=depset.override_flags,
|
|
name=depset.name,
|
|
output=depset.output,
|
|
include_setuptools=depset.include_setuptools,
|
|
)
|
|
elif depset.operation == "relax":
|
|
self.relax(
|
|
source_depset=depset.source_depset,
|
|
packages=depset.packages,
|
|
name=depset.name,
|
|
output=depset.output,
|
|
)
|
|
click.echo(f"Dependency set {depset.name} compiled successfully")
|
|
|
|
def compile(
|
|
self,
|
|
constraints: List[str],
|
|
name: str,
|
|
output: str,
|
|
append_flags: Optional[List[str]] = None,
|
|
override_flags: Optional[List[str]] = None,
|
|
packages: Optional[List[str]] = None,
|
|
requirements: Optional[List[str]] = None,
|
|
include_setuptools: Optional[bool] = False,
|
|
):
|
|
"""Compile a dependency set."""
|
|
args = DEFAULT_UV_FLAGS.copy()
|
|
stdin = None
|
|
if not include_setuptools:
|
|
args.extend(_flatten_flags(["--unsafe-package setuptools"]))
|
|
if self._uv_cache_dir:
|
|
args.extend(["--cache-dir", self._uv_cache_dir])
|
|
if override_flags:
|
|
args = _override_uv_flags(override_flags, args)
|
|
if append_flags:
|
|
args.extend(_flatten_flags(append_flags))
|
|
if constraints:
|
|
for constraint in sorted(constraints):
|
|
args.extend(["-c", constraint])
|
|
if requirements:
|
|
for requirement in sorted(requirements):
|
|
args.extend([requirement])
|
|
if packages:
|
|
# need to add a dash to process stdin
|
|
args.append("-")
|
|
stdin = _get_bytes(packages)
|
|
if output:
|
|
args.extend(["-o", output])
|
|
self.exec_uv_cmd("compile", args, stdin)
|
|
|
|
def subset(
|
|
self,
|
|
source_depset: str,
|
|
requirements: List[str],
|
|
name: str,
|
|
output: str = None,
|
|
append_flags: Optional[List[str]] = None,
|
|
override_flags: Optional[List[str]] = None,
|
|
include_setuptools: Optional[bool] = False,
|
|
):
|
|
"""Subset a dependency set."""
|
|
source_depset = _get_depset(self.config.depsets, source_depset)
|
|
self.check_subset_exists(source_depset, requirements)
|
|
self.compile(
|
|
constraints=[source_depset.output],
|
|
requirements=requirements,
|
|
name=name,
|
|
output=output,
|
|
append_flags=append_flags,
|
|
override_flags=override_flags,
|
|
include_setuptools=include_setuptools,
|
|
)
|
|
|
|
def expand(
|
|
self,
|
|
depsets: List[str],
|
|
requirements: List[str],
|
|
constraints: List[str],
|
|
name: str,
|
|
output: str = None,
|
|
append_flags: Optional[List[str]] = None,
|
|
override_flags: Optional[List[str]] = None,
|
|
include_setuptools: Optional[bool] = False,
|
|
):
|
|
"""Expand a dependency set."""
|
|
# handle both depsets and requirements
|
|
depset_req_list = []
|
|
for depset_name in depsets:
|
|
dep = _get_depset(self.config.depsets, depset_name)
|
|
if dep.operation == "relax":
|
|
depset_req_list.append(dep.output)
|
|
else:
|
|
depset_req_list.extend(
|
|
self.get_expanded_depset_requirements(depset_name, [])
|
|
)
|
|
if requirements:
|
|
depset_req_list.extend(requirements)
|
|
self.compile(
|
|
constraints=constraints,
|
|
requirements=depset_req_list,
|
|
name=name,
|
|
output=output,
|
|
append_flags=append_flags,
|
|
override_flags=override_flags,
|
|
include_setuptools=include_setuptools,
|
|
)
|
|
|
|
def relax(
|
|
self,
|
|
source_depset: str,
|
|
packages: List[str],
|
|
name: str,
|
|
output: str = None,
|
|
):
|
|
"""Relax a dependency set by removing specified packages from the lock file."""
|
|
source_depset = _get_depset(self.config.depsets, source_depset)
|
|
|
|
lock_file_path = self.get_path(source_depset.output)
|
|
requirements_file = parse_lock_file(str(lock_file_path))
|
|
requirements_list = [req.name for req in requirements_file.requirements]
|
|
for package in packages:
|
|
if package not in requirements_list:
|
|
raise RuntimeError(
|
|
f"Package {package} not found in lock file {source_depset.output}"
|
|
)
|
|
|
|
# Remove specified packages from requirements
|
|
requirements_file.requirements = [
|
|
req for req in requirements_file.requirements if req.name not in packages
|
|
]
|
|
|
|
# Write the modified lock file
|
|
output_path = self.get_path(output) if output else lock_file_path
|
|
write_lock_file(requirements_file, str(output_path))
|
|
click.echo(
|
|
f"Relaxed {source_depset.name} by removing packages {packages} and wrote to {output_path}"
|
|
)
|
|
|
|
def read_lock_file(self, file_path: Path) -> List[str]:
|
|
"""Read and return the contents of a lock file as a list of lines."""
|
|
if not file_path.exists():
|
|
raise RuntimeError(f"Lock file {file_path} does not exist")
|
|
with open(file_path, "r") as f:
|
|
return f.readlines()
|
|
|
|
def get_path(self, path: str) -> Path:
|
|
"""Convert a relative path to an absolute path within the workspace."""
|
|
return Path(self.workspace.dir) / path
|
|
|
|
def check_subset_exists(self, source_depset: Depset, requirements: List[str]):
|
|
"""Verify that all requirements exist in the source depset."""
|
|
for req in requirements:
|
|
if req not in self.get_expanded_depset_requirements(source_depset.name, []):
|
|
raise RuntimeError(
|
|
f"Requirement {req} is not a subset of {source_depset.name} in config {source_depset.config_name}"
|
|
)
|
|
|
|
def get_expanded_depset_requirements(
|
|
self, depset_name: str, requirements_list: List[str]
|
|
) -> List[str]:
|
|
"""Get all requirements for expanded depsets
|
|
|
|
Args:
|
|
depset_name: The name of the expanded depset to get the requirements for.
|
|
requirements_list: The list of requirements to extend.
|
|
|
|
Returns:
|
|
A list of requirements for the expanded depset.
|
|
"""
|
|
depset = _get_depset(self.config.depsets, depset_name)
|
|
requirements_list.extend(depset.requirements)
|
|
if depset.operation == "expand":
|
|
for dep in depset.depsets:
|
|
self.get_expanded_depset_requirements(dep, requirements_list)
|
|
return list(set(requirements_list))
|
|
|
|
def cleanup(self):
|
|
"""Remove the temporary directory used for lock file comparisons."""
|
|
if self.temp_dir:
|
|
shutil.rmtree(self.temp_dir)
|
|
|
|
|
|
def _get_bytes(packages: List[str]) -> bytes:
|
|
"""Convert a list of package names to newline-separated UTF-8 bytes."""
|
|
return ("\n".join(packages) + "\n").encode("utf-8")
|
|
|
|
|
|
def _get_depset(depsets: List[Depset], name: str) -> Depset:
|
|
"""Find and return a depset by name from a list of depsets."""
|
|
for depset in depsets:
|
|
if depset.name == name:
|
|
return depset
|
|
raise KeyError(f"Dependency set {name} not found")
|
|
|
|
|
|
def _flatten_flags(flags: List[str]) -> List[str]:
|
|
"""
|
|
Flatten a list of flags into a list of strings.
|
|
For example, ["--find-links https://pypi.org/simple"] will be flattened to
|
|
["--find-links", "https://pypi.org/simple"].
|
|
"""
|
|
flattened_flags = []
|
|
for flag in flags:
|
|
flattened_flags.extend(flag.split())
|
|
return flattened_flags
|
|
|
|
|
|
def _override_uv_flags(flags: List[str], args: List[str]) -> List[str]:
|
|
"""Override existing uv flags in args with new values from flags."""
|
|
flag_names = {f.split()[0] for f in flags if f.startswith("--")}
|
|
new_args = []
|
|
skip_next = False
|
|
for arg in args:
|
|
if skip_next:
|
|
skip_next = False
|
|
continue
|
|
if arg in flag_names:
|
|
skip_next = True
|
|
continue
|
|
new_args.append(arg)
|
|
|
|
return new_args + _flatten_flags(flags)
|
|
|
|
|
|
def parse_lock_file(lock_file_path: str) -> RequirementsFile:
|
|
"""
|
|
Parses a lock file and returns a RequirementsFile object, which contains
|
|
all information from the file, including requirements, options, and comments.
|
|
"""
|
|
return RequirementsFile.from_file(lock_file_path)
|
|
|
|
|
|
def write_lock_file(requirements_file: RequirementsFile, lock_file_path: str):
|
|
"""
|
|
Writes a RequirementsFile object to a lock file, preserving all its content.
|
|
"""
|
|
with open(lock_file_path, "w") as f:
|
|
f.write(requirements_file.dumps())
|
|
|
|
|
|
def _uv_binary():
|
|
"""Get the path to the uv binary for the current platform."""
|
|
r = runfiles.Create()
|
|
system = platform.system()
|
|
processor = platform.processor()
|
|
|
|
if system == "Linux" and processor == "x86_64":
|
|
return r.Rlocation("uv_x86_64-linux/uv-x86_64-unknown-linux-gnu/uv")
|
|
elif system == "Darwin" and (processor == "arm" or processor == "aarch64"):
|
|
return r.Rlocation("uv_aarch64-darwin/uv-aarch64-apple-darwin/uv")
|
|
else:
|
|
raise RuntimeError(f"Unsupported platform/processor: {system}/{processor}")
|