fix(accelerator/nvidia): add missing poller initialization

Signed-off-by: Gyuho Lee <[email protected]>
leptonai · Nov 15, 2024 · ee445ad · ee445ad
1 parent 9cc8610
commit ee445ad
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 0 deletions.
diff --git a/components/accelerator/nvidia/bad-envs/component.go b/components/accelerator/nvidia/bad-envs/component.go
@@ -17,6 +17,7 @@ func New(ctx context.Context, cfg Config) components.Component {
 	cfg.Query.SetDefaultsIfNotSet()
 
 	cctx, ccancel := context.WithCancel(ctx)
+	nvidia_query.SetDefaultPoller(cfg.Query.State.DB)
 	nvidia_query.GetDefaultPoller().Start(cctx, cfg.Query, bad_envs_id.Name)
 
 	return &component{

diff --git a/components/accelerator/nvidia/error-xid-sxid/component.go b/components/accelerator/nvidia/error-xid-sxid/component.go
@@ -24,6 +24,7 @@ func New(ctx context.Context, cfg Config) components.Component {
 
 	// this starts the Xid poller via "nvml.StartDefaultInstance"
 	cctx, ccancel := context.WithCancel(ctx)
+	nvidia_query.SetDefaultPoller(cfg.Query.State.DB)
 	nvidia_query.GetDefaultPoller().Start(cctx, cfg.Query, nvidia_error_xid_sxid_id.Name)
 
 	return &component{

diff --git a/components/accelerator/nvidia/gsp-firmware-mode/component.go b/components/accelerator/nvidia/gsp-firmware-mode/component.go
@@ -17,6 +17,7 @@ func New(ctx context.Context, cfg Config) components.Component {
 	cfg.Query.SetDefaultsIfNotSet()
 
 	cctx, ccancel := context.WithCancel(ctx)
+	nvidia_query.SetDefaultPoller(cfg.Query.State.DB)
 	nvidia_query.GetDefaultPoller().Start(cctx, cfg.Query, nvidia_gsp_firmware_mode_id.Name)
 
 	return &component{

diff --git a/components/accelerator/nvidia/nccl/component.go b/components/accelerator/nvidia/nccl/component.go
@@ -20,6 +20,7 @@ func New(ctx context.Context, cfg Config) components.Component {
 	cfg.Query.SetDefaultsIfNotSet()
 
 	cctx, ccancel := context.WithCancel(ctx)
+	nvidia_query.SetDefaultPoller(cfg.Query.State.DB)
 	nvidia_query.GetDefaultPoller().Start(cctx, cfg.Query, nvidia_nccl_id.Name)
 
 	return &component{

diff --git a/components/accelerator/nvidia/peermem/component.go b/components/accelerator/nvidia/peermem/component.go
@@ -20,6 +20,7 @@ func New(ctx context.Context, cfg Config) components.Component {
 	cfg.Query.SetDefaultsIfNotSet()
 
 	cctx, ccancel := context.WithCancel(ctx)
+	nvidia_query.SetDefaultPoller(cfg.Query.State.DB)
 	nvidia_query.GetDefaultPoller().Start(cctx, cfg.Query, nvidia_peermem_id.Name)
 
 	return &component{

diff --git a/components/accelerator/nvidia/persistence-mode/component.go b/components/accelerator/nvidia/persistence-mode/component.go
@@ -17,6 +17,7 @@ func New(ctx context.Context, cfg Config) components.Component {
 	cfg.Query.SetDefaultsIfNotSet()
 
 	cctx, ccancel := context.WithCancel(ctx)
+	nvidia_query.SetDefaultPoller(cfg.Query.State.DB)
 	nvidia_query.GetDefaultPoller().Start(cctx, cfg.Query, nvidia_persistence_mode_id.Name)
 
 	return &component{