COCOS-155 - Add python algo support (ultravioletrs#178)

* * feat(algorithm.go): add support for algorithm type context * feat(python.go): implement Python algorithm runtime * fix(cocos_defconfig): add IPTABLES package Signed-off-by: SammyOina <[email protected]> * update proto Signed-off-by: Sammy Oina <[email protected]> * small fixes Signed-off-by: Sammy Oina <[email protected]> * add metadata Signed-off-by: Sammy Oina <[email protected]> * debug Signed-off-by: Sammy Oina <[email protected]> * debug Signed-off-by: Sammy Oina <[email protected]> * chunk logger Signed-off-by: Sammy Oina <[email protected]> * debug logger Signed-off-by: Sammy Oina <[email protected]> * test lock Signed-off-by: Sammy Oina <[email protected]> * add req file Signed-off-by: SammyOina <[email protected]> * stream result Signed-off-by: SammyOina <[email protected]> * test with venv Signed-off-by: Sammy Oina <[email protected]> * fix missing requirements file Signed-off-by: Sammy Oina <[email protected]> * result stream Signed-off-by: Sammy Oina <[email protected]> * modify test server Signed-off-by: Sammy Oina <[email protected]> * remove debugging and cleaning up Signed-off-by: Sammy Oina <[email protected]> * original repo Signed-off-by: Sammy Oina <[email protected]> * add missing header Signed-off-by: Sammy Oina <[email protected]> * downgrade protoc Signed-off-by: Sammy Oina <[email protected]> --------- Signed-off-by: SammyOina <[email protected]> Signed-off-by: Sammy Oina <[email protected]>
WashingtonKK · Jul 23, 2024 · 67d01e3 · 67d01e3
1 parent 2f81098
commit 67d01e3
Show file tree

Hide file tree

Showing 24 changed files with 537 additions and 149 deletions.
diff --git a/.env b/.env
@@ -11,15 +11,11 @@ COCOS_JAEGER_OLTP_HTTP_PORT=4318
 ## Core Services
 
 ### Manager
-MANAGER_HTTP_HOST="cocos-manager"
-MANAGER_HTTP_PORT=9021
-MANAGER_HTTP_SERVER_CERT=""
-MANAGER_HTTP_SERVER_KEY=""
-MANAGER_GRPC_HOST="cocos-manager"
+MANAGER_GRPC_HOST=""
 MANAGER_GRPC_PORT=7003
 MANAGER_GRPC_SERVER_CERT=""
 MANAGER_GRPC_SERVER_KEY=""
-AGENT_GRPC_URL="192.168.100.4:7002"
+AGENT_GRPC_URL="localhost:7002"
 AGENT_GRPC_TIMEOUT=""
 AGENT_GRPC_CA_CERTS=""
 AGENT_GRPC_CLIENT_TLS=""
@@ -30,6 +26,3 @@ MANAGER_QEMU_ENABLE_SEV=false
 MANAGER_QEMU_SEV_CBITPOS=51
 MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/OVMF/OVMF_CODE.fd
 MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/OVMF_VARS.fd
-
-# Docker image tag
-COCOS_RELEASE_TAG=latest
diff --git a/.github/workflows/checkproto.yaml b/.github/workflows/checkproto.yaml
@@ -33,7 +33,7 @@ jobs:
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=27.2
+          PROTOC_VERSION=25.3
           PROTOC_GEN_VERSION=v1.34.2
           PROTOC_GRPC_VERSION=v1.4.0
 

diff --git a/agent/agent.pb.go b/agent/agent.pb.go
diff --git a/agent/agent.proto b/agent/agent.proto
@@ -10,12 +10,13 @@ option go_package = "./agent";
 service AgentService {
   rpc Algo(stream AlgoRequest) returns (AlgoResponse) {}
   rpc Data(stream DataRequest) returns (DataResponse) {}
-  rpc Result(ResultRequest) returns (ResultResponse) {}
+  rpc Result(ResultRequest) returns (stream ResultResponse) {}
   rpc Attestation(AttestationRequest) returns (AttestationResponse) {}
 }
 
 message AlgoRequest {
   bytes algorithm = 1;
+  bytes requirements = 2;
 }
 
 message AlgoResponse {}

diff --git a/agent/agent_grpc.pb.go b/agent/agent_grpc.pb.go
diff --git a/agent/algorithm/algorithm.go b/agent/algorithm/algorithm.go
@@ -2,8 +2,33 @@
 // SPDX-License-Identifier: Apache-2.0
 package algorithm
 
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+)
+
+type AlgorithType string
+
+const (
+	AlgoTypeBin    AlgorithType = "bin"
+	AlgoTypePython AlgorithType = "python"
+	AlgoTypeKey                 = "algo_type"
+)
+
+func AlgorithmTypeToContext(ctx context.Context, algoType string) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, AlgoTypeKey, algoType)
+}
+
+func AlgorithmTypeFromContext(ctx context.Context) string {
+	return metadata.ValueFromIncomingContext(ctx, AlgoTypeKey)[0]
+}
+
 // Algorithm is an interface that specifies the API for an algorithm.
 type Algorithm interface {
 	// Run executes the algorithm and returns the result.
 	Run() ([]byte, error)
+
+	// Add dataset to algorithm.
+	AddDataset(dataset string)
 }
diff --git a/agent/algorithm/binary/binary.go b/agent/algorithm/binary/binary.go
@@ -26,16 +26,19 @@ type binary struct {
 	stdout   io.Writer
 }
 
-func New(logger *slog.Logger, eventsSvc events.Service, algoFile string, datasets ...string) algorithm.Algorithm {
+func New(logger *slog.Logger, eventsSvc events.Service, algoFile string) algorithm.Algorithm {
 	return &binary{
 		algoFile: algoFile,
-		datasets: datasets,
 		logger:   logger,
 		stderr:   &algorithm.Stderr{Logger: logger, EventSvc: eventsSvc},
 		stdout:   &algorithm.Stdout{Logger: logger},
 	}
 }
 
+func (b *binary) AddDataset(dataset string) {
+	b.datasets = append(b.datasets, dataset)
+}
+
 func (b *binary) Run() ([]byte, error) {
 	defer os.Remove(b.algoFile)
 	defer func() {