Repository navigation
fix(train): discover Training Sessions models by capability keyword - #6506
Conversation
| #: before the Public Hub publishes the capability keyword above. Mirrors the | ||
| #: service's model configuration; a listed model that is absent from a | ||
| #: Region's Hub is simply not reported there. Remove when the keyword is live. | ||
| _TRAINABLE_MODELS_BY_REGION: Dict[str, FrozenSet[str]] = { |
There was a problem hiding this comment.
Does this mean we're going to now start maintaining this list of models in each region?
There was a problem hiding this comment.
Yes, until we have the Model Hub capability added.
| #: The Public Hub publishes each model's ``Capabilities`` as | ||
| #: ``@capability:{capability}`` (lowercase) keywords; ``TRAINING_SESSION`` is | ||
| #: set per region, so the keyword is present only where the model is served. | ||
| _TRAINABLE_CAPABILITY_KEYWORD = "@capability:training_session" |
There was a problem hiding this comment.
Do we still need list below or we can directly publish our list through internal pipeline to JumpStart?
There was a problem hiding this comment.
It's using the "OR" condition today, once JS new capability onboarding is done, we don't have to maintain the list below.
…bility list_base_models() and base_model resolution keep a Hub model when it carries @capability:training_session, or, until the Public Hub publishes that keyword, when it is in the client Region's static list of models served by Training Sessions. Replaces the MTRL recipe keyword filter, which listed only models that also have a managed MTRL recipe.
09a3610 to
784f6bb
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 AI code review
This PR changes Training Sessions model discovery from filtering Public Hub entries on the @recipe:finetuning_mtrl_ recipe keyword to filtering on the @capability:training_session capability keyword, with a transitional per-Region static allow-list (_TRAINABLE_MODELS_BY_REGION) for the interval before the Hub publishes the capability keyword. The change is well-contained, documented, and backed by new unit tests covering the static-list match, per-region behaviour, and the no-region case. No correctness, security, or backward-compatibility issues found; the implementation and tests look clean.
Reviewed commit 784f6bb. Automated review; verify before acting.
Why
list_base_models()andbase_modelresolution filtered Public Hub entries on the@recipe:finetuning_mtrl_keyword, so only models that also ship a managed MTRL recipe were discoverable. Training Sessions serves models independently of MTRL recipes (e.g.Qwen/Qwen3.8-27B).Change
@capability:training_session, or, until the Public Hub publishes that keyword, when it is in the client Region's static list of models served by Training Sessions (_TRAINABLE_MODELS_BY_REGION).Tests
tests/unit/train/training_session/{test_utils,test_service_client}.py(252 passed), including new cases for the static list (match, per-Region, no Region).Checklist