feat(llm): support cancel and resume for instant model import (#25519)

Allow deleting models to abort in-flight downloads, persist import input for resume-import after killed status, and clean up import cache on delete.

pkg/apis/llm/instantmodel.go

@@ -90,6 +90,9 @@ type MountedByLLMInfo struct {
 type InstantModelSyncstatusInput struct {
 }
 
+type InstantModelResumeImportInput struct {
+}
+
 // InstantModelVramRequirement reports the heuristic VRAM needed to run this
 // model, mirroring GPUStack's `estimate_model_vram`. Returns
 // `vram_required_mb=0` when `weight_size_bytes` is unknown — callers should

pkg/llm/models/instantmodel.go

@@ -811,6 +811,8 @@ func (model *SInstantModel) CustomizeDelete(ctx context.Context, userCred mcclie
 }
 
 func (model *SInstantModel) StartDeleteTask(ctx context.Context, userCred mcclient.TokenCredential, query jsonutils.JSONObject, parentTaskId string) error {
+	// Abort in-flight import/download before marking deleting so the worker exits promptly.
+	CancelInstantModelImport(model.GetId())
 	model.SetStatus(ctx, userCred, commonapis.STATUS_DELETING, "")
 	params := jsonutils.NewDict()
 	if query != nil && jsonutils.QueryBoolean(query, "purge", false) {
@@ -1117,7 +1119,13 @@ func (man *SInstantModelManager) DoImportWithParent(
 	return tempModel, nil
 }
 
+const instantModelImportInputMetadataKey = "import_input"
+
 func (model *SInstantModel) startImportTask(ctx context.Context, userCred mcclient.TokenCredential, input apis.InstantModelImportInput, parentTaskId string) error {
+	if err := model.SetMetadata(ctx, instantModelImportInputMetadataKey, jsonutils.Marshal(input), userCred); err != nil {
+		return errors.Wrap(err, "SetMetadata import_input")
+	}
+
 	params := jsonutils.NewDict()
 	params.Add(jsonutils.Marshal(input), "import_input")
 
@@ -1130,6 +1138,66 @@ func (model *SInstantModel) startImportTask(ctx context.Context, userCred mcclie
 	return nil
 }
 
+func (model *SInstantModel) resolveImportInput(ctx context.Context, userCred mcclient.TokenCredential) (apis.InstantModelImportInput, error) {
+	input := apis.InstantModelImportInput{}
+	if meta := model.GetMetadataJson(ctx, instantModelImportInputMetadataKey, userCred); meta != nil {
+		if err := meta.Unmarshal(&input); err != nil {
+			return input, errors.Wrap(err, "unmarshal import_input metadata")
+		}
+		if strings.TrimSpace(input.ModelName) != "" && input.LlmType != "" {
+			return input, nil
+		}
+	}
+	// Fallback for records imported before metadata persistence.
+	input = apis.InstantModelImportInput{
+		ModelName: model.ModelName,
+		ModelTag:  model.ModelTag,
+		LlmType:   apis.LLMContainerType(model.LlmType),
+		Source:    defaultInstantModelSource(model.Source),
+		RepoId:    model.ModelName,
+		Revision:  model.ModelTag,
+	}
+	if strings.TrimSpace(input.ModelName) == "" || input.LlmType == "" {
+		return input, httperrors.NewInvalidStatusError("cannot resume import: missing model_name or llm_type")
+	}
+	return input, nil
+}
+
+func (model *SInstantModel) PerformResumeImport(
+	ctx context.Context,
+	userCred mcclient.TokenCredential,
+	query jsonutils.JSONObject,
+	input apis.InstantModelResumeImportInput,
+) (jsonutils.JSONObject, error) {
+	if model.Status != imageapi.IMAGE_STATUS_KILLED {
+		return nil, httperrors.NewInvalidStatusError("can only resume import when status is killed, current: %s", model.Status)
+	}
+
+	importInput, err := model.resolveImportInput(ctx, userCred)
+	if err != nil {
+		return nil, errors.Wrap(err, "resolveImportInput")
+	}
+
+	if err := model.SetStatus(ctx, userCred, imageapi.IMAGE_STATUS_QUEUED, "resume import"); err != nil {
+		return nil, errors.Wrap(err, "SetStatus queued")
+	}
+	if err := model.SetProgress(0); err != nil {
+		log.Warningf("PerformResumeImport: reset progress for %s: %s", model.Id, err)
+	}
+
+	if err := model.startImportTask(ctx, userCred, importInput, ""); err != nil {
+		_ = model.SetStatus(ctx, userCred, imageapi.IMAGE_STATUS_KILLED, err.Error())
+		db.OpsLog.LogEvent(model, logclient.ACT_RESUME_IMPORT, err, userCred)
+		logclient.AddActionLogWithContext(ctx, model, logclient.ACT_RESUME_IMPORT, err, userCred, false)
+		return nil, errors.Wrap(err, "startImportTask")
+	}
+
+	notes := fmt.Sprintf("resume import %s:%s", importInput.ModelName, importInput.ModelTag)
+	db.OpsLog.LogEvent(model, logclient.ACT_RESUME_IMPORT, notes, userCred)
+	logclient.AddActionLogWithContext(ctx, model, logclient.ACT_RESUME_IMPORT, notes, userCred, true)
+	return nil, nil
+}
+
 func getInstantModelImportWorkDir(root string, input apis.InstantModelImportInput) (string, error) {
 	root = strings.TrimSpace(root)
 	if root == "" {
@@ -1182,6 +1250,12 @@ func sanitizeInstantModelImportCacheComponent(s string) string {
 }
 
 func (model *SInstantModel) updateImportStatus(ctx context.Context, userCred mcclient.TokenCredential, status string, reason string) error {
+	if err := ctx.Err(); err != nil {
+		return errors.Wrap(err, "import cancelled")
+	}
+	if model.Status == commonapis.STATUS_DELETING || model.Deleted || model.PendingDeleted {
+		return errors.Errorf("import cancelled: model is deleting")
+	}
 	if model.Status == status {
 		return nil
 	}
@@ -1205,6 +1279,9 @@ func (model *SInstantModel) updateImportStatus(ctx context.Context, userCred mcc
 }
 
 func (model *SInstantModel) DoImport(ctx context.Context, userCred mcclient.TokenCredential, s *mcclient.ClientSession, input apis.InstantModelImportInput) (tmpDir string, err error) {
+	ctx, endImport := model.beginImportContext(ctx)
+	defer endImport()
+
 	progress := newInstantModelImportProgressUpdater(model)
 	progress.set(0, true)
 
@@ -1235,6 +1312,10 @@ func (model *SInstantModel) DoImport(ctx context.Context, userCred mcclient.Toke
 		err = errors.Wrap(err, "DownloadModel")
 		return
 	}
+	if err = ctx.Err(); err != nil {
+		err = errors.Wrap(err, "import cancelled")
+		return
+	}
 	progress.set(apis.InstantModelImportDownloadProgressEnd, true)
 	log.Infof("Downloaded model %s:%s with modelId: %s to %s", input.ModelName, input.ModelTag, modelId, tmpDir)
 
@@ -1483,10 +1564,16 @@ func (model *SInstantModel) GetDetailsVramRequirement(
 
 func (model *SInstantModel) CleanupImportTmpDir(ctx context.Context, userCred mcclient.TokenCredential, tmpDir string) error {
 	// sync image status
-	err := model.syncImageStatus(ctx, userCred)
-	if err != nil {
-		return errors.Wrap(err, "syncImageStatus")
+	if len(model.ImageId) > 0 {
+		err := model.syncImageStatus(ctx, userCred)
+		if err != nil {
+			return errors.Wrap(err, "syncImageStatus")
+		}
 	}
+	return removeImportWorkDir(tmpDir)
+}
+
+func removeImportWorkDir(tmpDir string) error {
 	if tmpDir == "" {
 		return nil
 	}
@@ -1497,6 +1584,23 @@ func (model *SInstantModel) CleanupImportTmpDir(ctx context.Context, userCred mc
 	return nil
 }
 
+// CleanupImportCacheBestEffort removes the import work directory for this model if resolvable.
+func (model *SInstantModel) CleanupImportCacheBestEffort(ctx context.Context, userCred mcclient.TokenCredential) {
+	input, err := model.resolveImportInput(ctx, userCred)
+	if err != nil {
+		log.Warningf("CleanupImportCacheBestEffort %s: resolveImportInput: %s", model.Id, err)
+		return
+	}
+	tmpDir, err := getInstantModelImportWorkDir(options.Options.LLMWorkingDirectory, input)
+	if err != nil {
+		log.Warningf("CleanupImportCacheBestEffort %s: get work dir: %s", model.Id, err)
+		return
+	}
+	if err := removeImportWorkDir(tmpDir); err != nil {
+		log.Warningf("CleanupImportCacheBestEffort %s: %s", model.Id, err)
+	}
+}
+
 // GetOllamaRegistryYAML returns the Ollama registry YAML content
 func (man *SInstantModelManager) GetOllamaRegistryYAML() string {
 	return apis.OLLAMA_REGISTRY_YAML

pkg/llm/models/instantmodel_import_cancel.go

@@ -0,0 +1,31 @@
+package models
+
+import (
+	"context"
+	"sync"
+)
+
+var instantModelImportCancels sync.Map // map[string]context.CancelFunc
+
+// beginImportContext registers a cancel func so StartDeleteTask can abort an in-flight download.
+func (model *SInstantModel) beginImportContext(parent context.Context) (context.Context, func()) {
+	ctx, cancel := context.WithCancel(parent)
+	modelId := model.GetId()
+	instantModelImportCancels.Store(modelId, cancel)
+	return ctx, func() {
+		cancel()
+		instantModelImportCancels.Delete(modelId)
+	}
+}
+
+// CancelInstantModelImport cancels an in-flight DoImport for the given model id, if any.
+func CancelInstantModelImport(modelId string) {
+	if modelId == "" {
+		return
+	}
+	if v, ok := instantModelImportCancels.LoadAndDelete(modelId); ok {
+		if cancel, ok := v.(context.CancelFunc); ok && cancel != nil {
+			cancel()
+		}
+	}
+}

pkg/llm/tasks/llm/llm_instant_model_delete_task.go

@@ -136,6 +136,7 @@ func (task *LLMInstantModelDeleteTask) OnImageDeleteCompleteFailed(ctx context.C
 }
 
 func (task *LLMInstantModelDeleteTask) OnImageDeleteComplete(ctx context.Context, model *models.SInstantModel, body jsonutils.JSONObject) {
+	model.CleanupImportCacheBestEffort(ctx, task.GetUserCred())
 	err := model.RealDelete(ctx, task.GetUserCred())
 	if err != nil {
 		task.taskFailed(ctx, model, err)

pkg/llm/tasks/llm/llm_instant_model_import_task.go

@@ -6,12 +6,14 @@ import (
 	"yunion.io/x/jsonutils"
 	"yunion.io/x/log"
 
+	commonapis "yunion.io/x/onecloud/pkg/apis"
 	imageapi "yunion.io/x/onecloud/pkg/apis/image"
 	apis "yunion.io/x/onecloud/pkg/apis/llm"
 	"yunion.io/x/onecloud/pkg/cloudcommon/db"
 	"yunion.io/x/onecloud/pkg/cloudcommon/db/taskman"
 	"yunion.io/x/onecloud/pkg/llm/models"
 	"yunion.io/x/onecloud/pkg/llm/options"
+	"yunion.io/x/onecloud/pkg/llm/tasks/worker"
 	"yunion.io/x/onecloud/pkg/mcclient/auth"
 	"yunion.io/x/onecloud/pkg/util/logclient"
 )
@@ -25,9 +27,18 @@ func init() {
 }
 
 func (task *LLMInstantModelImportTask) taskFailed(ctx context.Context, model *models.SInstantModel, err string) {
-	model.SetStatus(ctx, task.UserCred, imageapi.IMAGE_STATUS_KILLED, err)
-	db.OpsLog.LogEvent(model, db.ACT_CREATE, err, task.UserCred)
-	logclient.AddActionLogWithStartable(task, model, logclient.ACT_CREATE, err, task.UserCred, false)
+	// Deleting: do not overwrite status with killed; download was cancelled by delete.
+	if model.Status == commonapis.STATUS_DELETING || model.Deleted || model.PendingDeleted {
+		task.SetStageFailed(ctx, jsonutils.NewString(err))
+		return
+	}
+	notes := err
+	if notes != "" {
+		notes = notes + "; resume with resume-import to continue download"
+	}
+	model.SetStatus(ctx, task.UserCred, imageapi.IMAGE_STATUS_KILLED, notes)
+	db.OpsLog.LogEvent(model, db.ACT_CREATE, notes, task.UserCred)
+	logclient.AddActionLogWithStartable(task, model, logclient.ACT_CREATE, notes, task.UserCred, false)
 	task.SetStageFailed(ctx, jsonutils.NewString(err))
 }
 
@@ -43,36 +54,40 @@ func (task *LLMInstantModelImportTask) OnInit(ctx context.Context, obj db.IStand
 
 	task.SetStage("OnImportComplete", nil)
 	s := auth.GetAdminSession(ctx, options.Options.Region)
-	var fileDir string
-	// err = s.WithTaskCallback(task.GetId(), func() error {
-	// 	fileDir, err = model.DoImport(ctx, task.UserCred, s, input)
-	// 	// 将 fileDir 存储到 task.Params 中，以便后续阶段可以访问
-	// 	if fileDir != "" {
-	// 		task.Params.Set("file_dir", jsonutils.NewString(fileDir))
-	// 	}
-	// 	return err
-	// })
-	fileDir, err = model.DoImport(ctx, task.UserCred, s, input)
-	// 将 fileDir 存储到 task.Params 中，以便后续阶段可以访问
-	if fileDir != "" {
-		task.Params.Set("file_dir", jsonutils.NewString(fileDir))
-	}
-	if err != nil {
-		task.OnImportCompleteFailed(ctx, model, jsonutils.NewString(err.Error()))
-		return
-	}
-	task.OnImportComplete(ctx, model, nil)
+	// Run DoImport off the object lock so Delete API is not blocked by long downloads.
+	worker.ImportTaskRun(task, func() (jsonutils.JSONObject, error) {
+		fileDir, err := model.DoImport(ctx, task.UserCred, s, input)
+		result := jsonutils.NewDict()
+		if fileDir != "" {
+			result.Set("file_dir", jsonutils.NewString(fileDir))
+		}
+		if err != nil {
+			return result, err
+		}
+		return result, nil
+	})
 }
 
 func (task *LLMInstantModelImportTask) OnImportComplete(ctx context.Context, obj db.IStandaloneModel, body jsonutils.JSONObject) {
 	model := obj.(*models.SInstantModel)
 
-	// 确保删除 fileDir
-	if fileDirObj, err := task.Params.Get("file_dir"); err == nil {
-		if fileDir, _ := fileDirObj.GetString(); fileDir != "" {
-			model.CleanupImportTmpDir(ctx, task.GetUserCred(), fileDir)
+	if model.Status == commonapis.STATUS_DELETING || model.Deleted || model.PendingDeleted {
+		task.SetStageComplete(ctx, nil)
+		return
+	}
+
+	fileDir := ""
+	if body != nil {
+		fileDir, _ = body.GetString("file_dir")
+	}
+	if fileDir == "" {
+		if fileDirObj, err := task.Params.Get("file_dir"); err == nil {
+			fileDir, _ = fileDirObj.GetString()
 		}
 	}
+	if fileDir != "" {
+		model.CleanupImportTmpDir(ctx, task.GetUserCred(), fileDir)
+	}
 
 	// Best-effort: estimate the model's weight-file size for downstream
 	// VRAM-claim calculation. Failure here is a warning, not a fatal — most
@@ -130,6 +145,9 @@ func fetchWeightSizeForImport(ctx context.Context, input apis.InstantModelImport
 
 func (task *LLMInstantModelImportTask) OnImportCompleteFailed(ctx context.Context, obj db.IStandaloneModel, err jsonutils.JSONObject) {
 	model := obj.(*models.SInstantModel)
-
-	task.taskFailed(ctx, model, err.String())
+	reason, _ := err.GetString("__reason__")
+	if reason == "" {
+		reason = err.String()
+	}
+	task.taskFailed(ctx, model, reason)
 }

pkg/util/logclient/consts.go

@@ -321,4 +321,6 @@ const (
 
 	ACT_REGISTER_AIPROXY   = "register_aiproxy"
 	ACT_UNREGISTER_AIPROXY = "unregister_aiproxy"
+
+	ACT_RESUME_IMPORT = "resume_import"
 )

pkg/util/logclient/consts_i18n.go

@@ -1477,6 +1477,11 @@ func init() {
 		CN("保存镜像"),
 	)
 
+	o.Set(ACT_RESUME_IMPORT, i18n.NewTableEntry().
+		EN("Resume Import").
+		CN("继续导入"),
+	)
+
 	o.Set(ACT_CLOUD_SYNC, i18n.NewTableEntry().
 		EN("Sync Cloud Resource").
 		CN("同步云资源"),