|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import json |
| 4 | +import os |
| 5 | +import subprocess |
| 6 | +import sys |
| 7 | +from typing import ( |
| 8 | + Any, |
| 9 | + Dict, |
| 10 | + Generator, |
| 11 | + List, |
| 12 | + Literal, |
| 13 | + NotRequired, |
| 14 | + Optional, |
| 15 | + Set, |
| 16 | + TypedDict, |
| 17 | +) |
| 18 | + |
| 19 | + |
| 20 | +class NixEvalJobsOutput(TypedDict): |
| 21 | + """Raw output from nix-eval-jobs command.""" |
| 22 | + |
| 23 | + attr: str |
| 24 | + attrPath: List[str] |
| 25 | + cacheStatus: Literal["notBuilt", "cached", "local"] |
| 26 | + drvPath: str |
| 27 | + isCached: bool |
| 28 | + name: str |
| 29 | + system: str |
| 30 | + neededBuilds: NotRequired[List[Any]] |
| 31 | + neededSubstitutes: NotRequired[List[Any]] |
| 32 | + outputs: NotRequired[Dict[str, str]] |
| 33 | + |
| 34 | + |
| 35 | +class RunsOnConfig(TypedDict): |
| 36 | + """GitHub Actions runs-on configuration.""" |
| 37 | + |
| 38 | + group: NotRequired[str] |
| 39 | + labels: List[str] |
| 40 | + |
| 41 | + |
| 42 | +class GitHubActionPackage(TypedDict): |
| 43 | + """Processed package for GitHub Actions matrix.""" |
| 44 | + |
| 45 | + attr: str |
| 46 | + name: str |
| 47 | + system: str |
| 48 | + already_cached: bool |
| 49 | + runs_on: RunsOnConfig |
| 50 | + |
| 51 | + |
| 52 | +BUILD_RUNNER_MAP: Dict[str, RunsOnConfig] = { |
| 53 | + "aarch64-linux": { |
| 54 | + "group": "self-hosted-runners-nix", |
| 55 | + "labels": ["aarch64-linux"], |
| 56 | + }, |
| 57 | + "aarch64-darwin": { |
| 58 | + "group": "self-hosted-runners-nix", |
| 59 | + "labels": ["aarch64-darwin"], |
| 60 | + }, |
| 61 | + "x86_64-linux": { |
| 62 | + "labels": ["blacksmith-32vcpu-ubuntu-2404"], |
| 63 | + }, |
| 64 | +} |
| 65 | + |
| 66 | + |
| 67 | +def get_worker_count() -> int: |
| 68 | + """Get optimal worker count based on CPU cores.""" |
| 69 | + try: |
| 70 | + return max(1, int(os.cpu_count())) |
| 71 | + except (OSError, AttributeError): |
| 72 | + print( |
| 73 | + "Warning: Unable to get CPU count, using default max_workers=1", |
| 74 | + file=sys.stderr, |
| 75 | + ) |
| 76 | + return 1 |
| 77 | + |
| 78 | + |
| 79 | +def build_nix_eval_command(max_workers: int) -> List[str]: |
| 80 | + """Build the nix-eval-jobs command with appropriate flags.""" |
| 81 | + return [ |
| 82 | + "nix-eval-jobs", |
| 83 | + "--flake", |
| 84 | + ".#checks", |
| 85 | + "--check-cache-status", |
| 86 | + "--force-recurse", |
| 87 | + "--quiet", |
| 88 | + "--workers", |
| 89 | + str(max_workers), |
| 90 | + ] |
| 91 | + |
| 92 | + |
| 93 | +def parse_nix_eval_line( |
| 94 | + line: str, drv_paths: Set[str] |
| 95 | +) -> Optional[GitHubActionPackage]: |
| 96 | + """Parse a single line of nix-eval-jobs output""" |
| 97 | + if not line.strip(): |
| 98 | + return None |
| 99 | + |
| 100 | + try: |
| 101 | + data: NixEvalJobsOutput = json.loads(line) |
| 102 | + if data["drvPath"] in drv_paths: |
| 103 | + return None |
| 104 | + drv_paths.add(data["drvPath"]) |
| 105 | + |
| 106 | + runs_on_config = BUILD_RUNNER_MAP[data["system"]] |
| 107 | + |
| 108 | + return { |
| 109 | + "attr": "checks." + data["attr"], |
| 110 | + "name": data["name"], |
| 111 | + "system": data["system"], |
| 112 | + "already_cached": data.get("cacheStatus") != "notBuilt", |
| 113 | + "runs_on": runs_on_config, |
| 114 | + } |
| 115 | + except json.JSONDecodeError: |
| 116 | + print(f"Skipping invalid JSON line: {line}", file=sys.stderr) |
| 117 | + return None |
| 118 | + |
| 119 | + |
| 120 | +def run_nix_eval_jobs(cmd: List[str]) -> Generator[GitHubActionPackage, None, None]: |
| 121 | + """Run nix-eval-jobs and yield parsed package data.""" |
| 122 | + print(f"Running command: {' '.join(cmd)}", file=sys.stderr) |
| 123 | + |
| 124 | + with subprocess.Popen( |
| 125 | + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True |
| 126 | + ) as process: |
| 127 | + drv_paths = set() |
| 128 | + |
| 129 | + for line in process.stdout: |
| 130 | + package = parse_nix_eval_line(line, drv_paths) |
| 131 | + if package: |
| 132 | + yield package |
| 133 | + |
| 134 | + if process.returncode and process.returncode != 0: |
| 135 | + print("Error: Evaluation failed", file=sys.stderr) |
| 136 | + sys.stderr.write(process.stderr.read()) |
| 137 | + sys.exit(process.returncode) |
| 138 | + |
| 139 | + |
| 140 | +def main() -> None: |
| 141 | + max_workers = get_worker_count() |
| 142 | + cmd = build_nix_eval_command(max_workers) |
| 143 | + |
| 144 | + gh_action_packages = list(run_nix_eval_jobs(cmd)) |
| 145 | + gh_output = {"include": gh_action_packages} |
| 146 | + print(json.dumps(gh_output)) |
| 147 | + |
| 148 | + |
| 149 | +if __name__ == "__main__": |
| 150 | + main() |
0 commit comments