[RFC PATCH bpf-next v4 0/4] BPF-based per-cgroup NUMA balancing mode selection

[RFC PATCH bpf-next v4 0/4] BPF-based per-cgroup NUMA balancing mode selection

[Yafang Shao]

Problem
=======

In our large fleet of Kubernetes-managed servers, NUMA balancing has
historically been disabled globally on each server. With the growing
deployment of multi-socket AMD EPYC servers, cross-NUMA access has become
a critical performance issue, which led us to reconsider enabling NUMA
balancing.

kernel.numa_balancing is a single system-wide sysctl. Latency-sensitive
workloads want NUMA balancing disabled, while NUMA-aware workloads want it
enabled. No single global setting serves both.

Design
======

Core structure
~~~~~~~~~~~~~~

This series follows Tejun's suggestion to store the existing NUMA
balancing mode bits in the cgroup's memcg and expose them through BPF
kfuncs.

A new field, mem_cgroup::numab_mode, holds the per-cgroup NUMA balancing
setting. Two kfuncs are added to set and read it:

  bpf_numab_set_mode(): set the mode for a specific memory cgroup. The
                        update is rejected with -EOPNOTSUPP if the cgroup
                        has no visible, online memcg. Valid
                        values are the NUMAB_* mode enum and
                        NUMAB_FOLLOW_SYSTEM (-1); NUMAB_FOLLOW_SYSTEM
                        uses the global sysctl state.

  bpf_numab_get_mode(): return the combined task-side request of the
                        cgroup's effective memcg and of all its ancestors.
                        NUMAB_FOLLOW_SYSTEM is returned when no request is
                        present. A cgroup without a memcg resolves to its
                        nearest ancestor memcg. The following section
                        describes how the request is combined across the
                        hierarchy.

Both kfuncs are registered for BPF_PROG_TYPE_SYSCALL. bpf_numab_set_mode()
is KF_SLEEPABLE because it takes the mode mutex and the jump-label locks;
bpf_numab_get_mode() does not sleep. The mode values are exposed as a BTF
enum, so a BPF program can use NUMAB_FOLLOW_SYSTEM and the NUMAB_* mode
names without hardcoding private integers.

Effective mode semantics
~~~~~~~~~~~~~~~~~~~~~~~~

In the current implementation, a memory cgroup can be set to the
following modes:

  NUMAB_FOLLOW_SYSTEM (-1):   use the global sysctl state
  NUMAB_DISABLED (0):         disabled
  NUMAB_NORMAL (0x1):         normal mode
  NUMAB_TIERING (0x2):        tiering mode
  NUMAB_NORMAL_TIERING (0x3): normal mode plus memory tiering

The diagrams below assume that each cgroup shown has its own memcg. A
cgroup without one resolves to its nearest ancestor memcg, as described
in Storage of numab_mode.

Except for the sentinel and the opt-out described below, the effective
mode is the OR of this memcg's mode and that of all its ancestors.

- An opt-out wins for the task side and binds the subtree

  A DISABLED mode anywhere in the cgroup's ancestry opts the cgroup's
  tasks out entirely. Given:
                     Pods (normal)
                     /
          Pod-A (disabled)
          /           \
   cgroup-A1(normal)   cgroup-A2(tiering)

  The effective task-side mode for tasks(*) in cgroup-A1 and cgroup-A2
  is:

                     Pods (normal)
                     /
          Pod-A (disabled)
          /           \
   cgroup-A1(disabled)   cgroup-A2(disabled)

  (*) The folio side is described below.

- NUMAB_FOLLOW_SYSTEM is skipped

  Given:

                     Pods (normal)
                     /
          Pod-A (-1)
          /           \
   cgroup-A1(none)   cgroup-A2(none)

  The effective mode for tasks in cgroup-A1 and cgroup-A2 is:

                     Pods (normal)
                     /
          Pod-A (normal)
          /           \
   cgroup-A1(normal)   cgroup-A2(normal)

- For all other modes, the effective value is the OR of this memcg's
  mode and that of its ancestors

  Given:

                     Pods (normal)
                     /
          Pod-A (-1)
          /           \
   cgroup-A1(none)   cgroup-A2(tiering)

  The effective mode for tasks in cgroup-A1 and cgroup-A2 is:

                     Pods (normal)
                     /
          Pod-A (-1)
          /           \
   cgroup-A1(normal)   cgroup-A2(normal|tiering)

- The rationale behind these semantics

  The cgroup hierarchy semantics invoked here are those of resource
  limits: an ancestor's limit binds every descendant. That is the right
  semantics for limits: a limit partitions a shared, finite resource,
  and what a descendant consumes comes out of its ancestors' pool, so an
  ancestor's limit must bind every level below it.

  A policy request is a different kind of control: it is a preference
  about how the kernel should treat the workload, not a quantity drawn
  from a parent's pool. Expressing it consumes nothing, so the limit
  semantics does not carry over; the question is only how the preferences
  of several levels should combine.

  We chose the additive interpretation: a descendant may ask for more on
  top of what its ancestors asked for, while an opt-out at any level
  still wins: DISABLED anywhere in the ancestry caps the whole subtree.

Task and folio keying
~~~~~~~~~~~~~~~~~~~~~

The NUMA balancing knob is consulted at some callsites with only a folio
and no task, and at others with only a task and no folio.

- Task only

  kernel/sched/fair.c: task_numa_fault()
  kernel/sched/fair.c: update_scan_period()
  kernel/sched/fair.c: migrate_degrades_locality()
  kernel/sched/fair.c: exceed_llc_capacity()
  kernel/sched/fair.c: get_pref_llc()
  kernel/sched/fair.c: get_scan_cpumasks()

- Task and folio

  kernel/sched/fair.c: should_numa_migrate_memory()

- Folio only

  mm/memory-tiers.c: folio_use_access_time()
  mm/mempolicy.c: folio_can_map_prot_numa()
  mm/migrate.c: folio_migrate_flags()
  mm/migrate.c: migrate_misplaced_folio_prepare()
  mm/migrate.c: migrate_misplaced_folio()

Therefore, the helpers are split by side:

  folio_numab_effective_mode()
  folio_numab_in_*_mode()
  task_numab_enabled()
  task_numab_get_modes()

Another concern is how to handle callsites that have both a task and a
folio. The following anchor resolves it. Both the user fault path and the
kswapd demotion path reach folio_migrate_flags():

kswapd demotion (no task):

  kswapd
  + shrink_folio_list
    + demote_folio_list
      + migrate_pages
        + folio_migrate_flags
          + folio_numab_in_tiering_mode

Fault path (user task):

  do_numa_page
  + mpol_misplaced
    + should_numa_migrate_memory(task-m, folio-n) (A)
      + task or folio?
  + migrate_misplaced_folio_prepare(folio-n)
  + migrate_misplaced_folio(folio-n)
    + migrate_pages
      + folio_migrate_flags(folio-n) (B)
        + folio_numab_in_tiering_mode

folio_migrate_flags() must key on the folio, since kswapd has no task to
consult. In the user fault path, both should_numa_migrate_memory() and
folio_numab_in_tiering_mode() are called. Since both refer to the same
folio in this context, keying on the folio at B but not at A would be
inconsistent. The rule is therefore simple: if a folio is available, key
on the folio; otherwise, key on the task. This also keeps the keying
robust to future refactoring: if either folio-keyed function is ever
reached by a kernel thread, the answer does not change.

Storage of numab_mode
~~~~~~~~~~~~~~~~~~~~~

Some cgroups do not have the memory controller enabled:

Given:
          Pods (+memory)
          /
       Pod-A ()
        /
      cgroup-A1 ()

cgroup-A1 has no memcg. Given the task/folio split, where should the
mode for such a cgroup live?

- If the mode were stored in struct cgroup and setting it on cgroup-A1
  were allowed, tasks in cgroup-A1 would use cgroup-A1's request while
  its folios would use Pod-A's memcg request, which is inconsistent.

- Storing the mode in struct mem_cgroup avoids that: cgroup-A1 cannot be
  set, and both its tasks and its folios resolve to the same ancestor
  memcg (Pod-A).

For consistency, the mode is stored in mem_cgroup, setting it on a
cgroup without the memory controller enabled is rejected with
-EOPNOTSUPP, and both task and folio lookups resolve a cgroup without a
memcg to its nearest ancestor memcg. A request is effective only while
its memcg is visible and online. Disabling the memory controller or
offlining the memcg clears the request, and lookups skip such a memcg, so
a pinned memcg cannot keep applying a stale request. A re-enabled memcg
starts transparent.

Shared memory semantics
~~~~~~~~~~~~~~~~~~~~~~~

With the mode stored in the memcg, the longstanding shared-memory
problem in memcg has to be addressed for per-memcg NUMA balancing as
well.

Given:

    cgroup-A     cgroup-B
           \       | (charge)
           Memory_shared_by_cgroup-A_and_cgroup-B

A choice is required when the NUMA modes of cgroup-A and cgroup-B
differ. Since the folio-side lookup starts at the folio's memcg, a task
in cgroup-A operating on this shared memory uses cgroup-B's folio-side
mode. If DISABLED were applied to memory, cgroup-B's DISABLED would
prevent cgroup-A from promoting the shared page. We avoid that by making
DISABLED task-side only.

In our experience, NUMA balancing is mostly disabled to avoid the
potential latency spikes it can cause. The intent of disabling it for
cgroup-B is to avoid latency spikes for tasks in cgroup-B; it does not
matter whether tasks in other cgroups perform NUMA balancing. We
therefore disable the task side only, not the memory side: DISABLED
applies to tasks, not to memory.

Demotion is global reclaim, not a per-cgroup balancing request.
DISABLED only opts the cgroup's tasks out of NUMA balancing; it must
not become a memory-side reclaim veto, or a cgroup could keep cold
pages in the fast tier, leaving less fast-tier capacity for other
cgroups' hot pages and defeating promotion for those pages.

So, in the end, DISABLED is a task-side property and does not gate
memory-side operations.

Shared fast tier
~~~~~~~~~~~~~~~~

A promotion competes for capacity in the shared fast tier, but it
cannot consume more memory than the cgroup may already consume: it moves
the cgroup's own folios, within its own memory limit, from the slow tier
to the fast tier. Each contributor's own limit bounds its footprint,
but the fast tier is still a shared resource; demotion is the global
release valve that keeps cold pages from occupying it indefinitely.

So a child can still enable tiering mode if its ancestor has only set
normal.

Given:

                     Pod (normal)
                     /        \
          cgroup-A (none)  cgroup-B (tiering)

the effective mode is:

                      Pod (normal)
                     /        \
          cgroup-A (normal)  cgroup-B (normal|tiering)

What an ancestor cannot do now is cap only tiering while leaving normal
balancing available below. If that targeted cap is wanted, the natural
extension is a negative request: NO_TIERING would propagate down the
tree like the positive requests and mask TIERING out of the effective
mode. It would cap promotion, not demotion; demotion remains global
reclaim. With:

                  Pod (normal | no_tiering)
                     /            \
            cgroup-A (none)   cgroup-B (tiering)

the effective modes would be:

                      Pod (normal)
                     /            \
            cgroup-A (normal)  cgroup-B (normal)

For this version, I would keep the modes minimal and add NO_TIERING as a
follow-up once there is a real user for it.

Cost
====

On the task side, an ancestor can opt out the whole subtree with
DISABLED, but cannot cap only tiering while leaving normal balancing
available below. On the memory side there is no opt-out at all:
DISABLED is task-side only, so a descendant's TIERING request is enough
to enable promotion for its own memory, and promotion is the mode that
consumes the shared fast tier.

The per-cgroup override was measured with schbench, 384 workers, on a
two-node 384-core machine. In that configuration, a cgroup mode put the
cgroup at the throughput of the equivalent system-wide setting, within a
fraction of a percent. This is a throughput measurement, not a
microbenchmark of the scan path: while the static key is on, a folio
lookup is O(cgroup depth), and callers that need more than one bit now
share a single lookup. An effective-mode cache may be worthwhile for
very deep hierarchies; that is left for follow-up measurements.

Limitations
===========

The mode is per-memcg, not per-task. A cgroup without its own memcg
cannot be set; its tasks and folios resolve to the nearest ancestor
memcg. If multiple tasks in a single cgroup need different modes, place
them in separate child cgroups with the memory controller enabled.

For threaded cgroups, the mode must be set on the nearest domain
ancestor that has the memory controller enabled, because the memory
controller is not threaded.

Changelog
=========

RFC v3 -> RFC v4:
- Store numab_mode in mem_cgroup instead of struct cgroup
- Refine the commit log

RFC v3: https://lwn.net/Articles/1101442/

RFC v2 -> RFC v3:
- Replace the struct_ops interface with two kfuncs, following Tejun's
  suggestion.
- Improve function naming and simplify the code (David)

RFC v2: https://lwn.net/Articles/1099962/

RFC v1 -> RFC v2:
- v1 keyed the hook on the memory cgroup and temporarily embedded the
  cgroup id into the struct_ops value, because cgroup struct_ops support
  was still under development.
- v2 attaches through the cgroup struct_ops support that has since landed
  in bpf-next, keys the task side on the task's cgroup and the memory side
  on the folio's memcg, and adds the memory-tiering mode alongside enable
  and disable.

RFC v1: https://lwn.net/Articles/1054030/

Yafang Shao (4):
  sched: add helpers for numa balancing
  sched: sample NUMA balancing mode once per path
  sched: add BPF-based per-cgroup NUMA balancing mode selection
  selftests/bpf: add tests for the numab kfuncs

 include/linux/memcontrol.h                    |   6 +
 include/linux/sched/numa_balancing.h          | 144 +++++++++
 init/Kconfig                                  |  23 ++
 kernel/sched/bpf_numab.c                      | 261 +++++++++++++++++
 kernel/sched/build_policy.c                   |   4 +
 kernel/sched/fair.c                           |  99 +++++--
 kernel/sched/sched.h                          |   1 -
 mm/memcontrol.c                               |  14 +
 mm/memory-tiers.c                             |   9 +-
 mm/mempolicy.c                                |   9 +-
 mm/migrate.c                                  |   9 +-
 mm/vmscan.c                                   |   6 +-
 tools/testing/selftests/bpf/config            |   2 +
 .../selftests/bpf/prog_tests/bpf_numab.c      | 273 ++++++++++++++++++
 tools/testing/selftests/bpf/progs/bpf_numab.c |  32 ++
 15 files changed, 844 insertions(+), 48 deletions(-)
 create mode 100644 kernel/sched/bpf_numab.c
 create mode 100644 tools/testing/selftests/bpf/prog_tests/bpf_numab.c
 create mode 100644 tools/testing/selftests/bpf/progs/bpf_numab.c

-- 
2.52.0