def register_model()

in source/controlplaneapi/runtime/app.py [0:0]

• 90 lines of code
• 7 McCabe index (conditional complexity)

def register_model():
    """
    Register a new Machine Learning (ML) model endpoint or publish a new version of an 
    existing model endpoint.
    
    Body:

    .. code-block:: python

        {
            "Name": string,
            "Description": string,
            "ContentGroups": list,
            "Endpoint": string,
            "PluginClass": ["Classifier"|"Optimizer"|"Featurer"|"Labeler"]
        }
    
    Parameters:

        - Name: Name of the Machine Learning (ML) Model
        - Description: Description of the Machine Learning (ML) Model
        - ContentGroups: List of Content Groups this Machine Learning (ML) Model is used for
        - Endpoint: ARN of the Machine Learning (ML) model endpoint. For example ARN of rekognition custom label project endpoint or Sagemaker Endpoint
        - PluginClass: One of "Classifier"|"Optimizer"|"Featurer"|"Labeler"

    Returns:

        A dict containing the Name and Version of the registered model

        .. code-block:: python

            {
                "Name": string,
                "Version": string
            }

    Raises:
        400 - BadRequestError
        500 - ChaliceViewError
    """
    try:
        model = json.loads(app.current_request.raw_body.decode())

        validate(instance=model, schema=API_SCHEMA["register_model"])

        print("Got a valid model schema")

        name = model["Name"]

        model_table = ddb_resource.Table(MODEL_TABLE_NAME)

        response = model_table.get_item(
            Key={
                "Name": name,
                "Version": "v0"
            },
            ConsistentRead=True
        )

        if "Item" not in response:
            print(f"Registering a new model endpoint '{name}'")
            latest_version = 0
            higher_version = 1

        else:
            print(f"Publishing a new version of the model endpoint '{name}'")
            latest_version = response["Item"]["Latest"]
            higher_version = int(latest_version) + 1

        model["Created"] = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")
        model["Enabled"] = True

        # Serialize Python object to DynamoDB object
        serialized_model = {k: serializer.serialize(v) for k, v in model.items()}

        ddb_client.transact_write_items(
            TransactItems=[
                {
                    "Update": {
                        "TableName": MODEL_TABLE_NAME,
                        "Key": {
                            "Name": {"S": name},
                            "Version": {"S": "v0"}
                        },
                        "ConditionExpression": "attribute_not_exists(#Latest) OR #Latest = :Latest",
                        "UpdateExpression": "SET #Latest = :Higher_version, #Description = :Description, #ContentGroups = :ContentGroups, #Endpoint = :Endpoint, #PluginClass = :PluginClass, #Created = :Created, #Enabled = :Enabled",
                        "ExpressionAttributeNames": {
                            "#Latest": "Latest",
                            "#Description": "Description",
                            "#ContentGroups": "ContentGroups",
                            "#Endpoint": "Endpoint",
                            "#PluginClass": "PluginClass",
                            "#Created": "Created",
                            "#Enabled": "Enabled"
                        },
                        "ExpressionAttributeValues": {
                            ":Latest": {"N": str(latest_version)},
                            ":Higher_version": {"N": str(higher_version)},
                            ":Description": serialized_model["Description"] if "Description" in serialized_model else {
                                "S": ""},
                            ":ContentGroups": serialized_model["ContentGroups"],
                            ":Endpoint": serialized_model["Endpoint"],
                            ":PluginClass": serialized_model["PluginClass"],
                            ":Created": serialized_model["Created"],
                            ":Enabled": serialized_model["Enabled"]
                        }
                    }
                },
                {
                    "Put": {
                        "TableName": MODEL_TABLE_NAME,
                        "Item": {
                            "Name": {"S": name},
                            "Version": {"S": "v" + str(higher_version)},
                            "Description": serialized_model["Description"] if "Description" in serialized_model else {
                                "S": ""},
                            "ContentGroups": serialized_model["ContentGroups"],
                            "Endpoint": serialized_model["Endpoint"],
                            "PluginClass": serialized_model["PluginClass"],
                            "Created": serialized_model["Created"],
                            "Enabled": serialized_model["Enabled"]
                        }
                    }
                }
            ]
        )

    except ValidationError as e:
        print(f"Got jsonschema ValidationError: {str(e)}")
        raise BadRequestError(e.message)

    except Exception as e:
        print(f"Unable to register or publish a new version of the Machine Learning model endpoint: {str(e)}")
        raise ChaliceViewError(
            f"Unable to register or publish a new version of the Machine Learning model endpoint: {str(e)}")

    else:
        print(
            f"Successfully registered or published a new version of the Machine Learning model endpoint: {json.dumps(model)}")

        return {
            "Name": model["Name"],
            "Version": "v" + str(higher_version)
        }