update

Signed-off-by: Gyuho Lee <[email protected]>
leptonai · Jan 13, 2025 · 4eaa783 · 4eaa783
1 parent 1bd689e
commit 4eaa783
Show file tree

Hide file tree

Showing 8 changed files with 62 additions and 0 deletions.
diff --git a/components/accelerator/nvidia/nccl/component.go b/components/accelerator/nvidia/nccl/component.go
@@ -75,6 +75,11 @@ func (c *component) Events(ctx context.Context, since time.Time) ([]components.E
 	if common_dmesg.GetDefaultLogPoller() == nil {
 		return nil, nil
 	}
+	select {
+	case <-common_dmesg.GetDefaultLogPoller().WaitStart():
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
 
 	logItems, err := common_dmesg.GetDefaultLogPoller().Find(since)
 	if err != nil {

diff --git a/components/accelerator/nvidia/peermem/component.go b/components/accelerator/nvidia/peermem/component.go
@@ -119,6 +119,11 @@ func (c *component) getEvents(ctx context.Context, since time.Time) ([]component
 	if common_dmesg.GetDefaultLogPoller() == nil {
 		return nil, nil
 	}
+	select {
+	case <-common_dmesg.GetDefaultLogPoller().WaitStart():
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
 
 	logItems, err := common_dmesg.GetDefaultLogPoller().Find(since)
 	if err != nil {

diff --git a/components/cpu/component.go b/components/cpu/component.go
@@ -107,6 +107,11 @@ func (c *component) Events(ctx context.Context, since time.Time) ([]components.E
 	if common_dmesg.GetDefaultLogPoller() == nil {
 		return nil, nil
 	}
+	select {
+	case <-common_dmesg.GetDefaultLogPoller().WaitStart():
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
 
 	logItems, err := common_dmesg.GetDefaultLogPoller().Find(since)
 	if err != nil {

diff --git a/components/cpu/component_test.go b/components/cpu/component_test.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	query_config "github.com/leptonai/gpud/components/query/config"
+	"github.com/leptonai/gpud/pkg/sqlite"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
@@ -16,11 +17,21 @@ func TestComponent(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
+	db, err := sqlite.Open(":memory:")
+	if err != nil {
+		t.Fatalf("failed to open db: %v", err)
+	}
+	defer db.Close()
+
 	component := New(
 		ctx,
 		Config{
 			Query: query_config.Config{
 				Interval: metav1.Duration{Duration: time.Second},
+				State: &query_config.State{
+					DBRO: db,
+					DBRW: db,
+				},
 			},
 		},
 	)

diff --git a/components/fd/component.go b/components/fd/component.go
@@ -107,6 +107,11 @@ func (c *component) Events(ctx context.Context, since time.Time) ([]components.E
 	if common_dmesg.GetDefaultLogPoller() == nil {
 		return nil, nil
 	}
+	select {
+	case <-common_dmesg.GetDefaultLogPoller().WaitStart():
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
 
 	logItems, err := common_dmesg.GetDefaultLogPoller().Find(since)
 	if err != nil {

diff --git a/components/fd/component_test.go b/components/fd/component_test.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	query_config "github.com/leptonai/gpud/components/query/config"
+	"github.com/leptonai/gpud/pkg/sqlite"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
@@ -16,11 +17,21 @@ func TestComponent(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
+	db, err := sqlite.Open(":memory:")
+	if err != nil {
+		t.Fatalf("failed to open db: %v", err)
+	}
+	defer db.Close()
+
 	component := New(
 		ctx,
 		Config{
 			Query: query_config.Config{
 				Interval: metav1.Duration{Duration: 5 * time.Second},
+				State: &query_config.State{
+					DBRO: db,
+					DBRW: db,
+				},
 			},
 		},
 	)

diff --git a/components/memory/component.go b/components/memory/component.go
@@ -107,6 +107,11 @@ func (c *component) Events(ctx context.Context, since time.Time) ([]components.E
 	if common_dmesg.GetDefaultLogPoller() == nil {
 		return nil, nil
 	}
+	select {
+	case <-common_dmesg.GetDefaultLogPoller().WaitStart():
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
 
 	logItems, err := common_dmesg.GetDefaultLogPoller().Find(since)
 	if err != nil {

diff --git a/components/query/poller.go b/components/query/poller.go
@@ -27,6 +27,8 @@ type Poller interface {
 	// Redundant calls will be skipped if there's an existing poller.
 	Start(ctx context.Context, cfg query_config.Config, componentName string)
 
+	WaitStart() <-chan any
+
 	// Config returns the config used to start the poller.
 	// This is useful for debugging and logging.
 	Config() query_config.Config
@@ -84,6 +86,8 @@ func New(id string, cfg query_config.Config, getFunc GetFunc, getErrHandler GetE
 		startPollFunc:      startPoll,
 		getFunc:            getFunc,
 		getErrHandler:      getErrHandler,
+		startedCloseOnce:   sync.Once{},
+		startedCh:          make(chan any),
 		cfg:                cfg,
 		inflightComponents: make(map[string]any),
 	}
@@ -102,6 +106,9 @@ type poller struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
+	startedCloseOnce sync.Once
+	startedCh        chan any
+
 	cfgMu sync.RWMutex
 	cfg   query_config.Config
 
@@ -213,9 +220,17 @@ func (pl *poller) Start(ctx context.Context, cfg query_config.Config, componentN
 		}
 	}()
 
+	pl.startedCloseOnce.Do(func() {
+		close(pl.startedCh)
+	})
+
 	log.Logger.Debugw("started poller", "caller", componentName, "inflightComponents", len(pl.inflightComponents))
 }
 
+func (pl *poller) WaitStart() <-chan any {
+	return pl.startedCh
+}
+
 func (pl *poller) Stop(componentName string) bool {
 	pl.ctxMu.Lock()
 	defer pl.ctxMu.Unlock()