Skip to content

Commit 2ccf232

Browse files
JohnN193JohnN193
authored
[RSDK-9677] add GetModelsFromModules to robot interface (viamrobotics#4676)
1 parent fe84eee commit 2ccf232

File tree

13 files changed

+318
-14
lines changed

13 files changed

+318
-14
lines changed

‎module/modmanager/manager.go

+23
Original file line numberDiff line numberDiff line change
@@ -612,6 +612,29 @@ func (mgr *Manager) Configs() []config.Module {
612612
return configs
613613
}
614614

615+
// AllModels returns a slice of resource.ModuleModelDiscovery representing the available models
616+
// from the currently managed modules.
617+
func (mgr *Manager) AllModels() []resource.ModuleModelDiscovery {
618+
moduleTypes := map[string]config.ModuleType{}
619+
models := []resource.ModuleModelDiscovery{}
620+
for _, moduleConfig := range mgr.Configs() {
621+
moduleName := moduleConfig.Name
622+
moduleTypes[moduleName] = moduleConfig.Type
623+
}
624+
for moduleName, handleMap := range mgr.Handles() {
625+
for api, handle := range handleMap {
626+
for _, model := range handle {
627+
modelModel := resource.ModuleModelDiscovery{
628+
ModuleName: moduleName, Model: model, API: api.API,
629+
FromLocalModule: moduleTypes[moduleName] == config.ModuleTypeLocal,
630+
}
631+
models = append(models, modelModel)
632+
}
633+
}
634+
}
635+
return models
636+
}
637+
615638
// Provides returns true if a component/service config WOULD be handled by a module.
616639
func (mgr *Manager) Provides(conf resource.Config) bool {
617640
_, ok := mgr.getModule(conf)

‎module/modmanager/manager_test.go

+26
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,32 @@ func TestModManagerFunctions(t *testing.T) {
206206
test.That(t, ok, test.ShouldBeTrue)
207207
test.That(t, reg.Constructor, test.ShouldNotBeNil)
208208

209+
t.Log("test AllModels")
210+
modCfg2 := config.Module{
211+
Name: "simple-module2",
212+
ExePath: modPath,
213+
Type: config.ModuleTypeLocal,
214+
}
215+
err = mgr.Add(ctx, modCfg2)
216+
test.That(t, err, test.ShouldBeNil)
217+
models := mgr.AllModels()
218+
for _, model := range models {
219+
test.That(t, model.Model, test.ShouldResemble, resource.NewModel("acme", "demo", "mycounter"))
220+
test.That(t, model.API, test.ShouldResemble, resource.NewAPI("rdk", "component", "generic"))
221+
switch model.ModuleName {
222+
case "simple-module":
223+
test.That(t, model.FromLocalModule, test.ShouldEqual, false)
224+
case "simple-module2":
225+
test.That(t, model.FromLocalModule, test.ShouldEqual, true)
226+
default:
227+
t.Fail()
228+
t.Logf("test AllModels failure: unrecoginzed moduleName %v", model.ModuleName)
229+
}
230+
}
231+
names, err := mgr.Remove(modCfg2.Name)
232+
test.That(t, names, test.ShouldBeEmpty)
233+
test.That(t, err, test.ShouldBeNil)
234+
209235
t.Log("test Provides")
210236
ok = mgr.Provides(cfgCounter1)
211237
test.That(t, ok, test.ShouldBeTrue)

‎module/modmaninterface/interface.go

+1
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ type ModuleManager interface {
2424
CleanModuleDataDirectory() error
2525

2626
Configs() []config.Module
27+
AllModels() []resource.ModuleModelDiscovery
2728
Provides(cfg resource.Config) bool
2829
Handles() map[string]module.HandlerMap
2930

‎resource/discovery.go

+18
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import (
44
"context"
55
"fmt"
66

7+
pb "go.viam.com/api/robot/v1"
8+
79
"go.viam.com/rdk/logging"
810
)
911

@@ -33,8 +35,24 @@ type (
3335
Query DiscoveryQuery
3436
Cause error
3537
}
38+
39+
// ModuleModelDiscovery holds the API and Model information of models within a module.
40+
ModuleModelDiscovery struct {
41+
ModuleName string
42+
API API
43+
Model Model
44+
FromLocalModule bool
45+
}
3646
)
3747

48+
// ToProto converts a ModuleModelDiscovery into the equivalent proto message.
49+
func (mm *ModuleModelDiscovery) ToProto() *pb.ModuleModel {
50+
return &pb.ModuleModel{
51+
Model: mm.Model.String(), Api: mm.API.String(), ModuleName: mm.ModuleName,
52+
FromLocalModule: mm.FromLocalModule,
53+
}
54+
}
55+
3856
func (e *DiscoverError) Error() string {
3957
return fmt.Sprintf("failed to get discovery for api %q and model %q error: %v", e.Query.API, e.Query.Model, e.Cause)
4058
}

‎robot/client/client.go

+26
Original file line numberDiff line numberDiff line change
@@ -922,6 +922,32 @@ func (rc *RobotClient) DiscoverComponents(ctx context.Context, qs []resource.Dis
922922
return discoveries, nil
923923
}
924924

925+
// GetModelsFromModules returns the available models from the configured modules on a given machine.
926+
func (rc *RobotClient) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) {
927+
resp, err := rc.client.GetModelsFromModules(ctx, &pb.GetModelsFromModulesRequest{})
928+
if err != nil {
929+
return nil, err
930+
}
931+
protoModels := resp.GetModels()
932+
models := []resource.ModuleModelDiscovery{}
933+
for _, protoModel := range protoModels {
934+
modelTriplet, err := resource.NewModelFromString(protoModel.Model)
935+
if err != nil {
936+
return nil, err
937+
}
938+
api, err := resource.NewAPIFromString(protoModel.Api)
939+
if err != nil {
940+
return nil, err
941+
}
942+
model := resource.ModuleModelDiscovery{
943+
ModuleName: protoModel.ModuleName, Model: modelTriplet, API: api,
944+
FromLocalModule: protoModel.FromLocalModule,
945+
}
946+
models = append(models, model)
947+
}
948+
return models, nil
949+
}
950+
925951
// FrameSystemConfig returns the configuration of the frame system of a given machine.
926952
//
927953
// frameSystem, err := machine.FrameSystemConfig(context.Background(), nil)

‎robot/client/client_test.go

+54
Original file line numberDiff line numberDiff line change
@@ -1355,6 +1355,60 @@ func TestClientDiscovery(t *testing.T) {
13551355
test.That(t, err, test.ShouldBeNil)
13561356
}
13571357

1358+
func TestClientGetModelsFromModules(t *testing.T) {
1359+
injectRobot := &inject.Robot{}
1360+
injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil }
1361+
injectRobot.ResourceNamesFunc = func() []resource.Name {
1362+
return finalResources
1363+
}
1364+
injectRobot.MachineStatusFunc = func(_ context.Context) (robot.MachineStatus, error) {
1365+
return robot.MachineStatus{State: robot.StateRunning}, nil
1366+
}
1367+
expectedModels := []resource.ModuleModelDiscovery{
1368+
{
1369+
ModuleName: "simple-module",
1370+
API: resource.NewAPI("rdk", "component", "generic"),
1371+
Model: resource.NewModel("acme", "demo", "mycounter"),
1372+
FromLocalModule: false,
1373+
},
1374+
{
1375+
ModuleName: "simple-module2",
1376+
API: resource.NewAPI("rdk", "component", "generic"),
1377+
Model: resource.NewModel("acme", "demo", "mycounter"),
1378+
FromLocalModule: true,
1379+
},
1380+
}
1381+
injectRobot.GetModelsFromModulesFunc = func(context.Context) ([]resource.ModuleModelDiscovery, error) {
1382+
return expectedModels, nil
1383+
}
1384+
1385+
gServer := grpc.NewServer()
1386+
pb.RegisterRobotServiceServer(gServer, server.New(injectRobot))
1387+
listener, err := net.Listen("tcp", "localhost:0")
1388+
test.That(t, err, test.ShouldBeNil)
1389+
logger := logging.NewTestLogger(t)
1390+
1391+
go gServer.Serve(listener)
1392+
defer gServer.Stop()
1393+
1394+
client, err := New(context.Background(), listener.Addr().String(), logger)
1395+
test.That(t, err, test.ShouldBeNil)
1396+
1397+
resp, err := client.GetModelsFromModules(context.Background())
1398+
test.That(t, err, test.ShouldBeNil)
1399+
test.That(t, len(resp), test.ShouldEqual, 2)
1400+
test.That(t, resp, test.ShouldResemble, expectedModels)
1401+
for index, model := range resp {
1402+
test.That(t, model.ModuleName, test.ShouldEqual, expectedModels[index].ModuleName)
1403+
test.That(t, model.Model, test.ShouldResemble, expectedModels[index].Model)
1404+
test.That(t, model.API, test.ShouldResemble, expectedModels[index].API)
1405+
test.That(t, model.FromLocalModule, test.ShouldEqual, expectedModels[index].FromLocalModule)
1406+
}
1407+
1408+
err = client.Close(context.Background())
1409+
test.That(t, err, test.ShouldBeNil)
1410+
}
1411+
13581412
func ensurePartsAreEqual(part, otherPart *referenceframe.FrameSystemPart) error {
13591413
if part.FrameConfig.Name() != otherPart.FrameConfig.Name() {
13601414
return fmt.Errorf("part had name %s while other part had name %s", part.FrameConfig.Name(), otherPart.FrameConfig.Name())

‎robot/impl/discovery_test.go

+72
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,19 @@ import (
88
modulepb "go.viam.com/api/module/v1"
99
"go.viam.com/test"
1010

11+
"go.viam.com/rdk/components/base"
12+
"go.viam.com/rdk/components/generic"
1113
"go.viam.com/rdk/config"
14+
"go.viam.com/rdk/examples/customresources/apis/gizmoapi"
15+
"go.viam.com/rdk/examples/customresources/apis/summationapi"
16+
"go.viam.com/rdk/examples/customresources/models/mybase"
17+
"go.viam.com/rdk/examples/customresources/models/mygizmo"
18+
"go.viam.com/rdk/examples/customresources/models/mynavigation"
19+
"go.viam.com/rdk/examples/customresources/models/mysum"
1220
"go.viam.com/rdk/logging"
1321
"go.viam.com/rdk/resource"
1422
"go.viam.com/rdk/robot"
23+
"go.viam.com/rdk/services/navigation"
1524
rtestutils "go.viam.com/rdk/testutils"
1625
)
1726

@@ -168,3 +177,66 @@ func TestDiscovery(t *testing.T) {
168177
test.That(t, len(complexHandles.Handlers), test.ShouldBeGreaterThan, 1)
169178
})
170179
}
180+
181+
func TestGetModelsFromModules(t *testing.T) {
182+
t.Run("no modules configured", func(t *testing.T) {
183+
r := setupLocalRobotWithFakeConfig(t)
184+
models, err := r.GetModelsFromModules(context.Background())
185+
test.That(t, err, test.ShouldBeNil)
186+
test.That(t, models, test.ShouldBeEmpty)
187+
})
188+
t.Run("local and registry modules are configured", func(t *testing.T) {
189+
r := setupLocalRobotWithFakeConfig(t)
190+
ctx := context.Background()
191+
192+
// add modules
193+
complexPath := rtestutils.BuildTempModule(t, "examples/customresources/demos/complexmodule")
194+
simplePath := rtestutils.BuildTempModule(t, "examples/customresources/demos/simplemodule")
195+
cfg := &config.Config{
196+
Modules: []config.Module{
197+
{
198+
Name: "simple",
199+
ExePath: simplePath,
200+
Type: config.ModuleTypeRegistry,
201+
},
202+
{
203+
Name: "complex",
204+
ExePath: complexPath,
205+
Type: config.ModuleTypeLocal,
206+
},
207+
},
208+
}
209+
r.Reconfigure(ctx, cfg)
210+
models, err := r.GetModelsFromModules(context.Background())
211+
test.That(t, err, test.ShouldBeNil)
212+
test.That(t, models, test.ShouldHaveLength, 5)
213+
214+
for _, model := range models {
215+
switch model.Model {
216+
case mygizmo.Model:
217+
test.That(t, model.FromLocalModule, test.ShouldEqual, true)
218+
test.That(t, model.ModuleName, test.ShouldEqual, "complex")
219+
test.That(t, model.API, test.ShouldResemble, gizmoapi.API)
220+
case mysum.Model:
221+
test.That(t, model.FromLocalModule, test.ShouldEqual, true)
222+
test.That(t, model.ModuleName, test.ShouldEqual, "complex")
223+
test.That(t, model.API, test.ShouldResemble, summationapi.API)
224+
case mybase.Model:
225+
test.That(t, model.FromLocalModule, test.ShouldEqual, true)
226+
test.That(t, model.ModuleName, test.ShouldEqual, "complex")
227+
test.That(t, model.API, test.ShouldResemble, base.API)
228+
case mynavigation.Model:
229+
test.That(t, model.FromLocalModule, test.ShouldEqual, true)
230+
test.That(t, model.ModuleName, test.ShouldEqual, "complex")
231+
test.That(t, model.API, test.ShouldResemble, navigation.API)
232+
case resource.NewModel("acme", "demo", "mycounter"):
233+
test.That(t, model.FromLocalModule, test.ShouldEqual, false)
234+
test.That(t, model.ModuleName, test.ShouldEqual, "simple")
235+
test.That(t, model.API, test.ShouldResemble, generic.API)
236+
default:
237+
t.Fail()
238+
t.Logf("test GetModelsFromModules failure: unrecoginzed model %v", model.Model)
239+
}
240+
}
241+
})
242+
}

‎robot/impl/local_robot.go

+4
Original file line numberDiff line numberDiff line change
@@ -1077,6 +1077,10 @@ func (r *localRobot) discoverRobotInternals(query resource.DiscoveryQuery) (inte
10771077
}
10781078
}
10791079

1080+
func (r *localRobot) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) {
1081+
return r.manager.moduleManager.AllModels(), nil
1082+
}
1083+
10801084
func dialRobotClient(
10811085
ctx context.Context,
10821086
config config.Remote,

‎robot/impl/resource_manager_test.go

+9
Original file line numberDiff line numberDiff line change
@@ -1886,6 +1886,15 @@ func (rr *dummyRobot) DiscoverComponents(ctx context.Context, qs []resource.Disc
18861886
return rr.robot.DiscoverComponents(ctx, qs)
18871887
}
18881888

1889+
func (rr *dummyRobot) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) {
1890+
rr.mu.Lock()
1891+
defer rr.mu.Unlock()
1892+
if rr.offline {
1893+
return nil, errors.New("offline")
1894+
}
1895+
return rr.robot.GetModelsFromModules(ctx)
1896+
}
1897+
18891898
func (rr *dummyRobot) RemoteNames() []string {
18901899
return nil
18911900
}

‎robot/robot.go

+4
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,10 @@ type Robot interface {
8888
// Only implemented for webcam cameras in builtin components.
8989
DiscoverComponents(ctx context.Context, qs []resource.DiscoveryQuery) ([]resource.Discovery, error)
9090

91+
// GetModelsFromModules returns a list of models supported by the configured modules,
92+
// and specifies whether the models are from a local or registry module.
93+
GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error)
94+
9195
// RemoteByName returns a remote robot by name.
9296
RemoteByName(name string) (Robot, bool)
9397

‎robot/server/server.go

+13
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,19 @@ func (s *Server) DiscoverComponents(ctx context.Context, req *pb.DiscoverCompone
228228
return &pb.DiscoverComponentsResponse{Discovery: pbDiscoveries}, nil
229229
}
230230

231+
// GetModelsFromModules returns all models from the currently managed modules.
232+
func (s *Server) GetModelsFromModules(ctx context.Context, req *pb.GetModelsFromModulesRequest) (*pb.GetModelsFromModulesResponse, error) {
233+
models, err := s.robot.GetModelsFromModules(ctx)
234+
if err != nil {
235+
return nil, err
236+
}
237+
resp := pb.GetModelsFromModulesResponse{}
238+
for _, mm := range models {
239+
resp.Models = append(resp.Models, mm.ToProto())
240+
}
241+
return &resp, nil
242+
}
243+
231244
// FrameSystemConfig returns the info of each individual part that makes up the frame system.
232245
func (s *Server) FrameSystemConfig(ctx context.Context, req *pb.FrameSystemConfigRequest) (*pb.FrameSystemConfigResponse, error) {
233246
fsCfg, err := s.robot.FrameSystemConfig(ctx)

‎robot/server/server_test.go

+43
Original file line numberDiff line numberDiff line change
@@ -448,6 +448,49 @@ func TestServer(t *testing.T) {
448448
})
449449
})
450450

451+
t.Run("GetModelsFromModules", func(t *testing.T) {
452+
injectRobot := &inject.Robot{}
453+
injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil }
454+
injectRobot.ResourceNamesFunc = func() []resource.Name { return []resource.Name{} }
455+
server := server.New(injectRobot)
456+
457+
expectedModels := []resource.ModuleModelDiscovery{
458+
{
459+
ModuleName: "simple-module",
460+
API: resource.NewAPI("rdk", "component", "generic"),
461+
Model: resource.NewModel("acme", "demo", "mycounter"),
462+
FromLocalModule: false,
463+
},
464+
{
465+
ModuleName: "simple-module2",
466+
API: resource.NewAPI("rdk", "component", "generic"),
467+
Model: resource.NewModel("acme", "demo", "mycounter"),
468+
FromLocalModule: true,
469+
},
470+
}
471+
injectRobot.GetModelsFromModulesFunc = func(context.Context) ([]resource.ModuleModelDiscovery, error) {
472+
return expectedModels, nil
473+
}
474+
expectedProto := []*pb.ModuleModel{expectedModels[0].ToProto(), expectedModels[1].ToProto()}
475+
476+
req := &pb.GetModelsFromModulesRequest{}
477+
resp, err := server.GetModelsFromModules(context.Background(), req)
478+
test.That(t, err, test.ShouldBeNil)
479+
protoModels := resp.GetModels()
480+
test.That(t, len(protoModels), test.ShouldEqual, 2)
481+
test.That(t, protoModels, test.ShouldResemble, expectedProto)
482+
for index, protoModel := range protoModels {
483+
test.That(t, protoModel.ModuleName, test.ShouldEqual, expectedProto[index].ModuleName)
484+
test.That(t, protoModel.ModuleName, test.ShouldEqual, expectedModels[index].ModuleName)
485+
test.That(t, protoModel.Model, test.ShouldEqual, expectedProto[index].Model)
486+
test.That(t, protoModel.Model, test.ShouldEqual, expectedModels[index].Model.String())
487+
test.That(t, protoModel.Api, test.ShouldEqual, expectedProto[index].Api)
488+
test.That(t, protoModel.Api, test.ShouldEqual, expectedModels[index].API.String())
489+
test.That(t, protoModel.FromLocalModule, test.ShouldEqual, expectedProto[index].FromLocalModule)
490+
test.That(t, protoModel.FromLocalModule, test.ShouldEqual, expectedModels[index].FromLocalModule)
491+
}
492+
})
493+
451494
t.Run("ResourceRPCSubtypes", func(t *testing.T) {
452495
injectRobot := &inject.Robot{}
453496
injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil }

0 commit comments

Comments
 (0)