From: Kefu Chai <kchai@redhat.com>
Date: Sat, 20 Feb 2021 05:15:58 +0000 (+0800)
Subject: mgr/diskprediction_local: add Predictor abstract class
X-Git-Tag: v17.1.0~2905^2~2
X-Git-Url: http://git-server-git.apps.pok.os.sepia.ceph.com/?a=commitdiff_plain;h=2bb6821f50c15d5303bb3833b32c16d69e251b6e;p=ceph.git

mgr/diskprediction_local: add Predictor abstract class

to make the interface more explicit, also pave the road to mypy.

Signed-off-by: Kefu Chai <kchai@redhat.com>
---

diff --git a/doc/_ext/ceph_commands.py b/doc/_ext/ceph_commands.py
index 897343cee214..eec031ee3d2a 100644
--- a/doc/_ext/ceph_commands.py
+++ b/doc/_ext/ceph_commands.py
@@ -275,13 +275,14 @@ class CephMgrCommands(Directive):
         mock_imports += ['OpenSSL',
                          'jwt',
                          'bcrypt',
-                         'scipy',
                          'jsonpatch',
                          'rook.rook_client',
                          'rook.rook_client.ceph',
                          'rook.rook_client._helper',
                          'cherrypy=3.2.3']
-
+        # make diskprediction_local happy
+        mock_imports += ['numpy',
+                         'scipy']
         # make restful happy
         mock_imports += ['pecan',
                          'pecan.rest',
diff --git a/src/pybind/mgr/diskprediction_local/module.py b/src/pybind/mgr/diskprediction_local/module.py
index 930fd1bbce4b..31b1ccf925e3 100644
--- a/src/pybind/mgr/diskprediction_local/module.py
+++ b/src/pybind/mgr/diskprediction_local/module.py
@@ -15,6 +15,7 @@ from mgr_module import CommandResult, MgrModule, Option
 #
 # in a command thread.  See https://tracker.ceph.com/issues/42764
 import scipy # noqa: ignore=F401
+from .predictor import Predictor, get_diskfailurepredictor_path
 
 
 TIME_FORMAT = '%Y%m%d-%H%M%S'
@@ -136,24 +137,14 @@ class Module(MgrModule):
             self.log.error('failed to get device %s health data due to %s', devid, str(e))
 
         # initialize appropriate disk failure predictor model
-        from .predictor import get_diskfailurepredictor_path
-        if self.predictor_model == 'prophetstor':
-            from .predictor import PSDiskFailurePredictor
-            obj_predictor = PSDiskFailurePredictor()
-            ret = obj_predictor.initialize("{}/models/{}".format(get_diskfailurepredictor_path(), self.predictor_model))
-            if ret is not None:
-                self.log.error('Error initializing predictor')
-                return predicted_result
-        elif self.predictor_model == 'redhat':
-            from .predictor import RHDiskFailurePredictor
-            obj_predictor = RHDiskFailurePredictor()
-            ret = obj_predictor.initialize("{}/models/{}".format(get_diskfailurepredictor_path(), self.predictor_model))
-            if ret is not None:
-                self.log.error('Error initializing predictor')
-                return predicted_result
-        else:
+        obj_predictor = Predictor.create(self.predictor_model)
+        if obj_predictor is None:
             self.log.error('invalid value received for MODULE_OPTIONS.predictor_model')
             return predicted_result
+        ret = obj_predictor.initialize("{}/models/{}".format(get_diskfailurepredictor_path(), self.predictor_model))
+        if ret is not None:
+            self.log.error('Error initializing predictor')
+            return predicted_result
 
         if len(health_data) >= 6:
             o_keys = sorted(health_data.keys(), reverse=True)
diff --git a/src/pybind/mgr/diskprediction_local/predictor.py b/src/pybind/mgr/diskprediction_local/predictor.py
index fbea9dc94989..e2b7fc4e720c 100644
--- a/src/pybind/mgr/diskprediction_local/predictor.py
+++ b/src/pybind/mgr/diskprediction_local/predictor.py
@@ -35,7 +35,24 @@ def get_diskfailurepredictor_path():
     return dir_path
 
 
-class RHDiskFailurePredictor(object):
+class Predictor:
+    @classmethod
+    def create(cls, name: str) -> Optional['Predictor']:
+        if name == 'prophetstor':
+            return PSDiskFailurePredictor()
+        elif name == 'redhat':
+            return RHDiskFailurePredictor()
+        else:
+            return None
+
+    def initialize(self, model_dir: str) -> str:
+        raise NotImplementedError()
+
+    def predict(self, dataset: Sequence[DevSmartT]) -> str:
+        raise NotImplementedError()
+
+
+class RHDiskFailurePredictor(Predictor):
     """Disk failure prediction module developed at Red Hat
 
     This class implements a disk failure prediction module.
@@ -229,7 +246,7 @@ class RHDiskFailurePredictor(object):
         return RHDiskFailurePredictor.PREDICTION_CLASSES[pred_class_id]
 
 
-class PSDiskFailurePredictor(object):
+class PSDiskFailurePredictor(Predictor):
     """Disk failure prediction developed at ProphetStor
 
     This class implements a disk failure prediction module.