Skip to content
94 changes: 94 additions & 0 deletions singlestoredb/fusion/handlers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from .files import ShowFilesHandler
from .utils import get_file_space
from .utils import get_inference_api
from .utils import get_inference_api_manager


class ShowCustomModelsHandler(ShowFilesHandler):
Expand Down Expand Up @@ -346,3 +347,96 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:


StopModelHandler.register(overwrite=True)


class ShowModelsHandler(SQLHandler):
"""
SHOW MODELS ;

Description
-----------
Displays the list of inference APIs in the current project.

Example
--------
The following command lists all inference APIs::

SHOW MODELS;

See Also
--------
* ``START MODEL model_name``
* ``STOP MODEL model_name``
* ``DROP MODEL 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


ShowModelsHandler.register(overwrite=True)


class DropModelHandler(SQLHandler):
"""
DROP MODEL 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 MODEL my_model;

See Also
--------
* ``START MODEL model_name``
* ``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.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


DropModelHandler.register(overwrite=True)
92 changes: 92 additions & 0 deletions singlestoredb/management/inference_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import os
from typing import Any
from typing import Dict
from typing import List
from typing import Optional

from .utils import vars_to_str
Expand Down Expand Up @@ -76,6 +77,48 @@ def from_stop_response(cls, response: Dict[str, Any]) -> 'ModelOperationResult':
hosting_platform=response.get('hostingPlatform', ''),
)

@classmethod
def from_drop_response(cls, response: Dict[str, Any]) -> 'ModelOperationResult':
"""
Create a ModelOperationResult from a drop operation response.

Parameters
----------
response : dict
Response from the drop endpoint

Returns
-------
ModelOperationResult

"""
return cls(
name=response.get('name', ''),
status=response.get('status', 'Deleted'),
hosting_platform=response.get('hostingPlatform', ''),
)

@classmethod
def from_show_response(cls, response: Dict[str, Any]) -> 'ModelOperationResult':
"""
Create a ModelOperationResult from a show operation response.

Parameters
----------
response : dict
Response from the show endpoint (single model info)

Returns
-------
ModelOperationResult

"""
return cls(
name=response.get('name', ''),
status=response.get('status', ''),
hosting_platform=response.get('hostingPlatform', ''),
)

def get_message(self) -> str:
"""
Get a human-readable message about the operation.
Expand Down Expand Up @@ -195,6 +238,20 @@ def stop(self) -> ModelOperationResult:
raise ManagementError(msg='No manager associated with this inference API')
return self._manager.stop(self.name)

def drop(self) -> ModelOperationResult:
"""
Drop this inference API model.

Returns
-------
ModelOperationResult
Result object containing status information about the dropped model

"""
if self._manager is None:
raise ManagementError(msg='No manager associated with this inference API')
return self._manager.drop(self.name)


class InferenceAPIManager(object):
"""
Expand Down Expand Up @@ -263,3 +320,38 @@ def stop(self, model_name: str) -> ModelOperationResult:
raise ManagementError(msg='Manager not initialized')
res = self._manager._post(f'models/{model_name}/stop')
return ModelOperationResult.from_stop_response(res.json())

def show(self) -> List[ModelOperationResult]:
"""
Show all inference APIs in the project.

Returns
-------
List[ModelOperationResult]
List of ModelOperationResult objects with status information

"""
if self._manager is None:
raise ManagementError(msg='Manager not initialized')
res = self._manager._get('models').json()
return [ModelOperationResult.from_show_response(api) for api in res]

def drop(self, model_name: str) -> ModelOperationResult:
"""
Drop an inference API model.

Parameters
----------
model_name : str
Name of the model to drop

Returns
-------
ModelOperationResult
Result object containing status information about the dropped model

"""
if self._manager is None:
raise ManagementError(msg='Manager not initialized')
res = self._manager._delete(f'models/{model_name}')
return ModelOperationResult.from_drop_response(res.json())