enpi_api.l2.client.api.ml_api
1from typing import Literal 2 3from loguru import logger 4 5from enpi_api.l1 import openapi_client 6from enpi_api.l2.client.api.file_api import FileApi 7from enpi_api.l2.events.workflow_execution_task_waitable import WorkflowExecutionTaskWaitable 8from enpi_api.l2.events.workflow_execution_waitable import WorkflowExecutionWaitable 9from enpi_api.l2.types.api_error import ApiError, ApiErrorContext 10from enpi_api.l2.types.clone import CloneId 11from enpi_api.l2.types.collection import CollectionId 12from enpi_api.l2.types.execution import Execution 13from enpi_api.l2.types.file import File, FileId 14from enpi_api.l2.types.log import LogLevel 15from enpi_api.l2.types.ml import ( 16 MlAwsEndpointConfig, 17 MlEndpoint, 18 MlEndpointId, 19 MlEndpointSignature, 20 MlflowModelUri, 21 MlInput, 22 MlInvocationId, 23 MlInvocationIntentConfig, 24 MlInvocationStats, 25 MlOutputIntent, 26 MlParam, 27) 28from enpi_api.l2.types.task import TaskState 29from enpi_api.l2.types.workflow import WorkflowExecutionId, WorkflowExecutionTaskId, WorkflowTaskTemplateName 30 31 32class InvocationFailed(Exception): 33 """Indicates that the ML invocation has failed.""" 34 35 def __init__(self, invocation_id: MlInvocationId): 36 """@private""" 37 super().__init__(f"ML invocation with ID `{invocation_id}` failed") 38 39 40class DeploymentFailed(Exception): 41 """Indicates that the ML deployment has failed.""" 42 43 def __init__(self, workflow_execution_task_id: WorkflowExecutionTaskId): 44 """@private""" 45 super().__init__(f"ML deployment task with ID `{workflow_execution_task_id}` failed") 46 47 48class MlApi: 49 _inner_api_client: openapi_client.ApiClient 50 _log_level: LogLevel 51 52 def __init__(self, inner_api_client: openapi_client.ApiClient, log_level: LogLevel): 53 """@private""" 54 self._inner_api_client = inner_api_client 55 self._log_level = log_level 56 57 def get_ml_endpoints(self) -> list[MlEndpoint]: 58 """Get all ML endpoints. 59 60 Returns: 61 list[enpi_api.l2.types.ml.MlEndpoint]: A list of ML endpoints. 62 """ 63 64 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 65 try: 66 return [MlEndpoint.from_raw(i) for i in ml_api_instance.get_ml_endpoints().ml_endpoints] 67 except openapi_client.ApiException as e: 68 raise ApiError(e) 69 70 def get_ml_invocation_stats(self) -> list[MlInvocationStats]: 71 """Get ML invocation statistics. 72 73 Returns: 74 list[enpi_api.l2.types.ml.MlInvocationStats]: A list of ML invocation statistics. 75 """ 76 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 77 try: 78 stats = ml_api_instance.get_ml_invocation_stats().stats 79 return [MlInvocationStats.from_raw(i) for i in stats] 80 except openapi_client.ApiException as e: 81 raise ApiError(e) 82 83 def register_ml_endpoint( 84 self, 85 display_name: str, 86 input_mapping: list[MlInput], 87 output_intents: list[MlOutputIntent], 88 vendor_config: MlAwsEndpointConfig, 89 signatures: list[MlEndpointSignature], 90 parameter_mapping: list[MlParam] | None = None, 91 ) -> MlEndpointId: 92 """Register a ML endpoint. 93 94 Args: 95 display_name (str): The display name of the ML endpoint. 96 input_mapping (list[MlInput]): The input mapping of the ML endpoint. 97 output_intents (list[MlOutputIntent]): The output intents of the ML endpoint. 98 vendor_config (MlAwsEndpointConfig): The AWS endpoint configuration of the ML endpoint. 99 signatures (list[MlEndpointSignature]): The signatures of the ML endpoint. 100 parameter_mapping (list[MlParam] | None): The parameter mapping of the ML endpoint. 101 102 Returns: 103 endpoint_id (str): The unique identifier of a ML endpoint. 104 """ 105 106 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 107 try: 108 parameter_mapping_items = ( 109 [ 110 openapi_client.MlParameterMapItem.model_validate( 111 { 112 "type": pm["type"], 113 "input_key": pm["input_key"], 114 "label": pm["label"] if "label" in pm and pm["label"] is not None else None, 115 "source": "parameter", 116 } 117 ) 118 for pm in parameter_mapping 119 ] 120 if parameter_mapping is not None 121 else None 122 ) 123 124 vendor_config_to_register = openapi_client.MlAwsEndpointConfig.from_dict( 125 {**vendor_config, "region": vendor_config.get("region", "eu-west-1"), "endpoint_type": "external"} 126 ) 127 assert vendor_config_to_register is not None 128 129 result = ml_api_instance.register_ml_endpoint( 130 register_ml_endpoint_request=openapi_client.RegisterMlEndpointRequest( 131 display_name=display_name, 132 input_mapping=[openapi_client.MlInput.model_validate(i) for i in input_mapping], 133 parameter_mapping=parameter_mapping_items, 134 output_intents=[openapi_client.MlOutputIntent.from_dict(dict(i)) for i in output_intents], 135 vendor_config=vendor_config_to_register, 136 signatures=[openapi_client.MlEndpointSignature.model_validate(i) for i in signatures], 137 ) 138 ) 139 assert result.endpoint_id is not None 140 return MlEndpointId(result.endpoint_id) 141 except openapi_client.ApiException as e: 142 raise ApiError(e) 143 144 def unregister_ml_endpoint(self, endpoint_id: MlEndpointId) -> None: 145 """Unregister a ML endpoint. 146 147 Args: 148 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 149 """ 150 151 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 152 try: 153 ml_api_instance.unregister_ml_endpoint(id=endpoint_id, body={}) 154 except openapi_client.ApiException as e: 155 raise ApiError(e) 156 157 def deploy_model( 158 self, 159 display_name: str, 160 input_mapping: list[MlInput], 161 output_intents: list[MlOutputIntent], 162 model_uri: MlflowModelUri, 163 signatures: list[MlEndpointSignature], 164 parameter_mapping: list[MlParam] | None = None, 165 base_image: Literal["cpu"] | Literal["gpu"] = "cpu", 166 ) -> Execution[MlEndpointId]: 167 """Deploy a registered MLFlow model as an endpoint on the KServe inference platform. 168 169 Args: 170 display_name (str): The display name of the ML endpoint. 171 input_mapping (list[MlInput]): The input mapping configuration for the ML endpoint. 172 Each MlInput defines a mapping between a platform tag and the input to the ML model 173 174 Example: 175 If the inference function of the model has input arguments "heavy_chain" and "light_chain" that map to the heavy and light chain of the amino acids, 176 the input mapping would specified as follows: 177 ``` 178 input_mapping=[ 179 MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="heavy_chain", chains=[Chain.HEAVY]), 180 MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="light_chain", chains=[Chain.KAPPA, Chain.LAMBDA]), 181 ] 182 ``` 183 output_intents (list[MlOutputIntent]): The output intent is a contract between the ML model and the platform. Each intent specifies 184 a schema that the model output must adhere to. This ensures the platform can faithfully process and represent the model output in the platform UI and databases. 185 A deployed model can potentially support multiple output intents, provided the model output adheres to the schemas outline below. 186 The platform currently supports 6 output intents: 187 1. Metadata intent 188 This output intent can be used for a scalar model output that updates (or creates) a tag for a sequence, collection, or clone. 189 An example of a model with this output intent would a model that predicts sequence perplexity: 190 ``` 191 output_intents=[ 192 MlMetadataIntent(type="metadata", tag_id=output_tag_id, output_key="perplexity", chains=[]), 193 ] 194 ``` 195 For this output intent, the model output must adhere to the following schema: 196 ``` 197 {"perplexity": list[float]} 198 199 ``` 200 The list mapped to the output key is a list of the tag values produced by the model. 201 202 2. Structure intent 203 This output intent can handle outputs from structure prediction models. 204 ``` 205 output_intents=[ 206 MlStructureIntent( 207 type="structure", 208 output_key="my_structure", 209 )] 210 ``` 211 For this output intent, the model output must adhere to the following schema: 212 ``` 213 {"my_structure": list[str]} 214 215 ``` 216 The list mapped to the output key is a list of the PDB file contents (1 list item per structure produced) 217 218 3. Cluster intent 219 This output intent can handle outputs clustering models. 220 ``` 221 output_intents=[ 222 MlClusterIntent(type="cluster", cluster_label_key="cluster_labels"), 223 ] 224 225 ``` 226 For this output intent, the model output must adhere to the following schema: 227 ``` 228 {"cluster_labels": list[int]} 229 ``` 230 The list mapped to the cluster label keys is a list of cluster labels 231 232 4. New sequence intent 233 This output intent can handle outputs from model that produces sequences (e.g. a humanization model). 234 In addition to a key value pair typical of other intents, this intent requires the specification of a 235 sequence template and quality control template to handle the sequences produced by the models. 236 ``` 237 output_intents=[ 238 MlNewSequenceIntent( 239 type="new_sequence", 240 new_sequences_key="humanized_sequences", 241 quality_control_template_id=quality_control_template.id, 242 quality_control_template_version=quality_control_template.version, 243 sequence_template_id=sequence_template.id, 244 sequence_template_version=sequence_template.version, 245 ) 246 ] 247 248 ``` 249 For this output intent, the model output must adhere to the following schema: 250 ``` 251 {"humanized_sequences":list[list[str]]} 252 253 ``` 254 Here, the inner list contains the (potential) multiple output sequences generated for a single input sequence. 255 If the output sequences are paired chain, the inner list should adhere to the following form: 256 ``` 257 [HC1, LC1, ..., HCn, LCn] 258 ``` 259 260 5. File intent 261 This output intent can be used to produce "raw" output i.e. the output will be made available as a dataframe without 262 any further processing or downstream updates in the platform. This intent allows the user to customize the output processing. 263 ``` 264 output_intents=[ 265 MlMetadataIntent(type="file", filetype="json", output_key="ab_property"), 266 ] 267 ``` 268 For this output intent, the model output must adhere to the following schema: 269 ``` 270 {"ab_property": list[float | str]} 271 ``` 272 273 6. Sequence position/ Liability intent 274 This output intent can be used by models that produce per-residue outputs (e.g. saliency maps or liability predictors). 275 The supported 'value_types' in the output intent specification are "number" and "text" 276 ``` 277 output_intents = [ 278 MlLiabilityIntent( 279 type="liability", 280 category="humanness", 281 value_type="number", 282 output_key="humanness" 283 ) 284 ] 285 ``` 286 For this output intent, the model output must adhere to the following schema: 287 ``` 288 {"humanness": 289 [ 290 { 291 # these fields are required! 292 "chain": "Heavy" | "Light", 293 "exposed": bool, 294 "length": int, 295 "position": int, 296 "value": float, 297 } 298 ] 299 } 300 301 ``` 302 303 304 model_uri (MlflowModelUri): The URI of a registered MLflow model in the format 305 'models:/model_name/model_version' 306 Example: 'models:/antibody_binding_model/1' 307 308 signatures (list[MlEndpointSignature]): The signature definitions for the ML endpoint. ** This argument will be deprecated. ** 309 Example: 310 ``` 311 signatures=[MlEndpointSignature(key="input", type="object", kind="input")], 312 313 ``` 314 315 parameter_mapping (list[MlParam] | None, optional): Specify runtime parameters for the ML model 316 317 Example: 318 Assume we have a model with the following runtime parameters: 319 1. Number of variants (int) 320 2. Number of diffusion steps (int) 321 3. Regions to mask (array[enum]) 322 ``` 323 parameter_mapping=[ 324 MlBaseParam(type="integer", input_key="num_samples", label="Number Of Variants per clone"), 325 MlBaseParam(type="integer", input_key="num_steps", label="Number of diffusion steps"), 326 MlEnumParam( 327 type="enum_array", 328 input_key="regions_to_mask", 329 label="Regions to mutate", 330 enum_options=[region.value for region in [Region.FR1, Region.FR2, Region.FR3, Region.FR4, Region.CDR1, Region.CDR2, Region.CDR3]], 331 ), 332 ] 333 334 ``` 335 The input_key must be indentical to the name used in the signature of the prediction/inference function of the model. 336 base_image (Literal["cpu"] | Literal["gpu"], optional): The base container image to use 337 for model deployment. Use "gpu" for models requiring GPU acceleration. Defaults to "cpu". 338 339 Returns: 340 MlEndpointId: Endpoint ID of the deployed model 341 342 Raises: 343 DeploymentFailed: If the model deployment task fails. 344 """ 345 346 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 347 parameter_mapping_items = ( 348 [ 349 openapi_client.MlParameterMapItem.model_validate( 350 { 351 "type": pm["type"], 352 "input_key": pm["input_key"], 353 "label": pm["label"] if "label" in pm and pm["label"] is not None else None, 354 "source": "parameter", 355 "enum_options": pm.get("enum_options"), 356 } 357 ) 358 for pm in parameter_mapping 359 ] 360 if parameter_mapping is not None 361 else None 362 ) 363 364 payload = openapi_client.DeployModelRequest( 365 display_name=display_name, 366 parameter_mapping=parameter_mapping_items, 367 input_mapping=[openapi_client.MlInput.model_validate(i) for i in input_mapping], 368 output_intents=[openapi_client.MlOutputIntent.from_dict(dict(i)) for i in output_intents], 369 model_uri=model_uri, 370 signatures=[openapi_client.MlEndpointSignature.model_validate(i) for i in signatures], 371 base_image=base_image, 372 ) 373 374 with ApiErrorContext(): 375 deploy_model_response = ml_api_instance.deploy_model(deploy_model_request=payload) 376 assert deploy_model_response.workflow_execution_id is not None 377 378 workflow_execution_id = WorkflowExecutionId(deploy_model_response.workflow_execution_id) 379 380 def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> MlEndpointId: 381 # If the task has succeeded, return the endpoint_id 382 match task_state: 383 case TaskState.SUCCEEDED: 384 result = ml_api_instance.get_endpoint_by_workflow_execution_task_id(workflow_execution_task_id=task_id) 385 assert result.endpoint_id is not None 386 return MlEndpointId(result.endpoint_id) 387 case _: 388 raise DeploymentFailed(task_id) 389 390 waitable = WorkflowExecutionTaskWaitable[MlEndpointId]( 391 workflow_execution_id=workflow_execution_id, 392 task_template_name=WorkflowTaskTemplateName.ENPI_APP_ML_MONITOR_INFERENCE_SERVICE, 393 on_complete=on_complete, 394 ) 395 396 return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state) 397 398 def undeploy_model(self, endpoint_id: MlEndpointId) -> None: 399 """Undeploys the model and delete the endpoint 400 401 Args: 402 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 403 """ 404 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 405 try: 406 ml_api_instance.undeploy_model(id=endpoint_id, body={}) 407 except openapi_client.ApiException as e: 408 raise ApiError(e) 409 410 def invoke_endpoint( 411 self, 412 endpoint_id: MlEndpointId, 413 clone_ids: list[CloneId] | None = None, 414 collection_ids: list[CollectionId] | None = None, 415 parameters: dict[str, str | int | float | bool | list[str]] | None = None, 416 intent_config: MlInvocationIntentConfig | None = None, 417 ) -> Execution[list[File]]: 418 """Invoke a ML endpoint for a set of clones or collections. 419 420 Args: 421 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 422 clone_ids (list[CloneId]): The unique identifiers of the clones. 423 collection_ids (list[CollectionIds]): The unique identifiers of the collections 424 parameters (dict | None): parameters passed to the invocation. 425 intent_config (MlInvocationIntentConfig | None): This argument is only used for invoking models with a new sequence output intent. 426 It is used to specify the name of the output collection under which all the output seqences are stored. 427 Example: 428 ``` 429 intent_config=MlInvocationIntentConfig(new_sequence=MlInvocationNewSequenceIntentConfig(output_collection_name="humanized_sequences")) 430 ``` 431 432 Returns: 433 output_files (list[File]): The list of output files generated by the File Intent of ML model. 434 """ 435 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 436 file_api = FileApi(self._inner_api_client, self._log_level) 437 if clone_ids is not None and collection_ids is not None: 438 raise ValueError("Either clone_ids or collection_ids must be provided, but not both") 439 440 with ApiErrorContext(): 441 invoke_ml_endpoint_request = openapi_client.InvokeMlEndpointRequest.from_dict( 442 dict( 443 clone_ids=clone_ids, 444 collection_ids=collection_ids, 445 parameters=parameters, 446 intent_config=intent_config.model_dump() if intent_config else None, 447 ) 448 ) 449 assert invoke_ml_endpoint_request is not None 450 result = ml_api_instance.invoke_endpoint(id=endpoint_id, invoke_ml_endpoint_request=invoke_ml_endpoint_request) 451 assert result.workflow_execution_id is not None 452 453 workflow_execution_id = WorkflowExecutionId(result.workflow_execution_id) 454 455 def on_complete(workflow_execution_id: WorkflowExecutionId, execution_state: TaskState) -> list[File]: 456 assert execution_state == TaskState.SUCCEEDED, ( 457 f"Workflow execution {workflow_execution_id} did not reach {TaskState.SUCCEEDED} state, got {execution_state} state instead" 458 ) 459 460 logger.success(f"Ml invocation with workflow execution ID: {workflow_execution_id} has successfully finished.") 461 462 # Get potential file intent output files 463 file_intent_output_ids = ml_api_instance.get_invocation_output(workflow_execution_id).file_ids 464 if len(file_intent_output_ids) > 0: 465 logger.success("Ml invocation File Intent output files retrieved successfully") 466 return [file_api.get_file_by_id(FileId(file_id)) for file_id in file_intent_output_ids] 467 else: 468 return [] 469 470 waitable = WorkflowExecutionWaitable( 471 workflow_execution_id=workflow_execution_id, 472 on_complete=on_complete, 473 ) 474 return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
class
InvocationFailed(builtins.Exception):
33class InvocationFailed(Exception): 34 """Indicates that the ML invocation has failed.""" 35 36 def __init__(self, invocation_id: MlInvocationId): 37 """@private""" 38 super().__init__(f"ML invocation with ID `{invocation_id}` failed")
Indicates that the ML invocation has failed.
class
DeploymentFailed(builtins.Exception):
41class DeploymentFailed(Exception): 42 """Indicates that the ML deployment has failed.""" 43 44 def __init__(self, workflow_execution_task_id: WorkflowExecutionTaskId): 45 """@private""" 46 super().__init__(f"ML deployment task with ID `{workflow_execution_task_id}` failed")
Indicates that the ML deployment has failed.
class
MlApi:
49class MlApi: 50 _inner_api_client: openapi_client.ApiClient 51 _log_level: LogLevel 52 53 def __init__(self, inner_api_client: openapi_client.ApiClient, log_level: LogLevel): 54 """@private""" 55 self._inner_api_client = inner_api_client 56 self._log_level = log_level 57 58 def get_ml_endpoints(self) -> list[MlEndpoint]: 59 """Get all ML endpoints. 60 61 Returns: 62 list[enpi_api.l2.types.ml.MlEndpoint]: A list of ML endpoints. 63 """ 64 65 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 66 try: 67 return [MlEndpoint.from_raw(i) for i in ml_api_instance.get_ml_endpoints().ml_endpoints] 68 except openapi_client.ApiException as e: 69 raise ApiError(e) 70 71 def get_ml_invocation_stats(self) -> list[MlInvocationStats]: 72 """Get ML invocation statistics. 73 74 Returns: 75 list[enpi_api.l2.types.ml.MlInvocationStats]: A list of ML invocation statistics. 76 """ 77 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 78 try: 79 stats = ml_api_instance.get_ml_invocation_stats().stats 80 return [MlInvocationStats.from_raw(i) for i in stats] 81 except openapi_client.ApiException as e: 82 raise ApiError(e) 83 84 def register_ml_endpoint( 85 self, 86 display_name: str, 87 input_mapping: list[MlInput], 88 output_intents: list[MlOutputIntent], 89 vendor_config: MlAwsEndpointConfig, 90 signatures: list[MlEndpointSignature], 91 parameter_mapping: list[MlParam] | None = None, 92 ) -> MlEndpointId: 93 """Register a ML endpoint. 94 95 Args: 96 display_name (str): The display name of the ML endpoint. 97 input_mapping (list[MlInput]): The input mapping of the ML endpoint. 98 output_intents (list[MlOutputIntent]): The output intents of the ML endpoint. 99 vendor_config (MlAwsEndpointConfig): The AWS endpoint configuration of the ML endpoint. 100 signatures (list[MlEndpointSignature]): The signatures of the ML endpoint. 101 parameter_mapping (list[MlParam] | None): The parameter mapping of the ML endpoint. 102 103 Returns: 104 endpoint_id (str): The unique identifier of a ML endpoint. 105 """ 106 107 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 108 try: 109 parameter_mapping_items = ( 110 [ 111 openapi_client.MlParameterMapItem.model_validate( 112 { 113 "type": pm["type"], 114 "input_key": pm["input_key"], 115 "label": pm["label"] if "label" in pm and pm["label"] is not None else None, 116 "source": "parameter", 117 } 118 ) 119 for pm in parameter_mapping 120 ] 121 if parameter_mapping is not None 122 else None 123 ) 124 125 vendor_config_to_register = openapi_client.MlAwsEndpointConfig.from_dict( 126 {**vendor_config, "region": vendor_config.get("region", "eu-west-1"), "endpoint_type": "external"} 127 ) 128 assert vendor_config_to_register is not None 129 130 result = ml_api_instance.register_ml_endpoint( 131 register_ml_endpoint_request=openapi_client.RegisterMlEndpointRequest( 132 display_name=display_name, 133 input_mapping=[openapi_client.MlInput.model_validate(i) for i in input_mapping], 134 parameter_mapping=parameter_mapping_items, 135 output_intents=[openapi_client.MlOutputIntent.from_dict(dict(i)) for i in output_intents], 136 vendor_config=vendor_config_to_register, 137 signatures=[openapi_client.MlEndpointSignature.model_validate(i) for i in signatures], 138 ) 139 ) 140 assert result.endpoint_id is not None 141 return MlEndpointId(result.endpoint_id) 142 except openapi_client.ApiException as e: 143 raise ApiError(e) 144 145 def unregister_ml_endpoint(self, endpoint_id: MlEndpointId) -> None: 146 """Unregister a ML endpoint. 147 148 Args: 149 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 150 """ 151 152 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 153 try: 154 ml_api_instance.unregister_ml_endpoint(id=endpoint_id, body={}) 155 except openapi_client.ApiException as e: 156 raise ApiError(e) 157 158 def deploy_model( 159 self, 160 display_name: str, 161 input_mapping: list[MlInput], 162 output_intents: list[MlOutputIntent], 163 model_uri: MlflowModelUri, 164 signatures: list[MlEndpointSignature], 165 parameter_mapping: list[MlParam] | None = None, 166 base_image: Literal["cpu"] | Literal["gpu"] = "cpu", 167 ) -> Execution[MlEndpointId]: 168 """Deploy a registered MLFlow model as an endpoint on the KServe inference platform. 169 170 Args: 171 display_name (str): The display name of the ML endpoint. 172 input_mapping (list[MlInput]): The input mapping configuration for the ML endpoint. 173 Each MlInput defines a mapping between a platform tag and the input to the ML model 174 175 Example: 176 If the inference function of the model has input arguments "heavy_chain" and "light_chain" that map to the heavy and light chain of the amino acids, 177 the input mapping would specified as follows: 178 ``` 179 input_mapping=[ 180 MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="heavy_chain", chains=[Chain.HEAVY]), 181 MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="light_chain", chains=[Chain.KAPPA, Chain.LAMBDA]), 182 ] 183 ``` 184 output_intents (list[MlOutputIntent]): The output intent is a contract between the ML model and the platform. Each intent specifies 185 a schema that the model output must adhere to. This ensures the platform can faithfully process and represent the model output in the platform UI and databases. 186 A deployed model can potentially support multiple output intents, provided the model output adheres to the schemas outline below. 187 The platform currently supports 6 output intents: 188 1. Metadata intent 189 This output intent can be used for a scalar model output that updates (or creates) a tag for a sequence, collection, or clone. 190 An example of a model with this output intent would a model that predicts sequence perplexity: 191 ``` 192 output_intents=[ 193 MlMetadataIntent(type="metadata", tag_id=output_tag_id, output_key="perplexity", chains=[]), 194 ] 195 ``` 196 For this output intent, the model output must adhere to the following schema: 197 ``` 198 {"perplexity": list[float]} 199 200 ``` 201 The list mapped to the output key is a list of the tag values produced by the model. 202 203 2. Structure intent 204 This output intent can handle outputs from structure prediction models. 205 ``` 206 output_intents=[ 207 MlStructureIntent( 208 type="structure", 209 output_key="my_structure", 210 )] 211 ``` 212 For this output intent, the model output must adhere to the following schema: 213 ``` 214 {"my_structure": list[str]} 215 216 ``` 217 The list mapped to the output key is a list of the PDB file contents (1 list item per structure produced) 218 219 3. Cluster intent 220 This output intent can handle outputs clustering models. 221 ``` 222 output_intents=[ 223 MlClusterIntent(type="cluster", cluster_label_key="cluster_labels"), 224 ] 225 226 ``` 227 For this output intent, the model output must adhere to the following schema: 228 ``` 229 {"cluster_labels": list[int]} 230 ``` 231 The list mapped to the cluster label keys is a list of cluster labels 232 233 4. New sequence intent 234 This output intent can handle outputs from model that produces sequences (e.g. a humanization model). 235 In addition to a key value pair typical of other intents, this intent requires the specification of a 236 sequence template and quality control template to handle the sequences produced by the models. 237 ``` 238 output_intents=[ 239 MlNewSequenceIntent( 240 type="new_sequence", 241 new_sequences_key="humanized_sequences", 242 quality_control_template_id=quality_control_template.id, 243 quality_control_template_version=quality_control_template.version, 244 sequence_template_id=sequence_template.id, 245 sequence_template_version=sequence_template.version, 246 ) 247 ] 248 249 ``` 250 For this output intent, the model output must adhere to the following schema: 251 ``` 252 {"humanized_sequences":list[list[str]]} 253 254 ``` 255 Here, the inner list contains the (potential) multiple output sequences generated for a single input sequence. 256 If the output sequences are paired chain, the inner list should adhere to the following form: 257 ``` 258 [HC1, LC1, ..., HCn, LCn] 259 ``` 260 261 5. File intent 262 This output intent can be used to produce "raw" output i.e. the output will be made available as a dataframe without 263 any further processing or downstream updates in the platform. This intent allows the user to customize the output processing. 264 ``` 265 output_intents=[ 266 MlMetadataIntent(type="file", filetype="json", output_key="ab_property"), 267 ] 268 ``` 269 For this output intent, the model output must adhere to the following schema: 270 ``` 271 {"ab_property": list[float | str]} 272 ``` 273 274 6. Sequence position/ Liability intent 275 This output intent can be used by models that produce per-residue outputs (e.g. saliency maps or liability predictors). 276 The supported 'value_types' in the output intent specification are "number" and "text" 277 ``` 278 output_intents = [ 279 MlLiabilityIntent( 280 type="liability", 281 category="humanness", 282 value_type="number", 283 output_key="humanness" 284 ) 285 ] 286 ``` 287 For this output intent, the model output must adhere to the following schema: 288 ``` 289 {"humanness": 290 [ 291 { 292 # these fields are required! 293 "chain": "Heavy" | "Light", 294 "exposed": bool, 295 "length": int, 296 "position": int, 297 "value": float, 298 } 299 ] 300 } 301 302 ``` 303 304 305 model_uri (MlflowModelUri): The URI of a registered MLflow model in the format 306 'models:/model_name/model_version' 307 Example: 'models:/antibody_binding_model/1' 308 309 signatures (list[MlEndpointSignature]): The signature definitions for the ML endpoint. ** This argument will be deprecated. ** 310 Example: 311 ``` 312 signatures=[MlEndpointSignature(key="input", type="object", kind="input")], 313 314 ``` 315 316 parameter_mapping (list[MlParam] | None, optional): Specify runtime parameters for the ML model 317 318 Example: 319 Assume we have a model with the following runtime parameters: 320 1. Number of variants (int) 321 2. Number of diffusion steps (int) 322 3. Regions to mask (array[enum]) 323 ``` 324 parameter_mapping=[ 325 MlBaseParam(type="integer", input_key="num_samples", label="Number Of Variants per clone"), 326 MlBaseParam(type="integer", input_key="num_steps", label="Number of diffusion steps"), 327 MlEnumParam( 328 type="enum_array", 329 input_key="regions_to_mask", 330 label="Regions to mutate", 331 enum_options=[region.value for region in [Region.FR1, Region.FR2, Region.FR3, Region.FR4, Region.CDR1, Region.CDR2, Region.CDR3]], 332 ), 333 ] 334 335 ``` 336 The input_key must be indentical to the name used in the signature of the prediction/inference function of the model. 337 base_image (Literal["cpu"] | Literal["gpu"], optional): The base container image to use 338 for model deployment. Use "gpu" for models requiring GPU acceleration. Defaults to "cpu". 339 340 Returns: 341 MlEndpointId: Endpoint ID of the deployed model 342 343 Raises: 344 DeploymentFailed: If the model deployment task fails. 345 """ 346 347 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 348 parameter_mapping_items = ( 349 [ 350 openapi_client.MlParameterMapItem.model_validate( 351 { 352 "type": pm["type"], 353 "input_key": pm["input_key"], 354 "label": pm["label"] if "label" in pm and pm["label"] is not None else None, 355 "source": "parameter", 356 "enum_options": pm.get("enum_options"), 357 } 358 ) 359 for pm in parameter_mapping 360 ] 361 if parameter_mapping is not None 362 else None 363 ) 364 365 payload = openapi_client.DeployModelRequest( 366 display_name=display_name, 367 parameter_mapping=parameter_mapping_items, 368 input_mapping=[openapi_client.MlInput.model_validate(i) for i in input_mapping], 369 output_intents=[openapi_client.MlOutputIntent.from_dict(dict(i)) for i in output_intents], 370 model_uri=model_uri, 371 signatures=[openapi_client.MlEndpointSignature.model_validate(i) for i in signatures], 372 base_image=base_image, 373 ) 374 375 with ApiErrorContext(): 376 deploy_model_response = ml_api_instance.deploy_model(deploy_model_request=payload) 377 assert deploy_model_response.workflow_execution_id is not None 378 379 workflow_execution_id = WorkflowExecutionId(deploy_model_response.workflow_execution_id) 380 381 def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> MlEndpointId: 382 # If the task has succeeded, return the endpoint_id 383 match task_state: 384 case TaskState.SUCCEEDED: 385 result = ml_api_instance.get_endpoint_by_workflow_execution_task_id(workflow_execution_task_id=task_id) 386 assert result.endpoint_id is not None 387 return MlEndpointId(result.endpoint_id) 388 case _: 389 raise DeploymentFailed(task_id) 390 391 waitable = WorkflowExecutionTaskWaitable[MlEndpointId]( 392 workflow_execution_id=workflow_execution_id, 393 task_template_name=WorkflowTaskTemplateName.ENPI_APP_ML_MONITOR_INFERENCE_SERVICE, 394 on_complete=on_complete, 395 ) 396 397 return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state) 398 399 def undeploy_model(self, endpoint_id: MlEndpointId) -> None: 400 """Undeploys the model and delete the endpoint 401 402 Args: 403 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 404 """ 405 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 406 try: 407 ml_api_instance.undeploy_model(id=endpoint_id, body={}) 408 except openapi_client.ApiException as e: 409 raise ApiError(e) 410 411 def invoke_endpoint( 412 self, 413 endpoint_id: MlEndpointId, 414 clone_ids: list[CloneId] | None = None, 415 collection_ids: list[CollectionId] | None = None, 416 parameters: dict[str, str | int | float | bool | list[str]] | None = None, 417 intent_config: MlInvocationIntentConfig | None = None, 418 ) -> Execution[list[File]]: 419 """Invoke a ML endpoint for a set of clones or collections. 420 421 Args: 422 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 423 clone_ids (list[CloneId]): The unique identifiers of the clones. 424 collection_ids (list[CollectionIds]): The unique identifiers of the collections 425 parameters (dict | None): parameters passed to the invocation. 426 intent_config (MlInvocationIntentConfig | None): This argument is only used for invoking models with a new sequence output intent. 427 It is used to specify the name of the output collection under which all the output seqences are stored. 428 Example: 429 ``` 430 intent_config=MlInvocationIntentConfig(new_sequence=MlInvocationNewSequenceIntentConfig(output_collection_name="humanized_sequences")) 431 ``` 432 433 Returns: 434 output_files (list[File]): The list of output files generated by the File Intent of ML model. 435 """ 436 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 437 file_api = FileApi(self._inner_api_client, self._log_level) 438 if clone_ids is not None and collection_ids is not None: 439 raise ValueError("Either clone_ids or collection_ids must be provided, but not both") 440 441 with ApiErrorContext(): 442 invoke_ml_endpoint_request = openapi_client.InvokeMlEndpointRequest.from_dict( 443 dict( 444 clone_ids=clone_ids, 445 collection_ids=collection_ids, 446 parameters=parameters, 447 intent_config=intent_config.model_dump() if intent_config else None, 448 ) 449 ) 450 assert invoke_ml_endpoint_request is not None 451 result = ml_api_instance.invoke_endpoint(id=endpoint_id, invoke_ml_endpoint_request=invoke_ml_endpoint_request) 452 assert result.workflow_execution_id is not None 453 454 workflow_execution_id = WorkflowExecutionId(result.workflow_execution_id) 455 456 def on_complete(workflow_execution_id: WorkflowExecutionId, execution_state: TaskState) -> list[File]: 457 assert execution_state == TaskState.SUCCEEDED, ( 458 f"Workflow execution {workflow_execution_id} did not reach {TaskState.SUCCEEDED} state, got {execution_state} state instead" 459 ) 460 461 logger.success(f"Ml invocation with workflow execution ID: {workflow_execution_id} has successfully finished.") 462 463 # Get potential file intent output files 464 file_intent_output_ids = ml_api_instance.get_invocation_output(workflow_execution_id).file_ids 465 if len(file_intent_output_ids) > 0: 466 logger.success("Ml invocation File Intent output files retrieved successfully") 467 return [file_api.get_file_by_id(FileId(file_id)) for file_id in file_intent_output_ids] 468 else: 469 return [] 470 471 waitable = WorkflowExecutionWaitable( 472 workflow_execution_id=workflow_execution_id, 473 on_complete=on_complete, 474 ) 475 return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
58 def get_ml_endpoints(self) -> list[MlEndpoint]: 59 """Get all ML endpoints. 60 61 Returns: 62 list[enpi_api.l2.types.ml.MlEndpoint]: A list of ML endpoints. 63 """ 64 65 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 66 try: 67 return [MlEndpoint.from_raw(i) for i in ml_api_instance.get_ml_endpoints().ml_endpoints] 68 except openapi_client.ApiException as e: 69 raise ApiError(e)
71 def get_ml_invocation_stats(self) -> list[MlInvocationStats]: 72 """Get ML invocation statistics. 73 74 Returns: 75 list[enpi_api.l2.types.ml.MlInvocationStats]: A list of ML invocation statistics. 76 """ 77 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 78 try: 79 stats = ml_api_instance.get_ml_invocation_stats().stats 80 return [MlInvocationStats.from_raw(i) for i in stats] 81 except openapi_client.ApiException as e: 82 raise ApiError(e)
Get ML invocation statistics.
Returns:
list[enpi_api.l2.types.ml.MlInvocationStats]: A list of ML invocation statistics.
def
register_ml_endpoint( self, display_name: str, input_mapping: list[enpi_api.l2.types.ml.MlInput], output_intents: list[enpi_api.l2.types.ml.MlClusterIntent | enpi_api.l2.types.ml.MlFileIntent | enpi_api.l2.types.ml.MlLiabilityIntent | enpi_api.l2.types.ml.MlMetadataIntent | enpi_api.l2.types.ml.MlNewSequenceIntent | enpi_api.l2.types.ml.MlStructureIntent], vendor_config: enpi_api.l2.types.ml.MlAwsEndpointConfig, signatures: list[enpi_api.l2.types.ml.MlEndpointSignature], parameter_mapping: list[enpi_api.l2.types.ml.MlBaseParam | enpi_api.l2.types.ml.MlEnumParam] | None = None) -> enpi_api.l2.types.ml.MlEndpointId:
84 def register_ml_endpoint( 85 self, 86 display_name: str, 87 input_mapping: list[MlInput], 88 output_intents: list[MlOutputIntent], 89 vendor_config: MlAwsEndpointConfig, 90 signatures: list[MlEndpointSignature], 91 parameter_mapping: list[MlParam] | None = None, 92 ) -> MlEndpointId: 93 """Register a ML endpoint. 94 95 Args: 96 display_name (str): The display name of the ML endpoint. 97 input_mapping (list[MlInput]): The input mapping of the ML endpoint. 98 output_intents (list[MlOutputIntent]): The output intents of the ML endpoint. 99 vendor_config (MlAwsEndpointConfig): The AWS endpoint configuration of the ML endpoint. 100 signatures (list[MlEndpointSignature]): The signatures of the ML endpoint. 101 parameter_mapping (list[MlParam] | None): The parameter mapping of the ML endpoint. 102 103 Returns: 104 endpoint_id (str): The unique identifier of a ML endpoint. 105 """ 106 107 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 108 try: 109 parameter_mapping_items = ( 110 [ 111 openapi_client.MlParameterMapItem.model_validate( 112 { 113 "type": pm["type"], 114 "input_key": pm["input_key"], 115 "label": pm["label"] if "label" in pm and pm["label"] is not None else None, 116 "source": "parameter", 117 } 118 ) 119 for pm in parameter_mapping 120 ] 121 if parameter_mapping is not None 122 else None 123 ) 124 125 vendor_config_to_register = openapi_client.MlAwsEndpointConfig.from_dict( 126 {**vendor_config, "region": vendor_config.get("region", "eu-west-1"), "endpoint_type": "external"} 127 ) 128 assert vendor_config_to_register is not None 129 130 result = ml_api_instance.register_ml_endpoint( 131 register_ml_endpoint_request=openapi_client.RegisterMlEndpointRequest( 132 display_name=display_name, 133 input_mapping=[openapi_client.MlInput.model_validate(i) for i in input_mapping], 134 parameter_mapping=parameter_mapping_items, 135 output_intents=[openapi_client.MlOutputIntent.from_dict(dict(i)) for i in output_intents], 136 vendor_config=vendor_config_to_register, 137 signatures=[openapi_client.MlEndpointSignature.model_validate(i) for i in signatures], 138 ) 139 ) 140 assert result.endpoint_id is not None 141 return MlEndpointId(result.endpoint_id) 142 except openapi_client.ApiException as e: 143 raise ApiError(e)
Register a ML endpoint.
Arguments:
- display_name (str): The display name of the ML endpoint.
- input_mapping (list[MlInput]): The input mapping of the ML endpoint.
- output_intents (list[MlOutputIntent]): The output intents of the ML endpoint.
- vendor_config (MlAwsEndpointConfig): The AWS endpoint configuration of the ML endpoint.
- signatures (list[MlEndpointSignature]): The signatures of the ML endpoint.
- parameter_mapping (list[MlParam] | None): The parameter mapping of the ML endpoint.
Returns:
endpoint_id (str): The unique identifier of a ML endpoint.
145 def unregister_ml_endpoint(self, endpoint_id: MlEndpointId) -> None: 146 """Unregister a ML endpoint. 147 148 Args: 149 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 150 """ 151 152 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 153 try: 154 ml_api_instance.unregister_ml_endpoint(id=endpoint_id, body={}) 155 except openapi_client.ApiException as e: 156 raise ApiError(e)
Unregister a ML endpoint.
Arguments:
- endpoint_id (MlEndpointId): The unique identifier of a ML endpoint.
def
deploy_model( self, display_name: str, input_mapping: list[enpi_api.l2.types.ml.MlInput], output_intents: list[enpi_api.l2.types.ml.MlClusterIntent | enpi_api.l2.types.ml.MlFileIntent | enpi_api.l2.types.ml.MlLiabilityIntent | enpi_api.l2.types.ml.MlMetadataIntent | enpi_api.l2.types.ml.MlNewSequenceIntent | enpi_api.l2.types.ml.MlStructureIntent], model_uri: enpi_api.l2.types.ml.MlflowModelUri, signatures: list[enpi_api.l2.types.ml.MlEndpointSignature], parameter_mapping: list[enpi_api.l2.types.ml.MlBaseParam | enpi_api.l2.types.ml.MlEnumParam] | None = None, base_image: Union[Literal['cpu'], Literal['gpu']] = 'cpu') -> enpi_api.l2.types.execution.Execution[NewType]:
158 def deploy_model( 159 self, 160 display_name: str, 161 input_mapping: list[MlInput], 162 output_intents: list[MlOutputIntent], 163 model_uri: MlflowModelUri, 164 signatures: list[MlEndpointSignature], 165 parameter_mapping: list[MlParam] | None = None, 166 base_image: Literal["cpu"] | Literal["gpu"] = "cpu", 167 ) -> Execution[MlEndpointId]: 168 """Deploy a registered MLFlow model as an endpoint on the KServe inference platform. 169 170 Args: 171 display_name (str): The display name of the ML endpoint. 172 input_mapping (list[MlInput]): The input mapping configuration for the ML endpoint. 173 Each MlInput defines a mapping between a platform tag and the input to the ML model 174 175 Example: 176 If the inference function of the model has input arguments "heavy_chain" and "light_chain" that map to the heavy and light chain of the amino acids, 177 the input mapping would specified as follows: 178 ``` 179 input_mapping=[ 180 MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="heavy_chain", chains=[Chain.HEAVY]), 181 MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="light_chain", chains=[Chain.KAPPA, Chain.LAMBDA]), 182 ] 183 ``` 184 output_intents (list[MlOutputIntent]): The output intent is a contract between the ML model and the platform. Each intent specifies 185 a schema that the model output must adhere to. This ensures the platform can faithfully process and represent the model output in the platform UI and databases. 186 A deployed model can potentially support multiple output intents, provided the model output adheres to the schemas outline below. 187 The platform currently supports 6 output intents: 188 1. Metadata intent 189 This output intent can be used for a scalar model output that updates (or creates) a tag for a sequence, collection, or clone. 190 An example of a model with this output intent would a model that predicts sequence perplexity: 191 ``` 192 output_intents=[ 193 MlMetadataIntent(type="metadata", tag_id=output_tag_id, output_key="perplexity", chains=[]), 194 ] 195 ``` 196 For this output intent, the model output must adhere to the following schema: 197 ``` 198 {"perplexity": list[float]} 199 200 ``` 201 The list mapped to the output key is a list of the tag values produced by the model. 202 203 2. Structure intent 204 This output intent can handle outputs from structure prediction models. 205 ``` 206 output_intents=[ 207 MlStructureIntent( 208 type="structure", 209 output_key="my_structure", 210 )] 211 ``` 212 For this output intent, the model output must adhere to the following schema: 213 ``` 214 {"my_structure": list[str]} 215 216 ``` 217 The list mapped to the output key is a list of the PDB file contents (1 list item per structure produced) 218 219 3. Cluster intent 220 This output intent can handle outputs clustering models. 221 ``` 222 output_intents=[ 223 MlClusterIntent(type="cluster", cluster_label_key="cluster_labels"), 224 ] 225 226 ``` 227 For this output intent, the model output must adhere to the following schema: 228 ``` 229 {"cluster_labels": list[int]} 230 ``` 231 The list mapped to the cluster label keys is a list of cluster labels 232 233 4. New sequence intent 234 This output intent can handle outputs from model that produces sequences (e.g. a humanization model). 235 In addition to a key value pair typical of other intents, this intent requires the specification of a 236 sequence template and quality control template to handle the sequences produced by the models. 237 ``` 238 output_intents=[ 239 MlNewSequenceIntent( 240 type="new_sequence", 241 new_sequences_key="humanized_sequences", 242 quality_control_template_id=quality_control_template.id, 243 quality_control_template_version=quality_control_template.version, 244 sequence_template_id=sequence_template.id, 245 sequence_template_version=sequence_template.version, 246 ) 247 ] 248 249 ``` 250 For this output intent, the model output must adhere to the following schema: 251 ``` 252 {"humanized_sequences":list[list[str]]} 253 254 ``` 255 Here, the inner list contains the (potential) multiple output sequences generated for a single input sequence. 256 If the output sequences are paired chain, the inner list should adhere to the following form: 257 ``` 258 [HC1, LC1, ..., HCn, LCn] 259 ``` 260 261 5. File intent 262 This output intent can be used to produce "raw" output i.e. the output will be made available as a dataframe without 263 any further processing or downstream updates in the platform. This intent allows the user to customize the output processing. 264 ``` 265 output_intents=[ 266 MlMetadataIntent(type="file", filetype="json", output_key="ab_property"), 267 ] 268 ``` 269 For this output intent, the model output must adhere to the following schema: 270 ``` 271 {"ab_property": list[float | str]} 272 ``` 273 274 6. Sequence position/ Liability intent 275 This output intent can be used by models that produce per-residue outputs (e.g. saliency maps or liability predictors). 276 The supported 'value_types' in the output intent specification are "number" and "text" 277 ``` 278 output_intents = [ 279 MlLiabilityIntent( 280 type="liability", 281 category="humanness", 282 value_type="number", 283 output_key="humanness" 284 ) 285 ] 286 ``` 287 For this output intent, the model output must adhere to the following schema: 288 ``` 289 {"humanness": 290 [ 291 { 292 # these fields are required! 293 "chain": "Heavy" | "Light", 294 "exposed": bool, 295 "length": int, 296 "position": int, 297 "value": float, 298 } 299 ] 300 } 301 302 ``` 303 304 305 model_uri (MlflowModelUri): The URI of a registered MLflow model in the format 306 'models:/model_name/model_version' 307 Example: 'models:/antibody_binding_model/1' 308 309 signatures (list[MlEndpointSignature]): The signature definitions for the ML endpoint. ** This argument will be deprecated. ** 310 Example: 311 ``` 312 signatures=[MlEndpointSignature(key="input", type="object", kind="input")], 313 314 ``` 315 316 parameter_mapping (list[MlParam] | None, optional): Specify runtime parameters for the ML model 317 318 Example: 319 Assume we have a model with the following runtime parameters: 320 1. Number of variants (int) 321 2. Number of diffusion steps (int) 322 3. Regions to mask (array[enum]) 323 ``` 324 parameter_mapping=[ 325 MlBaseParam(type="integer", input_key="num_samples", label="Number Of Variants per clone"), 326 MlBaseParam(type="integer", input_key="num_steps", label="Number of diffusion steps"), 327 MlEnumParam( 328 type="enum_array", 329 input_key="regions_to_mask", 330 label="Regions to mutate", 331 enum_options=[region.value for region in [Region.FR1, Region.FR2, Region.FR3, Region.FR4, Region.CDR1, Region.CDR2, Region.CDR3]], 332 ), 333 ] 334 335 ``` 336 The input_key must be indentical to the name used in the signature of the prediction/inference function of the model. 337 base_image (Literal["cpu"] | Literal["gpu"], optional): The base container image to use 338 for model deployment. Use "gpu" for models requiring GPU acceleration. Defaults to "cpu". 339 340 Returns: 341 MlEndpointId: Endpoint ID of the deployed model 342 343 Raises: 344 DeploymentFailed: If the model deployment task fails. 345 """ 346 347 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 348 parameter_mapping_items = ( 349 [ 350 openapi_client.MlParameterMapItem.model_validate( 351 { 352 "type": pm["type"], 353 "input_key": pm["input_key"], 354 "label": pm["label"] if "label" in pm and pm["label"] is not None else None, 355 "source": "parameter", 356 "enum_options": pm.get("enum_options"), 357 } 358 ) 359 for pm in parameter_mapping 360 ] 361 if parameter_mapping is not None 362 else None 363 ) 364 365 payload = openapi_client.DeployModelRequest( 366 display_name=display_name, 367 parameter_mapping=parameter_mapping_items, 368 input_mapping=[openapi_client.MlInput.model_validate(i) for i in input_mapping], 369 output_intents=[openapi_client.MlOutputIntent.from_dict(dict(i)) for i in output_intents], 370 model_uri=model_uri, 371 signatures=[openapi_client.MlEndpointSignature.model_validate(i) for i in signatures], 372 base_image=base_image, 373 ) 374 375 with ApiErrorContext(): 376 deploy_model_response = ml_api_instance.deploy_model(deploy_model_request=payload) 377 assert deploy_model_response.workflow_execution_id is not None 378 379 workflow_execution_id = WorkflowExecutionId(deploy_model_response.workflow_execution_id) 380 381 def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> MlEndpointId: 382 # If the task has succeeded, return the endpoint_id 383 match task_state: 384 case TaskState.SUCCEEDED: 385 result = ml_api_instance.get_endpoint_by_workflow_execution_task_id(workflow_execution_task_id=task_id) 386 assert result.endpoint_id is not None 387 return MlEndpointId(result.endpoint_id) 388 case _: 389 raise DeploymentFailed(task_id) 390 391 waitable = WorkflowExecutionTaskWaitable[MlEndpointId]( 392 workflow_execution_id=workflow_execution_id, 393 task_template_name=WorkflowTaskTemplateName.ENPI_APP_ML_MONITOR_INFERENCE_SERVICE, 394 on_complete=on_complete, 395 ) 396 397 return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
Deploy a registered MLFlow model as an endpoint on the KServe inference platform.
Arguments:
- display_name (str): The display name of the ML endpoint.
input_mapping (list[MlInput]): The input mapping configuration for the ML endpoint. Each MlInput defines a mapping between a platform tag and the input to the ML model
Example: If the inference function of the model has input arguments "heavy_chain" and "light_chain" that map to the heavy and light chain of the amino acids, the input mapping would specified as follows:
input_mapping=[ MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="heavy_chain", chains=[Chain.HEAVY]), MlInput(tag_id=SequenceTags.ReceptorAminoAcids, input_key="light_chain", chains=[Chain.KAPPA, Chain.LAMBDA]), ]output_intents (list[MlOutputIntent]): The output intent is a contract between the ML model and the platform. Each intent specifies
- a schema that the model output must adhere to. This ensures the platform can faithfully process and represent the model output in the platform UI and databases.
- A deployed model can potentially support multiple output intents, provided the model output adheres to the schemas outline below.
The platform currently supports 6 output intents: 1. Metadata intent This output intent can be used for a scalar model output that updates (or creates) a tag for a sequence, collection, or clone. An example of a model with this output intent would a model that predicts sequence perplexity:
output_intents=[ MlMetadataIntent(type="metadata", tag_id=output_tag_id, output_key="perplexity", chains=[]), ]For this output intent, the model output must adhere to the following schema:
{"perplexity": list[float]}The list mapped to the output key is a list of the tag values produced by the model.
- Structure intent This output intent can handle outputs from structure prediction models.
output_intents=[ MlStructureIntent( type="structure", output_key="my_structure", )]For this output intent, the model output must adhere to the following schema:
{"my_structure": list[str]}The list mapped to the output key is a list of the PDB file contents (1 list item per structure produced)
- Cluster intent This output intent can handle outputs clustering models.
output_intents=[ MlClusterIntent(type="cluster", cluster_label_key="cluster_labels"), ]For this output intent, the model output must adhere to the following schema:
{"cluster_labels": list[int]}The list mapped to the cluster label keys is a list of cluster labels
- New sequence intent This output intent can handle outputs from model that produces sequences (e.g. a humanization model). In addition to a key value pair typical of other intents, this intent requires the specification of a sequence template and quality control template to handle the sequences produced by the models.
output_intents=[ MlNewSequenceIntent( type="new_sequence", new_sequences_key="humanized_sequences", quality_control_template_id=quality_control_template.id, quality_control_template_version=quality_control_template.version, sequence_template_id=sequence_template.id, sequence_template_version=sequence_template.version, ) ]For this output intent, the model output must adhere to the following schema:
{"humanized_sequences":list[list[str]]}Here, the inner list contains the (potential) multiple output sequences generated for a single input sequence. If the output sequences are paired chain, the inner list should adhere to the following form:
[HC1, LC1, ..., HCn, LCn]File intent This output intent can be used to produce "raw" output i.e. the output will be made available as a dataframe without any further processing or downstream updates in the platform. This intent allows the user to customize the output processing.
output_intents=[ MlMetadataIntent(type="file", filetype="json", output_key="ab_property"), ]For this output intent, the model output must adhere to the following schema:
{"ab_property": list[float | str]}Sequence position/ Liability intent This output intent can be used by models that produce per-residue outputs (e.g. saliency maps or liability predictors). The supported 'value_types' in the output intent specification are "number" and "text"
output_intents = [ MlLiabilityIntent( type="liability", category="humanness", value_type="number", output_key="humanness" ) ]For this output intent, the model output must adhere to the following schema:
{"humanness": [ { # these fields are required! "chain": "Heavy" | "Light", "exposed": bool, "length": int, "position": int, "value": float, } ] }
model_uri (MlflowModelUri): The URI of a registered MLflow model in the format 'models:/model_name/model_version' Example: 'models:/antibody_binding_model/1'
- signatures (list[MlEndpointSignature]): The signature definitions for the ML endpoint. * This argument will be deprecated. *
- Example:
- ``` signatures=[MlEndpointSignature(key="input", type="object", kind="input")],
- ```
parameter_mapping (list[MlParam] | None, optional): Specify runtime parameters for the ML model
Example: Assume we have a model with the following runtime parameters: 1. Number of variants (int) 2. Number of diffusion steps (int) 3. Regions to mask (array[enum])
parameter_mapping=[ MlBaseParam(type="integer", input_key="num_samples", label="Number Of Variants per clone"), MlBaseParam(type="integer", input_key="num_steps", label="Number of diffusion steps"), MlEnumParam( type="enum_array", input_key="regions_to_mask", label="Regions to mutate", enum_options=[region.value for region in [Region.FR1, Region.FR2, Region.FR3, Region.FR4, Region.CDR1, Region.CDR2, Region.CDR3]], ), ]The input_key must be indentical to the name used in the signature of the prediction/inference function of the model.
- base_image (Literal["cpu"] | Literal["gpu"], optional): The base container image to use for model deployment. Use "gpu" for models requiring GPU acceleration. Defaults to "cpu".
Returns:
MlEndpointId: Endpoint ID of the deployed model
Raises:
- DeploymentFailed: If the model deployment task fails.
399 def undeploy_model(self, endpoint_id: MlEndpointId) -> None: 400 """Undeploys the model and delete the endpoint 401 402 Args: 403 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 404 """ 405 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 406 try: 407 ml_api_instance.undeploy_model(id=endpoint_id, body={}) 408 except openapi_client.ApiException as e: 409 raise ApiError(e)
Undeploys the model and delete the endpoint
Arguments:
- endpoint_id (MlEndpointId): The unique identifier of a ML endpoint.
def
invoke_endpoint( self, endpoint_id: enpi_api.l2.types.ml.MlEndpointId, clone_ids: list[enpi_api.l2.types.clone.CloneId] | None = None, collection_ids: list[enpi_api.l2.types.collection.CollectionId] | None = None, parameters: dict[str, str | int | float | bool | list[str]] | None = None, intent_config: enpi_api.l2.types.ml.MlInvocationIntentConfig | None = None) -> enpi_api.l2.types.execution.Execution[list[File]]:
411 def invoke_endpoint( 412 self, 413 endpoint_id: MlEndpointId, 414 clone_ids: list[CloneId] | None = None, 415 collection_ids: list[CollectionId] | None = None, 416 parameters: dict[str, str | int | float | bool | list[str]] | None = None, 417 intent_config: MlInvocationIntentConfig | None = None, 418 ) -> Execution[list[File]]: 419 """Invoke a ML endpoint for a set of clones or collections. 420 421 Args: 422 endpoint_id (MlEndpointId): The unique identifier of a ML endpoint. 423 clone_ids (list[CloneId]): The unique identifiers of the clones. 424 collection_ids (list[CollectionIds]): The unique identifiers of the collections 425 parameters (dict | None): parameters passed to the invocation. 426 intent_config (MlInvocationIntentConfig | None): This argument is only used for invoking models with a new sequence output intent. 427 It is used to specify the name of the output collection under which all the output seqences are stored. 428 Example: 429 ``` 430 intent_config=MlInvocationIntentConfig(new_sequence=MlInvocationNewSequenceIntentConfig(output_collection_name="humanized_sequences")) 431 ``` 432 433 Returns: 434 output_files (list[File]): The list of output files generated by the File Intent of ML model. 435 """ 436 ml_api_instance = openapi_client.MlApi(self._inner_api_client) 437 file_api = FileApi(self._inner_api_client, self._log_level) 438 if clone_ids is not None and collection_ids is not None: 439 raise ValueError("Either clone_ids or collection_ids must be provided, but not both") 440 441 with ApiErrorContext(): 442 invoke_ml_endpoint_request = openapi_client.InvokeMlEndpointRequest.from_dict( 443 dict( 444 clone_ids=clone_ids, 445 collection_ids=collection_ids, 446 parameters=parameters, 447 intent_config=intent_config.model_dump() if intent_config else None, 448 ) 449 ) 450 assert invoke_ml_endpoint_request is not None 451 result = ml_api_instance.invoke_endpoint(id=endpoint_id, invoke_ml_endpoint_request=invoke_ml_endpoint_request) 452 assert result.workflow_execution_id is not None 453 454 workflow_execution_id = WorkflowExecutionId(result.workflow_execution_id) 455 456 def on_complete(workflow_execution_id: WorkflowExecutionId, execution_state: TaskState) -> list[File]: 457 assert execution_state == TaskState.SUCCEEDED, ( 458 f"Workflow execution {workflow_execution_id} did not reach {TaskState.SUCCEEDED} state, got {execution_state} state instead" 459 ) 460 461 logger.success(f"Ml invocation with workflow execution ID: {workflow_execution_id} has successfully finished.") 462 463 # Get potential file intent output files 464 file_intent_output_ids = ml_api_instance.get_invocation_output(workflow_execution_id).file_ids 465 if len(file_intent_output_ids) > 0: 466 logger.success("Ml invocation File Intent output files retrieved successfully") 467 return [file_api.get_file_by_id(FileId(file_id)) for file_id in file_intent_output_ids] 468 else: 469 return [] 470 471 waitable = WorkflowExecutionWaitable( 472 workflow_execution_id=workflow_execution_id, 473 on_complete=on_complete, 474 ) 475 return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
Invoke a ML endpoint for a set of clones or collections.
Arguments:
- endpoint_id (MlEndpointId): The unique identifier of a ML endpoint.
- clone_ids (list[CloneId]): The unique identifiers of the clones.
- collection_ids (list[CollectionIds]): The unique identifiers of the collections
- parameters (dict | None): parameters passed to the invocation.
- intent_config (MlInvocationIntentConfig | None): This argument is only used for invoking models with a new sequence output intent.
- It is used to specify the name of the output collection under which all the output seqences are stored.
- Example:
- ``` intent_config=MlInvocationIntentConfig(new_sequence=MlInvocationNewSequenceIntentConfig(output_collection_name="humanized_sequences"))
- ```
Returns:
output_files (list[File]): The list of output files generated by the File Intent of ML model.