Model show drop

singlestoredb/fusion/handlers/models.py

@@ -4,10 +4,13 @@
 from typing import Dict
 from typing import Optional
 
+from .. import result
 from ..handler import SQLHandler
 from ..result import FusionSQLResult
 from .files import ShowFilesHandler
 from .utils import get_file_space
+from .utils import get_inference_api
+from .utils import get_inference_api_manager
 
 
 class ShowModelsHandler(ShowFilesHandler):
@@ -248,3 +251,192 @@ 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)
+        operation_result = inference_api.start()
+
+        res = FusionSQLResult()
+        res.add_field('Status', result.STRING)
+        res.add_field('Message', result.STRING)
+        res.set_rows([
+            (
+                operation_result.status,
+                operation_result.get_message(),
+            ),
+        ])
+
+        return res
+
+
+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)
+        operation_result = inference_api.stop()
+
+        res = FusionSQLResult()
+        res.add_field('Status', result.STRING)
+        res.add_field('Message', result.STRING)
+        res.set_rows([
+            (
+                operation_result.status,
+                operation_result.get_message(),
+            ),
+        ])
+
+        return res
+
+
+StopModelHandler.register(overwrite=True)
+
+
+class ShowInferenceAPIsHandler(SQLHandler):
+    """
+    SHOW INFERENCE APIS ;
+
+    Description
+    -----------
+    Displays the list of inference APIs in the current project.
+
+    Example
+    --------
+    The following command lists all inference APIs::
+
+        SHOW INFERENCE APIS;
+
+    See Also
+    --------
+    * ``START MODEL model_name``
+    * ``STOP MODEL model_name``
+    * ``DROP INFERENCE API model_name``
+
+    """  # noqa: E501
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        inference_api_manager = get_inference_api_manager()
+        models = inference_api_manager.show()
+
+        res = FusionSQLResult()
+        res.add_field('Model Name', result.STRING)
+        res.add_field('Status', result.STRING)
+
+        rows = []
+        for model in models:
+            rows.append((
+                model.name,
+                model.status,
+            ))
+
+        res.set_rows(rows)
+        return res
+
+
+ShowInferenceAPIsHandler.register(overwrite=True)
+
+
+class DropInferenceAPIHandler(SQLHandler):
+    """
+    DROP INFERENCE API model_name ;
+
+    # Model Name
+    model_name = '<model-name>'
+
+    Description
+    -----------
+    Drops (deletes) an inference API model.
+
+    Arguments
+    ---------
+    * ``<model-name>``: Name of the model to drop.
+
+    Example
+    --------
+    The following command drops an inference API::
+
+        DROP INFERENCE API my_model;
+
+    See Also
+    --------
+    * ``START MODEL model_name``
+    * ``STOP MODEL model_name``
+    * ``SHOW INFERENCE APIS``
+
+    """  # noqa: E501
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        inference_api = get_inference_api(params)
+        operation_result = inference_api.drop()
+
+        res = FusionSQLResult()
+        res.add_field('Status', result.STRING)
+        res.add_field('Message', result.STRING)
+        res.set_rows([
+            (
+                operation_result.status,
+                operation_result.get_message(),
+            ),
+        ])
+
+        return res
+
+
+DropInferenceAPIHandler.register(overwrite=True)

singlestoredb/fusion/handlers/utils.py

@@ -12,6 +12,8 @@
 from ...management.files import FilesManager
 from ...management.files import FileSpace
 from ...management.files import manage_files
+from ...management.inference_api import InferenceAPIInfo
+from ...management.inference_api import InferenceAPIManager
 from ...management.workspace import StarterWorkspace
 from ...management.workspace import Workspace
 from ...management.workspace import WorkspaceGroup
@@ -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)