From 453f74626adcdaf8a8524162aaeca42a83a0a578 Mon Sep 17 00:00:00 2001 From: Zexi Li Date: Thu, 3 Sep 2026 20:14:23 +0800 Subject: [PATCH] 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 | 3 + pkg/llm/models/instantmodel.go | 110 +++++++++++++++++- pkg/llm/models/instantmodel_import_cancel.go | 31 +++++ .../llm/llm_instant_model_delete_task.go | 1 + .../llm/llm_instant_model_import_task.go | 74 +++++++----- pkg/util/logclient/consts.go | 2 + pkg/util/logclient/consts_i18n.go | 5 + 7 files changed, 195 insertions(+), 31 deletions(-) create mode 100644 pkg/llm/models/instantmodel_import_cancel.go diff --git a/pkg/apis/llm/instantmodel.go b/pkg/apis/llm/instantmodel.go index 26851172e5..5fffcf8bf5 100644 --- a/pkg/apis/llm/instantmodel.go +++ b/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 diff --git a/pkg/llm/models/instantmodel.go b/pkg/llm/models/instantmodel.go index 477d508044..e5c0a4951f 100644 --- a/pkg/llm/models/instantmodel.go +++ b/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 diff --git a/pkg/llm/models/instantmodel_import_cancel.go b/pkg/llm/models/instantmodel_import_cancel.go new file mode 100644 index 0000000000..ba6ebc6dd0 --- /dev/null +++ b/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() + } + } +} diff --git a/pkg/llm/tasks/llm/llm_instant_model_delete_task.go b/pkg/llm/tasks/llm/llm_instant_model_delete_task.go index 059b6f2a55..fba15d664c 100644 --- a/pkg/llm/tasks/llm/llm_instant_model_delete_task.go +++ b/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) diff --git a/pkg/llm/tasks/llm/llm_instant_model_import_task.go b/pkg/llm/tasks/llm/llm_instant_model_import_task.go index 3c16ce5048..51b2b267cc 100644 --- a/pkg/llm/tasks/llm/llm_instant_model_import_task.go +++ b/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) } diff --git a/pkg/util/logclient/consts.go b/pkg/util/logclient/consts.go index 3d0dfbfaa5..65a35d1241 100644 --- a/pkg/util/logclient/consts.go +++ b/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" ) diff --git a/pkg/util/logclient/consts_i18n.go b/pkg/util/logclient/consts_i18n.go index 937c17d582..3d70cff376 100644 --- a/pkg/util/logclient/consts_i18n.go +++ b/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("同步云资源"),