Add sync logic

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

Dockerfile.agent

@@ -32,4 +32,7 @@ WORKDIR /
 COPY --from=builder /workspace/manager .
 USER 65532:65532
 
+# Expose the http server.
+EXPOSE 9090
+
 ENTRYPOINT ["/manager"]

README.md

@@ -6,7 +6,7 @@
 </p>
 
 <h3 align="center">
-A lightweight P2P-based cache system for model distributions.
+A lightweight P2P-based cache system for model distributions on Kubernetes.
 </h3>
 
 [![stability-alpha](https://img.shields.io/badge/stability-alpha-f4d03f.svg)](https://github.com/mkenney/software-guides/blob/master/STABILITY-BADGES.md#alpha)
@@ -39,7 +39,7 @@ _Name Story: the inspiration of the name `Manta` is coming from Dota2, called [M
 
 Read the [Installation](./docs//installation.md) for guidance.
 
-### Preload Models
+### Preheat Models
 
 A toy sample to preload the `Qwen/Qwen2.5-0.5B-Instruct` model:
 
@@ -69,6 +69,8 @@ spec:
     zone: zone-a
 ```
 
+### Delete Models
+
 If you want to remove the model weights once `Torrent` is deleted, set the `ReclaimPolicy=Delete`, default to `Retain`:
 
 ```yaml
@@ -80,8 +82,6 @@ spec:
   replicas: 1
   hub:
     repoID: Qwen/Qwen2.5-0.5B-Instruct
-  nodeSelector:
-    zone: zone-a
   reclaimPolicy: Delete
 ```
 

agent/deploy/daemonset.yaml

@@ -14,6 +14,7 @@ spec:
       labels:
         app: manta-agent
     spec:
+      hostNetwork: true
       serviceAccountName: manta-agent
       initContainers:
       - name: init-permissions
@@ -24,9 +25,10 @@ spec:
             mountPath: /workspace/models
       containers:
       - name: agent
-        image: inftyai/manta-agent:v0.0.1
+        # image: inftyai/manta-agent:v0.0.1
+        image: inftyai/test:manta-agent-110810
         ports:
-        - containerPort: 8080
+        - containerPort: 9090
         resources:
           limits:
             memory: 200Mi

agent/pkg/controller/replication_controller.go

@@ -83,6 +83,7 @@ func (r *ReplicationReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 
 	// This may take a long time, the concurrency is controlled by the MaxConcurrentReconciles.
 	if err := handler.HandleReplication(ctx, replication); err != nil {
+		logger.Error(err, "error to handle replication", "Replication", klog.KObj(replication))
 		return ctrl.Result{}, err
 	} else {
 		if err := r.updateNodeTracker(ctx, replication); err != nil {

agent/pkg/handler/chunk_handler.go

@@ -21,6 +21,8 @@ import (
 	"io"
 	"net/http"
 	"os"
+
+	"github.com/inftyai/manta/api"
 )
 
 const (
@@ -30,6 +32,7 @@ const (
 // SendChunk will send the chunk content via http request.
 func SendChunk(w http.ResponseWriter, r *http.Request) {
 	path := r.URL.Query().Get("path")
+
 	if path == "" {
 		http.Error(w, "path is required", http.StatusBadRequest)
 		return
@@ -45,25 +48,29 @@ func SendChunk(w http.ResponseWriter, r *http.Request) {
 	buffer := make([]byte, buffSize)
 	for {
 		n, err := file.Read(buffer)
+		if err != nil {
+			if err == io.EOF {
+				break
+			} else {
+				fmt.Println("Error reading file")
+				http.Error(w, "Error reading file", http.StatusInternalServerError)
+				return
+			}
+		}
+
 		if n > 0 {
 			_, writeErr := w.Write(buffer[:n])
 			if writeErr != nil {
 				fmt.Println("Error writing to response:", writeErr)
+				http.Error(w, "Error writing to response", http.StatusInternalServerError)
 				return
 			}
 		}
-		if err == io.EOF {
-			break
-		}
-		if err != nil {
-			http.Error(w, "Error reading file", http.StatusInternalServerError)
-			return
-		}
 	}
 }
 
 func recvChunk(blobPath, snapshotPath, peerName string) error {
-	url := fmt.Sprintf("http://%s:8080/sync?path=%s", peerName, blobPath)
+	url := fmt.Sprintf("http://%s:%s/sync?path=%s", peerName, api.HttpPort, blobPath)
 
 	resp, err := http.Get(url)
 	if err != nil {
@@ -87,6 +94,5 @@ func recvChunk(blobPath, snapshotPath, peerName string) error {
 		return err
 	}
 
-	fmt.Println("Chunk synced successfully")
 	return nil
 }

agent/pkg/handler/replication_handler.go

@@ -42,7 +42,8 @@ func HandleReplication(ctx context.Context, replication *api.Replication) error
 
 	if replication.Spec.Source.URI != nil {
 		host, _ := parseURI(*replication.Spec.Source.URI)
-		if host == api.URI_REMOTE {
+		// TODO: handel other uris.
+		if host != api.URI_LOCALHOST {
 			return syncChunk(ctx, replication)
 		}
 		// TODO: handle uri with object store.
@@ -77,18 +78,17 @@ func downloadChunk(ctx context.Context, replication *api.Replication) error {
 			// TODO: handle modelScope
 		}
 		// TODO: Handle address
-		logger.Info("download file successfully", "file", filename)
 	}
 
-	// symlink can helps to validate the file is downloaded successfully.
+	// symlink can help to validate the file is downloaded successfully.
 	// TODO: once we support split a file to several chunks, the targetPath should be
 	// changed here, such as targetPath-0001.
 	if err := createSymlink(blobPath, targetPath); err != nil {
 		logger.Error(err, "failed to create symlink")
 		return err
 	}
 
-	logger.Info("create symlink successfully", "file", filename)
+	logger.Info("download chunk successfully", "file", filename)
 	return nil
 }
 
@@ -97,12 +97,19 @@ func syncChunk(ctx context.Context, replication *api.Replication) error {
 
 	logger.Info("start to sync chunks", "Replication", klog.KObj(replication))
 
+	// The source URI looks like remote://node@<path-to-your-file>
 	sourceSplits := strings.Split(*replication.Spec.Source.URI, "://")
+	addresses := strings.Split(sourceSplits[1], "@")
+	nodeName, blobPath := addresses[0], addresses[1]
+	// The destination URI looks like localhost://<path-to-your-file>
 	destSplits := strings.Split(*replication.Spec.Destination.URI, "://")
-	if err := recvChunk(sourceSplits[1], destSplits[1], replication.Spec.NodeName); err != nil {
+
+	if err := recvChunk(blobPath, destSplits[1], nodeName); err != nil {
 		logger.Error(err, "failed to sync chunk")
 		return err
 	}
+
+	logger.Info("sync chunk successfully", "Replication", klog.KObj(replication), "target", destSplits[0])
 	return nil
 }
 

agent/pkg/server/server.go

@@ -18,11 +18,12 @@ package server
 
 import (
 	"context"
-	"fmt"
 	"net/http"
 	"time"
 
 	"github.com/inftyai/manta/agent/pkg/handler"
+	"github.com/inftyai/manta/api"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 )
 
 func Run(ctx context.Context) {
@@ -31,13 +32,14 @@ func Run(ctx context.Context) {
 
 	mux := http.NewServeMux()
 	mux.HandleFunc("/sync", handler.SendChunk)
-	server := &http.Server{Addr: ":8080", Handler: mux}
+	server := &http.Server{Addr: ":" + api.HttpPort, Handler: mux}
 
 	go func() {
-		fmt.Println("Server started on port 8080")
+		logger := log.FromContext(ctx)
+		logger.Info("Server started on port 9090")
 
 		if err := server.ListenAndServe(); err != nil {
-			fmt.Printf("ListenAndServe error: %s\n", err)
+			logger.Error(err, "listen and server error")
 			cancel()
 		}
 	}()
@@ -47,9 +49,11 @@ func Run(ctx context.Context) {
 	ctxShutdown, cancelShutdown := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancelShutdown()
 
+	logger := log.FromContext(ctx)
+
 	if err := server.Shutdown(ctxShutdown); err != nil {
-		fmt.Printf("Server shutdown error: %s\n", err)
+		logger.Error(err, "server shutdown error")
 	}
 
-	fmt.Println("Server shutdown successfully")
+	logger.Info("server shutdown successfully")
 }

agent/pkg/task/task.go

@@ -33,14 +33,14 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 
-	"github.com/inftyai/manta/agent/pkg/util"
+	cons "github.com/inftyai/manta/api"
 	api "github.com/inftyai/manta/api/v1alpha1"
 )
 
 const (
 	syncDuration = 5 * time.Minute
 
-	workspace = util.DefaultWorkspace
+	workspace = cons.DefaultWorkspace
 )
 
 var (

agent/pkg/util/consts.go → api/consts.go

@@ -14,8 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package util
+package api
+
+// Shared by both controller plane and agent.
 
 const (
 	DefaultWorkspace = "/workspace/models/"
+	HttpPort         = "9090"
 )

api/v1alpha1/replication_types.go

@@ -32,6 +32,7 @@ type Target struct {
 	// URI represents the file address with different storages, e.g.:
 	// 	 - oss://<bucket>.<endpoint>/<path-to-your-file>
 	// 	 - localhost://<path-to-your-file>
+	// 	 - remote://<node-name>@<path-to-your-file>
 	// Note: if it's a folder, all the files under the folder will be considered,
 	// otherwise, only one file will be replicated.
 	URI *string `json:"uri,omitempty"`

api/v1alpha1/torrent_types.go

@@ -84,6 +84,8 @@ type TorrentSpec struct {
 	// +optional
 	Replicas *int32 `json:"replicas,omitempty"`
 	// ReclaimPolicy represents how to handle the file replicas when Torrent is deleted.
+	// Be careful to use the Delete policy because once two Torrents refer to the same
+	// repo, delete one Torrent will remove the whole files.
 	// +kubebuilder:default=Retain
 	// +kubebuilder:validation:Enum={Retain,Delete}
 	// +optional
@@ -97,8 +99,11 @@ type TorrentSpec struct {
 type TrackerState string
 
 const (
-	PendingTrackerState  TrackerState = "Pending"
-	ReadyTrackerState    TrackerState = "Ready"
+	// Pending means the chunk is waiting for downloading.
+	PendingTrackerState TrackerState = "Pending"
+	// Ready means the chunk is ready for downloading or downloaded.
+	ReadyTrackerState TrackerState = "Ready"
+	// Deleting means the chunk is being removed.
 	DeletingTrackerState TrackerState = "Deleting"
 )
 
@@ -167,6 +172,7 @@ type TorrentStatus struct {
 //+kubebuilder:subresource:status
 //+kubebuilder:resource:scope=Cluster
 //+kubebuilder:printcolumn:name="Phase",type=string,JSONPath=".status.phase"
+//+kubebuilder:printcolumn:name="Age",type=date,JSONPath=".metadata.creationTimestamp"
 
 // Torrent is the Schema for the torrents API
 type Torrent struct {

cmd/main.go

@@ -22,6 +22,7 @@ import (
 
 	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
 	// to ensure that exec-entrypoint and run can make use of them.
+
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
 
 	"k8s.io/apimachinery/pkg/runtime"
@@ -60,14 +61,13 @@ func main() {
 	var metricsAddr string
 	var enableLeaderElection bool
 	var probeAddr string
+
 	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
 	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
 	flag.BoolVar(&enableLeaderElection, "leader-elect", false,
 		"Enable leader election for controller manager. "+
 			"Enabling this will ensure there is only one active controller manager.")
-	opts := zap.Options{
-		Development: true,
-	}
+	opts := zap.Options{}
 	opts.BindFlags(flag.CommandLine)
 	flag.Parse()
 
@@ -133,7 +133,7 @@ func setupControllers(mgr ctrl.Manager, certsReady chan struct{}) {
 	<-certsReady
 	setupLog.Info("certs ready")
 
-	dispatcher, err := dispatcher.NewDispatcher([]framework.RegisterFunc{nodeselector.New, diskaware.New}, []framework.RegisterFunc{gnumber.New})
+	dispatcher, err := dispatcher.NewDispatcher([]framework.RegisterFunc{nodeselector.New, diskaware.New, gnumber.New})
 	if err != nil {
 		setupLog.Error(err, "unable to create dispatcher")
 		os.Exit(1)

config/crd/bases/manta.io_replications.yaml

@@ -88,9 +88,9 @@ spec:
                   uri:
                     description: "URI represents the file address with different storages,
                       e.g.:\n\t - oss://<bucket>.<endpoint>/<path-to-your-file>\n\t
-                      - localhost://<path-to-your-file>\nNote: if it's a folder, all
-                      the files under the folder will be considered,\notherwise, only
-                      one file will be replicated."
+                      - localhost://<path-to-your-file>\n\t - remote://<node-name>@<path-to-your-file>\nNote:
+                      if it's a folder, all the files under the folder will be considered,\notherwise,
+                      only one file will be replicated."
                     type: string
                 type: object
               nodeName:
@@ -139,9 +139,9 @@ spec:
                   uri:
                     description: "URI represents the file address with different storages,
                       e.g.:\n\t - oss://<bucket>.<endpoint>/<path-to-your-file>\n\t
-                      - localhost://<path-to-your-file>\nNote: if it's a folder, all
-                      the files under the folder will be considered,\notherwise, only
-                      one file will be replicated."
+                      - localhost://<path-to-your-file>\n\t - remote://<node-name>@<path-to-your-file>\nNote:
+                      if it's a folder, all the files under the folder will be considered,\notherwise,
+                      only one file will be replicated."
                     type: string
                 type: object
             required:

config/crd/bases/manta.io_torrents.yaml

@@ -18,6 +18,9 @@ spec:
     - jsonPath: .status.phase
       name: Phase
       type: string
+    - jsonPath: .metadata.creationTimestamp
+      name: Age
+      type: date
     name: v1alpha1
     schema:
       openAPIV3Schema:
@@ -83,8 +86,10 @@ spec:
                 type: object
               reclaimPolicy:
                 default: Retain
-                description: ReclaimPolicy represents how to handle the file replicas
-                  when Torrent is deleted.
+                description: |-
+                  ReclaimPolicy represents how to handle the file replicas when Torrent is deleted.
+                  Be careful to use the Delete policy because once two Torrents refer to the same
+                  repo, delete one Torrent will remove the whole files.
                 enum:
                 - Retain
                 - Delete

config/default/manager_auth_proxy_patch.yaml

@@ -37,3 +37,4 @@ spec:
         - "--health-probe-bind-address=:8081"
         - "--metrics-bind-address=127.0.0.1:8080"
         - "--leader-elect"
+        - "--zap-log-level=2"

config/manager/kustomization.yaml

@@ -5,4 +5,4 @@ kind: Kustomization
 images:
 - name: controller
   newName: inftyai/manta
-  newTag: v0.0.1
+  newTag: "110905"