homey/modules/services/mlflow.nix

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

# MLflow — ML experiment tracking server.
#
# Auth model: MLflow's built-in basic-auth plugin (--app-name basic-auth).
#   - Web UI: login form (MLflow's own user database — separate from LDAP).
#   - Python SDK: set MLFLOW_TRACKING_TOKEN=<token> for bearer-token auth,
#                 or MLFLOW_TRACKING_USERNAME/PASSWORD for basic auth.
#   - No Authelia forward_auth — MLflow manages its own user database.
#
# On first boot, basic-auth.ini seeds the initial admin account. After that,
# credentials live in /mlflow/data/basic_auth.db and the ini is ignored for auth.
#
# Secrets consumed from sops:
#   mlflow/secret_key       — Flask CSRF secret (generate: openssl rand -hex 32)
#   mlflow/admin_password   — initial admin password (used once on first boot)
#
# Volume layout:
#   <dataDir>/mlflow/data/       → /mlflow/data  (SQLite tracking DB, auth DB, basic-auth.ini)
#   <dataDir>/mlflow/artifacts/  → /mlflow/artifacts  (model files, plots, etc.)

let
  cfg     = config.homey.mlflow;
  dataDir = config.homey.storage.mountPoint;
  domain  = homeyConfig.domain;
in
{
  options.homey.mlflow = {
    enable = lib.mkEnableOption "MLflow experiment tracking server";

    image = lib.mkOption {
      type    = lib.types.str;
      default = "ghcr.io/mlflow/mlflow:latest";
    };

    port = lib.mkOption {
      type        = lib.types.port;
      default     = 5050;
      description = "Host port MLflow listens on (bound to 127.0.0.1, Caddy proxy).";
    };
  };

  config = lib.mkIf cfg.enable {
    # -------------------------------------------------------------------------
    # Secrets
    # -------------------------------------------------------------------------
    sops.secrets."mlflow/secret_key"     = { owner = "root"; };
    sops.secrets."mlflow/admin_password" = { owner = "root"; };

    # -------------------------------------------------------------------------
    # Container
    # -------------------------------------------------------------------------
    virtualisation.oci-containers.containers.mlflow = {
      image = cfg.image;
      ports = [ "127.0.0.1:${toString cfg.port}:5000" ];
      volumes = [
        "${dataDir}/mlflow/data:/mlflow/data"
        "${dataDir}/mlflow/artifacts:/mlflow/artifacts"
      ];
      extraOptions = [ "--network=homey" ];
      environment = {
        # v3.x: auth config path is env var, not a CLI flag
        MLFLOW_AUTH_CONFIG_PATH = "/mlflow/data/basic-auth.ini";
      };
      environmentFiles = [ "/run/mlflow-secrets.env" ];
      cmd = [
        "mlflow" "server"
        "--host" "0.0.0.0"
        "--port" "5000"
        "--backend-store-uri" "sqlite:////mlflow/data/mlflow.db"
        "--default-artifact-root" "/mlflow/artifacts"
        "--app-name" "basic-auth"
        # v3.x security middleware: must explicitly allow the public hostname
        "--allowed-hosts" "mlflow.${domain}"
        # Allow browser UI (ajax-api) requests from the public origin.
        # Without this, fastapi_security blocks all cross-origin requests with 403,
        # breaking chart data and UI telemetry (SDK api/2.0 calls are unaffected).
        "--cors-allowed-origins" "https://mlflow.${domain}"
      ];
    };

    # -------------------------------------------------------------------------
    # ExecStartPre: write secrets env file and seed basic-auth.ini on first boot
    # -------------------------------------------------------------------------
    systemd.services."podman-mlflow" = {
      serviceConfig.ExecStartPre = [
        (pkgs.writeShellScript "mlflow-write-secrets" ''
          set -euo pipefail

          install -m 600 /dev/null /run/mlflow-secrets.env
          printf 'MLFLOW_FLASK_SERVER_SECRET_KEY=%s\n' \
            "$(cat ${config.sops.secrets."mlflow/secret_key".path})" \
            >> /run/mlflow-secrets.env

          # Seed basic-auth.ini on first boot only.
          # After first run MLflow stores credentials in basic_auth.db.
          if [ ! -f "${dataDir}/mlflow/data/basic-auth.ini" ]; then
            printf '[mlflow]\ndefault_permission = NO_PERMISSIONS\nadmin_username = admin\nadmin_password = %s\ndatabase_uri = sqlite:////mlflow/data/basic_auth.db\n' \
              "$(cat ${config.sops.secrets."mlflow/admin_password".path})" \
              > "${dataDir}/mlflow/data/basic-auth.ini"
          fi
        '')
      ];
      postStop = "rm -f /run/mlflow-secrets.env";
      after    = lib.mkAfter [ "mnt-data.mount" "podman-homey-network.service" ];
      requires = lib.mkAfter [ "mnt-data.mount" "podman-homey-network.service" ];
    };

    # -------------------------------------------------------------------------
    # Caddy virtual host — auth=false, MLflow handles its own login
    # -------------------------------------------------------------------------
    homey.caddy.virtualHosts = [{
      subdomain = "mlflow";
      port      = cfg.port;
      auth      = false;
    }];

    # -------------------------------------------------------------------------
    # Storage directories on external HD
    # -------------------------------------------------------------------------
    homey.storage.extraDirs = [
      { path = "mlflow"; }
      { path = "mlflow/data";      mode = "0750"; }
      { path = "mlflow/artifacts"; mode = "0750"; }
    ];

    # -------------------------------------------------------------------------
    # Backup
    # -------------------------------------------------------------------------
    homey.backup.extraPaths = [
      "${dataDir}/mlflow/data"
      "${dataDir}/mlflow/artifacts"
    ];

    # -------------------------------------------------------------------------
    # Monitoring
    # -------------------------------------------------------------------------
    homey.monitoring.monitors = [{
      name     = "MLflow";
      url      = "https://mlflow.${domain}";
      interval = 60;
    }];
  };
}