feat!: separate Prom trainer and Druid trainer

Signed-off-by: Avik Basu <[email protected]>

numalogic/connectors/_config.py

@@ -16,8 +16,9 @@ class ConnectorConf:
 
 @dataclass
 class PrometheusConf(ConnectorConf):
- pushgateway: str
+ pushgateway: str = ""
  scrape_interval: int = 30
+ return_labels: list[str] = field(default_factory=list)
 
 
 @dataclass

numalogic/udfs/__init__.py

@@ -9,7 +9,7 @@
 from numalogic.udfs.inference import InferenceUDF
 from numalogic.udfs.postprocess import PostprocessUDF
 from numalogic.udfs.preprocess import PreprocessUDF
-from numalogic.udfs.trainer import TrainerUDF
+from numalogic.udfs.trainer import TrainerUDF, PromTrainerUDF, DruidTrainerUDF
 
 
 def set_logger() -> None:
@@ -27,6 +27,8 @@ def set_logger() -> None:
  "PreprocessUDF",
  "InferenceUDF",
  "TrainerUDF",
+ "PromTrainerUDF",
+ "DruidTrainerUDF",
  "PostprocessUDF",
  "UDFFactory",
  "StreamConf",

numalogic/udfs/factory.py

@@ -10,27 +10,28 @@
 # limitations under the License.
 
 import logging
+import warnings
 from typing import ClassVar
 
-from pynumaflow.mapper import Mapper, MultiProcMapper, AsyncMapper
-
-from numalogic.udfs import NumalogicUDF
-from numalogic.udfs.inference import InferenceUDF
-from numalogic.udfs.postprocess import PostprocessUDF
-from numalogic.udfs.preprocess import PreprocessUDF
-from numalogic.udfs.trainer import TrainerUDF
 
 _LOGGER = logging.getLogger(__name__)
 
 
 class UDFFactory:
  """Factory class to fetch the right UDF."""
 
+ from numalogic.udfs import NumalogicUDF
+ from numalogic.udfs.inference import InferenceUDF
+ from numalogic.udfs.postprocess import PostprocessUDF
+ from numalogic.udfs.preprocess import PreprocessUDF
+ from numalogic.udfs.trainer import DruidTrainerUDF, PromTrainerUDF
+
  _UDF_MAP: ClassVar[dict[str, type[NumalogicUDF]]] = {
  "preprocess": PreprocessUDF,
  "inference": InferenceUDF,
  "postprocess": PostprocessUDF,
- "trainer": TrainerUDF,
+ "druidtrainer": DruidTrainerUDF,
+ "promtrainer": PromTrainerUDF,
  }
 
  @classmethod
@@ -50,6 +51,12 @@ def get_udf_cls(cls, udf_name: str) -> type[NumalogicUDF]:
  ------
  ValueError: If the UDF name is invalid
  """
+ if udf_name == "trainer":
+ warnings.warn(
+ "UDF name 'trainer' is deprecated. Use 'druidtrainer' or 'promtrainer' instead."
+ )
+ udf_name = "druidtrainer"
+
  try:
  return cls._UDF_MAP[udf_name]
  except KeyError as err:
@@ -81,6 +88,8 @@ def get_udf_instance(cls, udf_name: str, **kwargs) -> NumalogicUDF:
 class ServerFactory:
  """Factory class to fetch the right pynumaflow function server/mapper."""
 
+ from pynumaflow.mapper import Mapper, MultiProcMapper, AsyncMapper
+
  _SERVER_MAP: ClassVar[dict] = {
  "sync": Mapper,
  "async": AsyncMapper,

numalogic/udfs/trainer/__init__.py

@@ -0,0 +1,5 @@
+from numalogic.udfs.trainer._base import TrainerUDF
+from numalogic.udfs.trainer._prom import PromTrainerUDF
+from numalogic.udfs.trainer._druid import DruidTrainerUDF
+
+__all__ = ["TrainerUDF", "PromTrainerUDF", "DruidTrainerUDF"]

numalogic/udfs/trainer.py → numalogic/udfs/trainer/_base.py

@@ -14,8 +14,6 @@
 from numalogic.base import StatelessTransformer
 from numalogic.config import PreprocessFactory, ModelFactory, ThresholdFactory, RegistryFactory
 from numalogic.config._config import NumalogicConf
-from numalogic.config.factory import ConnectorFactory
-from numalogic.connectors import DruidFetcherConf
 from numalogic.models.autoencoder import TimeseriesTrainer
 from numalogic.tools.data import StreamingDataset
 from numalogic.tools.exceptions import ConfigNotFoundError, RedisRegistryError
@@ -55,52 +53,12 @@ def __init__(
  jitter_sec=jitter_sec,
  jitter_steps_sec=jitter_steps_sec,
  )
- self.druid_conf = self.pl_conf.druid_conf
-
- data_fetcher_cls = ConnectorFactory.get_cls("DruidFetcher")
- try:
- self.data_fetcher = data_fetcher_cls(
- url=self.druid_conf.url, endpoint=self.druid_conf.endpoint
- )
- except AttributeError:
- _LOGGER.warning("Druid config not found, data fetcher will not be initialized!")
- self.data_fetcher = None
 
  self._model_factory = ModelFactory()
  self._preproc_factory = PreprocessFactory()
  self._thresh_factory = ThresholdFactory()
  self.train_msg_deduplicator = TrainMsgDeduplicator(r_client)
 
- def register_druid_fetcher_conf(self, config_id: str, conf: DruidFetcherConf) -> None:
- """
- Register DruidFetcherConf with the UDF.
-
- Args:
- config_id: Config ID
- conf: DruidFetcherConf object
- """
- self.pl_conf.druid_conf.id_fetcher[config_id] = conf
-
- def get_druid_fetcher_conf(self, config_id: str) -> DruidFetcherConf:
- """
- Get DruidFetcherConf with the given ID.
-
- Args:
- config_id: Config ID
-
- Returns
- -------
- DruidFetcherConf object
-
- Raises
- ------
- ConfigNotFoundError: If config with the given ID is not found
- """
- try:
- return self.pl_conf.druid_conf.id_fetcher[config_id]
- except KeyError as err:
- raise ConfigNotFoundError(f"Config with ID {config_id} not found!") from err
-
  @classmethod
  def compute(
  cls,
@@ -312,7 +270,7 @@ def get_feature_arr(
 
  def fetch_data(self, payload: TrainerPayload) -> pd.DataFrame:
  """
- Fetch data from druid.
+ Fetch data from a data connector.
 
  Args:
  payload: TrainerPayload object
@@ -321,37 +279,4 @@ def fetch_data(self, payload: TrainerPayload) -> pd.DataFrame:
  -------
  Dataframe
  """
- _start_time = time.perf_counter()
- _conf = self.get_conf(payload.config_id)
- _fetcher_conf = self.druid_conf.fetcher or (
- self.get_druid_fetcher_conf(payload.config_id) if self.druid_conf.id_fetcher else None
- )
- if not _fetcher_conf:
- raise ConfigNotFoundError(
- f"Druid fetcher config not found for config_id: {payload.config_id}!"
- )
-
- try:
- _df = self.data_fetcher.fetch(
- datasource=_fetcher_conf.datasource,
- filter_keys=_conf.composite_keys,
- filter_values=payload.composite_keys,
- dimensions=list(_fetcher_conf.dimensions),
- delay=self.druid_conf.delay_hrs,
- granularity=_fetcher_conf.granularity,
- aggregations=dict(_fetcher_conf.aggregations),
- group_by=list(_fetcher_conf.group_by),
- pivot=_fetcher_conf.pivot,
- hours=_conf.numalogic_conf.trainer.train_hours,
- )
- except Exception:
- _LOGGER.exception("%s - Error while fetching data from druid", payload.uuid)
- return pd.DataFrame()
-
- _LOGGER.debug(
- "%s - Time taken to fetch data: %.3f sec, df shape: %s",
- payload.uuid,
- time.perf_counter() - _start_time,
- _df.shape,
- )
- return _df
+ raise NotImplementedError("fetch_data method not implemented")

numalogic/udfs/trainer/_druid.py

@@ -0,0 +1,118 @@
+import logging
+import time
+from typing import Optional
+
+import pandas as pd
+
+from numalogic.config.factory import ConnectorFactory
+from numalogic.connectors import DruidFetcherConf
+from numalogic.tools.exceptions import ConfigNotFoundError
+from numalogic.tools.types import redis_client_t
+from numalogic.udfs._config import PipelineConf
+from numalogic.udfs.entities import TrainerPayload
+from numalogic.udfs.trainer._base import TrainerUDF
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class DruidTrainerUDF(TrainerUDF):
+ """
+ Trainer UDF using Druid as data source.
+
+ Args:
+ r_client: Redis client
+ pl_conf: Pipeline config
+ """
+
+ def __init__(
+ self,
+ r_client: redis_client_t,
+ pl_conf: Optional[PipelineConf] = None,
+ ):
+ super().__init__(r_client=r_client, pl_conf=pl_conf)
+ self.dataconn_conf = self.pl_conf.druid_conf
+ data_fetcher_cls = ConnectorFactory.get_cls("DruidFetcher")
+ try:
+ self.data_fetcher = data_fetcher_cls(
+ url=self.dataconn_conf.url, endpoint=self.dataconn_conf.endpoint
+ )
+ except AttributeError as err:
+ raise ConfigNotFoundError("Druid config not found!") from err
+
+ def register_druid_fetcher_conf(self, config_id: str, conf: DruidFetcherConf) -> None:
+ """
+ Register DruidFetcherConf with the UDF.
+
+ Args:
+ config_id: Config ID
+ conf: DruidFetcherConf object
+ """
+ self.pl_conf.druid_conf.id_fetcher[config_id] = conf
+
+ def get_druid_fetcher_conf(self, config_id: str) -> DruidFetcherConf:
+ """
+ Get DruidFetcherConf with the given ID.
+
+ Args:
+ config_id: Config ID
+
+ Returns
+ -------
+ DruidFetcherConf object
+
+ Raises
+ ------
+ ConfigNotFoundError: If config with the given ID is not found
+ """
+ try:
+ return self.pl_conf.druid_conf.id_fetcher[config_id]
+ except KeyError as err:
+ raise ConfigNotFoundError(f"Config with ID {config_id} not found!") from err
+
+ def fetch_data(self, payload: TrainerPayload) -> pd.DataFrame:
+ """
+ Fetch data from druid.
+
+ Args:
+ payload: TrainerPayload object
+
+ Returns
+ -------
+ Dataframe
+ """
+ _start_time = time.perf_counter()
+ _conf = self.get_conf(payload.config_id)
+ _fetcher_conf = self.dataconn_conf.fetcher or (
+ self.get_druid_fetcher_conf(payload.config_id)
+ if self.dataconn_conf.id_fetcher
+ else None
+ )
+ if not _fetcher_conf:
+ raise ConfigNotFoundError(
+ f"Druid fetcher config not found for config_id: {payload.config_id}!"
+ )
+
+ try:
+ _df = self.data_fetcher.fetch(
+ datasource=_fetcher_conf.datasource,
+ filter_keys=_conf.composite_keys,
+ filter_values=payload.composite_keys,
+ dimensions=list(_fetcher_conf.dimensions),
+ delay=self.dataconn_conf.delay_hrs,
+ granularity=_fetcher_conf.granularity,
+ aggregations=dict(_fetcher_conf.aggregations),
+ group_by=list(_fetcher_conf.group_by),
+ pivot=_fetcher_conf.pivot,
+ hours=_conf.numalogic_conf.trainer.train_hours,
+ )
+ except Exception:
+ _LOGGER.exception("%s - Error while fetching data from druid", payload.uuid)
+ return pd.DataFrame()
+
+ _LOGGER.debug(
+ "%s - Time taken to fetch data: %.3f sec, df shape: %s",
+ payload.uuid,
+ time.perf_counter() - _start_time,
+ _df.shape,
+ )
+ return _df