Add option to include child processes in observation

internal/mock/psutil/psutil.go

@@ -89,7 +89,7 @@ func (u *MockPSUtil) GPU() ([]psutil.GPUInfo, error) {
 	return gpu, nil
 }
 
-func (u *MockPSUtil) Process(pid int32) (psutil.Process, error) {
+func (u *MockPSUtil) Process(pid int32, children bool) (psutil.Process, error) {
 	return &mockPSUtilProcess{}, nil
 }
 

resources/psutil/process.go

@@ -36,6 +36,7 @@ type process struct {
 	cpuLimit  uint64
 	ncpu      float64
 	proc      *psprocess.Process
+	children  bool
 	procfs    Procfs
 
 	stopTicker context.CancelFunc
@@ -51,13 +52,14 @@ type process struct {
 	gpu gpu.GPU
 }
 
-func (u *util) Process(pid int32) (Process, error) {
+func (u *util) Process(pid int32, children bool) (Process, error) {
 	p := &process{
 		pid:       pid,
 		hasCgroup: u.hasCgroup,
 		cpuLimit:  u.cpuLimit,
 		ncpu:      u.ncpu,
 		gpu:       u.gpu,
+		children:  children,
 		procfs:    u.procfs,
 	}
 
@@ -105,13 +107,15 @@ func (p *process) collectCPU() cpuTimesStat {
 		}
 	}
 
-	cstat := p.collectCPUFromChildren(p.proc)
+	if p.children {
+		cstat := p.collectCPUFromChildren(p.proc)
 
-	stat.total += cstat.total
-	stat.system += cstat.system
-	stat.user += cstat.user
-	stat.idle += cstat.idle
-	stat.other += cstat.other
+		stat.total += cstat.total
+		stat.system += cstat.system
+		stat.user += cstat.user
+		stat.idle += cstat.idle
+		stat.other += cstat.other
+	}
 
 	return *stat
 }
@@ -163,7 +167,9 @@ func (p *process) collectMemory() uint64 {
 
 	rss := info.RSS
 
-	rss += p.collectMemoryFromChildren(p.proc)
+	if p.children {
+		rss += p.collectMemoryFromChildren(p.proc)
+	}
 
 	return rss
 }

resources/psutil/psutil.go

@@ -113,7 +113,7 @@ type Util interface {
 	GPU() ([]GPUInfo, error)
 
 	// Process returns a process observer for a process with the given pid.
-	Process(pid int32) (Process, error)
+	Process(pid int32, children bool) (Process, error)
 }
 
 type util struct {

resources/resources.go

@@ -238,7 +238,7 @@ func New(config Config) (Resources, error) {
 		"max_gpu_memory": r.maxGPUMemory,
 	})
 
-	r.self, err = r.psutil.Process(int32(os.Getpid()))
+	r.self, err = r.psutil.Process(int32(os.Getpid()), false)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create process observer for self: %w", err)
 	}
@@ -639,7 +639,7 @@ func (r *resources) Network() ([]NetworkInfo, error) {
 }
 
 func (r *resources) Process(pid int32) (Process, error) {
-	proc, err := r.psutil.Process(pid)
+	proc, err := r.psutil.Process(pid, true)
 	if err != nil {
 		return nil, err
 	}