Files
cocos/pkg/sdk/agent.go
T
Sammy Kerata Oina 226704cf0d COCOS-122 - Enable streaming RPCs for Algo and Data services (#123)
* Enable streaming RPCs for Algo and Data services

Modified the gRPC service definitions for Algo and Data methods to use stream processing, enabling the handling of larger datasets and algorithms without being limited by memory restrictions. This allows client and server to send chunks of data sequentially rather than requiring the entire payload to be loaded into memory at once.

Updated server implementations to accumulate data from multiple chunks, allowing for more efficient processing and communication when dealing with large files. Client implementations have been adjusted to segment and send data in a streaming fashion.

Removed previously existing synchronous client code as it became redundant with the new streaming approach, streamlining the client's communication patterns with the gRPC backend.

This change allows for better resource management, especially in systems with constraints on memory, improving overall scalability and performance of the data and algorithm processing pipeline.

Signed-off-by: SammyOina <sammyoina@gmail.com>

* Refactor algorithm ID check logic

Simplify the algorithm validation logic in the agent service by replacing the previous containment check with direct ID comparison. This change streamlines the error handling for undeclared algorithms and hash mismatches, while also ensuring clear and direct provider validation. The modifications enhance the readability and maintainability of the code without altering functionality.

Signed-off-by: SammyOina <sammyoina@gmail.com>

* Updated README to build single-file executable with PyInstaller

Modified the PyInstaller command in the manual testing README to bundle the linear regression script into a single executable file. This simplifies distribution and execution of the script by eliminating the need for multiple dependency files.

Ref: Optimization of deployment process
Signed-off-by: SammyOina <sammyoina@gmail.com>

---------

Signed-off-by: SammyOina <sammyoina@gmail.com>
2024-05-13 13:14:50 +02:00

122 lines
2.3 KiB
Go

// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package sdk
import (
"bytes"
"context"
"io"
"log/slog"
"github.com/ultravioletrs/cocos/agent"
)
var _ agent.Service = (*agentSDK)(nil)
const (
size64 = 64
bufferSize = 1024 * 1024
)
type agentSDK struct {
client agent.AgentServiceClient
logger *slog.Logger
}
func NewAgentSDK(log *slog.Logger, agentClient agent.AgentServiceClient) *agentSDK {
return &agentSDK{
client: agentClient,
logger: log,
}
}
func (sdk *agentSDK) Algo(ctx context.Context, algorithm agent.Algorithm) error {
stream, err := sdk.client.Algo(ctx)
if err != nil {
sdk.logger.Error("Failed to call Algo RPC")
return err
}
algoBuffer := bytes.NewBuffer(algorithm.Algorithm)
buf := make([]byte, bufferSize)
for {
n, err := algoBuffer.Read(buf)
if err == io.EOF {
break
}
if err != nil {
return err
}
err = stream.Send(&agent.AlgoRequest{Id: algorithm.ID, Provider: algorithm.Provider, Algorithm: buf[:n]})
if err != nil {
return err
}
}
if _, err := stream.CloseAndRecv(); err != nil {
return err
}
return nil
}
func (sdk *agentSDK) Data(ctx context.Context, dataset agent.Dataset) error {
stream, err := sdk.client.Data(ctx)
if err != nil {
sdk.logger.Error("Failed to call Algo RPC")
return err
}
dataBuffer := bytes.NewBuffer(dataset.Dataset)
buf := make([]byte, bufferSize)
for {
n, err := dataBuffer.Read(buf)
if err == io.EOF {
break
}
if err != nil {
return err
}
err = stream.Send(&agent.DataRequest{Id: dataset.ID, Provider: dataset.Provider, Dataset: buf[:n]})
if err != nil {
return err
}
}
if _, err := stream.CloseAndRecv(); err != nil {
return err
}
return nil
}
func (sdk *agentSDK) Result(ctx context.Context, consumer string) ([]byte, error) {
request := &agent.ResultRequest{
Consumer: consumer,
}
response, err := sdk.client.Result(ctx, request)
if err != nil {
sdk.logger.Error("Failed to call Result RPC")
return nil, err
}
return response.File, nil
}
func (sdk *agentSDK) Attestation(ctx context.Context, reportData [size64]byte) ([]byte, error) {
request := &agent.AttestationRequest{
ReportData: reportData[:],
}
response, err := sdk.client.Attestation(ctx, request)
if err != nil {
sdk.logger.Error("Failed to call Attestation RPC")
return nil, err
}
return response.File, nil
}