mirror of
https://github.com/ultravioletrs/cocos.git
synced 2026-06-23 04:10:25 +00:00
COCOS-344 - New agent structure (#350)
* new agent structure Signed-off-by: Sammy Oina <sammyoina@gmail.com> * minor fixes and testing Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix lint Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> * cvm tests fix Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix test Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix cli test Signed-off-by: Sammy Oina <sammyoina@gmail.com> * rename Signed-off-by: Sammy Oina <sammyoina@gmail.com> * rename cvm to cvms plural Signed-off-by: Sammy Oina <sammyoina@gmail.com> * rename service Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> * remove context Signed-off-by: Sammy Oina <sammyoina@gmail.com> * refactor: reorder parameters in NewAlgorithm functions and update CVMClient to CVMSClient Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix(tests): update SendEvent mock to include an additional parameter Signed-off-by: Sammy Oina <sammyoina@gmail.com> * move expectations Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix(tests): move event initialization to the correct scope in service tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix(tests): update SendEvent mock to use EXPECT instead of On in service tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> --------- Signed-off-by: Sammy Oina <sammyoina@gmail.com>
This commit is contained in:
Sammy Kerata Oina
committed by
GitHub
GitHub
parent
59b8057e5c
commit
ecad6514f3
@@ -13,18 +13,18 @@ import (
|
||||
|
||||
mglog "github.com/absmach/magistrala/logger"
|
||||
"github.com/caarlos0/env/v11"
|
||||
"github.com/ultravioletrs/cocos/agent/cvms"
|
||||
cvmgrpc "github.com/ultravioletrs/cocos/agent/cvms/api/grpc"
|
||||
"github.com/ultravioletrs/cocos/internal"
|
||||
"github.com/ultravioletrs/cocos/internal/server"
|
||||
grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
managergrpc "github.com/ultravioletrs/cocos/manager/api/grpc"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials"
|
||||
"google.golang.org/grpc/reflection"
|
||||
)
|
||||
|
||||
var _ managergrpc.Service = (*svc)(nil)
|
||||
var _ cvmgrpc.Service = (*svc)(nil)
|
||||
|
||||
const (
|
||||
svcName = "computations_test_server"
|
||||
@@ -42,7 +42,7 @@ type svc struct {
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
func (s *svc) Run(ctx context.Context, ipAddress string, sendMessage managergrpc.SendFunc, authInfo credentials.AuthInfo) {
|
||||
func (s *svc) Run(ctx context.Context, ipAddress string, sendMessage cvmgrpc.SendFunc, authInfo credentials.AuthInfo) {
|
||||
s.logger.Debug(fmt.Sprintf("received who am on ip address %s", ipAddress))
|
||||
|
||||
pubKey, err := os.ReadFile(pubKeyFile)
|
||||
@@ -52,7 +52,7 @@ func (s *svc) Run(ctx context.Context, ipAddress string, sendMessage managergrpc
|
||||
}
|
||||
pubPem, _ := pem.Decode(pubKey)
|
||||
|
||||
var datasets []*manager.Dataset
|
||||
var datasets []*cvms.Dataset
|
||||
for _, dataPath := range dataPaths {
|
||||
if _, err := os.Stat(dataPath); os.IsNotExist(err) {
|
||||
s.logger.Error(fmt.Sprintf("data file does not exist: %s", dataPath))
|
||||
@@ -64,7 +64,7 @@ func (s *svc) Run(ctx context.Context, ipAddress string, sendMessage managergrpc
|
||||
return
|
||||
}
|
||||
|
||||
datasets = append(datasets, &manager.Dataset{Hash: dataHash[:], UserKey: pubPem.Bytes})
|
||||
datasets = append(datasets, &cvms.Dataset{Hash: dataHash[:], UserKey: pubPem.Bytes})
|
||||
}
|
||||
|
||||
algoHash, err := internal.Checksum(algoPath)
|
||||
@@ -73,16 +73,16 @@ func (s *svc) Run(ctx context.Context, ipAddress string, sendMessage managergrpc
|
||||
return
|
||||
}
|
||||
|
||||
if err := sendMessage(&manager.ServerStreamMessage{
|
||||
Message: &manager.ServerStreamMessage_RunReq{
|
||||
RunReq: &manager.ComputationRunReq{
|
||||
if err := sendMessage(&cvms.ServerStreamMessage{
|
||||
Message: &cvms.ServerStreamMessage_RunReq{
|
||||
RunReq: &cvms.ComputationRunReq{
|
||||
Id: "1",
|
||||
Name: "sample computation",
|
||||
Description: "sample descrption",
|
||||
Datasets: datasets,
|
||||
Algorithm: &manager.Algorithm{Hash: algoHash[:], UserKey: pubPem.Bytes},
|
||||
ResultConsumers: []*manager.ResultConsumer{{UserKey: pubPem.Bytes}},
|
||||
AgentConfig: &manager.AgentConfig{
|
||||
Algorithm: &cvms.Algorithm{Hash: algoHash[:], UserKey: pubPem.Bytes},
|
||||
ResultConsumers: []*cvms.ResultConsumer{{UserKey: pubPem.Bytes}},
|
||||
AgentConfig: &cvms.AgentConfig{
|
||||
Port: "7002",
|
||||
LogLevel: "debug",
|
||||
AttestedTls: attestedTLS,
|
||||
@@ -113,7 +113,7 @@ func main() {
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
g, ctx := errgroup.WithContext(ctx)
|
||||
incomingChan := make(chan *manager.ClientStreamMessage)
|
||||
incomingChan := make(chan *cvms.ClientStreamMessage)
|
||||
|
||||
logger, err := mglog.New(os.Stdout, "debug")
|
||||
if err != nil {
|
||||
@@ -128,7 +128,7 @@ func main() {
|
||||
|
||||
registerAgentServiceServer := func(srv *grpc.Server) {
|
||||
reflection.Register(srv)
|
||||
manager.RegisterManagerServiceServer(srv, managergrpc.NewServer(incomingChan, &svc{logger: logger}))
|
||||
cvms.RegisterServiceServer(srv, cvmgrpc.NewServer(incomingChan, &svc{logger: logger}))
|
||||
}
|
||||
grpcServerConfig := server.ServerConfig{
|
||||
BaseConfig: server.BaseConfig{
|
||||
@@ -1,127 +0,0 @@
|
||||
// Copyright (c) Ultraviolet
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Simplified script to pass configs to agent without manager and read logs and events for manager.
|
||||
// This tool is meant for testing purposes.
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"log"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
|
||||
"github.com/mdlayher/vsock"
|
||||
"github.com/ultravioletrs/cocos/agent"
|
||||
"github.com/ultravioletrs/cocos/internal"
|
||||
internalvsock "github.com/ultravioletrs/cocos/internal/vsock"
|
||||
"github.com/ultravioletrs/cocos/manager"
|
||||
"github.com/ultravioletrs/cocos/manager/events"
|
||||
"github.com/ultravioletrs/cocos/manager/qemu"
|
||||
)
|
||||
|
||||
const (
|
||||
managerVsockPort = events.ManagerVsockPort
|
||||
vsockConfigPort = qemu.VsockConfigPort
|
||||
)
|
||||
|
||||
func main() {
|
||||
if len(os.Args) < 5 {
|
||||
log.Fatalf("usage: %s <data-path> <algo-path> <public-key-path> <attested-tls-bool>", os.Args[0])
|
||||
}
|
||||
dataPath := os.Args[1]
|
||||
algoPath := os.Args[2]
|
||||
pubKeyFile := os.Args[3]
|
||||
attestedTLSParam, err := strconv.ParseBool(os.Args[4])
|
||||
if err != nil {
|
||||
log.Fatalf("usage: %s <data-path> <algo-path> <public-key-path> <attested-tls-bool>, <attested-tls-bool> must be a bool value", os.Args[0])
|
||||
}
|
||||
attestedTLS := attestedTLSParam
|
||||
|
||||
pubKey, err := os.ReadFile(pubKeyFile)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to read public key file: %s", err)
|
||||
}
|
||||
pubPem, _ := pem.Decode(pubKey)
|
||||
algoHash, err := internal.Checksum(algoPath)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to calculate checksum: %s", err)
|
||||
}
|
||||
dataHash, err := internal.Checksum(dataPath)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to calculate checksum: %s", err)
|
||||
}
|
||||
|
||||
ac := agent.Computation{
|
||||
ID: "123",
|
||||
Datasets: agent.Datasets{agent.Dataset{Hash: [32]byte(dataHash), UserKey: pubPem.Bytes}},
|
||||
Algorithm: agent.Algorithm{Hash: [32]byte(algoHash), UserKey: pubPem.Bytes},
|
||||
ResultConsumers: []agent.ResultConsumer{{UserKey: pubPem.Bytes}},
|
||||
AgentConfig: agent.AgentConfig{
|
||||
LogLevel: "debug",
|
||||
Port: "7002",
|
||||
AttestedTls: attestedTLS,
|
||||
},
|
||||
}
|
||||
if err := sendAgentConfig(3, ac); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
listener, err := vsock.Listen(managerVsockPort, nil)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to listen on vsock: %s", err)
|
||||
}
|
||||
defer listener.Close()
|
||||
|
||||
log.Printf("Listening on vsock port %d", managerVsockPort)
|
||||
|
||||
for {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
log.Printf("failed to accept connection: %s", err)
|
||||
continue
|
||||
}
|
||||
|
||||
go handleConnection(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func sendAgentConfig(cid uint32, ac agent.Computation) error {
|
||||
conn, err := vsock.Dial(cid, qemu.VsockConfigPort, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close()
|
||||
payload, err := json.Marshal(ac)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var ac2 agent.Computation
|
||||
if err := json.Unmarshal(payload, &ac2); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := conn.Write(payload); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func handleConnection(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
|
||||
ackReader := internalvsock.NewAckReader(conn)
|
||||
|
||||
for {
|
||||
var message manager.ClientStreamMessage
|
||||
err := ackReader.ReadProto(&message)
|
||||
if err != nil {
|
||||
log.Printf("Error reading message: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("Received message: %s", message.String())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user