Initial commit

Use pickle for model2wasm

Fix TinyGo entrypoint and m2cgen wrapping workflow

Add FL task mode and metadata fields

Signed-off-by: JeffMboya <jangina.mboya@gmail.com>

Include FL mode and spec in start payload

Signed-off-by: JeffMboya <jangina.mboya@gmail.com>

Add FL update envelope type for results

Signed-off-by: JeffMboya <jangina.mboya@gmail.com>

Deduplicate FL helpers; update host and wazero

Signed-off-by: JeffMboya <jangina.mboya@gmail.com>

Handle train results as FL envelope updates

Signed-off-by: JeffMboya <jangina.mboya@gmail.com>

Author: JeffMboya

examples/ml-wasm/main.go

@@ -0,0 +1,14 @@
+package main
+
+func predict(x0 int32, x1 int32) int32 {
+    in := []float64{
+        float64(x0) / 100.0,
+        float64(x1) / 100.0,
+    }
+
+    y := score(in)
+
+    return int32(y * 100.0)
+}
+
+func main() {}

examples/ml-wasm/model_gen.go

@@ -0,0 +1,6 @@
+package main
+
+func score(input []float64) float64 {
+    return 0.0000000000000008881784197001252 + input[0] * 1.9999999999999996 + input[1] * 2.999999999999999
+}
+

examples/ml-wasm/mymodel.pkl

@@ -0,0 +1,28 @@
+from sklearn.linear_model import LinearRegression
+import numpy as np
+import os
+import pickle
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+
+model_path = os.path.join(script_dir, "mymodel.pkl")
+
+# Fake training data: y = 2*x1 + 3*x2
+X = np.array([
+    [0, 0],
+    [1, 0],
+    [0, 1],
+    [1, 1],
+    [2, 1],
+    [1, 2],
+], dtype=float)
+
+y = 2 * X[:, 0] + 3 * X[:, 1]
+
+model = LinearRegression()
+model.fit(X, y)
+
+with open(model_path, "wb") as f:
+    pickle.dump(model, f)
+
+print(f"Saved model to {model_path}")

examples/ml-wasm/mymodel.wasm

@@ -2,13 +2,15 @@ package manager
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"time"
 
 	"github.com/0x6flab/namegenerator"
 	pkgerrors "github.com/absmach/propeller/pkg/errors"
+	flpkg "github.com/absmach/propeller/pkg/fl"
 	"github.com/absmach/propeller/pkg/mqtt"
 	"github.com/absmach/propeller/pkg/proplet"
 	"github.com/absmach/propeller/pkg/scheduler"
@@ -435,6 +437,70 @@ func (svc *service) updateResultsHandler(ctx context.Context, msg map[string]any
 	t.UpdatedAt = time.Now()
 	t.FinishTime = time.Now()
 
+	if errMsg, ok := msg["error"].(string); ok && errMsg != "" {
+		t.Error = errMsg
+	}
+
+		if err := svc.tasksDB.Update(ctx, t.ID, t); err != nil {
+			return err
+		}
+
+		const scanLimit = 10_000
+		data, _, lErr := svc.tasksDB.List(ctx, 0, scanLimit)
+		if lErr != nil {
+			return lErr
+		}
+
+		var (
+			pendingOrRunning bool
+			numUpdates       uint64
+			totalSamples     uint64
+		)
+
+		for i := range data {
+			tt, ok := data[i].(task.Task)
+			if !ok {
+				return pkgerrors.ErrInvalidData
+			}
+			if tt.Mode != task.ModeTrain || tt.FL == nil {
+				continue
+			}
+			if tt.FL.JobID != envlp.JobID || tt.FL.RoundID != envlp.RoundID {
+				continue
+			}
+
+			switch tt.State {
+			case task.Pending, task.Scheduled, task.Running:
+				pendingOrRunning = true
+			}
+
+			if tt.State == task.Completed && tt.Error == "" {
+				if u, ok := tt.Results.(flpkg.UpdateEnvelope); ok {
+					numUpdates++
+					totalSamples += u.NumSamples
+				} else {
+					// If storage rehydrates as map[string]any, ignore for aggregation stub.
+				}
+			}
+		}
+
+		if !pendingOrRunning {
+			svc.logger.InfoContext(ctx, "FL round complete (aggregation stub)",
+				slog.String("job_id", envlp.JobID),
+				slog.Uint64("round_id", envlp.RoundID),
+				slog.Uint64("num_updates", numUpdates),
+				slog.Uint64("total_samples", totalSamples),
+			)
+		}
+
+		return nil
+	}
+
+	t.Results = msg["results"]
+	t.State = task.Completed
+	t.UpdatedAt = time.Now()
+	t.FinishTime = time.Now()
+
 	if errMsg, ok := msg["error"].(string); ok && errMsg != "" {
 		t.Error = errMsg
 	}
@@ -534,6 +600,94 @@ func (svc *service) handlePropletMetrics(ctx context.Context, msg map[string]any
 	return nil
 }
 
+func (svc *service) handleTaskMetrics(ctx context.Context, msg map[string]any) error {
+	taskID, ok := msg["task_id"].(string)
+	if !ok {
+		return errors.New("invalid task_id")
+	}
+	if taskID == "" {
+		return errors.New("task id is empty")
+	}
+
+	propletID, ok := msg["proplet_id"].(string)
+	if !ok {
+		return errors.New("invalid proplet_id")
+	}
+
+	taskMetrics := TaskMetrics{
+		TaskID:    taskID,
+		PropletID: propletID,
+	}
+
+	if ts, ok := msg["timestamp"].(string); ok {
+		if t, err := time.Parse(time.RFC3339Nano, ts); err == nil {
+			taskMetrics.Timestamp = t
+		}
+	}
+	if taskMetrics.Timestamp.IsZero() {
+		taskMetrics.Timestamp = time.Now()
+	}
+
+	if metricsData, ok := msg["metrics"].(map[string]any); ok {
+		taskMetrics.Metrics = svc.parseProcessMetrics(metricsData)
+	}
+
+	if aggData, ok := msg["aggregated"].(map[string]any); ok {
+		taskMetrics.Aggregated = svc.parseAggregatedMetrics(aggData)
+	}
+
+	key := fmt.Sprintf("%s:%d", taskID, taskMetrics.Timestamp.UnixNano())
+	if err := svc.metricsDB.Create(ctx, key, taskMetrics); err != nil {
+		svc.logger.WarnContext(ctx, "failed to store task metrics", "error", err, "task_id", taskID)
+
+		return err
+	}
+
+	return nil
+}
+
+func (svc *service) handlePropletMetrics(ctx context.Context, msg map[string]any) error {
+	propletID, ok := msg["proplet_id"].(string)
+	if !ok {
+		return errors.New("invalid proplet_id")
+	}
+	if propletID == "" {
+		return errors.New("proplet id is empty")
+	}
+	namespace, _ := msg["namespace"].(string)
+
+	propletMetrics := PropletMetrics{
+		PropletID: propletID,
+		Namespace: namespace,
+	}
+
+	if ts, ok := msg["timestamp"].(string); ok {
+		if t, err := time.Parse(time.RFC3339Nano, ts); err == nil {
+			propletMetrics.Timestamp = t
+		}
+	}
+	if propletMetrics.Timestamp.IsZero() {
+		propletMetrics.Timestamp = time.Now()
+	}
+
+	if cpuData, ok := msg["cpu_metrics"].(map[string]any); ok {
+		propletMetrics.CPU = svc.parseCPUMetrics(cpuData)
+	}
+
+	if memData, ok := msg["memory_metrics"].(map[string]any); ok {
+		propletMetrics.Memory = svc.parseMemoryMetrics(memData)
+	}
+
+	key := fmt.Sprintf("%s:%d", propletID, propletMetrics.Timestamp.UnixNano())
+	if err := svc.metricsDB.Create(ctx, key, propletMetrics); err != nil {
+		svc.logger.WarnContext(ctx, "failed to store proplet metrics", "error", err, "proplet_id", propletID)
+
+		return err
+	}
+
+	return nil
+}
+
 func (svc *service) parseProcessMetrics(data map[string]any) proplet.ProcessMetrics {
 	metrics := proplet.ProcessMetrics{}
 

examples/ml-wasm/train_model.py

@@ -0,0 +1,13 @@
+package fl
+
+type UpdateEnvelope struct {
+	TaskID        string         `json:"task_id,omitempty"`
+	JobID         string         `json:"job_id"`
+	RoundID       uint64         `json:"round_id"`
+	GlobalVersion string         `json:"global_version"`
+	PropletID     string         `json:"proplet_id"`
+	NumSamples    uint64         `json:"num_samples"`
+	UpdateB64     string         `json:"update_b64"`
+	Metrics       map[string]any `json:"metrics,omitempty"`
+	Format        string         `json:"format,omitempty"`
+}

manager/service.go

@@ -0,0 +1,79 @@
+package runtimes
+
+import (
+	"encoding/base64"
+	"fmt"
+	"strconv"
+
+	flpkg "github.com/absmach/propeller/pkg/fl"
+)
+
+func buildFLPayloadFromString(taskID, mode, propletID string, env map[string]string, rawOut string) map[string]any {
+	// Backward-compatible inference behavior.
+	if mode != "train" {
+		return map[string]any{
+			"task_id": taskID,
+			"results": rawOut,
+		}
+	}
+
+	envlp := flpkg.UpdateEnvelope{
+		TaskID:        taskID,
+		JobID:         env["FL_JOB_ID"],
+		RoundID:       parseUint64(env, "FL_ROUND_ID"),
+		GlobalVersion: env["FL_GLOBAL_VERSION"],
+		PropletID:     propletID,
+		NumSamples:    parseUint64(env, "FL_NUM_SAMPLES"),
+		UpdateB64:     base64.StdEncoding.EncodeToString([]byte(rawOut)),
+		Metrics:       nil,
+		Format:        env["FL_FORMAT"],
+	}
+
+	return map[string]any{
+		"task_id": taskID,
+		"results": envlp,
+	}
+}
+
+func buildFLPayloadFromUint64Slice(taskID, mode, propletID string, env map[string]string, results []uint64) map[string]any {
+	if mode != "train" {
+		return map[string]any{
+			"task_id": taskID,
+			"results": results,
+		}
+	}
+
+	raw := fmt.Sprint(results)
+
+	envlp := flpkg.UpdateEnvelope{
+		TaskID:        taskID,
+		JobID:         env["FL_JOB_ID"],
+		RoundID:       parseUint64(env, "FL_ROUND_ID"),
+		GlobalVersion: env["FL_GLOBAL_VERSION"],
+		PropletID:     propletID,
+		NumSamples:    parseUint64(env, "FL_NUM_SAMPLES"),
+		UpdateB64:     base64.StdEncoding.EncodeToString([]byte(raw)),
+		Metrics:       nil,
+		Format:        env["FL_FORMAT"],
+	}
+
+	return map[string]any{
+		"task_id": taskID,
+		"results": envlp,
+	}
+}
+
+func parseUint64(env map[string]string, key string) uint64 {
+	if env == nil {
+		return 0
+	}
+	s, ok := env[key]
+	if !ok || s == "" {
+		return 0
+	}
+	v, err := strconv.ParseUint(s, 10, 64)
+	if err != nil {
+		return 0
+	}
+	return v
+}