DRAFT: Make embedding model switchable - #1752
Conversation
|
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds selectable torch and Triton embedding backends, configurable model variants, chunked Triton inference, runtime PIL-image checks, optional dependency handling, and related tests and documentation. ChangesEmbedding backend integration
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
lightly_studio/.env.example (1)
17-22: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winClarify
LIGHTLY_STUDIO_EMBEDDINGS_MODEL_NAMEsemantics for thetritontype.Line 18 documents this variable via "# Embeddings. The torch backend supports MobileCLIP and Perception Encoder model names." but doesn't state what value is expected when
LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE=triton. Inembedding_manager.py, the same env var is forwarded verbatim as the Triton model name (TritonEmbeddingGenerator(url=..., model_name=model_name)), so with the defaultmobileclip_s0a user switching to Triton without also overridingMODEL_NAMEwill send a Triton model name that likely doesn't match their actual deployed model repository name.📝 Suggested doc clarification
-# LIGHTLY_STUDIO_EMBEDDINGS_MODEL_NAME=mobileclip_s0 +# For torch: name of a MobileCLIP or Perception Encoder variant. +# For triton: must match the model name in your Triton model repository. +# LIGHTLY_STUDIO_EMBEDDINGS_MODEL_NAME=mobileclip_s0🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lightly_studio/.env.example` around lines 17 - 22, Update the embedding environment variable comments to state that LIGHTLY_STUDIO_EMBEDDINGS_MODEL_NAME must match the deployed Triton model repository name when LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE=triton, while retaining the existing torch model-name guidance.lightly_studio/src/lightly_studio/dataset/embedding_manager.py (1)
247-281: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winChunked
embed_imageslogic verified correct; consider adding test coverage.The chunk-relative
kept_indicesre-offsetting, the "raise only if no chunk succeeded"AllInputFilesFailedErrorhandling, and the empty-embeddings fallback are all consistent. However, no test intest_embedding_manager.pyexercises this chunking/failure path (only_load_torch_embedding_generatoris covered by the new tests), despite this being flagged as high complexity.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lightly_studio/src/lightly_studio/dataset/embedding_manager.py` around lines 247 - 281, Add tests in test_embedding_manager.py covering the chunked embed_images flow: verify kept_indices are offset to global sample IDs across chunks, successful chunks are stored when another raises AllInputFilesFailedError, and the error is re-raised only when every chunk fails with an empty-embeddings fallback. Use mocks for model.embed_images and _store_embeddings.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@lightly_studio/pyproject.toml`:
- Line 68: The Triton client should be optional rather than installed for every
user. In lightly_studio/pyproject.toml lines 68-68, move
tritonclient[grpc]>=2.41.1 from the main dependencies into a Triton
optional-dependencies/extras group; in
lightly_studio/src/lightly_studio/dataset/embedding_manager.py lines 606-618,
retain the existing try/except ImportError guard around the Triton image
embedding generator loader with no direct change required.
In `@lightly_studio/src/lightly_studio/dataset/embedding_manager.py`:
- Around line 600-626: Update _load_image_embedding_generator_from_env so a
failed TritonEmbeddingGenerator import logs the disabled message once and
returns None immediately. Prevent the Triton path from falling through to the
unsupported-model-type warning, while preserving the existing behavior for
successfully loaded Triton and torch generators.
In
`@lightly_studio/src/lightly_studio/dataset/triton_mobileclip_embedding_generator.py`:
- Line 42: Update the embedding generator’s Triton RPCs to pass the configured
finite client_timeout to both get_model_metadata() during initialization and
infer() during batch embedding, ensuring neither call can wait indefinitely
while preserving the existing request behavior.
---
Nitpick comments:
In `@lightly_studio/.env.example`:
- Around line 17-22: Update the embedding environment variable comments to state
that LIGHTLY_STUDIO_EMBEDDINGS_MODEL_NAME must match the deployed Triton model
repository name when LIGHTLY_STUDIO_EMBEDDINGS_MODEL_TYPE=triton, while
retaining the existing torch model-name guidance.
In `@lightly_studio/src/lightly_studio/dataset/embedding_manager.py`:
- Around line 247-281: Add tests in test_embedding_manager.py covering the
chunked embed_images flow: verify kept_indices are offset to global sample IDs
across chunks, successful chunks are stored when another raises
AllInputFilesFailedError, and the error is re-raised only when every chunk fails
with an empty-embeddings fallback. Use mocks for model.embed_images and
_store_embeddings.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 71447df4-a94f-429a-955f-f914ac402434
⛔ Files ignored due to path filters (3)
lightly_studio/uv.lockis excluded by!**/*.locktritonserver/server_trt/__pycache__/dali_pipeline.cpython-37.pycis excluded by!**/*.pyctritonserver/server_trt/model_repository/mobileclip_s0/1/__pycache__/model.cpython-37.pycis excluded by!**/*.pyc
📒 Files selected for processing (22)
.gitignorelightly_studio/.env.examplelightly_studio/docs/docs/tutorials/yolo26-model-evaluation.mdlightly_studio/pyproject.tomllightly_studio/src/lightly_studio/api/routes/api/enterprise.pylightly_studio/src/lightly_studio/dataset/embedding_generator.pylightly_studio/src/lightly_studio/dataset/embedding_manager.pylightly_studio/src/lightly_studio/dataset/env.pylightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.pylightly_studio/src/lightly_studio/dataset/perception_encoder_embedding_generator.pylightly_studio/src/lightly_studio/dataset/triton_mobileclip_embedding_generator.pylightly_studio/src/lightly_studio/examples/example_custom_embedding_model.pylightly_studio/tests/api/routes/api/test_enterprise.pylightly_studio/tests/dataset/test_embedding_manager.pylightly_studio/tests/dataset/test_fsspec_lister.pylightly_studio/tests/dataset/test_image_crop_embedding.pylightly_studio/tests/dataset/test_image_embedding.pylightly_studio/tests/dataset/test_mobileclip_embedding_generator.pylightly_studio/tests/dataset/test_perception_encoder_embedding_generator.pylightly_studio/tests/dataset/test_triton_mobileclip_embedding_generator.pylightly_studio_view/src/lib/components/PlotPanel/SearchMarkerOverlay.tslightly_studio_view/src/lib/hooks/useEmbeddingProjection/useEmbeddingProjection.ts
liopeer
changed the title
|
/review |
2 participants
What has changed and why?
NOTE: Not 100% finished, but would like early feedback, since this is my first PR here! Especially the tests still need some focus.
How has it been tested?
Did you update CHANGELOG.md?
Summary by CodeRabbit