Set timeout default to None (#10487)

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
mlflow · Nov 22, 2023 · 8ca8429 · 8ca8429
1 parent fdb2779
commit 8ca8429
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/mlflow/deployments/databricks/__init__.py b/mlflow/deployments/databricks/__init__.py
@@ -70,7 +70,7 @@ def _call_endpoint(
         prefix: str = "/api/2.0",
         route: Optional[str] = None,
         json_body: Optional[Dict[str, Any]] = None,
-        timeout: int = MLFLOW_HTTP_REQUEST_TIMEOUT.get(),
+        timeout: Optional[int] = None,
     ):
         call_kwargs = {}
         if method.lower() == "get":
@@ -82,7 +82,7 @@ def _call_endpoint(
             host_creds=get_databricks_host_creds(self.target_uri),
             endpoint=posixpath.join(prefix, "serving-endpoints", route or ""),
             method=method,
-            timeout=timeout,
+            timeout=MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout,
             raise_on_status=False,
             retry_codes=MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
             **call_kwargs,

diff --git a/mlflow/deployments/mlflow/__init__.py b/mlflow/deployments/mlflow/__init__.py
@@ -93,7 +93,7 @@ def _call_endpoint(
         method: str,
         route: str,
         json_body: Optional[str] = None,
-        timeout: int = MLFLOW_HTTP_REQUEST_TIMEOUT.get(),
+        timeout: Optional[int] = None,
     ):
         call_kwargs = {}
         if method.lower() == "get":
@@ -105,7 +105,7 @@ def _call_endpoint(
             host_creds=_get_default_host_creds(self.target_uri),
             endpoint=route,
             method=method,
-            timeout=timeout,
+            timeout=MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout,
             retry_codes=MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
             raise_on_status=False,
             **call_kwargs,