Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RSDK-9677] add GetModelsFromModules to robot interface #4676

Merged
merged 18 commits into from
Jan 27, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions module/modmanager/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
26 changes: 26 additions & 0 deletions module/modmanager/manager_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions module/modmaninterface/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
18 changes: 18 additions & 0 deletions resource/discovery.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ import (
"context"
"fmt"

pb "go.viam.com/api/robot/v1"

"go.viam.com/rdk/logging"
)

Expand Down Expand Up @@ -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)
}
Expand Down
26 changes: 26 additions & 0 deletions robot/client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -917,6 +917,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)
Expand Down
54 changes: 54 additions & 0 deletions robot/client/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1354,6 +1354,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())
Expand Down
72 changes: 72 additions & 0 deletions robot/impl/discovery_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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)
}
}
})
}
4 changes: 4 additions & 0 deletions robot/impl/local_robot.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
9 changes: 9 additions & 0 deletions robot/impl/resource_manager_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
4 changes: 4 additions & 0 deletions robot/robot.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
13 changes: 13 additions & 0 deletions robot/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,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)
Expand Down
43 changes: 43 additions & 0 deletions robot/server/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -446,6 +446,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 }
Expand Down
Loading
Loading