diff --git a/src/huggingface_hub/_inference_endpoints.py b/src/huggingface_hub/_inference_endpoints.py index 651efc97b0..282b627b54 100644 --- a/src/huggingface_hub/_inference_endpoints.py +++ b/src/huggingface_hub/_inference_endpoints.py @@ -245,6 +245,7 @@ def update( revision: Optional[str] = None, task: Optional[str] = None, custom_image: Optional[Dict] = None, + secrets: Optional[Dict[str, str]] = None, ) -> "InferenceEndpoint": """Update the Inference Endpoint. @@ -279,7 +280,8 @@ def update( custom_image (`Dict`, *optional*): A custom Docker image to use for the Inference Endpoint. This is useful if you want to deploy an Inference Endpoint running on the `text-generation-inference` (TGI) framework (see examples). - + secrets (`Dict[str, str]`, *optional*): + Secret values to inject in the container environment. Returns: [`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data. """ @@ -298,6 +300,7 @@ def update( revision=revision, task=task, custom_image=custom_image, + secrets=secrets, token=self._token, # type: ignore [arg-type] ) diff --git a/src/huggingface_hub/hf_api.py b/src/huggingface_hub/hf_api.py index 17fa4889a7..3b1e601be5 100644 --- a/src/huggingface_hub/hf_api.py +++ b/src/huggingface_hub/hf_api.py @@ -136,9 +136,7 @@ ) from .utils import tqdm as hf_tqdm from .utils._typing import CallableT -from .utils.endpoint_helpers import ( - _is_emission_within_threshold, -) +from .utils.endpoint_helpers import _is_emission_within_threshold R = TypeVar("R") # Return type @@ -7418,6 +7416,7 @@ def create_inference_endpoint( revision: Optional[str] = None, task: Optional[str] = None, custom_image: Optional[Dict] = None, + secrets: Optional[Dict[str, str]] = None, type: InferenceEndpointType = InferenceEndpointType.PROTECTED, namespace: Optional[str] = None, token: Union[bool, str, None] = None, @@ -7456,6 +7455,8 @@ def create_inference_endpoint( custom_image (`Dict`, *optional*): A custom Docker image to use for the Inference Endpoint. This is useful if you want to deploy an Inference Endpoint running on the `text-generation-inference` (TGI) framework (see examples). + secrets (`Dict[str, str]`, *optional*): + Secret values to inject in the container environment. type ([`InferenceEndpointType]`, *optional*): The type of the Inference Endpoint, which can be `"protected"` (default), `"public"` or `"private"`. namespace (`str`, *optional*): @@ -7518,6 +7519,7 @@ def create_inference_endpoint( ... }, ... "url": "ghcr.io/huggingface/text-generation-inference:1.1.0", ... }, + ... secrets={"MY_SECRET_KEY": "secret_value"}, ... ) ``` @@ -7543,6 +7545,7 @@ def create_inference_endpoint( "revision": revision, "task": task, "image": image, + "secrets": secrets, }, "name": name, "provider": { @@ -7625,6 +7628,7 @@ def update_inference_endpoint( revision: Optional[str] = None, task: Optional[str] = None, custom_image: Optional[Dict] = None, + secrets: Optional[Dict[str, str]] = None, # Other namespace: Optional[str] = None, token: Union[bool, str, None] = None, @@ -7664,7 +7668,8 @@ def update_inference_endpoint( custom_image (`Dict`, *optional*): A custom Docker image to use for the Inference Endpoint. This is useful if you want to deploy an Inference Endpoint running on the `text-generation-inference` (TGI) framework (see examples). - + secrets (`Dict[str, str]`, *optional*): + Secret values to inject in the container environment. namespace (`str`, *optional*): The namespace where the Inference Endpoint will be updated. Defaults to the current user's namespace. token (Union[bool, str, None], optional): @@ -7702,6 +7707,8 @@ def update_inference_endpoint( payload["model"]["task"] = task if custom_image is not None: payload["model"]["image"] = {"custom": custom_image} + if secrets is not None: + payload["model"]["secrets"] = secrets response = get_session().put( f"{constants.INFERENCE_ENDPOINTS_ENDPOINT}/endpoint/{namespace}/{name}", diff --git a/tests/test_inference_endpoints.py b/tests/test_inference_endpoints.py index 019d04f57e..67ec4cab0d 100644 --- a/tests/test_inference_endpoints.py +++ b/tests/test_inference_endpoints.py @@ -30,6 +30,7 @@ "task": "text-generation", "framework": "pytorch", "image": {"huggingface": {}}, + "secret": {"token": "my-token"}, }, "status": { "createdAt": "2023-10-26T12:41:53.263078506Z", @@ -61,6 +62,7 @@ "task": "text-generation", "framework": "pytorch", "image": {"huggingface": {}}, + "secrets": {"token": "my-token"}, }, "status": { "createdAt": "2023-10-26T12:41:53.263Z", @@ -93,6 +95,7 @@ "task": "text-generation", "framework": "pytorch", "image": {"huggingface": {}}, + "secrets": {"token": "my-token"}, }, "status": { "createdAt": "2023-10-26T12:41:53.263Z",