nix-builder-autoscaler/flake.nix

{
  description = "nix-builder-autoscaler - autoscaler daemon for Nix remote builders on EC2 Spot";

  inputs = {
    nixpkgs.url = "https://flakehub.com/f/NixOS/nixpkgs/0.1";
    treefmt-nix = {
      url = "github:numtide/treefmt-nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    pyproject-nix = {
      url = "github:pyproject-nix/pyproject.nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    uv2nix = {
      url = "github:pyproject-nix/uv2nix";
      inputs.pyproject-nix.follows = "pyproject-nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    pyproject-build-systems = {
      url = "github:pyproject-nix/build-system-pkgs";
      inputs.pyproject-nix.follows = "pyproject-nix";
      inputs.uv2nix.follows = "uv2nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
  };

  outputs =
    {
      self,
      nixpkgs,
      treefmt-nix,
      pyproject-nix,
      uv2nix,
      pyproject-build-systems,
      ...
    }:
    let
      systems = [ "x86_64-linux" ];
      forAllSystems = fn: nixpkgs.lib.genAttrs systems (system: fn nixpkgs.legacyPackages.${system});

      agentWorkspace = uv2nix.lib.workspace.loadWorkspace { workspaceRoot = ./agent; };
      agentOverlay = agentWorkspace.mkPyprojectOverlay { sourcePreference = "wheel"; };

      buildbotExtWorkspace = uv2nix.lib.workspace.loadWorkspace { workspaceRoot = ./buildbot-ext; };
      buildbotExtOverlay = buildbotExtWorkspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
      pyprojectOverrides = final: prev: {
        py-ubjson = prev.py-ubjson.overrideAttrs (old: {
          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ];
        });
      };
      mkBuildbotExtPythonSet =
        pkgs: python:
        (pkgs.callPackage pyproject-nix.build.packages {
          inherit python;
        }).overrideScope
          (
            pkgs.lib.composeManyExtensions [
              pyproject-build-systems.overlays.default
              buildbotExtOverlay
              pyprojectOverrides
            ]
          );
      mkBuildbotAutoscaleExtPythonPackagesExtension = pkgs: pyFinal: _pyPrev: {
        buildbot-autoscale-ext =
          pyFinal.toPythonModule (mkBuildbotExtPythonSet pkgs pyFinal.python)."buildbot-autoscale-ext";
      };
    in
    {
      lib = {
        inherit mkBuildbotAutoscaleExtPythonPackagesExtension;
      };

      formatter = forAllSystems (
        pkgs: (treefmt-nix.lib.evalModule pkgs ./treefmt.nix).config.build.wrapper
      );

      nixosModules = {
        default = import ./nix/modules/nixos;
        nix-builder-autoscaler = import ./nix/modules/nixos/services/nix-builder-autoscaler.nix;
        buildbot-nix-autoscaler = import ./nix/modules/nixos/services/buildbot-nix-autoscaler.nix;
      };

      packages = forAllSystems (
        pkgs:
        let
          agentPythonSet =
            (pkgs.callPackage pyproject-nix.build.packages {
              python = pkgs.python312;
            }).overrideScope
              (
                pkgs.lib.composeManyExtensions [
                  pyproject-build-systems.overlays.default
                  agentOverlay
                  pyprojectOverrides
                ]
              );
          buildbotExtPythonSet =
            mkBuildbotExtPythonSet pkgs pkgs.python3;
          venv = agentPythonSet.mkVirtualEnv "nix-builder-autoscaler-env" agentWorkspace.deps.default;
          buildbotExtVenv = buildbotExtPythonSet.mkVirtualEnv "buildbot-autoscale-ext-env" buildbotExtWorkspace.deps.default;
        in
        {
          nix-builder-autoscaler = venv;
          buildbot-autoscale-ext = buildbotExtPythonSet."buildbot-autoscale-ext";
          buildbot-autoscale-ext-venv = buildbotExtVenv;
          default = venv;
        }
      );

      apps = forAllSystems (
        pkgs:
        let
          venv = self.packages.${pkgs.stdenv.hostPlatform.system}.nix-builder-autoscaler;
        in
        {
          nix-builder-autoscaler = {
            type = "app";
            program = "${venv}/bin/python";
            meta.description = "Nix builder autoscaler daemon";
          };
          autoscalerctl = {
            type = "app";
            program = "${venv}/bin/autoscalerctl";
            meta.description = "Autoscaler CLI";
          };
          default = {
            type = "app";
            program = "${venv}/bin/autoscalerctl";
            meta.description = "Autoscaler CLI";
          };
        }
      );

      checks = forAllSystems (
        pkgs:
        let
          agentPythonSet =
            (pkgs.callPackage pyproject-nix.build.packages {
              python = pkgs.python312;
            }).overrideScope
              (
                pkgs.lib.composeManyExtensions [
                  pyproject-build-systems.overlays.default
                  agentOverlay
                  pyprojectOverrides
                ]
              );
          buildbotExtPythonSet =
            mkBuildbotExtPythonSet pkgs pkgs.python3;
          testVenv = agentPythonSet.mkVirtualEnv "nix-builder-autoscaler-test-env" {
            nix-builder-autoscaler = [ "dev" ];
          };
          buildbotExtTestVenv = buildbotExtPythonSet.mkVirtualEnv "buildbot-autoscale-ext-test-env" {
            buildbot-autoscale-ext = [ "dev" ];
          };
          src = ./agent;
          buildbotExtSrc = ./buildbot-ext;
        in
        {
          devShell = self.devShells.${pkgs.stdenv.hostPlatform.system}.default;

          nix-builder-autoscaler-unit-tests = pkgs.stdenv.mkDerivation {
            name = "nix-builder-autoscaler-unit-tests";
            inherit src;
            dontConfigure = true;
            dontBuild = true;
            nativeBuildInputs = [ testVenv ];
            checkPhase = ''
              runHook preCheck
              export HOME=$(mktemp -d)
              pytest nix_builder_autoscaler/tests/ --ignore=nix_builder_autoscaler/tests/integration/ -v
              runHook postCheck
            '';
            doCheck = true;
            installPhase = ''
              mkdir -p $out
              touch $out/passed
            '';
          };

          nix-builder-autoscaler-integration-tests = pkgs.stdenv.mkDerivation {
            name = "nix-builder-autoscaler-integration-tests";
            inherit src;
            dontConfigure = true;
            dontBuild = true;
            nativeBuildInputs = [ testVenv ];
            checkPhase = ''
              runHook preCheck
              export HOME=$(mktemp -d)
              pytest nix_builder_autoscaler/tests/integration/ -v
              runHook postCheck
            '';
            doCheck = true;
            installPhase = ''
              mkdir -p $out
              touch $out/passed
            '';
          };

          nix-builder-autoscaler-ruff = pkgs.stdenv.mkDerivation {
            name = "nix-builder-autoscaler-ruff";
            inherit src;
            dontConfigure = true;
            dontBuild = true;
            nativeBuildInputs = [ testVenv ];
            checkPhase = ''
              runHook preCheck
              ruff check nix_builder_autoscaler/
              ruff format --check nix_builder_autoscaler/
              runHook postCheck
            '';
            doCheck = true;
            installPhase = ''
              mkdir -p $out
              touch $out/passed
            '';
          };

          nix-builder-autoscaler-pyright = pkgs.stdenv.mkDerivation {
            name = "nix-builder-autoscaler-pyright";
            inherit src;
            dontConfigure = true;
            dontBuild = true;
            nativeBuildInputs = [
              testVenv
              pkgs.nodejs
            ];
            checkPhase = ''
              runHook preCheck
              export HOME=$(mktemp -d)
              pyright nix_builder_autoscaler/
              runHook postCheck
            '';
            doCheck = true;
            installPhase = ''
              mkdir -p $out
              touch $out/passed
            '';
          };

          buildbot-autoscale-ext-tests = pkgs.stdenv.mkDerivation {
            name = "buildbot-autoscale-ext-tests";
            src = buildbotExtSrc;
            dontConfigure = true;
            dontBuild = true;
            nativeBuildInputs = [ buildbotExtTestVenv ];
            checkPhase = ''
              runHook preCheck
              export HOME=$(mktemp -d)
              pytest buildbot_autoscale_ext/tests/ -v
              runHook postCheck
            '';
            doCheck = true;
            installPhase = ''
              mkdir -p $out
              touch $out/passed
            '';
          };

          buildbot-autoscale-ext-ruff = pkgs.stdenv.mkDerivation {
            name = "buildbot-autoscale-ext-ruff";
            src = buildbotExtSrc;
            dontConfigure = true;
            dontBuild = true;
            nativeBuildInputs = [ buildbotExtTestVenv ];
            checkPhase = ''
              runHook preCheck
              ruff check buildbot_autoscale_ext/
              ruff format --check buildbot_autoscale_ext/
              runHook postCheck
            '';
            doCheck = true;
            installPhase = ''
              mkdir -p $out
              touch $out/passed
            '';
          };

          buildbot-autoscale-ext-pyright = pkgs.stdenv.mkDerivation {
            name = "buildbot-autoscale-ext-pyright";
            src = buildbotExtSrc;
            dontConfigure = true;
            dontBuild = true;
            nativeBuildInputs = [
              buildbotExtTestVenv
              pkgs.nodejs
            ];
            checkPhase = ''
              runHook preCheck
              export HOME=$(mktemp -d)
              pyright buildbot_autoscale_ext/
              runHook postCheck
            '';
            doCheck = true;
            installPhase = ''
              mkdir -p $out
              touch $out/passed
            '';
          };
        }
      );

      devShells = forAllSystems (pkgs: {
        default = pkgs.mkShell {
          packages = with pkgs; [
            uv
            ruff
            pyright
          ];
        };
      });
    };
}