feat(sdk): hacky way to set extra headers for the gql client (wandb#5237)

Author: dmitryduev

wandb/sdk/internal/internal_api.py

@@ -87,6 +87,7 @@ class DefaultSettings(TypedDict):
         api_key: Optional[str]
         entity: Optional[str]
         project: Optional[str]
+        _extra_http_headers: Optional[Mapping[str, str]]
 
     _Response = MutableMapping
     SweepState = Literal["RUNNING", "PAUSED", "CANCELED", "FINISHED"]
@@ -178,6 +179,7 @@ def __init__(
             "api_key": None,
             "entity": None,
             "project": None,
+            "_extra_http_headers": None,
         }
         self.retry_timedelta = retry_timedelta
         # todo: Old Settings do not follow the SupportsKeysAndGetItem Protocol
@@ -194,12 +196,22 @@ def __init__(
             "system_samples": 15,
             "heartbeat_seconds": 30,
         }
+
+        # todo: remove this hacky hack after settings refactor is complete
+        #  keeping this code here to limit scope and so that it is easy to remove later
+        extra_http_headers = self.settings(
+            "_extra_http_headers"
+        ) or wandb.sdk.wandb_settings._str_as_dict(
+            self._environ.get("WANDB__EXTRA_HTTP_HEADERS", {})
+        )
+
         self.client = Client(
             transport=GraphQLSession(
                 headers={
                     "User-Agent": self.user_agent,
                     "X-WANDB-USERNAME": env.get_username(env=self._environ),
                     "X-WANDB-USER-EMAIL": env.get_user_email(env=self._environ),
+                    **extra_http_headers,
                 },
                 use_json=True,
                 # this timeout won't apply when the DNS lookup fails. in that case, it will be 60s

wandb/sdk/internal/settings_static.py

@@ -62,6 +62,7 @@ class SettingsStatic:
     _flow_control_custom: bool
     disable_job_creation: bool
     _async_upload_concurrency_limit: Optional[int]
+    _extra_http_headers: Optional[Mapping[str, str]]
 
     # TODO(jhr): clean this up, it is only in SettingsStatic and not in Settings
     _log_level: int

wandb/sdk/lib/_settings_toposort_generated.py

@@ -395,6 +395,7 @@ class Settings:
     _disable_viewer: bool  # Prevent early viewer query
     _except_exit: bool
     _executable: str
+    _extra_http_headers: Mapping[str, str]
     _flow_control_custom: bool
     _flow_control_disabled: bool
     _internal_check_process: Union[int, float]
@@ -540,6 +541,7 @@ def _default_props(self) -> Dict[str, Dict[str, Any]]:
             },
             _disable_stats={"preprocessor": _str_as_bool},
             _disable_viewer={"preprocessor": _str_as_bool},
+            _extra_http_headers={"preprocessor": _str_as_dict},
             _network_buffer={"preprocessor": int},
             _colab={
                 "hook": lambda _: "google.colab" in sys.modules,