Re: [PATCH 5/6] sched_ext: Add a kfunc to read a task's preferred NUMA node

[Andrea Righi <arighi@nvidia.com>]

On Tue, Oct 27, 2026 at 11:50:56AM +0100, Peter Zijlstra wrote:
> On Tue, Oct 27, 2026 at 07:57:36AM +0100, Andrea Righi wrote:
> > Hi Prateek,
> > 
> > On Tue, Oct 27, 2026 at 11:42:22AM +0530, K Prateek Nayak wrote:
> > > Hello Andrea,
> > > 
> > > On 10/27/2026 9:43 AM, Andrea Righi wrote:
> > > > +__bpf_kfunc s32 scx_bpf_task_numa_nid(const struct task_struct *p)
> > > > +{
> > > > +#ifdef CONFIG_NUMA_BALANCING
> > > > +	if (task_on_scx(p)) {
> > > > +		struct scx_sched *sch;
> > > > +
> > > > +		guard(rcu)();
> > > > +		sch = scx_task_sched_rcu(p);
> > > > +		if (!sch || !SCX_HAS_OP(sch, set_numa_node))
> > > > +			return NUMA_NO_NODE;
> > > > +	}
> > > > +
> > > > +	return READ_ONCE(p->numa_preferred_nid);
> > > 
> > > Would the ext scheduler need the raw numa_preferred_nid when
> > > sched_numa_balancing is disabled?
> > > 
> > > It may just be pointing to some stale value if NUMA balancing
> > > was disabled at runtime right?
> > 
> > Right, if NUMA balancing is disabled at runtime the value stays at the
> > last preference.
> > 
> > This is consistent with fair: disabling NUMA balancing at runtime only
> > flips the sched_numa_balancing static key, p->numa_preferred_nid is
> > never reset, fair stops acting on it in most places (e.g.,
> > migrate_degrades_locality()), but not everywhere, find_idlest_group()
> > still uses it in the slow wakeup path. So a BPF scheduler sees the same
> > value that fair keeps and can decide whether to use it or not.
> > 
> > A more accurate approach could be to notify the sched classes when
> > sched_numa_balancing changes (e.g., resetting the preferred node of all
> > tasks when it's disabled and recomputing it when it's re-enabled), but
> > that would be a significant rework, involving fair as well, and it's
> > outside the scope of this series. I'd rather start by making ext
> > consistent with fair, and then rework both in a more accurate and
> > complete way as a follow-up, if we want.
> 
> But should we not cleanup first and then add shiny on top?

Fair enough. Looking more closely, I think there're only two places in
fair that still use p->numa_preferred_nid when NUMA balancing is
disabled at runtime: find_idlest_group() and the fbq classification in
the load balancer (via rq->nr_numa_running / nr_preferred_running).
Everything else is already gated by sched_numa_balancing.

Assuming I'm not missing anything, we can probably add a preparatory
patch that gates those two on sched_numa_balancing as well, then make
scx_bpf_task_numa_nid() return NUMA_NO_NODE when NUMA balancing is
disabled. In this way the stale value is never used.

Thanks,
-Andrea