feat(p2p): wire P2P distribution into model-agent main

Integrate P2P model distribution into the model-agent entry point:
- Add P2P configuration fields (ports, rates, encryption, timeout)
- Read P2P settings from environment variables
- Create ModelDistributor when P2P_ENABLED=true
- Create P2PLeaseManager for download coordination
- Call gopher.EnableP2P() to activate P2P download flow
- Start MetainfoServer for peer discovery
- Add graceful shutdown for P2P resources

Environment variables:
- P2P_ENABLED: Enable/disable P2P (default: false)
- PEERS_SERVICE: Headless service DNS for peer discovery
- P2P_TORRENT_PORT: BitTorrent port (default: 6881)
- P2P_METAINFO_PORT: Metainfo HTTP port (default: 8081)
- P2P_MAX_DOWNLOAD_RATE: Max download rate in bytes/s
- P2P_MAX_UPLOAD_RATE: Max upload rate in bytes/s
- P2P_ENCRYPTION_ENABLED: Enable BitTorrent encryption
- P2P_DOWNLOAD_TIMEOUT: P2P download timeout in seconds

Author: slin1237

cmd/model-agent/main.go

@@ -23,6 +23,7 @@ import (
 
 	omev1beta1client "github.com/sgl-project/ome/pkg/client/clientset/versioned"
 	omev1beta1informers "github.com/sgl-project/ome/pkg/client/informers/externalversions"
+	"github.com/sgl-project/ome/pkg/distributor"
 	"github.com/sgl-project/ome/pkg/logging"
 	"github.com/sgl-project/ome/pkg/modelagent"
 	"github.com/sgl-project/ome/pkg/version"
@@ -43,6 +44,16 @@ type config struct {
 	numDownloadWorker    int
 	namespace            string
 	logLevel             string
+	// P2P configuration
+	p2pEnabled           bool
+	p2pPeersService      string
+	p2pTorrentPort       int
+	p2pMetainfoPort      int
+	p2pMaxDownloadRate   int64
+	p2pMaxUploadRate     int64
+	p2pEnableEncryption  bool
+	p2pRequireEncryption bool
+	p2pDownloadTimeout   time.Duration
 }
 
 // Logger type alias for zap.SugaredLogger
@@ -89,6 +100,40 @@ func initConfig(_ *cobra.Command, _ []string) {
 		panic("NODE_NAME environment variable is empty")
 	}
 	cfg.nodeName = nodeName
+
+	// P2P configuration from environment variables
+	cfg.p2pEnabled = os.Getenv("P2P_ENABLED") == "true"
+	cfg.p2pPeersService = os.Getenv("PEERS_SERVICE")
+	if cfg.p2pPeersService == "" {
+		cfg.p2pPeersService = "ome-peers.ome.svc.cluster.local"
+	}
+	cfg.p2pTorrentPort = getEnvInt("P2P_TORRENT_PORT", 6881)
+	cfg.p2pMetainfoPort = getEnvInt("P2P_METAINFO_PORT", 8081)
+	cfg.p2pMaxDownloadRate = getEnvInt64("P2P_MAX_DOWNLOAD_RATE", 524288000) // 500 MB/s
+	cfg.p2pMaxUploadRate = getEnvInt64("P2P_MAX_UPLOAD_RATE", 524288000)     // 500 MB/s
+	cfg.p2pEnableEncryption = os.Getenv("P2P_ENCRYPTION_ENABLED") == "true"
+	cfg.p2pRequireEncryption = os.Getenv("P2P_ENCRYPTION_REQUIRED") == "true"
+	cfg.p2pDownloadTimeout = time.Duration(getEnvInt("P2P_DOWNLOAD_TIMEOUT", 3600)) * time.Second // 1 hour default
+}
+
+// getEnvInt reads an integer from environment variable with a default value
+func getEnvInt(key string, defaultVal int) int {
+	if val := os.Getenv(key); val != "" {
+		if i, err := fmt.Sscanf(val, "%d", &defaultVal); err == nil && i == 1 {
+			return defaultVal
+		}
+	}
+	return defaultVal
+}
+
+// getEnvInt64 reads an int64 from environment variable with a default value
+func getEnvInt64(key string, defaultVal int64) int64 {
+	if val := os.Getenv(key); val != "" {
+		if i, err := fmt.Sscanf(val, "%d", &defaultVal); err == nil && i == 1 {
+			return defaultVal
+		}
+	}
+	return defaultVal
 }
 
 // initializeLogger creates and configures a zap logger with the specified settings
@@ -334,6 +379,75 @@ func runCommand(cmd *cobra.Command, args []string) {
 		logger.Fatalf("Failed to initialize components: %v", err)
 	}
 
+	// Initialize P2P distribution if enabled
+	var p2pDistributor *distributor.ModelDistributor
+	var metainfoServer *distributor.MetainfoServer
+	if cfg.p2pEnabled {
+		logger.Info("P2P model distribution is enabled, initializing...")
+
+		// Get POD_NAME and POD_IP for P2P peer identification
+		podName := os.Getenv("POD_NAME")
+		if podName == "" {
+			logger.Warn("POD_NAME not set, P2P coordination may not work correctly")
+		}
+		podIP := os.Getenv("POD_IP")
+		if podIP == "" {
+			logger.Warn("POD_IP not set, P2P peer discovery may not work correctly")
+		}
+
+		// Create distributor configuration
+		distCfg := distributor.Config{
+			DataDir:                   cfg.modelsRootDir,
+			PodName:                   podName,
+			PodIP:                     podIP,
+			PeersService:              cfg.p2pPeersService,
+			TorrentPort:               cfg.p2pTorrentPort,
+			MetainfoPort:              cfg.p2pMetainfoPort,
+			MaxDownloadRate:           cfg.p2pMaxDownloadRate,
+			MaxUploadRate:             cfg.p2pMaxUploadRate,
+			EnableEncryption:          cfg.p2pEnableEncryption,
+			RequireEncryption:         cfg.p2pRequireEncryption,
+			Namespace:                 cfg.namespace,
+			LeaseDurationSeconds:      120, // 2 minutes
+			LeaseRenewIntervalSeconds: 30,  // renew every 30 seconds
+			P2PTimeoutSeconds:         int(cfg.p2pDownloadTimeout.Seconds()),
+			EnableP2P:                 true,
+		}
+
+		// Create the P2P distributor
+		p2pDistributor, err = distributor.New(distCfg, logger)
+		if err != nil {
+			logger.Errorf("Failed to create P2P distributor: %v", err)
+			logger.Warn("Continuing without P2P support")
+		} else {
+			// Create lease manager for P2P coordination
+			leaseManager := modelagent.NewP2PLeaseManager(kubeClient, cfg.namespace, cfg.nodeName, logger)
+
+			// Enable P2P on the gopher
+			gopher.EnableP2P(p2pDistributor, leaseManager)
+			gopher.SetP2PTimeout(cfg.p2pDownloadTimeout)
+
+			// Create and start metainfo server
+			metainfoServer = distributor.NewMetainfoServer(
+				cfg.modelsRootDir,
+				cfg.p2pMetainfoPort,
+				p2pDistributor,
+				logger,
+			)
+
+			go func() {
+				logger.Infof("Starting P2P metainfo server on port %d", cfg.p2pMetainfoPort)
+				if err := metainfoServer.ServeWithContext(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) {
+					logger.Errorf("Metainfo server error: %v", err)
+				}
+			}()
+
+			logger.Info("P2P model distribution initialized successfully")
+		}
+	} else {
+		logger.Info("P2P model distribution is disabled")
+	}
+
 	// Set up a health check server
 	server := setupServer(cfg.port, cfg.modelsRootDir, logger)
 	go func() {
@@ -350,6 +464,13 @@ func runCommand(cmd *cobra.Command, args []string) {
 	if err := scout.Run(stopCh); err != nil {
 		logger.Fatalf("Error running scout: %v", err)
 	}
+
+	// Cleanup P2P resources on shutdown
+	if p2pDistributor != nil {
+		logger.Info("Shutting down P2P distributor...")
+		p2pDistributor.Close()
+	}
+	_ = metainfoServer // Suppress unused warning - shutdown handled via context
 }
 
 // createKubeClient creates a Kubernetes client from the provided config

pkg/constants/constants.go

@@ -383,10 +383,11 @@ const (
 	P2PDefaultPieceLength          = 4 * 1024 * 1024 // 4MB pieces
 
 	// P2P wait configuration for nodes waiting for P2P availability
-	P2PDefaultWaitMaxAttempts    = 60    // Maximum attempts before giving up on P2P
-	P2PDefaultWaitBaseDelayMs    = 2000  // Base delay between attempts (2 seconds)
-	P2PDefaultWaitMaxDelayMs     = 30000 // Maximum delay between attempts (30 seconds)
-	P2PDefaultWaitBackoffDivisor = 10    // Backoff increases every N attempts
+	P2PDefaultWaitMaxAttempts    = 60    // Deprecated: kept for compatibility
+	P2PDefaultWaitBaseDelayMs    = 2000  // Interval between lease/peer checks (2 seconds)
+	P2PDefaultWaitMaxDelayMs     = 30000 // Deprecated: kept for compatibility
+	P2PDefaultWaitBackoffDivisor = 10    // Deprecated: kept for compatibility
+	P2PMaxWaitTimeMinutes        = 240   // Maximum absolute wait time (4 hours) for very large models
 )
 
 // P2P environment variable keys

pkg/distributor/distributor.go

@@ -184,11 +184,17 @@ func (d *ModelDistributor) SeedModel(path, modelHash string) error {
 		return nil
 	}
 
+	d.logger.Infof("Creating metainfo for model %s at path %s (this may take several minutes for large models)...", modelHash, path)
+	startTime := time.Now()
+
 	mi, err := d.createMetainfo(path, modelHash)
 	if err != nil {
 		return fmt.Errorf("failed to create metainfo: %w", err)
 	}
 
+	metainfoTime := time.Since(startTime)
+	d.logger.Infof("Metainfo created for model %s in %v", modelHash, metainfoTime.Round(time.Second))
+
 	t, err := d.torrentClient.AddTorrent(mi)
 	if err != nil {
 		return fmt.Errorf("failed to add torrent: %w", err)
@@ -198,7 +204,8 @@ func (d *ModelDistributor) SeedModel(path, modelHash string) error {
 	d.activeTorrents[modelHash] = t
 	d.metrics.RecordSeeding(modelHash)
 
-	d.logger.Infof("Started seeding model %s", modelHash)
+	totalTime := time.Since(startTime)
+	d.logger.Infof("Started seeding model %s (total setup time: %v)", modelHash, totalTime.Round(time.Second))
 	return nil
 }
 

pkg/distributor/metainfo_server.go

@@ -46,7 +46,7 @@ func (s *MetainfoServer) Start() error {
 		Addr:         fmt.Sprintf(":%d", s.port),
 		Handler:      mux,
 		ReadTimeout:  30 * time.Second,
-		WriteTimeout: 30 * time.Second,
+		WriteTimeout: 10 * time.Minute, // Large metainfo files (30+ MB for 1TB models) need more time
 		IdleTimeout:  60 * time.Second,
 	}
 

pkg/modelagent/gopher.go

@@ -1138,55 +1138,85 @@ func (s *Gopher) downloadFromHuggingFace(ctx context.Context, task *GopherTask,
 
 // waitForP2PAvailability waits for the model to become available via P2P.
 // This is used when another node holds the download lease.
-// It uses exponential backoff to avoid overwhelming the cluster with DNS lookups.
-func (s *Gopher) waitForP2PAvailability(ctx context.Context, modelHash, modelInfo string) error {
+// It checks the lease status to determine whether to keep waiting:
+// - If lease is complete: P2P should be available, try download
+// - If lease exists and not expired: keep waiting (lease holder still downloading)
+// - If lease expired or not found: give up (lease holder crashed)
+// - If model is being deleted: abort early
+func (s *Gopher) waitForP2PAvailability(ctx context.Context, task *GopherTask, modelHash, modelInfo, leaseName string) error {
 	if s.p2pDistributor == nil {
 		return fmt.Errorf("P2P distributor not configured")
 	}
 
 	// Use constants for configurable wait behavior
-	maxAttempts := constants.P2PDefaultWaitMaxAttempts
-	baseDelay := time.Duration(constants.P2PDefaultWaitBaseDelayMs) * time.Millisecond
-	maxDelay := time.Duration(constants.P2PDefaultWaitMaxDelayMs) * time.Millisecond
-	backoffDivisor := constants.P2PDefaultWaitBackoffDivisor
+	checkInterval := time.Duration(constants.P2PDefaultWaitBaseDelayMs) * time.Millisecond
+	maxWaitTime := time.Duration(constants.P2PMaxWaitTimeMinutes) * time.Minute
+	startTime := time.Now()
+
+	for {
+		elapsed := time.Since(startTime)
 
-	for attempt := 0; attempt < maxAttempts; attempt++ {
 		// Check context cancellation first
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		default:
 		}
 
+		// Check absolute maximum wait time
+		if elapsed > maxWaitTime {
+			return fmt.Errorf("absolute timeout waiting for P2P availability for model %s after %v", modelInfo, elapsed)
+		}
+
+		// Check if the model is being deleted - abort early to allow cleanup
+		if s.isModelBeingDeleted(task) {
+			s.logger.Infof("Model %s is being deleted, aborting P2P wait", modelInfo)
+			return fmt.Errorf("model %s is being deleted, aborting P2P wait", modelInfo)
+		}
+
+		// Check lease status to decide whether to keep waiting
+		lease, err := s.p2pLeaseManager.Get(ctx, leaseName)
+		if err != nil {
+			s.logger.Debugf("Failed to get lease %s: %v, will retry", leaseName, err)
+		} else if lease != nil {
+			// Check if lease is complete (download finished, seeding started)
+			if s.p2pLeaseManager.IsComplete(lease) {
+				s.logger.Infof("Lease %s is complete, P2P should be available for model %s", leaseName, modelInfo)
+				// Give a short delay for seeding to fully start
+				time.Sleep(2 * time.Second)
+			} else if s.p2pLeaseManager.IsExpired(lease) {
+				// Lease expired - holder might have crashed
+				s.logger.Warnf("Lease %s expired for model %s, giving up on P2P wait", leaseName, modelInfo)
+				return fmt.Errorf("lease expired while waiting for P2P availability for model %s", modelInfo)
+			} else {
+				// Lease is active but not complete - holder still downloading
+				s.logger.Debugf("Lease %s still active (holder: %s) for model %s, waiting... (elapsed: %v)",
+					leaseName, *lease.Spec.HolderIdentity, modelInfo, elapsed.Round(time.Second))
+			}
+		}
+
 		// Check if model is available via P2P
 		if s.p2pDistributor.HasPeers(ctx, modelHash) {
 			s.logger.Infof("P2P peers now available for model %s, attempting download", modelInfo)
 			if err := s.p2pDistributor.TryP2PDownload(ctx, modelHash, s.p2pTimeout); err == nil {
-				s.logger.Infof("Successfully downloaded model %s via P2P after waiting", modelInfo)
+				s.logger.Infof("Successfully downloaded model %s via P2P after waiting %v", modelInfo, elapsed.Round(time.Second))
 				return nil
 			} else {
 				s.logger.Warnf("P2P download attempt failed for model %s: %v", modelInfo, err)
 			}
 		}
 
-		// Calculate delay with exponential backoff
-		// Backoff increases every backoffDivisor attempts to avoid rapid polling
-		delay := baseDelay * time.Duration(1<<uint(attempt/backoffDivisor))
-		if delay > maxDelay {
-			delay = maxDelay
+		// Log progress periodically (every 30 seconds)
+		if int(elapsed.Seconds())%30 == 0 && elapsed.Seconds() > 0 {
+			s.logger.Infof("Still waiting for P2P availability for model %s (elapsed: %v)", modelInfo, elapsed.Round(time.Second))
 		}
 
-		s.logger.Debugf("Waiting %v before next P2P check for model %s (attempt %d/%d)",
-			delay, modelInfo, attempt+1, maxAttempts)
-
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-time.After(delay):
+		case <-time.After(checkInterval):
 		}
 	}
-
-	return fmt.Errorf("timeout waiting for P2P availability for model %s after %d attempts", modelInfo, maxAttempts)
 }
 
 // downloadWithP2P orchestrates the model download with P2P support.
@@ -1240,7 +1270,7 @@ func (s *Gopher) downloadWithP2P(ctx context.Context, task *GopherTask, baseMode
 
 	// Lease held by another node - wait for P2P availability
 	s.logger.Infof("Lease held by another node for model %s, waiting for P2P availability", modelInfo)
-	if err := s.waitForP2PAvailability(ctx, modelHash, modelInfo); err == nil {
+	if err := s.waitForP2PAvailability(ctx, task, modelHash, modelInfo, leaseName); err == nil {
 		s.logger.Infof("Model %s now available via P2P", modelInfo)
 		return nil
 	} else {
@@ -1286,6 +1316,28 @@ func (s *Gopher) downloadWithLeaseHeld(ctx context.Context, task *GopherTask, ba
 	return nil
 }
 
+// isModelBeingDeleted checks if the model resource is being deleted (has deletionTimestamp).
+// This is used to abort long-running operations early when the resource is deleted.
+func (s *Gopher) isModelBeingDeleted(task *GopherTask) bool {
+	if task.BaseModel != nil {
+		bm, err := s.baseModelLister.BaseModels(task.BaseModel.Namespace).Get(task.BaseModel.Name)
+		if err != nil {
+			// If we can't get the resource, it might be deleted
+			return true
+		}
+		return !bm.ObjectMeta.DeletionTimestamp.IsZero()
+	}
+	if task.ClusterBaseModel != nil {
+		cbm, err := s.clusterBaseModelLister.Get(task.ClusterBaseModel.Name)
+		if err != nil {
+			// If we can't get the resource, it might be deleted
+			return true
+		}
+		return !cbm.ObjectMeta.DeletionTimestamp.IsZero()
+	}
+	return false
+}
+
 // startSeeding begins seeding the model to peers. Errors are logged but not returned
 // since seeding failure shouldn't fail the overall download operation.
 func (s *Gopher) startSeeding(destPath, modelHash, modelInfo string) {