Update Torrent API and add tests

Signed-off-by: kerthcet <[email protected]>

Author: kerthcet

.github/workflows/golang-workflow.yaml

@@ -0,0 +1,12 @@
+name: CI Workflow
+
+on:
+  pull_request:
+    types:
+      - opened
+      - synchronize
+
+jobs:
+  golang-ci:
+    # Use a released version for stable.
+    uses: kerthcet/github-workflow-as-kube/.github/workflows/[email protected]

api/v1alpha1/nodetracker_types.go

@@ -23,7 +23,7 @@ import (
 // NodeTrackerSpec defines the desired state of NodeTracker
 // It acts like a cache.
 type NodeTrackerSpec struct {
-	NodeName string `json:"nodeName"`
+	// TODO: GC disk rate.
 }
 
 // NodeTrackerStatus defines the observed state of NodeTracker

api/v1alpha1/torrent_types.go

@@ -23,9 +23,10 @@ import (
 // This is inspired by https://github.com/InftyAI/llmaz.
 // ModelHub represents the model registry for model downloads.
 type ModelHub struct {
+	// TODO: support ModelScope
 	// Name refers to the model registry, such as huggingface.
 	// +kubebuilder:default=Huggingface
-	// +kubebuilder:validation:Enum={Huggingface,ModelScope}
+	// +kubebuilder:validation:Enum={Huggingface}
 	// +optional
 	Name *string `json:"name,omitempty"`
 	// ModelID refers to the model identifier on model hub,
@@ -37,15 +38,26 @@ type ModelHub struct {
 	// TODO: this is only supported with Huggingface, add support for ModelScope
 	// in the near future.
 	Filename *string `json:"filename,omitempty"`
+
+	// TODO: not supported
 	// Revision refers to a Git revision id which can be a branch name, a tag, or a commit hash.
 	// Most of the time, you don't need to specify it.
 	// +optional
-	Revision *string `json:"revision,omitempty"`
+	// Revision *string `json:"revision,omitempty"`
 }
 
 // URIProtocol represents the protocol of the URI.
 type URIProtocol string
 
+type ReclaimPolicy string
+
+const (
+	// RetainReclaimPolicy represents keep the files when Torrent is deleted.
+	RetainReclaimPolicy ReclaimPolicy = "Retain"
+	// DeleteReclaimPolicy represents delete the files when Torrent is deleted.
+	DeleteReclaimPolicy ReclaimPolicy = "Delete"
+)
+
 // TorrentSpec defines the desired state of Torrent
 type TorrentSpec struct {
 	// ModelHub represents the model registry for model downloads.
@@ -58,29 +70,45 @@ type TorrentSpec struct {
 	// 	- Image: img://nginx:1.14.2
 	// 	- OSS: oss://<bucket>.<endpoint>/<path-to-your-files>
 	// +optional
-	// URI *URIProtocol `json:"uriuri,omitempty"`
+	// URI *URIProtocol `json:"uri,omitempty"`
 
 	// Replicas represents the replication number of each file.
 	// +kubebuilder:default=1
 	// +optional
-	Replicas *int32 ` json:"replicas,omitempty"`
+	Replicas *int32 `json:"replicas,omitempty"`
+	// ReclaimPolicy represents how to handle the file replicas when Torrent is deleted.
+	// +kubebuilder:default=Retain
+	// +kubebuilder:validation:Enum={Retain,Delete}
+	// +optional
+	ReclaimPolicy *ReclaimPolicy `json:"reclaimPolicy,omitempty"`
+	// NodeSelector represents the node constraint to download the chunks.
+	// It can be used to download the model to one node for special case.
+	// +optional
+	NodeSelector map[string]string `json:"nodeSelector,omitempty"`
 }
 
 type TrackerState string
 
 const (
+	PendingTrackerState  TrackerState = "Pending"
 	DownloadTrackerState TrackerState = "Downloading"
 	ReadyTrackerState    TrackerState = "Ready"
 )
 
 type FileTracker struct {
 	// Name represents the name of the file.
 	Name string `json:"name"`
-	// State represents the state of the file, whether Pending
-	// for download or downloaded ready.
+	// State represents the state of the file, whether in downloading
+	// or downloaded ready.
 	State TrackerState `json:"State"`
 	// SizeBytes represents the file size.
-	SizeBytes int64 `json:"sizeBytes"`
+	// This is only used in nodeTracker.
+	// +optional
+	SizeBytes *int64 `json:"sizeBytes,omitempty"`
+	// Path represents the absolute path of the file in the node.
+	// This is only used in nodeTracker.
+	// +optional
+	Path *string `json:"path,omitempty"`
 }
 
 const (

api/v1alpha1/zz_generated.deepcopy.go

@@ -40,11 +40,6 @@ spec:
             description: |-
               NodeTrackerSpec defines the desired state of NodeTracker
               It acts like a cache.
-            properties:
-              nodeName:
-                type: string
-            required:
-            - nodeName
             type: object
           status:
             description: NodeTrackerStatus defines the observed state of NodeTracker
@@ -54,20 +49,26 @@ spec:
                   properties:
                     State:
                       description: |-
-                        State represents the state of the file, whether Pending
-                        for download or downloaded ready.
+                        State represents the state of the file, whether in downloading
+                        or downloaded ready.
                       type: string
                     name:
                       description: Name represents the name of the file.
                       type: string
+                    path:
+                      description: |-
+                        Path represents the absolute path of the file in the node.
+                        This is only used in nodeTracker.
+                      type: string
                     sizeBytes:
-                      description: SizeBytes represents the file size.
+                      description: |-
+                        SizeBytes represents the file size.
+                        This is only used in nodeTracker.
                       format: int64
                       type: integer
                   required:
                   - State
                   - name
-                  - sizeBytes
                   type: object
                 type: array
             type: object

config/crd/bases/manta.io_nodetrackers.yaml

@@ -61,16 +61,25 @@ spec:
                     description: Name refers to the model registry, such as huggingface.
                     enum:
                     - Huggingface
-                    - ModelScope
-                    type: string
-                  revision:
-                    description: |-
-                      Revision refers to a Git revision id which can be a branch name, a tag, or a commit hash.
-                      Most of the time, you don't need to specify it.
                     type: string
                 required:
                 - modelID
                 type: object
+              nodeSelector:
+                additionalProperties:
+                  type: string
+                description: |-
+                  NodeSelector represents the node constraint to download the chunks.
+                  It can be used to download the model to one node for special case.
+                type: object
+              reclaimPolicy:
+                default: Retain
+                description: ReclaimPolicy represents how to handle the file replicas
+                  when Torrent is deleted.
+                enum:
+                - Retain
+                - Delete
+                type: string
               replicas:
                 default: 1
                 description: Replicas represents the replication number of each file.
@@ -143,20 +152,26 @@ spec:
                   properties:
                     State:
                       description: |-
-                        State represents the state of the file, whether Pending
-                        for download or downloaded ready.
+                        State represents the state of the file, whether in downloading
+                        or downloaded ready.
                       type: string
                     name:
                       description: Name represents the name of the file.
                       type: string
+                    path:
+                      description: |-
+                        Path represents the absolute path of the file in the node.
+                        This is only used in nodeTracker.
+                      type: string
                     sizeBytes:
-                      description: SizeBytes represents the file size.
+                      description: |-
+                        SizeBytes represents the file size.
+                        This is only used in nodeTracker.
                       format: int64
                       type: integer
                   required:
                   - State
                   - name
-                  - sizeBytes
                   type: object
                 type: array
             type: object

config/crd/bases/manta.io_torrents.yaml

@@ -21,29 +21,24 @@ import (
 
 	"k8s.io/apimachinery/pkg/runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/cache"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
-	inftyaicomv1alpha1 "github.com/inftyai/manta/api/v1alpha1"
+	api "github.com/inftyai/manta/api/v1alpha1"
 )
 
 // NodeTrackerReconciler reconciles a NodeTracker object
 type NodeTrackerReconciler struct {
 	client.Client
 	Scheme *runtime.Scheme
+	Cache  cache.Cache
 }
 
 //+kubebuilder:rbac:groups=manta.io,resources=nodetrackers,verbs=get;list;watch;create;update;patch;delete
 //+kubebuilder:rbac:groups=manta.io,resources=nodetrackers/status,verbs=get;update;patch
 //+kubebuilder:rbac:groups=manta.io,resources=nodetrackers/finalizers,verbs=update
 
-// Reconcile is part of the main kubernetes reconciliation loop which aims to
-// move the current state of the cluster closer to the desired state.
-// TODO(user): Modify the Reconcile function to compare the state specified by
-// the NodeTracker object against the actual cluster state, and then
-// perform operations to make the cluster state reflect the state specified by
-// the user.
-//
 // For more details, check Reconcile and its Result here:
 // - https://pkg.go.dev/sigs.k8s.io/[email protected]/pkg/reconcile
 func (r *NodeTrackerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
@@ -57,6 +52,6 @@ func (r *NodeTrackerReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 // SetupWithManager sets up the controller with the Manager.
 func (r *NodeTrackerReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
-		For(&inftyaicomv1alpha1.NodeTracker{}).
+		For(&api.NodeTracker{}).
 		Complete(r)
 }

pkg/controller/nodetracker_controller.go

@@ -22,6 +22,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	inftyaicomv1alpha1 "github.com/inftyai/manta/api/v1alpha1"
@@ -58,5 +59,6 @@ func (r *ReplicationReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 func (r *ReplicationReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&inftyaicomv1alpha1.Replication{}).
+		WithOptions(controller.Options{MaxConcurrentReconciles: 5}).
 		Complete(r)
 }