fix(trainer): extract runtime metadata from SDK Runtime in get_runtime - #192
Conversation
Extract framework, image, num_nodes, device, device_count, and trainer_type from the SDK Runtime.trainer object in get_runtime(), while keeping backward-compatible spec fallback. Signed-off-by: priyank <priyank8445@gmail.com>
| if image: | ||
| data["image"] = str(image) | ||
| num_nodes = getattr(trainer, "num_nodes", None) | ||
| if num_nodes: |
There was a problem hiding this comment.
| if num_nodes: | |
| if num_nodes is not None: |
if num_nodes skips when num_nodes is 0, because 0 is falsy in Python.
But 0 is a valid value so the response drops num_nodes entirely.
| class TestGetRuntime: | ||
| @patch("kubeflow_mcp.trainer.api.discovery.get_trainer_client") | ||
| def test_get_runtime_extracts_sdk_trainer_metadata(self, mock_client_fn): | ||
| mock_trainer = MagicMock() |
There was a problem hiding this comment.
Thinking would it be good to add a test with real Runtime / RuntimeTrainer objects instead of MagicMock so this doesn't slip through again : https://github.com/kubeflow/mcp-server/pull/192/changes#r3901070952
| return data | ||
|
|
||
|
|
||
| def get_runtime(name: str, include_packages: bool = False) -> dict[str, Any]: |
There was a problem hiding this comment.
include_packages=True triggers a second client.get_runtime() via _get_runtime_image(). Thinking .. we already have rt in get_runtime().. can we pass trainer.image through (or the whole rt) to avoid duplicate API calls on the slow path. WDYT?
|
|
||
|
|
||
| def get_runtime(name: str, include_packages: bool = False) -> dict[str, Any]: | ||
| """Get ClusterTrainingRuntime configuration. |
There was a problem hiding this comment.
Docstring still says "ClusterTrainingRuntime" only. SDK resolves both cluster-scoped and namespaced runtimes right?.. worth updating while this is not merged yet..
…date docstrings Pass the already resolved runtime object to _fetch_packages_via_pod() and _get_runtime_image() to prevent duplicate client.get_runtime() calls on include_packages=True. Also update docstrings to reflect that get_runtime and list_runtimes resolve both namespaced TrainingRuntimes and cluster-scoped ClusterTrainingRuntimes. Signed-off-by: priyank <priyank8445@gmail.com>
… Runtime types Update num_nodes check from truthy check to is not None so 0 is preserved. Add unit tests using real SDK Runtime and RuntimeTrainer objects. Signed-off-by: priyank <priyank8445@gmail.com>
|
🚀 Thanks for the great feedback and suggestions @abhijeet-dhumal! Updated:
|
|
/ok-to-test |
|
Thanks @priyank766 ! |
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: abhijeet-dhumal The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
What this PR does
get_runtime()to extract runtime metadata (framework,image,num_nodes,device,device_count,trainer_type) fromrt.trainerandrt.pretrained_model, aligning with the Kubeflow SDK'stypes.Runtimedataclass..spec._get_runtime_imageto checkclient.get_runtime(name)first (which resolves both namespaced and cluster-scoped runtimes) before falling back toCustomObjectsApi._extract_runtime_metadatato keep cyclomatic complexity low.kubeflow_mcp/trainer/api/discovery_test.py.Fixes #191
How it was tested
uv run pytest kubeflow_mcp/trainer/api/discovery_test.py(14 passed)uv run pytest(516 passed, all unit & conformance schema snapshot tests pass)uv run ruff check .anduv run ruff format --check .(clean)