diff --git a/src/sagemaker/huggingface/model.py b/src/sagemaker/huggingface/model.py index 107e7407cc..48e3681563 100644 --- a/src/sagemaker/huggingface/model.py +++ b/src/sagemaker/huggingface/model.py @@ -304,6 +304,7 @@ def register( approval_status=None, description=None, drift_check_baselines=None, + customer_metadata_properties=None, domain=None, ): """Creates a model package for creating SageMaker models or listing on Marketplace. @@ -332,6 +333,8 @@ def register( or "PendingManualApproval". Defaults to ``PendingManualApproval``. description (str): Model Package description. Defaults to ``None``. drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None). + customer_metadata_properties (dict[str, str]): A dictionary of key-value paired + metadata properties (default: None). domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING", "MACHINE_LEARNING" (default: None). @@ -362,6 +365,7 @@ def register( approval_status, description, drift_check_baselines=drift_check_baselines, + customer_metadata_properties=customer_metadata_properties, domain=domain, ) diff --git a/src/sagemaker/sklearn/model.py b/src/sagemaker/sklearn/model.py index 91380721ef..e0ae5e2c3d 100644 --- a/src/sagemaker/sklearn/model.py +++ b/src/sagemaker/sklearn/model.py @@ -151,6 +151,8 @@ def register( marketplace_cert=False, approval_status=None, description=None, + drift_check_baselines=None, + customer_metadata_properties=None, domain=None, ): """Creates a model package for creating SageMaker models or listing on Marketplace. @@ -176,6 +178,9 @@ def register( approval_status (str): Model Approval Status, values can be "Approved", "Rejected", or "PendingManualApproval" (default: "PendingManualApproval"). description (str): Model Package description (default: None). + drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None). + customer_metadata_properties (dict[str, str]): A dictionary of key-value paired + metadata properties (default: None). domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING", "MACHINE_LEARNING" (default: None). @@ -205,6 +210,8 @@ def register( marketplace_cert, approval_status, description, + drift_check_baselines=drift_check_baselines, + customer_metadata_properties=customer_metadata_properties, domain=domain, )