|
24 | 24 | import os
|
25 | 25 | import re
|
26 | 26 | import time
|
27 |
| -from enum import Enum |
28 | 27 | from functools import lru_cache
|
29 | 28 | from pathlib import Path
|
30 | 29 | from threading import RLock
|
@@ -163,18 +162,6 @@ def make_client(self) -> ClientLibrary:
|
163 | 162 | return client
|
164 | 163 |
|
165 | 164 |
|
166 |
| -class DiagnosticCategory(Enum): |
167 |
| - COMPUTES = "computes" |
168 |
| - LABS = "labs" |
169 |
| - LAB_EVENTS = "lab_events" |
170 |
| - NODE_LAUNCH_QUEUE = "node_launch_queue" |
171 |
| - SERVICES = "services" |
172 |
| - NODE_DEFINITIONS = "node_definitions" |
173 |
| - USER_LIST = "user_list" |
174 |
| - LICENSING = "licensing" |
175 |
| - STARTUP_SCHEDULER = "startup_scheduler" |
176 |
| - |
177 |
| - |
178 | 165 | class ClientLibrary:
|
179 | 166 | """Python bindings for the REST API of a CML controller."""
|
180 | 167 |
|
@@ -207,7 +194,7 @@ class ClientLibrary:
|
207 | 194 | "labs": "labs",
|
208 | 195 | "lab": "labs/{lab_id}",
|
209 | 196 | "lab_topology": "labs/{lab_id}/topology",
|
210 |
| - "diagnostics": "diagnostics/{category}", |
| 197 | + "diagnostics": "diagnostics", |
211 | 198 | "system_health": "system_health",
|
212 | 199 | "system_stats": "system_stats",
|
213 | 200 | "populate_lab_tiles": "populate_lab_tiles",
|
@@ -639,11 +626,16 @@ def all_labs(self, show_all: bool = False) -> list[Lab]:
|
639 | 626 | :param show_all: Whether to get only labs owned by the admin or all user labs.
|
640 | 627 | :returns: A list of Lab objects.
|
641 | 628 | """
|
642 |
| - lab_ids = self.get_lab_list(show_all=show_all) |
| 629 | + url = {"url": self._url_for("labs")} |
| 630 | + if show_all: |
| 631 | + url["params"] = {"show_all": True} |
| 632 | + lab_ids = self._session.get(**url).json() |
| 633 | + |
643 | 634 | result = []
|
644 | 635 | for lab_id in lab_ids:
|
645 | 636 | lab = self.join_existing_lab(lab_id)
|
646 | 637 | result.append(lab)
|
| 638 | + |
647 | 639 | return result
|
648 | 640 |
|
649 | 641 | @locked
|
@@ -818,31 +810,14 @@ def join_existing_lab(self, lab_id: str, sync_lab: bool = True) -> Lab:
|
818 | 810 | self._labs[lab_id] = lab
|
819 | 811 | return lab
|
820 | 812 |
|
821 |
| - def get_diagnostics( |
822 |
| - self, categories: list[DiagnosticCategory] | None = None |
823 |
| - ) -> dict: |
| 813 | + def get_diagnostics(self) -> dict: |
824 | 814 | """
|
825 |
| - Return selected diagnostic data as a JSON object. |
| 815 | + Return the controller diagnostic data as a JSON object. |
826 | 816 |
|
827 |
| - :param categories: List of diagnostic categories to fetch. If None, fetch all. |
828 | 817 | :returns: The diagnostic data.
|
829 | 818 | """
|
830 |
| - if categories is None: |
831 |
| - categories = list(DiagnosticCategory) |
832 |
| - |
833 |
| - diagnostics_data = {} |
834 |
| - for category in categories: |
835 |
| - value = category.value |
836 |
| - url = self._url_for("diagnostics", category=value) |
837 |
| - try: |
838 |
| - response = self._session.get(url) |
839 |
| - response.raise_for_status() |
840 |
| - diagnostics_data[value] = response.json() |
841 |
| - except httpx.HTTPStatusError: |
842 |
| - diagnostics_data[value] = { |
843 |
| - "error": f"Failed to fetch {value} diagnostics" |
844 |
| - } |
845 |
| - return diagnostics_data |
| 819 | + url = self._url_for("diagnostics") |
| 820 | + return self._session.get(url).json() |
846 | 821 |
|
847 | 822 | def get_system_health(self) -> dict:
|
848 | 823 | """
|
|
0 commit comments