diff --git a/module/modmanager/manager.go b/module/modmanager/manager.go index 96a7048d782..8e21558a97b 100644 --- a/module/modmanager/manager.go +++ b/module/modmanager/manager.go @@ -612,6 +612,29 @@ func (mgr *Manager) Configs() []config.Module { return configs } +// AllModels returns a slice of resource.ModuleModelDiscovery representing the available models +// from the currently managed modules. +func (mgr *Manager) AllModels() []resource.ModuleModelDiscovery { + moduleTypes := map[string]config.ModuleType{} + models := []resource.ModuleModelDiscovery{} + for _, moduleConfig := range mgr.Configs() { + moduleName := moduleConfig.Name + moduleTypes[moduleName] = moduleConfig.Type + } + for moduleName, handleMap := range mgr.Handles() { + for api, handle := range handleMap { + for _, model := range handle { + modelModel := resource.ModuleModelDiscovery{ + ModuleName: moduleName, Model: model, API: api.API, + FromLocalModule: moduleTypes[moduleName] == config.ModuleTypeLocal, + } + models = append(models, modelModel) + } + } + } + return models +} + // Provides returns true if a component/service config WOULD be handled by a module. func (mgr *Manager) Provides(conf resource.Config) bool { _, ok := mgr.getModule(conf) diff --git a/module/modmanager/manager_test.go b/module/modmanager/manager_test.go index 5ab749732f6..73029ea1846 100644 --- a/module/modmanager/manager_test.go +++ b/module/modmanager/manager_test.go @@ -206,6 +206,32 @@ func TestModManagerFunctions(t *testing.T) { test.That(t, ok, test.ShouldBeTrue) test.That(t, reg.Constructor, test.ShouldNotBeNil) + t.Log("test AllModels") + modCfg2 := config.Module{ + Name: "simple-module2", + ExePath: modPath, + Type: config.ModuleTypeLocal, + } + err = mgr.Add(ctx, modCfg2) + test.That(t, err, test.ShouldBeNil) + models := mgr.AllModels() + for _, model := range models { + test.That(t, model.Model, test.ShouldResemble, resource.NewModel("acme", "demo", "mycounter")) + test.That(t, model.API, test.ShouldResemble, resource.NewAPI("rdk", "component", "generic")) + switch model.ModuleName { + case "simple-module": + test.That(t, model.FromLocalModule, test.ShouldEqual, false) + case "simple-module2": + test.That(t, model.FromLocalModule, test.ShouldEqual, true) + default: + t.Fail() + t.Logf("test AllModels failure: unrecoginzed moduleName %v", model.ModuleName) + } + } + names, err := mgr.Remove(modCfg2.Name) + test.That(t, names, test.ShouldBeEmpty) + test.That(t, err, test.ShouldBeNil) + t.Log("test Provides") ok = mgr.Provides(cfgCounter1) test.That(t, ok, test.ShouldBeTrue) diff --git a/module/modmaninterface/interface.go b/module/modmaninterface/interface.go index 605910d2ba8..4763f923e25 100644 --- a/module/modmaninterface/interface.go +++ b/module/modmaninterface/interface.go @@ -24,6 +24,7 @@ type ModuleManager interface { CleanModuleDataDirectory() error Configs() []config.Module + AllModels() []resource.ModuleModelDiscovery Provides(cfg resource.Config) bool Handles() map[string]module.HandlerMap diff --git a/resource/discovery.go b/resource/discovery.go index 9b2a11f6cba..f4b6c8e686b 100644 --- a/resource/discovery.go +++ b/resource/discovery.go @@ -4,6 +4,8 @@ import ( "context" "fmt" + pb "go.viam.com/api/robot/v1" + "go.viam.com/rdk/logging" ) @@ -33,8 +35,24 @@ type ( Query DiscoveryQuery Cause error } + + // ModuleModelDiscovery holds the API and Model information of models within a module. + ModuleModelDiscovery struct { + ModuleName string + API API + Model Model + FromLocalModule bool + } ) +// ToProto converts a ModuleModelDiscovery into the equivalent proto message. +func (mm *ModuleModelDiscovery) ToProto() *pb.ModuleModel { + return &pb.ModuleModel{ + Model: mm.Model.String(), Api: mm.API.String(), ModuleName: mm.ModuleName, + FromLocalModule: mm.FromLocalModule, + } +} + func (e *DiscoverError) Error() string { return fmt.Sprintf("failed to get discovery for api %q and model %q error: %v", e.Query.API, e.Query.Model, e.Cause) } diff --git a/robot/client/client.go b/robot/client/client.go index 2dc5941e5a6..94461e0744c 100644 --- a/robot/client/client.go +++ b/robot/client/client.go @@ -922,6 +922,32 @@ func (rc *RobotClient) DiscoverComponents(ctx context.Context, qs []resource.Dis return discoveries, nil } +// GetModelsFromModules returns the available models from the configured modules on a given machine. +func (rc *RobotClient) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) { + resp, err := rc.client.GetModelsFromModules(ctx, &pb.GetModelsFromModulesRequest{}) + if err != nil { + return nil, err + } + protoModels := resp.GetModels() + models := []resource.ModuleModelDiscovery{} + for _, protoModel := range protoModels { + modelTriplet, err := resource.NewModelFromString(protoModel.Model) + if err != nil { + return nil, err + } + api, err := resource.NewAPIFromString(protoModel.Api) + if err != nil { + return nil, err + } + model := resource.ModuleModelDiscovery{ + ModuleName: protoModel.ModuleName, Model: modelTriplet, API: api, + FromLocalModule: protoModel.FromLocalModule, + } + models = append(models, model) + } + return models, nil +} + // FrameSystemConfig returns the configuration of the frame system of a given machine. // // frameSystem, err := machine.FrameSystemConfig(context.Background(), nil) diff --git a/robot/client/client_test.go b/robot/client/client_test.go index 76490b392d2..9e2e882cca9 100644 --- a/robot/client/client_test.go +++ b/robot/client/client_test.go @@ -1355,6 +1355,60 @@ func TestClientDiscovery(t *testing.T) { test.That(t, err, test.ShouldBeNil) } +func TestClientGetModelsFromModules(t *testing.T) { + injectRobot := &inject.Robot{} + injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil } + injectRobot.ResourceNamesFunc = func() []resource.Name { + return finalResources + } + injectRobot.MachineStatusFunc = func(_ context.Context) (robot.MachineStatus, error) { + return robot.MachineStatus{State: robot.StateRunning}, nil + } + expectedModels := []resource.ModuleModelDiscovery{ + { + ModuleName: "simple-module", + API: resource.NewAPI("rdk", "component", "generic"), + Model: resource.NewModel("acme", "demo", "mycounter"), + FromLocalModule: false, + }, + { + ModuleName: "simple-module2", + API: resource.NewAPI("rdk", "component", "generic"), + Model: resource.NewModel("acme", "demo", "mycounter"), + FromLocalModule: true, + }, + } + injectRobot.GetModelsFromModulesFunc = func(context.Context) ([]resource.ModuleModelDiscovery, error) { + return expectedModels, nil + } + + gServer := grpc.NewServer() + pb.RegisterRobotServiceServer(gServer, server.New(injectRobot)) + listener, err := net.Listen("tcp", "localhost:0") + test.That(t, err, test.ShouldBeNil) + logger := logging.NewTestLogger(t) + + go gServer.Serve(listener) + defer gServer.Stop() + + client, err := New(context.Background(), listener.Addr().String(), logger) + test.That(t, err, test.ShouldBeNil) + + resp, err := client.GetModelsFromModules(context.Background()) + test.That(t, err, test.ShouldBeNil) + test.That(t, len(resp), test.ShouldEqual, 2) + test.That(t, resp, test.ShouldResemble, expectedModels) + for index, model := range resp { + test.That(t, model.ModuleName, test.ShouldEqual, expectedModels[index].ModuleName) + test.That(t, model.Model, test.ShouldResemble, expectedModels[index].Model) + test.That(t, model.API, test.ShouldResemble, expectedModels[index].API) + test.That(t, model.FromLocalModule, test.ShouldEqual, expectedModels[index].FromLocalModule) + } + + err = client.Close(context.Background()) + test.That(t, err, test.ShouldBeNil) +} + func ensurePartsAreEqual(part, otherPart *referenceframe.FrameSystemPart) error { if part.FrameConfig.Name() != otherPart.FrameConfig.Name() { return fmt.Errorf("part had name %s while other part had name %s", part.FrameConfig.Name(), otherPart.FrameConfig.Name()) diff --git a/robot/impl/discovery_test.go b/robot/impl/discovery_test.go index 31c68be3044..f959b7104db 100644 --- a/robot/impl/discovery_test.go +++ b/robot/impl/discovery_test.go @@ -8,10 +8,19 @@ import ( modulepb "go.viam.com/api/module/v1" "go.viam.com/test" + "go.viam.com/rdk/components/base" + "go.viam.com/rdk/components/generic" "go.viam.com/rdk/config" + "go.viam.com/rdk/examples/customresources/apis/gizmoapi" + "go.viam.com/rdk/examples/customresources/apis/summationapi" + "go.viam.com/rdk/examples/customresources/models/mybase" + "go.viam.com/rdk/examples/customresources/models/mygizmo" + "go.viam.com/rdk/examples/customresources/models/mynavigation" + "go.viam.com/rdk/examples/customresources/models/mysum" "go.viam.com/rdk/logging" "go.viam.com/rdk/resource" "go.viam.com/rdk/robot" + "go.viam.com/rdk/services/navigation" rtestutils "go.viam.com/rdk/testutils" ) @@ -168,3 +177,66 @@ func TestDiscovery(t *testing.T) { test.That(t, len(complexHandles.Handlers), test.ShouldBeGreaterThan, 1) }) } + +func TestGetModelsFromModules(t *testing.T) { + t.Run("no modules configured", func(t *testing.T) { + r := setupLocalRobotWithFakeConfig(t) + models, err := r.GetModelsFromModules(context.Background()) + test.That(t, err, test.ShouldBeNil) + test.That(t, models, test.ShouldBeEmpty) + }) + t.Run("local and registry modules are configured", func(t *testing.T) { + r := setupLocalRobotWithFakeConfig(t) + ctx := context.Background() + + // add modules + complexPath := rtestutils.BuildTempModule(t, "examples/customresources/demos/complexmodule") + simplePath := rtestutils.BuildTempModule(t, "examples/customresources/demos/simplemodule") + cfg := &config.Config{ + Modules: []config.Module{ + { + Name: "simple", + ExePath: simplePath, + Type: config.ModuleTypeRegistry, + }, + { + Name: "complex", + ExePath: complexPath, + Type: config.ModuleTypeLocal, + }, + }, + } + r.Reconfigure(ctx, cfg) + models, err := r.GetModelsFromModules(context.Background()) + test.That(t, err, test.ShouldBeNil) + test.That(t, models, test.ShouldHaveLength, 5) + + for _, model := range models { + switch model.Model { + case mygizmo.Model: + test.That(t, model.FromLocalModule, test.ShouldEqual, true) + test.That(t, model.ModuleName, test.ShouldEqual, "complex") + test.That(t, model.API, test.ShouldResemble, gizmoapi.API) + case mysum.Model: + test.That(t, model.FromLocalModule, test.ShouldEqual, true) + test.That(t, model.ModuleName, test.ShouldEqual, "complex") + test.That(t, model.API, test.ShouldResemble, summationapi.API) + case mybase.Model: + test.That(t, model.FromLocalModule, test.ShouldEqual, true) + test.That(t, model.ModuleName, test.ShouldEqual, "complex") + test.That(t, model.API, test.ShouldResemble, base.API) + case mynavigation.Model: + test.That(t, model.FromLocalModule, test.ShouldEqual, true) + test.That(t, model.ModuleName, test.ShouldEqual, "complex") + test.That(t, model.API, test.ShouldResemble, navigation.API) + case resource.NewModel("acme", "demo", "mycounter"): + test.That(t, model.FromLocalModule, test.ShouldEqual, false) + test.That(t, model.ModuleName, test.ShouldEqual, "simple") + test.That(t, model.API, test.ShouldResemble, generic.API) + default: + t.Fail() + t.Logf("test GetModelsFromModules failure: unrecoginzed model %v", model.Model) + } + } + }) +} diff --git a/robot/impl/local_robot.go b/robot/impl/local_robot.go index 92269b31450..858544c542e 100644 --- a/robot/impl/local_robot.go +++ b/robot/impl/local_robot.go @@ -1077,6 +1077,10 @@ func (r *localRobot) discoverRobotInternals(query resource.DiscoveryQuery) (inte } } +func (r *localRobot) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) { + return r.manager.moduleManager.AllModels(), nil +} + func dialRobotClient( ctx context.Context, config config.Remote, diff --git a/robot/impl/resource_manager_test.go b/robot/impl/resource_manager_test.go index 6da6b4dc75d..771783a1c41 100644 --- a/robot/impl/resource_manager_test.go +++ b/robot/impl/resource_manager_test.go @@ -1886,6 +1886,15 @@ func (rr *dummyRobot) DiscoverComponents(ctx context.Context, qs []resource.Disc return rr.robot.DiscoverComponents(ctx, qs) } +func (rr *dummyRobot) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) { + rr.mu.Lock() + defer rr.mu.Unlock() + if rr.offline { + return nil, errors.New("offline") + } + return rr.robot.GetModelsFromModules(ctx) +} + func (rr *dummyRobot) RemoteNames() []string { return nil } diff --git a/robot/robot.go b/robot/robot.go index 652ff32abec..fc9e2f94b86 100644 --- a/robot/robot.go +++ b/robot/robot.go @@ -88,6 +88,10 @@ type Robot interface { // Only implemented for webcam cameras in builtin components. DiscoverComponents(ctx context.Context, qs []resource.DiscoveryQuery) ([]resource.Discovery, error) + // GetModelsFromModules returns a list of models supported by the configured modules, + // and specifies whether the models are from a local or registry module. + GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) + // RemoteByName returns a remote robot by name. RemoteByName(name string) (Robot, bool) diff --git a/robot/server/server.go b/robot/server/server.go index 93b02e1b827..fc391036e3c 100644 --- a/robot/server/server.go +++ b/robot/server/server.go @@ -228,6 +228,19 @@ func (s *Server) DiscoverComponents(ctx context.Context, req *pb.DiscoverCompone return &pb.DiscoverComponentsResponse{Discovery: pbDiscoveries}, nil } +// GetModelsFromModules returns all models from the currently managed modules. +func (s *Server) GetModelsFromModules(ctx context.Context, req *pb.GetModelsFromModulesRequest) (*pb.GetModelsFromModulesResponse, error) { + models, err := s.robot.GetModelsFromModules(ctx) + if err != nil { + return nil, err + } + resp := pb.GetModelsFromModulesResponse{} + for _, mm := range models { + resp.Models = append(resp.Models, mm.ToProto()) + } + return &resp, nil +} + // FrameSystemConfig returns the info of each individual part that makes up the frame system. func (s *Server) FrameSystemConfig(ctx context.Context, req *pb.FrameSystemConfigRequest) (*pb.FrameSystemConfigResponse, error) { fsCfg, err := s.robot.FrameSystemConfig(ctx) diff --git a/robot/server/server_test.go b/robot/server/server_test.go index fd8ca3a4598..dae6edd0f33 100644 --- a/robot/server/server_test.go +++ b/robot/server/server_test.go @@ -448,6 +448,49 @@ func TestServer(t *testing.T) { }) }) + t.Run("GetModelsFromModules", func(t *testing.T) { + injectRobot := &inject.Robot{} + injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil } + injectRobot.ResourceNamesFunc = func() []resource.Name { return []resource.Name{} } + server := server.New(injectRobot) + + expectedModels := []resource.ModuleModelDiscovery{ + { + ModuleName: "simple-module", + API: resource.NewAPI("rdk", "component", "generic"), + Model: resource.NewModel("acme", "demo", "mycounter"), + FromLocalModule: false, + }, + { + ModuleName: "simple-module2", + API: resource.NewAPI("rdk", "component", "generic"), + Model: resource.NewModel("acme", "demo", "mycounter"), + FromLocalModule: true, + }, + } + injectRobot.GetModelsFromModulesFunc = func(context.Context) ([]resource.ModuleModelDiscovery, error) { + return expectedModels, nil + } + expectedProto := []*pb.ModuleModel{expectedModels[0].ToProto(), expectedModels[1].ToProto()} + + req := &pb.GetModelsFromModulesRequest{} + resp, err := server.GetModelsFromModules(context.Background(), req) + test.That(t, err, test.ShouldBeNil) + protoModels := resp.GetModels() + test.That(t, len(protoModels), test.ShouldEqual, 2) + test.That(t, protoModels, test.ShouldResemble, expectedProto) + for index, protoModel := range protoModels { + test.That(t, protoModel.ModuleName, test.ShouldEqual, expectedProto[index].ModuleName) + test.That(t, protoModel.ModuleName, test.ShouldEqual, expectedModels[index].ModuleName) + test.That(t, protoModel.Model, test.ShouldEqual, expectedProto[index].Model) + test.That(t, protoModel.Model, test.ShouldEqual, expectedModels[index].Model.String()) + test.That(t, protoModel.Api, test.ShouldEqual, expectedProto[index].Api) + test.That(t, protoModel.Api, test.ShouldEqual, expectedModels[index].API.String()) + test.That(t, protoModel.FromLocalModule, test.ShouldEqual, expectedProto[index].FromLocalModule) + test.That(t, protoModel.FromLocalModule, test.ShouldEqual, expectedModels[index].FromLocalModule) + } + }) + t.Run("ResourceRPCSubtypes", func(t *testing.T) { injectRobot := &inject.Robot{} injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil } diff --git a/testutils/inject/robot.go b/testutils/inject/robot.go index 570cc8d1c2b..8b42d50c45b 100644 --- a/testutils/inject/robot.go +++ b/testutils/inject/robot.go @@ -26,20 +26,21 @@ import ( // Robot is an injected robot. type Robot struct { robot.LocalRobot - Mu sync.RWMutex // Ugly, has to be manually locked if a test means to swap funcs on an in-use robot. - DiscoverComponentsFunc func(ctx context.Context, keys []resource.DiscoveryQuery) ([]resource.Discovery, error) - RemoteByNameFunc func(name string) (robot.Robot, bool) - ResourceByNameFunc func(name resource.Name) (resource.Resource, error) - RemoteNamesFunc func() []string - ResourceNamesFunc func() []resource.Name - ResourceRPCAPIsFunc func() []resource.RPCAPI - ProcessManagerFunc func() pexec.ProcessManager - ConfigFunc func() *config.Config - LoggerFunc func() logging.Logger - CloseFunc func(ctx context.Context) error - StopAllFunc func(ctx context.Context, extra map[resource.Name]map[string]interface{}) error - FrameSystemConfigFunc func(ctx context.Context) (*framesystem.Config, error) - TransformPoseFunc func( + Mu sync.RWMutex // Ugly, has to be manually locked if a test means to swap funcs on an in-use robot. + DiscoverComponentsFunc func(ctx context.Context, keys []resource.DiscoveryQuery) ([]resource.Discovery, error) + GetModelsFromModulesFunc func(ctx context.Context) ([]resource.ModuleModelDiscovery, error) + RemoteByNameFunc func(name string) (robot.Robot, bool) + ResourceByNameFunc func(name resource.Name) (resource.Resource, error) + RemoteNamesFunc func() []string + ResourceNamesFunc func() []resource.Name + ResourceRPCAPIsFunc func() []resource.RPCAPI + ProcessManagerFunc func() pexec.ProcessManager + ConfigFunc func() *config.Config + LoggerFunc func() logging.Logger + CloseFunc func(ctx context.Context) error + StopAllFunc func(ctx context.Context, extra map[resource.Name]map[string]interface{}) error + FrameSystemConfigFunc func(ctx context.Context) (*framesystem.Config, error) + TransformPoseFunc func( ctx context.Context, pose *referenceframe.PoseInFrame, dst string, @@ -228,6 +229,16 @@ func (r *Robot) DiscoverComponents(ctx context.Context, keys []resource.Discover return r.DiscoverComponentsFunc(ctx, keys) } +// GetModelsFromModules calls the injected GetModelsFromModules or the real one. +func (r *Robot) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) { + r.Mu.RLock() + defer r.Mu.RUnlock() + if r.GetModelsFromModulesFunc == nil { + return r.LocalRobot.GetModelsFromModules(ctx) + } + return r.GetModelsFromModulesFunc(ctx) +} + // FrameSystemConfig calls the injected FrameSystemConfig or the real version. func (r *Robot) FrameSystemConfig(ctx context.Context) (*framesystem.Config, error) { r.Mu.RLock()