[PATCH 2/6] sched/numa: Leave the placement of a BPF-scheduled task to its scheduler

[Andrea Righi <arighi@nvidia.com>]

A task under a BPF scheduler can take NUMA hinting faults via
task_numa_fault() and then reach numa_migrate_preferred(), which may
migrate the task to its preferred node through task_numa_migrate().

The migration of such a task should always be driven by its BPF
scheduler, which owns the placement of its tasks.

Keep collecting the statistics, which is what makes
p->numa_preferred_nid worth reading, but skip the migration for a task
owned by a BPF scheduler. The scheduler can read the preferred node and
decide for itself whether moving the task is worth the cost.

For the same reason, never pick a task owned by a BPF scheduler as the
swap partner of another task that NUMA balancing migrates, e.g., a fair
task when the BPF scheduler runs in partial mode. Such a task is
treated like a kthread on the destination CPU: the migrating task can
still move there alone, but the BPF-scheduled task is left in place.

Reported-by: Vladimir Vdovin <deliran@verdict.gg>
Link: https://lore.kernel.org/r/20261002124559.10367-1-deliran@verdict.gg
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/fair.c | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 36df9a4dde534..5064c7b707f4d 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -3267,8 +3267,13 @@ static bool task_numa_compare(struct task_numa_env *env,
 
 	rcu_read_lock();
 	cur = rcu_dereference_all(dst_rq->curr);
+	/*
+	 * Never pick a task under a BPF scheduler as the swap partner: its
+	 * placement is left to that scheduler, see numa_migrate_preferred().
+	 * Moving the faulting task alone to the destination is still allowed.
+	 */
 	if (cur && ((cur->flags & (PF_EXITING | PF_KTHREAD)) ||
-		    !cur->mm))
+		    !cur->mm || task_on_scx(cur)))
 		cur = NULL;
 
 	/*
@@ -3645,6 +3650,14 @@ static void numa_migrate_preferred(struct task_struct *p)
 	if (task_node(p) == p->numa_preferred_nid)
 		return;
 
+	/*
+	 * A task under a BPF scheduler is placed by that scheduler. Keep the
+	 * statistics coming, which is what makes p->numa_preferred_nid worth
+	 * reading, but leave the placement alone.
+	 */
+	if (task_on_scx(p))
+		return;
+
 	/* Otherwise, try migrate to a CPU on the preferred node */
 	task_numa_migrate(p);
 }
-- 
2.56.0