Skip to content

Commit

Permalink
Add telemetry headers
Browse files Browse the repository at this point in the history
- Capture telemetry via a header at SDK init
- Enable users to opt-out of telemetry via setting env_var
  • Loading branch information
marwan116 committed Jun 25, 2024
1 parent ca782bd commit e171c72
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion anyscale_provider/hooks/anyscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,14 @@ def client(self) -> Anyscale:
if not token:
raise AirflowException(f"Missing API token for connection id {self.conn_id}")

return Anyscale(auth_token=token)
# Add custom headers if telemetry is enabled - by default telemetry is enabled.
headers = {}
telemetry_env = os.getenv("ANYSCALE__AIRFLOW_TELEMETRY_ENABLED", "true")
telemetry_enabled = telemetry_env.lower() in ["true", "1", "yes", "on"]
if telemetry_enabled:
headers["X-Anyscale-Source"] = "airflow"

return Anyscale(auth_token=token, headers=headers)

@classmethod
def get_ui_field_behaviour(cls) -> dict[str, Any]:
Expand Down

0 comments on commit e171c72

Please sign in to comment.