refactor: expose a public custom-weights packaging helper#503
Conversation
|
CI note: the |
|
|
4e5a23f to
0f536de
Compare
|
Follow-up hardening from a review pass on the packaging helper (all with tests):
Also rebased onto latest |
|
Ready for review. This pulls the custom-weights packaging into a public, non-interactive Since the first version I hardened a handful of edge cases that showed up while testing against real checkpoints: a 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. |
|
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. |
|
Thanks for the joint review. Both points are addressed on the MCP side (roboflow/roboflow-mcp#65):
🤖 Addressed by Claude Code |
…eights # Conflicts: # roboflow/util/model_processor.py # tests/util/test_model_processor.py
Mike-Medvedev
left a comment
There was a problem hiding this comment.
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:
-
filenameis not constrained undermodel_path. Joins likemodel_path / filenamefollow absolute paths and../. MCP forwards callerfilenameinto 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 absolutefilename. Add tests for../and absolute paths. -
YOLO happy path still leaks non-
ModelPackagingErrorexceptions. Accesses likecheckpoint["train_args"],model_instance.yaml["nc"], andmodel_instance.nccan raise rawKeyError/AttributeError. MCP maps onlyModelPackagingErrorto 400s, so atypical checkpoints become opaque 500s. RF-DETR / YOLO-NAS already got completeness guards; YOLO should match (at least missingtrain_argson the yolov10/11/12/26/-clsbranch).
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
left a comment
There was a problem hiding this comment.
Inline notes on the two contract blockers.
8730b60 to
7ae5856
Compare
Mike-Medvedev
left a comment
There was a problem hiding this comment.
Looks good. The earlier blockers are addressed:
_resolve_within_sourcerejects absolute/../filenames (cross-platform) and directory targets- YOLO completeness goes through
_require_*intoModelPackagingError - HF packaging tests are in place
CI is green. LGTM.
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 withinput(), can callsys.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.New public API:
roboflow/util/model_processor.pypackage_custom_weights(model_type, model_path, filename="weights/best.pt", *, build_dir=None, allow_dependency_mismatch=False, allow_size_mismatch=False) -> ModelUploadBundlemodel_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_dir=Nonepackages into a freshmkdtemp()owned by the returned bundle (bundle.cleanup()removes it). Passing an explicitbuild_diruses it as-is and leaves ownership to the caller. The old "write intomodel_path" behavior lives only behind the compatibilityprocess().ModelUploadBundlewitharchive_path,build_dir,owns_build_dir, the resolvedmodel_type, and collectedwarnings.Typed error contract:
ModelPackagingError(base), withUnsupportedModelError,TaskMismatchError,MissingFileError,MissingDependencyError,DependencyMismatchError, andSizeMismatchError. EveryModelPackagingErroris 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 existingexceptblocks keep working.Validations (some new, some already in this module):
yolov8,yolov11, ...). The size is inferred from the model yaml'sscaleletter or thedepth_multiple/width_multiplepair; a missing suffix is filled in (yolov8becomesyolov8n, with a warning), and a declared size that conflicts with the checkpoint raisesSizeMismatchErrornaming the size that fits, instead of failing server-side later.positional_encoding_size,resolution / patch_size, or the backboneposition_embeddingstensor) is reconciled against the requested variant; a known mislabel raises pointing at the matching variant, an unknown or custom grid warns and proceeds.torch/ultralytics.Prefix-based family gate:
model_type.startswith(SUPPORTED_MODELS)rather than substring containment, so an invalid type likefoo-yolov8nis 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 topackage_custom_weights_interactive(build_dir=model_path): it packages intomodel_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.deployandWorkspace.deploy_modelkeep their signatures. They callpackage_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, andmaybe_prepend_dummy_classkeep their signatures.yolo26-semand semantic-segmentation handling are unchanged.Hardening
Edge cases found while testing against real checkpoints, each fixed with a test:
SameFileError. In the legacyVersion.deploy/Workspace.deploy_modelflowbuild_dir == model_path, so a top-levelweights.ptwasshutil.copy'd onto itself and raisedSameFileError(anOSError, outside the packaging error contract). The self-copy is now skipped when source and destination resolve to the same path.yolov10) had no inferable size,_resolve_yolo_sizeraisedSizeMismatchErrorunconditionally, even underallow_mismatch; the interactive retry then re-prompted forever. It now converges (packages the bare name with a warning onceallow_size_mismatchis set), matching the old deploy behavior.filenamethat doesn't exist no longer falls back to discovering some other.pt/.pthin the directory (which would upload the wrong weights); it raisesMissingFileError. The default path still discovers, as before.args. A checkpoint storingargsas a bare scalar/list raised a rawTypeErrorfromvars(); a small_checkpoint_args_as_dicthelper 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 likeKeyError: 'args'. Packaging now rejects these up front with an actionable message:args(class names, class count, model config); a bare state_dict is rejected instead of failing server-side.processing_params.class_names; a bare state_dict is rejected._filtered_argsnow routes through_checkpoint_args_as_dict, so a corruptmodel.argscoerces to{}rather than raising a rawTypeError.Deferred (pre-existing behavior, not introduced here): the
_process_yolobranches dereferencecheckpoint["train_args"],model_instance.nc, andmodel_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)
sys.exit(1).MissingFileErrorinstead of prompting y/n and callingexit(1).filenamestill falls back to the first top-level.pt/.pth(sorted for determinism) with a warning, preserving historical behavior; an explicit non-defaultfilenamethat is missing is now a hard error rather than a silent fallback (see hardening above).Tests
tests/util/test_model_processor.pycover size/variant resolution, the error contract, and RF-DETR class names.package_custom_weightsfails the test ifinput()orsys.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 intomodel_path, keeps the interactive confirm, and returns the old tuple.python -m unittest(832 tests),ruff format --check,ruff check, andmypyall green on the CI matrix (Ubuntu + Windows, Python 3.10 to 3.13).Follow-ups and coordination
package_custom_weightsonce this lands in a PyPI release. For the MCP to install the SDK, the exactopencv-python-headless==4.10.0.84pin inrequirements.txtalso needs to relax; fix(deps): relax opencv-python-headless pin to >=4.10.0 #501 (still open) does that. Until both land, the MCP PR pins this branch by git ref.SUPPORTED_RFDETR_TYPES, ...) and the per-family_process_*function names were kept to make those rebases mechanical.