[DO NOT MERGE] Refactor agent scaler

joshua/joshua_model.py

@@ -96,6 +96,13 @@ def open(c_file=None, dir_path=("joshua",)):
 
     cluster_file = c_file
     db = fdb.open(cluster_file)
+    # Timeout applies to all transactions. This should prevent FDB to hang forever.
+    transaction_timeout = int(os.environ.get("TRANSACTION_TIMEOUT_MS", "10000"))
+    # Retry limit applies to all transactions. This should prevent FDB to hang forever.
+    transaction_retry_limit = int(os.environ.get("TRANSACTION_RETRY_LIMIT", "5"))
+    logger.info(f"transaction_timeout: {transaction_timeout} and transaction_retry_limit: {transaction_retry_limit}")
+    db.options.set_transaction_timeout(transaction_timeout)
+    db.options.set_transaction_retry_limit(transaction_retry_limit)
     dir_top = create_or_open_top_path(db, dir_path)
     dir_ensembles = dir_top.create_or_open(db, "ensembles")
     dir_active = dir_ensembles.create_or_open(db, "active")

k8s/agent-scaler/Dockerfile

@@ -1,11 +1,37 @@
+FROM golang:1.25.10-bookworm AS go-builder
+
+ARG FDB_VERSION="7.1.57"
+ARG FDB_WEBSITE
+ARG TAG="latest"
+
+# TODO(johscheuer): Add support for arm64.
+RUN set -eux && \
+    curl --fail -L "https://github.com/apple/foundationdb/releases/download/${FDB_VERSION}/foundationdb-clients_${FDB_VERSION}-1_amd64.deb" -o foundationdb-clients_${FDB_VERSION}-1_amd64.deb && \
+    curl --fail -L "https://github.com/apple/foundationdb/releases/download/${FDB_VERSION}/foundationdb-clients_${FDB_VERSION}-1_amd64.deb.sha256" -o foundationdb-clients_${FDB_VERSION}-1_amd64.deb.sha256 && \
+    sha256sum -c foundationdb-clients_${FDB_VERSION}-1_amd64.deb.sha256 && \
+    dpkg -i foundationdb-clients_${FDB_VERSION}-1_amd64.deb && \
+    rm foundationdb-clients_${FDB_VERSION}-1_amd64.deb foundationdb-clients_${FDB_VERSION}-1_amd64.deb.sha256
+
+ENV CGO_ENABLED=1
+
+COPY go.mod go.sum /src/
+WORKDIR /src
+RUN go mod download
+
+COPY cmd /src/cmd
+COPY internal /src/internal
+
+RUN CGO_ENABLED=1 GOOS=linux GO111MODULE=on go build -ldflags="-s -w" -o /agent-scaler ./cmd
+
 FROM rockylinux/rockylinux:9.6
 ARG AGENT_TAG=joshua-agent:latest
 
 RUN dnf -y update && \
     dnf install -y \
         epel-release \
         scl-utils \
-        yum-utils && \
+        yum-utils \
+        pkg-config && \
     dnf -y install \
         gettext \
         procps-ng \
@@ -31,16 +57,16 @@ RUN chmod +x \
 
 # libfdb_c.so
 ARG FDB_VERSION="7.1.57"
-RUN curl -Ls https://github.com/apple/foundationdb/releases/download/${FDB_VERSION}/libfdb_c.x86_64.so \
-         -o /lib64/libfdb_c.so && \
-    chmod +x /lib64/libfdb_c.so
-
-ENV LD_LIBRARY_PATH="/lib64:$LD_LIBRARY_PATH"
+RUN rpm -ivh --excludepath=/usr/bin --excludepath=/usr/lib/foundationdb/backup_agent https://github.com/apple/foundationdb/releases/download/${FDB_VERSION}/foundationdb-clients-${FDB_VERSION}-1.el7.x86_64.rpm
 
 # FDB python binding
 RUN pip3 install \
-        foundationdb==7.1.57 \
-        boto3
+         foundationdb==7.1.57 \
+         boto3
+
+COPY --from=go-builder /agent-scaler /usr/local/bin
+COPY entrypoint.sh /tools/
+RUN chmod +x /tools/entrypoint.sh
 
 ENV BATCH_SIZE=1
 ENV MAX_JOBS=10
@@ -49,6 +75,6 @@ ENV AGENT_TIMEOUT=300
 ENV AGENT_TAG=${AGENT_TAG}
 ENV NAMESPACE=joshua
 ENV AGENT_NAME=joshua-agent
+ENV SCALER_IMPL=bash
 
-# Entry point
-ENTRYPOINT /tools/agent-scaler.sh
+ENTRYPOINT ["/tools/entrypoint.sh"]

k8s/agent-scaler/cmd/main.go

@@ -0,0 +1,110 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+
+	batchv1 "k8s.io/api/batch/v1"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/healthz"
+	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
+
+	"github.com/apple/fdb-joshua/k8s/agent-scaler/internal/config"
+	"github.com/apple/fdb-joshua/k8s/agent-scaler/internal/joshua"
+	"github.com/apple/fdb-joshua/k8s/agent-scaler/internal/scaler"
+)
+
+func main() {
+	ctrl.SetLogger(zap.New())
+	log := ctrl.Log.WithName("agent-scaler")
+
+	cfg, err := config.Load()
+	if err != nil {
+		log.Error(err, "loading config")
+		os.Exit(1)
+	}
+
+	scheme := runtime.NewScheme()
+	if err := clientgoscheme.AddToScheme(scheme); err != nil {
+		log.Error(err, "adding client-go scheme")
+		os.Exit(1)
+	}
+	if err := batchv1.AddToScheme(scheme); err != nil {
+		log.Error(err, "adding batchv1 scheme")
+		os.Exit(1)
+	}
+	if err := corev1.AddToScheme(scheme); err != nil {
+		log.Error(err, "adding corev1 scheme")
+		os.Exit(1)
+	}
+
+	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
+		Scheme: scheme,
+		Metrics: metricsserver.Options{
+			BindAddress: cfg.MetricsAddr,
+		},
+		HealthProbeBindAddress: cfg.HealthAddr,
+		LeaderElection:         false,
+	})
+	if err != nil {
+		log.Error(err, "creating manager")
+		os.Exit(1)
+	}
+
+	counter, err := joshua.NewEnsembleCounter(cfg.FDBClusterFile, cfg.JoshuaNamespace)
+	if err != nil {
+		log.Error(err, "opening FDB connection")
+		os.Exit(1)
+	}
+	s := scaler.New(cfg, mgr.GetClient(), counter, log.WithName("scaler"))
+
+	if err := mgr.Add(s); err != nil {
+		log.Error(err, "adding scaler to manager")
+		os.Exit(1)
+	}
+
+	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
+		log.Error(err, "adding healthz check")
+		os.Exit(1)
+	}
+	if err := mgr.AddReadyzCheck("readyz", func(_ *http.Request) error {
+		if !s.IsReady() {
+			return fmt.Errorf("not ready")
+		}
+		return nil
+	}); err != nil {
+		log.Error(err, "adding readyz check")
+		os.Exit(1)
+	}
+
+	log.Info("starting agent-scaler",
+		"agentName", cfg.AgentName,
+		"maxJobs", cfg.MaxJobs,
+		"batchSize", cfg.BatchSize,
+		"checkDelay", cfg.CheckDelay,
+		"namespace", cfg.Namespace,
+	)
+
+	// SIGKILL cannot be intercepted — the OS delivers it unconditionally.
+	// We handle SIGTERM (Kubernetes graceful stop) and SIGINT (Ctrl-C).
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
+	defer stop()
+
+	go func() {
+		<-ctx.Done()
+		log.Info("shutdown signal received, draining")
+	}()
+
+	if err := mgr.Start(ctx); err != nil {
+		log.Error(err, "manager exited")
+		os.Exit(1)
+	}
+}

k8s/agent-scaler/entrypoint.sh

@@ -0,0 +1,5 @@
+#!/bin/sh
+case "${SCALER_IMPL:-bash}" in
+  go) exec /usr/local/bin/agent-scaler ;;
+  *)  exec /tools/agent-scaler.sh ;;
+esac

k8s/agent-scaler/go.mod

@@ -0,0 +1,69 @@
+module github.com/apple/fdb-joshua/k8s/agent-scaler
+
+go 1.25
+
+require (
+	github.com/apple/foundationdb/bindings/go v0.0.0-20240218203134-05afc09840a7
+	github.com/go-logr/logr v1.4.2
+	github.com/prometheus/client_golang v1.20.5
+	github.com/spf13/pflag v1.0.5
+	k8s.io/api v0.31.3
+	k8s.io/apimachinery v0.31.3
+	k8s.io/client-go v0.31.3
+	sigs.k8s.io/controller-runtime v0.19.3
+	sigs.k8s.io/yaml v1.4.0
+)
+
+require github.com/x448/float16 v0.8.4 // indirect
+
+require (
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/emicklei/go-restful/v3 v3.11.0 // indirect
+	github.com/evanphx/json-patch/v5 v5.9.0 // indirect
+	github.com/fxamacker/cbor/v2 v2.7.0 // indirect
+	github.com/go-logr/zapr v1.3.0 // indirect
+	github.com/go-openapi/jsonpointer v0.19.6 // indirect
+	github.com/go-openapi/jsonreference v0.20.2 // indirect
+	github.com/go-openapi/swag v0.22.4 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
+	github.com/google/gnostic-models v0.6.8 // indirect
+	github.com/google/go-cmp v0.6.0 // indirect
+	github.com/google/gofuzz v1.2.0 // indirect
+	github.com/google/uuid v1.6.0 // indirect
+	github.com/imdario/mergo v0.3.6 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/klauspost/compress v1.17.9 // indirect
+	github.com/mailru/easyjson v0.7.7 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/prometheus/client_model v0.6.1 // indirect
+	github.com/prometheus/common v0.55.0 // indirect
+	github.com/prometheus/procfs v0.15.1 // indirect
+	go.uber.org/multierr v1.11.0 // indirect
+	go.uber.org/zap v1.26.0 // indirect
+	golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc // indirect
+	golang.org/x/net v0.26.0 // indirect
+	golang.org/x/oauth2 v0.21.0 // indirect
+	golang.org/x/sys v0.22.0 // indirect
+	golang.org/x/term v0.21.0 // indirect
+	golang.org/x/text v0.16.0 // indirect
+	golang.org/x/time v0.3.0 // indirect
+	gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
+	google.golang.org/protobuf v1.34.2 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+	k8s.io/apiextensions-apiserver v0.31.0 // indirect
+	k8s.io/klog/v2 v2.130.1 // indirect
+	k8s.io/kube-openapi v0.0.0-20240228011516-70dd3763d340 // indirect
+	k8s.io/utils v0.0.0-20240711033017-18e509b52bc8 // indirect
+	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
+)