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

[Andrea Righi <arighi@nvidia.com>]

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.

> 
> > +#else
> > +	return NUMA_NO_NODE;
> > +#endif
> > +}
> > +
> >  /**
> >   * scx_bpf_locked_rq - Return the rq currently locked by SCX
> >   * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
> > @@ -11836,6 +11877,7 @@ BTF_ID_FLAGS(func, scx_bpf_put_cpumask, KF_RELEASE)
> >  BTF_ID_FLAGS(func, scx_bpf_task_running, KF_RCU)
> 
> On a tangential note, scx_bpf_task_running() just does a
> rcu_access_pointer() (resolves to a READ_ONCE()) and doesn't take a
> rcu_read_lock() so does it really need KF_RCU?

Ah yes, that's tricky. KF_RCU doesn't mean that the kfunc takes
rcu_read_lock(), it's a verifier flag that allows the kfunc to accept
RCU-protected pointers as arguments, in addition to trusted (refcounted)
ones.

scx_bpf_task_running() still dereferences p (task_rq(p)->donor), so p
needs to be valid for the duration of the call, and KF_RCU is the
weakest guarantee that the verifier can enforce for that. Same for
scx_bpf_task_numa_nid().

> >  BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
> >  BTF_ID_FLAGS(func, scx_bpf_task_cid, KF_RCU)
> > +BTF_ID_FLAGS(func, scx_bpf_task_numa_nid, KF_RCU)
> >  BTF_ID_FLAGS(func, scx_bpf_locked_rq, KF_IMPLICIT_ARGS | KF_RET_NULL)
> >  BTF_ID_FLAGS(func, scx_bpf_cpu_curr, KF_IMPLICIT_ARGS | KF_RET_NULL | KF_RCU_PROTECTED)
> >  BTF_ID_FLAGS(func, scx_bpf_cid_curr, KF_IMPLICIT_ARGS | KF_RET_NULL | KF_RCU_PROTECTED)

Thanks for taking a look at this series!
-Andrea