Fix based on reviews

Author: ono-max

launchable/commands/subset.py

@@ -8,12 +8,9 @@
 from typing import Any, Callable, Dict, List, Optional, Sequence, TextIO, Tuple, Union
 
 import click
-from requests import Session
-from requests.adapters import HTTPAdapter
 from tabulate import tabulate
 
 from launchable.utils.authentication import get_org_workspace
-from launchable.utils.http_client import MAX_RETRIES, default_retry_strategy
 from launchable.utils.session import parse_session
 from launchable.utils.tracking import Tracking, TrackingClient
 
@@ -508,19 +505,8 @@ def run(self):
             else:
                 try:
                     test_runner = context.invoked_subcommand
-
-                    strategy = default_retry_strategy()
-                    # Requesting a subset may require retrying the requests because it takes time to load
-                    # models on the Launchable side.
-                    # Therefore, we make an exception to allow retries.
-                    strategy.read = MAX_RETRIES
-                    adapter = HTTPAdapter(max_retries=strategy)
-                    s = Session()
-                    s.mount("http://", adapter)
-                    s.mount("https://", adapter)
                     client = LaunchableClient(
                         test_runner=test_runner,
-                        session=s,
                         app=app,
                         tracking_client=tracking_client)
 

launchable/utils/env_keys.py

@@ -3,3 +3,4 @@
 ORGANIZATION_KEY = "LAUNCHABLE_ORGANIZATION"
 WORKSPACE_KEY = "LAUNCHABLE_WORKSPACE"
 BASE_URL_KEY = "LAUNCHABLE_BASE_URL"
+SKIP_TIMEOUT_RETRY = "LAUNCHABLE_SKIP_TIMEOUT_RETRY"

launchable/utils/http_client.py

@@ -8,13 +8,13 @@
 from click import Context
 from requests import Session
 from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
+from requests.packages.urllib3.util.retry import Retry  # type: ignore
 
 from launchable.version import __version__
 
 from ..app import Application
 from .authentication import authentication_headers
-from .env_keys import BASE_URL_KEY
+from .env_keys import BASE_URL_KEY, SKIP_TIMEOUT_RETRY
 from .gzipgen import compress as gzipgen_compress
 from .logger import AUDIT_LOG_FORMAT, Logger
 
@@ -31,18 +31,6 @@ def get_base_url():
     return os.getenv(BASE_URL_KEY) or DEFAULT_BASE_URL
 
 
-def default_retry_strategy():
-    return Retry(
-        total=MAX_RETRIES,
-        # When Launchable server is unstable, ReadTimeout can occur.
-        # To prevent the execution from slowing down, we disable retrying the request in this case.
-        read=0,
-        allowed_methods=["GET", "PUT", "PATCH", "DELETE"],
-        status_forcelist=[429, 500, 502, 503, 504],
-        backoff_factor=2
-    )
-
-
 class DryRunResponse:
     def __init__(self, status_code, payload):
         self.status_code = status_code
@@ -63,7 +51,17 @@ def __init__(self, base_url: str = "", session: Optional[Session] = None,
         self.skip_cert_verification = bool(app and app.skip_cert_verification)
 
         if session is None:
-            strategy = default_retry_strategy()
+            read = MAX_RETRIES
+            if os.getenv(SKIP_TIMEOUT_RETRY):
+                read = 0
+            strategy = Retry(
+                total=MAX_RETRIES,
+                read=read,
+                allowed_methods=["GET", "PUT", "PATCH", "DELETE"],
+                status_forcelist=[429, 500, 502, 503, 504],
+                backoff_factor=2
+            )
+
             adapter = HTTPAdapter(max_retries=strategy)
             s = Session()
             s.mount("http://", adapter)