diff --git a/singlestoredb/fusion/handlers/models.py b/singlestoredb/fusion/handlers/models.py index 5bb5cefa..8dc53ffc 100644 --- a/singlestoredb/fusion/handlers/models.py +++ b/singlestoredb/fusion/handlers/models.py @@ -4,10 +4,12 @@ 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 class ShowModelsHandler(ShowFilesHandler): @@ -248,3 +250,99 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: DropModelsHandler.register(overwrite=True) + + +class StartModelHandler(SQLHandler): + """ + START MODEL model_name ; + + # Model Name + model_name = '' + + Description + ----------- + Starts an inference API model. + + Arguments + --------- + * ````: 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 = '' + + Description + ----------- + Stops an inference API model. + + Arguments + --------- + * ````: 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) diff --git a/singlestoredb/fusion/handlers/utils.py b/singlestoredb/fusion/handlers/utils.py index 68af2e58..f5e82b03 100644 --- a/singlestoredb/fusion/handlers/utils.py +++ b/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) diff --git a/singlestoredb/management/inference_api.py b/singlestoredb/management/inference_api.py index 13b6413c..775ccf6b 100644 --- a/singlestoredb/management/inference_api.py +++ b/singlestoredb/management/inference_api.py @@ -10,6 +10,93 @@ from singlestoredb.management.manager import Manager +class ModelOperationResult(object): + """ + Result of a model start or stop operation. + + Attributes + ---------- + name : str + Name of the model + status : str + Current status of the model (e.g., 'Active', 'Initializing', 'Suspended') + hosting_platform : str + Hosting platform (e.g., 'Nova', 'Amazon', 'Azure') + """ + + def __init__( + self, + name: str, + status: str, + hosting_platform: str, + ): + self.name = name + self.status = status + self.hosting_platform = hosting_platform + + @classmethod + def from_start_response(cls, response: Dict[str, Any]) -> 'ModelOperationResult': + """ + Create a ModelOperationResult from a start operation response. + + Parameters + ---------- + response : dict + Response from the start endpoint + + Returns + ------- + ModelOperationResult + + """ + return cls( + name=response.get('modelName', ''), + status='Initializing', + hosting_platform=response.get('hostingPlatform', ''), + ) + + @classmethod + def from_stop_response(cls, response: Dict[str, Any]) -> 'ModelOperationResult': + """ + Create a ModelOperationResult from a stop operation response. + + Parameters + ---------- + response : dict + Response from the stop endpoint + + Returns + ------- + ModelOperationResult + + """ + return cls( + name=response.get('name', ''), + status=response.get('status', 'Suspended'), + hosting_platform=response.get('hostingPlatform', ''), + ) + + def get_message(self) -> str: + """ + Get a human-readable message about the operation. + + Returns + ------- + str + Message describing the operation result + + """ + return f'Model is {self.status}' + + def __str__(self) -> str: + """Return string representation.""" + return vars_to_str(self) + + def __repr__(self) -> str: + """Return string representation.""" + return str(self) + + class InferenceAPIInfo(object): """ Inference API definition. @@ -24,6 +111,7 @@ class InferenceAPIInfo(object): connection_url: str project_id: str hosting_platform: str + _manager: Optional['InferenceAPIManager'] def __init__( self, @@ -33,6 +121,7 @@ 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 @@ -40,6 +129,7 @@ def __init__( self.name = name self.project_id = project_id self.hosting_platform = hosting_platform + self._manager = manager @classmethod def from_dict( @@ -77,6 +167,34 @@ def __repr__(self) -> str: """Return string representation.""" return str(self) + def start(self) -> ModelOperationResult: + """ + Start this inference API model. + + Returns + ------- + ModelOperationResult + Result object containing status information about the started model + + """ + if self._manager is None: + raise ManagementError(msg='No manager associated with this inference API') + return self._manager.start(self.name) + + def stop(self) -> ModelOperationResult: + """ + Stop this inference API model. + + Returns + ------- + ModelOperationResult + Result object containing status information about the stopped model + + """ + 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 +220,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) -> ModelOperationResult: + """ + Start an inference API model. + + Parameters + ---------- + model_name : str + Name of the model to start + + Returns + ------- + ModelOperationResult + Result object containing status information about the started model + + """ + if self._manager is None: + raise ManagementError(msg='Manager not initialized') + res = self._manager._post(f'models/{model_name}/start') + return ModelOperationResult.from_start_response(res.json()) + + def stop(self, model_name: str) -> ModelOperationResult: + """ + Stop an inference API model. + + Parameters + ---------- + model_name : str + Name of the model to stop + + Returns + ------- + ModelOperationResult + Result object containing status information about the stopped model + + """ + if self._manager is None: + raise ManagementError(msg='Manager not initialized') + res = self._manager._post(f'models/{model_name}/stop') + return ModelOperationResult.from_stop_response(res.json())