Skip to content

Commit

Permalink
fix: env vars to http client in sagemaker (#35)
Browse files Browse the repository at this point in the history
* fix: env vars to http

* fix: env vars to http
  • Loading branch information
asafgardin authored Jan 3, 2024
1 parent 4d4ef71 commit 4336c46
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion ai21/services/sagemaker.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from typing import List

from ai21 import AI21EnvConfig
from ai21.ai21_http_client import AI21HTTPClient
from ai21.clients.sagemaker.constants import (
SAGEMAKER_MODEL_PACKAGE_NAMES,
Expand Down Expand Up @@ -56,7 +57,13 @@ def list_model_package_versions(cls, model_name: str, region: str) -> List[str]:

@classmethod
def _create_ai21_http_client(cls) -> AI21HTTPClient:
return AI21HTTPClient()
return AI21HTTPClient(
api_key=AI21EnvConfig.api_key,
api_host=AI21EnvConfig.api_host,
api_version=AI21EnvConfig.api_version,
timeout_sec=AI21EnvConfig.timeout_sec,
num_retries=AI21EnvConfig.num_retries,
)


def _assert_model_package_exists(model_name, region):
Expand Down

0 comments on commit 4336c46

Please sign in to comment.