Add model download management, job progress tracking, and generative model switching#764
Open
Irozuku wants to merge 90 commits into
Open
Add model download management, job progress tracking, and generative model switching#764Irozuku wants to merge 90 commits into
Irozuku wants to merge 90 commits into
Conversation
Add progress and progress_message columns to task_copy (with a migration for existing databases), a report_progress method on the queue that refreshes last_update so changes surface through polling, and a safe no-op report_progress helper on BaseJob for jobs to opt in.
…r jobs Emit phase checkpoints from ModelJob, DatasetJob, PredictJob and ConverterJob so the queue exposes a real completion signal. The converter loop reports determinate progress across its converters.
Render a LinearProgress bar (determinate when a percent is known, indeterminate otherwise) plus the current phase message for running jobs in the queue widget and the job details dialog. Add the progress label to the five locale files.
Add queue tests for progress writes, completion setting 100, changes payload, the old-table migration and the BaseJob no-op, plus widget tests for the determinate and indeterminate progress bars.
An undownloaded, download-required model now renders disabled in the ModelsRightBar list and cannot open the config dialog; a download control is shown beneath it and the list refreshes on completion. The in-dialog download control is removed (the list is the download surface); the disabled-create-button guard remains as a safety net.
…nline download Cards for a download-required, not-yet-downloaded component are disabled (unclickable, dimmed) and render an inline download control; an onDownloadChange callback lets parents refresh. Non-downloadable components are unchanged.
Add BaseGenerativeModel.get_metadata exposing requires_download and download_size_bytes, a 409/422 download gate in upload_generative_session, and a defensive JobError in GenerativeJob.run for undownloaded models.
Replace the four enum-based model classes (QwenModel, SmolLMModel, LlamaModel, MistralModel) with nine thin subclasses of GGUFTextGenerationModel (HFDownloadableMixin). Each subclass represents one checkpoint and carries REPO_ID, GGUF_PATTERN, and DOWNLOAD_SIZE_BYTES so the download machinery can manage it independently. The shared schema and __init__/generate logic live in gguf_text_generation_base.py.
… GGUF-split tests Reverts the generative session-creation guard to the endpoint's existing 400 "not registered" behavior for unknown models (dropping the added 422 that broke existing tests); the 409 download gate is unchanged. Updates the session fixtures to use a non-download model after the GGUF split removed the old QwenModel component.
The delete affordance in the models side bar is now a Material IconButton (hover and ripple feedback) instead of a bare clickable icon; the download hint stays a plain icon.
A run whose model still needs downloading cannot be trained, so the Train/Retrain button is disabled with a tooltip until the model is downloaded. The state comes from the shared live download state so an inline download re-enables it.
Run All trained every not-started run, including ones whose model still needs downloading. It now trains only runs whose model is ready (downloaded or download-free), reading the live download state, and warns when any were skipped.
The chat input used a one-time reconciled download status that reports true mid-download (partial files), so a still-downloading model was usable. Both the input gate and the model switcher labels now read the shared live download state, blocking while downloading and updating the moment a download finishes.
The job poller's final flush reported every non-finished watched job as an error when the queue briefly looked empty, so a download queued right after a delete showed a false failure. It now only errors a job that vanished or is actually in an error state, and keeps watching pending or running jobs.
DOWNLOAD_SIZE_BYTES were rough estimates far below what is actually fetched: snapshot_download pulls the whole HF repo (every weight format and, for diffusion, all fp16/fp32 variants). Update each model's declared size to the exact repo byte total, verified against an on-disk download.
These two Opus-MT models were merged in still describing weights as downloading on first use and inheriting the generic 300 MB size. Correct their descriptions to state weights must be downloaded before use, and set their real download sizes (1.17 GB / 1.21 GB).
HF repos ship TensorFlow, Flax, Rust, ONNX, OpenVINO and CoreML copies of the weights that from_pretrained never uses here, so downloads were several times larger than needed. Add HF_IGNORE_PATTERNS (applied to every snapshot_download) to skip them, keeping both .bin and .safetensors so fine-tuning and inference still load, and update DOWNLOAD_SIZE_BYTES to the slimmer real sizes (e.g. Bert 3.4GB->882MB, T5 2.2GB->486MB, opus ~1.2GB->~300MB, SDXL 77GB->35GB).
The Xet transfer backend returns a 404 on its read-token endpoint for some repos (e.g. bert-base-uncased), aborting the whole snapshot download. Disable Xet before downloading so it uses the reliable HTTP/LFS path.
PixArtSigma __init__ discarded the _pretrained_source(None) result and then read self.model_name, which was never set, raising AttributeError on generation. Assign the resolved source to self.model_name like the other diffusion models.
The 512 checkpoint has no model_index.json (it ships only the transformer), so PixArtSigmaPipeline.from_pretrained on it failed with 'no file named model_index.json'. Download both the 512 (transformer) and 1024 (full pipeline) repos and build the pipeline from the 1024 scaffold with the 512 transformer injected, the standard PixArt-Sigma 512 recipe. Size updated to the combined download.
The 512 checkpoint has no standalone pipeline and needs the 1024 scaffold, so keeping them as separate components meant downloading the shared T5/VAE twice. Replace PixArtSigma512 and PixArtSigma1024 with a single PixArtSigma that downloads both repos once and picks the checkpoint (1024 or 512) via a schema parameter, injecting the 512 transformer into the 1024 pipeline when selected.
Deleting a downloaded component from the selector button or the models side bar icon now opens a confirmation modal instead of deleting immediately, reusing DeleteConfirmationModal with a per-component message (all five locales).
React bubbles events through the component tree even for portaled modals, so confirming a component deletion from within a clickable model row also fired the row's onClick and opened the configure dialog. Stop propagation on the modal content.
The delete button in the component download control now shows the download size (Delete download ({{size}})), keeping the size visible after a component is downloaded. Added deleteWithSize to all five locales.
Clicking the delete icon opened the confirmation modal without the mouse leaving the row, so the hover popover stayed anchored and visible after deletion. Clear the hover on any row click via a capture phase handler, which fires even for the action icon that stops propagation.
…loaded The per row train button in the model comparison table ignored download status. Disable it (with a tooltip) when the run's model still needs downloading, using the shared live download state so it enables once the download finishes.
Update the snapshot_download assertions for the new ignore_patterns, add save_pretrained to the opus DummyTokenizer (the base save persists the tokenizer now), and make the jobs test registry module scoped so DummyModel is registered before the module scoped run fixtures create runs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an explicit download system for every weight backed ML component (text classification, translation, image classification, GGUF LLMs, diffusion, ControlNet) and a per job progress bar in the job queue. Before this change these models silently pulled weights from HuggingFace on first use, with no visibility, no size information, and no way to manage them, and runs or sessions could start against weights that were not present. Now each component that requires a download is fetched on demand with progress, an accurate size, and a delete action, and runs, trainings, and generative sessions are blocked until the required weights are present, including components nested inside another component's parameters. The job queue also shows real progress (a determinate bar plus a phase message) for long running jobs instead of a static "running" label. Generative sessions can also change their model on the fly: the switcher lets you pick any compatible model (blocking input and offering a download when it is not present), records each model change and its parameters in the chat history, and restores the parameters last used for that model in the session or its defaults.
Component download in the model list and Component Selector (generative)
Download gating on train and run
Per job progress bar in the job queue
Generative session with model selector
Type of Change
Changes (by file)
Download infrastructure (backend)
DashAI/back/dependencies/downloads/downloadable.py:DownloadableMixinplusHFDownloadableMixin,HFPretrainedDownloadMixin, andTorchvisionDownloadMixin. Downloads each repo into the component folder, usesallow_patternsandignore_patternsto skip alternate framework weights (TensorFlow, Flax, Rust, ONNX, OpenVINO, CoreML), and forces the classic transfer path (Xet disabled).DashAI/back/dependencies/downloads/nested.py: walks a parameters dict for components chosen as another component's parameter at any depth and reports which ones still need downloading.DashAI/back/job/component_download_job.py,DashAI/back/job/base_job.py: the download job plus areport_progresshelper.DashAI/back/api/api_v1/endpoints/components.py: download, status, and delete endpoints,POST /components/downloads/required, and reconciliation of download status against the filesystem on read.DashAI/back/api/api_v1/endpoints/runs.py,generative_session.py,job/model_job.py: gate run creation, session creation, and training on the model and its nested components being downloaded, returning 409 or JobError with the list of missing components.Generative session model switching (backend)
DashAI/back/api/api_v1/endpoints/generative_session.py: allow changing a session's model, including to one that is not yet downloaded, record the change in the parameter history with the model name, emit a model change event, and restore the parameters last used for that model in the session or its defaults.DashAI/back/dependencies/database/models.pyand migrationa7d2c9e4f1b0: addmodel_nametoGenerativeSessionParameterHistory.Models made downloadable (backend)
REQUIRES_DOWNLOAD,DOWNLOAD_SIZE_BYTES, and its repos, and descriptions were corrected to state that weights must be downloaded before use.DashAI/back/models/hugging_face/pixart_sigma_model.py: merged the 512 and 1024 variants into onePixArtSigmawith acheckpointselector, because the 512 repo ships only a transformer and reuses the 1024 pipeline.DashAI/back/initial_components.py: registrations updated for the split and merged classes.Per job progress bar
DashAI/back/dependencies/job_queues/huey_job_queue.py,base_job_queue.py,job/base_job.py:progressandprogress_messagecolumns ontask_copy(with a migration for existing databases), areport_progressmethod, and pass through in/jobs/changes.DashAI/front/src/components/jobs/JobQueueWidget.jsx: a determinate or indeterminate bar plus a phase message for each job.Download UI (frontend)
api/component.ts: download, delete, and status calls plusgetRequiredDownloads.components/models/model/ComponentDownloadControl.jsx: the download and delete control, with a module level pub/sub so every control for the same component stays in sync.ModelDownloadStatusIcon.jsxfor the models side bar.custom/ComponentSelector.jsx,ComponentDetailsPanel.jsx: inline download control and the size shown on the delete button.shared/FormSchema*selectors andcontexts/schema.js: inline download for nested components, a fix for selecting a component nested two levels deep, and a root breadcrumb.models/ModelsRightBar.jsx,RunCard.jsx,ModelComparisonTable.jsx,SessionVisualization.jsx: disable Train, Retrain, and Run All until the model is downloaded (using live state), and update the list in place to avoid scroll resets.generative/ModelSwitcher.jsx,GenerativeChat.jsx,CreateSession*: a session level switcher to change the model mid session, model change events shown in the chat history, and blocking of input and switching until the model is downloaded.threeSectionLayout/DeleteConfirmationModal.jsx: delete confirmation that stops click propagation so confirming does not trigger the row underneath.utils/jobPoller.js: stop reporting freshly queued jobs as failed.common.jsonandcustom.json(en, es, pt, de, zh): keys for download, size, must download, and confirm delete.Testing
uv run pytest tests/back/covers the download mixin, the nested collector and/downloads/required, per model downloadable metadata, the run, session, and training gates, generative session model switching, and job progress reporting.yarn testcovers the download control, the side bar icon, and the job queue progress bar.Notes
PixArtSigma512andPixArtSigma1024becamePixArtSigma, and the GGUF models split per quantization), so generative sessions that reference old names need to be recreated.