Merge pull request #20 from gabrielmscampos/feat/add-ml-clients

cmsdials/clients/ml_bad_lumisection/client.py

@@ -0,0 +1,29 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import MLBadLumisection, MLBadLumisectionFilters, PaginatedMLBadLumisectionList
+
+
+class MLBadLumisectionClient(BaseAuthorizedAPIClient):
+    data_model = MLBadLumisection
+    pagination_model = PaginatedMLBadLumisectionList
+    filter_class = MLBadLumisectionFilters
+    lookup_url = "ml-bad-lumisection/"
+
+    def get(self, model_id: int, dataset_id: int, run_number: int, ls_number: int, me_id: int, **kwargs):
+        edp = f"{model_id}/{dataset_id}/{run_number}/{ls_number}/{me_id}/"
+        return super().get(edp, **kwargs)
+
+    def cert_json(self, model_id__in: list[int], dataset_id__in: list[int], run_number__in: list[int], **kwargs):
+        edp = "cert-json/"
+        midin = ",".join(str(v) for v in model_id__in)
+        didin = ",".join(str(v) for v in dataset_id__in)
+        ridin = ",".join(str(v) for v in run_number__in)
+        params = {"model_id__in": midin, "dataset_id__in": didin, "run_number__in": ridin}
+        return super().get(edp, params=params, return_raw_json=True, **kwargs)
+
+    def golden_json(self, model_id__in: list[int], dataset_id__in: list[int], run_number__in: list[int], **kwargs):
+        edp = "golden-json/"
+        midin = ",".join(str(v) for v in model_id__in)
+        didin = ",".join(str(v) for v in dataset_id__in)
+        ridin = ",".join(str(v) for v in run_number__in)
+        params = {"model_id__in": midin, "dataset_id__in": didin, "run_number__in": ridin}
+        return super().get(edp, params=params, return_raw_json=True, **kwargs)

cmsdials/clients/ml_bad_lumisection/models.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class MLBadLumisection(BaseModel):
+    model_id: int
+    dataset_id: int
+    file_id: int
+    run_number: int
+    ls_number: int
+    me_id: int
+    mse: float
+
+
+class PaginatedMLBadLumisectionList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[MLBadLumisection]
+
+
+class MLBadLumisectionFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    model_id: Optional[int] = None
+    model_id__in: Optional[list[int]] = None
+    dataset_id: Optional[int] = None
+    dataset_id__in: Optional[list[int]] = None
+    dataset: Optional[str] = None
+    dataset__regex: Optional[str] = None
+    me_id: Optional[int] = None
+    me: Optional[str] = None
+    me__regex: Optional[str] = None
+    run_number: Optional[int] = None
+    run_number__in: Optional[list[int]] = None
+    ls_number: Optional[int] = None

cmsdials/clients/ml_models_index/client.py

@@ -0,0 +1,13 @@
+from ...utils.api_client import BaseAuthorizedAPIClient
+from .models import MLModelsIndex, MLModelsIndexFilters, PaginatedMLModelsIndexList
+
+
+class MLModelsIndexClient(BaseAuthorizedAPIClient):
+    data_model = MLModelsIndex
+    pagination_model = PaginatedMLModelsIndexList
+    filter_class = MLModelsIndexFilters
+    lookup_url = "ml-models-index/"
+
+    def get(self, model_id: int, **kwargs):
+        edp = f"{model_id}/"
+        return super().get(edp, **kwargs)

cmsdials/clients/ml_models_index/models.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel
+
+from ...utils.base_model import OBaseModel, PaginatedBaseModel
+
+
+class MLModelsIndex(BaseModel):
+    model_id: int
+    filename: str
+    target_me: str
+    active: bool
+
+
+class PaginatedMLModelsIndexList(PaginatedBaseModel):
+    next: Optional[AnyUrl]
+    previous: Optional[AnyUrl]
+    results: list[MLModelsIndex]
+
+
+class MLModelsIndexFilters(OBaseModel):
+    next_token: Optional[str] = None
+    page_size: Optional[int] = None
+    model_id: Optional[int] = None
+    model_id__in: Optional[list[int]] = None
+    filename: Optional[str] = None
+    filename__regex: Optional[str] = None
+    target_me: Optional[str] = None
+    target_me__regex: Optional[str] = None
+    active: Optional[bool] = None

cmsdials/cmsdials.py

@@ -6,6 +6,8 @@
 from .clients.h2d.client import LumisectionHistogram2DClient
 from .clients.lumisection.client import LumisectionClient
 from .clients.mes.client import MonitoringElementClient
+from .clients.ml_bad_lumisection.client import MLBadLumisectionClient
+from .clients.ml_models_index.client import MLModelsIndexClient
 from .clients.oms_proxy.client import OMSProxyClient
 from .clients.run.client import RunClient
 
@@ -19,3 +21,5 @@ def __init__(self, creds: BaseCredentials, workspace: Optional[str] = None, *arg
         self.run = RunClient(creds, workspace, *args, **kwargs)
         self.mes = MonitoringElementClient(creds, workspace, *args, **kwargs)
         self.oms = OMSProxyClient(creds, *args, **kwargs)
+        self.ml_bad_lumis = MLBadLumisectionClient(creds, workspace, *args, **kwargs)
+        self.ml_models_index = MLModelsIndexClient(creds, workspace, *args, **kwargs)

cmsdials/filters.py

@@ -3,6 +3,8 @@
 from .clients.h2d.models import LumisectionHistogram2DFilters
 from .clients.lumisection.models import LumisectionFilters
 from .clients.mes.models import MEFilters
+from .clients.ml_bad_lumisection.models import MLBadLumisectionFilters
+from .clients.ml_models_index.models import MLModelsIndexFilters
 from .clients.oms_proxy.models import OMSFilter, OMSPage
 from .clients.run.models import RunFilters
 
@@ -14,6 +16,8 @@
     "LumisectionFilters",
     "RunFilters",
     "MEFilters",
+    "MLBadLumisectionFilters",
+    "MLModelsIndexFilters",
     "OMSFilter",
     "OMSPage",
 ]

cmsdials/utils/api_client.py

@@ -81,12 +81,13 @@ def _build_headers(self) -> dict:
         self.creds.before_request(base)
         return base
 
-    def get(self, edp: str, retries=DEFAULT_RETRIES):  # noqa: A002
+    def get(self, edp: str, retries=DEFAULT_RETRIES, params: Optional[dict] = None, return_raw_json: bool = False):
         endpoint_url = self.api_url + self.lookup_url + edp
         headers = self._build_headers()
         response = self._requests_get_retriable(
             endpoint_url,
             headers=headers,
+            params=params,
             timeout=self.default_timeout,
             retries=retries,
         )
@@ -98,6 +99,8 @@ def get(self, edp: str, retries=DEFAULT_RETRIES):  # noqa: A002
             raise err
 
         response = response.json()
+        if return_raw_json:
+            return response
         return self.data_model(**response)
 
     def list(self, filters=None, retries=DEFAULT_RETRIES):

cmsdials/utils/base_model.py

@@ -14,7 +14,11 @@
 
 class OBaseModel(BaseModel):
     def cleandict(self):
-        return {key: value for key, value in self.dict().items() if value is not None}
+        return {
+            key: ",".join(str(v) for v in value) if isinstance(value, (list, tuple)) else value
+            for key, value in self.dict().items()
+            if value is not None
+        }
 
 
 class PaginatedBaseModel(BaseModel):