Skip to content

refactor: expose a public custom-weights packaging helper#503

Merged
lrosemberg merged 12 commits into
mainfrom
lean/package-custom-weights
Jul 9, 2026
Merged

refactor: expose a public custom-weights packaging helper#503
lrosemberg merged 12 commits into
mainfrom
lean/package-custom-weights

Conversation

@lrosemberg

@lrosemberg lrosemberg commented Jul 2, 2026

Copy link
Copy Markdown
Contributor

What this PR does

Makes the SDK the single owner of custom-weights packaging. It exposes a public, non-interactive helper that both the SDK deploy flows and external callers (such as the Roboflow MCP server) use, so the packaging logic lives in one place instead of being copied per caller.

The old packaging entry point could not be reused outside the CLI: process() prompts with input(), can call sys.exit(1), prints to stdout, and writes artifacts into the caller's model directory. This PR extracts a clean, side-effect-free helper and keeps the CLI behavior behind a thin compatibility layer.

Before:                          After:
Version.deploy(...)              package_custom_weights(...)   (only packages, no I/O side effects)
Workspace.deploy_model(...)        ├─ Version.deploy(...)            package + SDK upload
  packaging logic inlined          ├─ Workspace.deploy_model(...)    package + SDK upload
  in each deploy flow              └─ external callers               package + their own upload

New public API: roboflow/util/model_processor.py

package_custom_weights(model_type, model_path, filename="weights/best.pt", *, build_dir=None, allow_dependency_mismatch=False, allow_size_mismatch=False) -> ModelUploadBundle

  • Non-interactive and side-effect free on model_path. It never prompts, prints, exits, or writes into the source directory. Heavy deps (torch, ultralytics) are imported lazily, only for the families that need them.
  • Build directory ownership. build_dir=None packages into a fresh mkdtemp() owned by the returned bundle (bundle.cleanup() removes it). Passing an explicit build_dir uses it as-is and leaves ownership to the caller. The old "write into model_path" behavior lives only behind the compatibility process().
  • Returns ModelUploadBundle with archive_path, build_dir, owns_build_dir, the resolved model_type, and collected warnings.

Typed error contract: ModelPackagingError (base), with UnsupportedModelError, TaskMismatchError, MissingFileError, MissingDependencyError, DependencyMismatchError, and SizeMismatchError. Every ModelPackagingError is a user-correctable input problem; anything else escaping the helper is a bug. Each subclass also inherits the builtin type the old code raised (ValueError / FileNotFoundError / RuntimeError), so existing except blocks keep working.

Validations (some new, some already in this module):

  • YOLO size-suffix inference and checking. Roboflow rejects bare family names (yolov8, yolov11, ...). The size is inferred from the model yaml's scale letter or the depth_multiple/width_multiple pair; a missing suffix is filled in (yolov8 becomes yolov8n, with a warning), and a declared size that conflicts with the checkpoint raises SizeMismatchError naming the size that fits, instead of failing server-side later.
  • RF-DETR variant checking. The checkpoint's position-encoding grid (from positional_encoding_size, resolution / patch_size, or the backbone position_embeddings tensor) is reconciled against the requested variant; a known mislabel raises pointing at the matching variant, an unknown or custom grid warns and proceeds.
  • Cleaner errors for missing files, malformed checkpoints, and missing torch/ultralytics.

Prefix-based family gate: model_type.startswith(SUPPORTED_MODELS) rather than substring containment, so an invalid type like foo-yolov8n is rejected up front instead of being half-recognized and rejected by the backend after upload.

Backwards compatibility

  • process(model_type, model_path, filename) is kept as a thin wrapper that delegates to package_custom_weights_interactive(build_dir=model_path): it packages into model_path (same intermediate artifacts and archive landing there as before), prints warnings, keeps the historical print-and-confirm on dependency/size mismatches, and returns (zip_file_name, model_type).
  • Version.deploy and Workspace.deploy_model keep their signatures. They call package_custom_weights_interactive, which preserves the historical CLI UX: warnings are printed, and dependency/size mismatches ask for confirmation before retrying with the matching override.
  • validate_model_type_for_project, task_of_model_type, _detect_yolo_task, _detect_rfdetr_task, get_classnames_txt_for_rfdetr, and maybe_prepend_dummy_class keep their signatures. yolo26-sem and semantic-segmentation handling are unchanged.

Hardening

Edge cases found while testing against real checkpoints, each fixed with a test:

  • RF-DETR re-deploy SameFileError. In the legacy Version.deploy / Workspace.deploy_model flow build_dir == model_path, so a top-level weights.pt was shutil.copy'd onto itself and raised SameFileError (an OSError, outside the packaging error contract). The self-copy is now skipped when source and destination resolve to the same path.
  • Bare-family infinite retry loop. When a bare family name (e.g. yolov10) had no inferable size, _resolve_yolo_size raised SizeMismatchError unconditionally, even under allow_mismatch; the interactive retry then re-prompted forever. It now converges (packages the bare name with a warning once allow_size_mismatch is set), matching the old deploy behavior.
  • Silent wrong checkpoint on a typo'd filename. An explicit, non-default filename that doesn't exist no longer falls back to discovering some other .pt/.pth in the directory (which would upload the wrong weights); it raises MissingFileError. The default path still discovers, as before.
  • Corrupt checkpoint args. A checkpoint storing args as a bare scalar/list raised a raw TypeError from vars(); a small _checkpoint_args_as_dict helper coerces those to {} so the error contract holds.

Checkpoint-completeness guards. Roboflow's server-side conversion reads metadata that a bare inference state_dict (e.g. {"model": <tensors>} with nothing else) does not carry, so such a file used to package and upload fine and then fail conversion with an opaque error like KeyError: 'args'. Packaging now rejects these up front with an actionable message:

  • RF-DETR: the checkpoint must contain args (class names, class count, model config); a bare state_dict is rejected instead of failing server-side.
  • YOLO-NAS: the checkpoint must contain processing_params.class_names; a bare state_dict is rejected.
  • YOLO args coercion: _filtered_args now routes through _checkpoint_args_as_dict, so a corrupt model.args coerces to {} rather than raising a raw TypeError.

Deferred (pre-existing behavior, not introduced here): the _process_yolo branches dereference checkpoint["train_args"], model_instance.nc, and model_instance.yaml["nc"] directly, which can raise raw exceptions on non-standard or reconstructed Ultralytics checkpoints. These only bite atypical inputs and guarding them touches the happy path across Ultralytics versions, so they are left for a separate hardening pass.

Intentional behavior changes (all previously interactive dead ends)

  • Declining a dependency-version prompt in the deploy flows now re-raises a typed error instead of sys.exit(1).
  • Hugging Face uploads with missing tokenizer/preprocessor sidecar files now fail with a clear MissingFileError instead of prompting y/n and calling exit(1).
  • RF-DETR checkpoint discovery: the default filename still falls back to the first top-level .pt/.pth (sorted for determinism) with a warning, preserving historical behavior; an explicit non-default filename that is missing is now a hard error rather than a silent fallback (see hardening above).

Tests

  • Packaging tests in tests/util/test_model_processor.py cover size/variant resolution, the error contract, and RF-DETR class names.
  • Contract tests: package_custom_weights fails the test if input() or sys.exit() is ever called; the source directory is byte-identical before and after packaging with a temp build dir; the owned temp dir is removed on failure; process() still writes the archive into model_path, keeps the interactive confirm, and returns the old tuple.
  • Regression tests for each hardening fix above.
  • python -m unittest (832 tests), ruff format --check, ruff check, and mypy all green on the CI matrix (Ubuntu + Windows, Python 3.10 to 3.13).

Follow-ups and coordination

@lrosemberg lrosemberg self-assigned this Jul 2, 2026
@lrosemberg lrosemberg marked this pull request as draft July 2, 2026 02:16
@lrosemberg

Copy link
Copy Markdown
Contributor Author

CI note: the build (ubuntu-latest, 3.13) failure was the pre-existing numpy 2.5 / mypy stub issue (#498), not this change; the matrix fail-fast then canceled the other jobs (3.11 had already passed the quality gate). I cherry-picked the same dev-extra numpy<2.5 cap that #495 carries in commit a4c0891 so this PR's CI can run green in the meantime. If #495 merges first, this line rebases away cleanly.

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ lrosemberg
❌ pre-commit-ci[bot]
You have signed the CLA already but the status is still pending? Let us recheck it.

@lrosemberg lrosemberg force-pushed the lean/package-custom-weights branch from 4e5a23f to 0f536de Compare July 6, 2026 15:04
@lrosemberg

Copy link
Copy Markdown
Contributor Author

Follow-up hardening from a review pass on the packaging helper (all with tests):

  • rf-detr re-deploy SameFileError: in the legacy Version.deploy / Workspace.deploy_model flow build_dir == model_path, so a top-level weights.pt was copied onto itself. The self-copy is now skipped.
  • Bare YOLO family infinite loop: _resolve_yolo_size raised unconditionally when a bare family name (e.g. yolov10) had no inferable size, even under allow_mismatch; the interactive retry then looped forever. It now converges (packages the bare name with a warning) once allow_size_mismatch is set, matching the old deploy behavior.
  • Silent wrong checkpoint: an explicit non-default filename that doesn't exist no longer falls back to discovering some other .pt in the directory (a typo would upload the wrong weights); it raises MissingFileError. The default path still discovers, as before.
  • Corrupt args: vars() on a scalar/list args raised a bare TypeError outside the ModelPackagingError contract; a small _checkpoint_args_as_dict coerces those to {}.

Also rebased onto latest main and dropped the redundant numpy<2.5 dev-extra commit (main already caps it in requirements.txt via #495) and the auto-format commit, so the history is single-author for the CLA check. python -m unittest green (832 tests).

@lrosemberg

Copy link
Copy Markdown
Contributor Author

Ready for review. This pulls the custom-weights packaging into a public, non-interactive package_custom_weights helper so the logic lives in one place, with a typed error contract, and keeps the old CLI behavior behind a thin compatibility wrapper (process(), Version.deploy, and Workspace.deploy_model are unchanged for callers).

Since the first version I hardened a handful of edge cases that showed up while testing against real checkpoints: a SameFileError on rf-detr re-deploy, an infinite retry loop on bare YOLO family names, a silent wrong-checkpoint on a typo'd filename, and a few bare-state-dict cases that used to fail server-side with an opaque KeyError (now caught up front with an actionable message). Each has a test.

CI is green across the matrix (832 tests, ruff, mypy), and the history is single-author so the CLA check passes. roboflow/roboflow-mcp#65 depends on this shipping to PyPI. Would appreciate a review when you have a minute.

@lrosemberg lrosemberg requested a review from Mike-Medvedev July 6, 2026 23:07
@lrosemberg lrosemberg marked this pull request as ready for review July 6, 2026 23:15
@jarbas-roboflow

Copy link
Copy Markdown

Status: ❌ Request Changes — Gist: https://gist.github.com/jarbasrf/f79a5d5a098e4b0508b9d203902dbf37. Resumo: The SDK/MCP integration shape is sound, but the MCP PR still declares the SDK dependency as a mutable git branch. Pin it to a release or immutable SHA before merging.

@lrosemberg

Copy link
Copy Markdown
Contributor Author

Thanks for the joint review. Both points are addressed on the MCP side (roboflow/roboflow-mcp#65):

  1. Mutable SDK branch dependencypyproject.toml now pins the SDK at the immutable commit 29d1aec3fcf7d8f75ebfb97275030f86e88ea135 instead of the lean/package-custom-weights branch, so the reviewed MCP behavior stays tied to the reviewed SDK behavior and can't drift on a rebase/delete. It will still be swapped for a released roboflow>=x.y.z before merge, per the pre-merge checklist.
  2. Versioned upload preflight — for version_number is not None, the tool now calls rf_api.versions_get(...) before packaging, so a typo in project_id/version_number fails fast instead of after torch.load + zip on a large checkpoint. Post-packaging task validation stays, since it needs the SDK-resolved bundle.model_type. Added a regression test asserting a missing version never reaches prepare_custom_weights_archive.

./dev/test.sh is green (327 passed).

🤖 Addressed by Claude Code

@lrosemberg lrosemberg requested a review from tonylampada July 8, 2026 18:33

@Mike-Medvedev Mike-Medvedev left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good extraction. package_custom_weights() is the right public surface for MCP and the deploy flows, and the interactive/compat split looks clean. CI is green.

Requesting changes on two contract issues:

  1. filename is not constrained under model_path. Joins like model_path / filename follow absolute paths and ../. MCP forwards caller filename into this helper, so on a hosted MCP host this is a path-read / wrong-weights risk, and it breaks the documented "relative to model_path" contract. After join, resolve() and require the checkpoint stays under the source dir; reject absolute filename. Add tests for ../ and absolute paths.

  2. YOLO happy path still leaks non-ModelPackagingError exceptions. Accesses like checkpoint["train_args"], model_instance.yaml["nc"], and model_instance.nc can raise raw KeyError / AttributeError. MCP maps only ModelPackagingError to 400s, so atypical checkpoints become opaque 500s. RF-DETR / YOLO-NAS already got completeness guards; YOLO should match (at least missing train_args on the yolov10/11/12/26/-cls branch).

Also missing HF packaging tests for the prompt→MissingFileError behavior change. Minor follow-ups: resolve explicit build_dir, tighten RF-DETR multi-checkpoint discovery.

Once this ships to PyPI with #501, MCP #65 can drop the git pin.

@Mike-Medvedev Mike-Medvedev left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inline notes on the two contract blockers.

Comment thread roboflow/util/model_processor.py
Comment thread roboflow/util/model_processor.py Outdated
@lrosemberg lrosemberg force-pushed the lean/package-custom-weights branch from 8730b60 to 7ae5856 Compare July 8, 2026 19:54
Mike-Medvedev
Mike-Medvedev previously approved these changes Jul 8, 2026

@Mike-Medvedev Mike-Medvedev left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. The earlier blockers are addressed:

  • _resolve_within_source rejects absolute/../ filenames (cross-platform) and directory targets
  • YOLO completeness goes through _require_* into ModelPackagingError
  • HF packaging tests are in place

CI is green. LGTM.

@digaobarbosa digaobarbosa left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@lrosemberg lrosemberg merged commit 07cb3d1 into main Jul 9, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants