gistfile1.nix

{ config, lib, pkgs, ... }:

with lib;
let
  cfg = config.services.image-classifier;
  worker-cpus = builtins.genList (i: {
    # not sure why it's not overriding the template unit file
    # https://github.com/NixOS/nixpkgs/issues/80933
    systemd.services."image-classifier-worker-cpu@${builtins.toString (i + 1)}" = {
      after = [ "network-online.target" ];
      wants = [ "network-online.target" ];
      wantedBy = [ "multi-user.target" ];
      serviceConfig = rec {
        Restart = "always";
        RuntimeDirectory = "image-classifier-worker-cpu";
        ExecStart = [
          ""
          ''
          ${pkgs.image-classifier}/bin/image-classifier-worker \
            --scheduler-ip='${cfg.worker-cpu.scheduler-ip}' \
            --scheduler-port='${builtins.toString cfg.worker-cpu.scheduler-port}' \
            --nthreads='${builtins.toString cfg.worker-cpu.nthreads}' \
            --memory-limit='${cfg.worker-cpu.memory-limit}' \
            --local-directory=/run/${RuntimeDirectory} \
            --name='image-classifier-worker-cpu-%i'
          ''
        ];
      };
    };
  }) cfg.worker-cpu.instances;
in
  {

    options.services.image-classifier = {
      scheduler = {
        enable = mkEnableOption "Image Classifier Scheduler";
        host = mkOption {
          type = types.str;
          default = "127.0.0.1";
        };
        port = mkOption {
          type = types.int;
          default = 3201;
        };
        dashboard-address = mkOption {
          type = types.str;
          default = ":3202";
        };
      };
      worker-cpu = {
        enable = mkEnableOption "Image Classifier Worker CPU";
        instances = mkOption {
          type = types.int;
          default = 1;
        };
        scheduler-ip = mkOption {
          type = types.str;
          default = "127.0.0.1";
        };
        scheduler-port = mkOption {
          type = types.int;
          default = 3201;
        };
        nthreads = mkOption {
          type = types.int;
          default = 1;
        };
        memory-limit = mkOption {
          type = types.str;
          default = "0";
        };
      };
      worker-gpu = {
        enable = mkEnableOption "Image Classifier Worker GPU";
        instances = mkOption {
          type = types.int;
          default = 1;
        };
        scheduler-ip = mkOption {
          type = types.str;
          default = "127.0.0.1";
        };
        scheduler-port = mkOption {
          type = types.int;
          default = 3201;
        };
        nthreads = mkOption {
          type = types.int;
          default = 1;
        };
        memory-limit = mkOption {
          type = types.str;
          default = "0";
        };
      };
    };

    config = mkMerge [
      (mkIf (cfg.scheduler.enable || cfg.worker-cpu.enable) {
        environment.systemPackages = [ pkgs.image-classifier ];
        systemd.packages = [ pkgs.image-classifier ];
      })
      (mkIf cfg.scheduler.enable {
        systemd.services.image-classifier-scheduler = {
          serviceConfig = {
            ExecStart = [
              ""
              ''
              ${pkgs.image-classifier}/bin/image-classifier-scheduler \
                --host='${cfg.scheduler.host}' \
                --port='${builtins.toString cfg.scheduler.port}' \
                --dashboard-address='${cfg.scheduler.dashboard-address}'
              ''
            ];
          };
        };
      })
    ] ++ (if cfg.worker-cpu.enable then worker-cpus else []);
  }