experiment with proto storage

Author: MiguelAAe

apis/go/mlops/scheduler/storage.pb.go

@@ -20,3 +20,9 @@ message ExperimentSnapshot {
   // on restart, which would guard against lost events in communication.
   bool deleted = 2;
 }
+
+
+message ModelSnapshot{
+  repeated ModelVersionStatus versions = 1;
+  bool deleted = 2;
+}

apis/mlops/scheduler/storage.proto

@@ -67,7 +67,7 @@ type ModelStatus struct {
 	ModelGwReason       string
 	AvailableReplicas   uint32
 	UnavailableReplicas uint32
-	DrainingReplicas    uint32
+	DrainingReplicas    uint32 //this field is never used and was never added to proto
 	Timestamp           time.Time
 }
 
@@ -90,6 +90,66 @@ func NewDefaultModelVersion(model *pb.Model, version uint32) *ModelVersion {
 	}
 }
 
+func NewModelSnapshot(model *pb.Model) *pb.ModelSnapshot {
+
+	generation := model.GetMeta().GetKubernetesMeta().GetGeneration()
+	version := max(uint32(1), uint32(generation))
+
+	var versions []*pb.ModelVersionStatus
+	versions = append(versions, &pb.ModelVersionStatus{
+		Version:           version,
+		ServerName:        "",
+		KubernetesMeta:    nil,
+		ModelReplicaState: make(map[int32]*pb.ModelReplicaStatus),
+		State: &pb.ModelStatus{
+			State:               pb.ModelStatus_ModelStateUnknown,
+			Reason:              "",
+			AvailableReplicas:   0,
+			UnavailableReplicas: 0,
+			LastChangeTimestamp: nil,
+			ModelGwState:        pb.ModelStatus_ModelCreate,
+			ModelGwReason:       "",
+		},
+		ModelDefn: model,
+	})
+
+	return &pb.ModelSnapshot{
+		Versions: versions,
+		Deleted:  false,
+	}
+}
+
+func CreateNextModelVersion(snapshot *pb.ModelSnapshot, model *pb.Model) *pb.ModelSnapshot {
+	if snapshot == nil {
+		return NewModelSnapshot(model)
+	}
+
+	if len(snapshot.Versions) < 1 {
+		return NewModelSnapshot(model)
+	}
+
+	version := snapshot.Versions[len(snapshot.Versions)-1].Version
+
+	snapshot.Versions = append(snapshot.Versions, &pb.ModelVersionStatus{
+		Version:           version,
+		ServerName:        "",
+		KubernetesMeta:    nil,
+		ModelReplicaState: make(map[int32]*pb.ModelReplicaStatus),
+		State: &pb.ModelStatus{
+			State:               pb.ModelStatus_ModelStateUnknown,
+			Reason:              "",
+			AvailableReplicas:   0,
+			UnavailableReplicas: 0,
+			LastChangeTimestamp: nil,
+			ModelGwState:        pb.ModelStatus_ModelCreate,
+			ModelGwReason:       "",
+		},
+		ModelDefn: model,
+	})
+
+	return snapshot
+}
+
 // TODO: remove deleted from here and reflect in callers
 // This is only used in tests, thus we don't need to worry about modelGWState
 func NewModelVersion(model *pb.Model, version uint32, server string, replicas map[int]ReplicaStatus, deleted bool, state ModelState) *ModelVersion {
@@ -396,6 +456,18 @@ func (m *Model) Inactive() bool {
 	return m.Latest().Inactive()
 }
 
+func ModelInactive(snap *pb.ModelSnapshot) bool {
+
+}
+
+func ModelLatest(snap *pb.ModelSnapshot) *pb.ModelVersionStatus {
+	if len(snap.Versions) > 0 {
+		return snap.Versions[len(snap.Versions)-1]
+	} else {
+		return nil
+	}
+}
+
 func (m *Model) IsDeleted() bool {
 	return m.deleted.Load()
 }