fusion commands to start and stop model

Author: mohit-singlestore

singlestoredb/fusion/handlers/models.py

@@ -8,6 +8,7 @@
 from ..result import FusionSQLResult
 from .files import ShowFilesHandler
 from .utils import get_file_space
+from .utils import get_inference_api
 
 
 class ShowModelsHandler(ShowFilesHandler):
@@ -248,3 +249,77 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
 
 
 DropModelsHandler.register(overwrite=True)
+
+
+class StartModelHandler(SQLHandler):
+    """
+    START MODEL model_name ;
+
+    # Model Name
+    model_name = '<model-name>'
+
+    Description
+    -----------
+    Starts an inference API model.
+
+    Arguments
+    ---------
+    * ``<model-name>``: Name of the model to start.
+
+    Example
+    --------
+    The following command starts a model::
+
+        START MODEL my_model;
+
+    See Also
+    --------
+    * ``STOP MODEL model_name``
+    * ``SHOW MODELS``
+
+    """  # noqa: E501
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        inference_api = get_inference_api(params)
+        inference_api.start()
+        return None
+
+
+StartModelHandler.register(overwrite=True)
+
+
+class StopModelHandler(SQLHandler):
+    """
+    STOP MODEL model_name ;
+
+    # Model Name
+    model_name = '<model-name>'
+
+    Description
+    -----------
+    Stops an inference API model.
+
+    Arguments
+    ---------
+    * ``<model-name>``: Name of the model to stop.
+
+    Example
+    --------
+    The following command stops a model::
+
+        STOP MODEL my_model;
+
+    See Also
+    --------
+    * ``START MODEL model_name``
+    * ``SHOW MODELS``
+
+    """  # noqa: E501
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        inference_api = get_inference_api(params)
+        inference_api.stop()
+        return None
+
+
+StopModelHandler.register(overwrite=True)

singlestoredb/fusion/handlers/utils.py

@@ -16,6 +16,8 @@
 from ...management.workspace import Workspace
 from ...management.workspace import WorkspaceGroup
 from ...management.workspace import WorkspaceManager
+from ...management.inference_api import InferenceAPIInfo
+from ...management.inference_api import InferenceAPIManager
 
 
 def get_workspace_manager() -> WorkspaceManager:
@@ -322,3 +324,16 @@ def get_file_space(params: Dict[str, Any]) -> FileSpace:
             raise ValueError(f'invalid file location: {file_location}')
 
     raise KeyError('no file space was specified')
+
+
+def get_inference_api_manager() -> InferenceAPIManager:
+    """Return the inference API manager for the current project."""
+    wm = get_workspace_manager()
+    return wm.organization.inference_apis
+
+
+def get_inference_api(params: Dict[str, Any]) -> InferenceAPIInfo:
+    """Return an inference API based on model name in params."""
+    inference_apis = get_inference_api_manager()
+    model_name = params['model_name']
+    return inference_apis.get(model_name)

singlestoredb/management/inference_api.py

@@ -24,6 +24,7 @@ class InferenceAPIInfo(object):
     connection_url: str
     project_id: str
     hosting_platform: str
+    _manager: Optional['InferenceAPIManager']
 
     def __init__(
         self,
@@ -33,13 +34,15 @@ def __init__(
         connection_url: str,
         project_id: str,
         hosting_platform: str,
+        manager: Optional['InferenceAPIManager'] = None,
     ):
         self.service_id = service_id
         self.connection_url = connection_url
         self.model_name = model_name
         self.name = name
         self.project_id = project_id
         self.hosting_platform = hosting_platform
+        self._manager = manager
 
     @classmethod
     def from_dict(
@@ -77,6 +80,34 @@ def __repr__(self) -> str:
         """Return string representation."""
         return str(self)
 
+    def start(self) -> Dict[str, Any]:
+        """
+        Start this inference API model.
+
+        Returns
+        -------
+        dict
+            Response from the start operation
+
+        """
+        if self._manager is None:
+            raise ManagementError(msg='No manager associated with this inference API')
+        return self._manager.start(self.name)
+
+    def stop(self) -> Dict[str, Any]:
+        """
+        Stop this inference API model.
+
+        Returns
+        -------
+        dict
+            Response from the stop operation
+
+        """
+        if self._manager is None:
+            raise ManagementError(msg='No manager associated with this inference API')
+        return self._manager.stop(self.name)
+
 
 class InferenceAPIManager(object):
     """
@@ -102,4 +133,46 @@ def get(self, model_name: str) -> InferenceAPIInfo:
         if self._manager is None:
             raise ManagementError(msg='Manager not initialized')
         res = self._manager._get(f'inferenceapis/{self.project_id}/{model_name}').json()
-        return InferenceAPIInfo.from_dict(res)
+        inference_api = InferenceAPIInfo.from_dict(res)
+        inference_api._manager = self  # Associate the manager
+        return inference_api
+
+    def start(self, model_name: str) -> Dict[str, Any]:
+        """
+        Start an inference API model.
+
+        Parameters
+        ----------
+        model_name : str
+            Name of the model to start
+
+        Returns
+        -------
+        dict
+            Response from the start operation
+
+        """
+        if self._manager is None:
+            raise ManagementError(msg='Manager not initialized')
+        res = self._manager._post(f'inferenceapis/{self.project_id}/{model_name}/start')
+        return res.json()
+
+    def stop(self, model_name: str) -> Dict[str, Any]:
+        """
+        Stop an inference API model.
+
+        Parameters
+        ----------
+        model_name : str
+            Name of the model to stop
+
+        Returns
+        -------
+        dict
+            Response from the stop operation
+
+        """
+        if self._manager is None:
+            raise ManagementError(msg='Manager not initialized')
+        res = self._manager._post(f'inferenceapis/{self.project_id}/{model_name}/stop')
+        return res.json()