Skip to content

Commit

Permalink
Use an Airflow config to manage provider configuration (#41)
Browse files Browse the repository at this point in the history
Before this change, the provider used standard operating system environment variables to receive user configuration.

This PR changes these to Airflow config, which already gracefully handles booleans and supports setting the configuration via environment variables and via `airflow.cfg`.

This PR introduces these changes:
* Replace `ANYSCALE__AIRFLOW_TELEMETRY_ENABLED` by `AIRFLOW__ANYSCALE__CLI_TOKEN` (or equivalent `airflow.cfg`)
* Support `AIRFLOW__ANYSCALE__CLI_TOKEN` (or equivalent `airflow.cfg`) as an alternative to `ANYSCALE_CLI_TOKEN`

Note: the following is discussed in a separate PR (#42), so I removed changing the default value for telemetry from the current PR.

~~The current documentation needs to provide more information about the telemetry and data collection, which is a crucial aspect that needs to be addressed.~~

~~Until the following questions are addressed, I believe it should be disabled by default:~~
~~* How does enabling telemetry help users?~~
~~* Where is the data being sent?~~
~~* What data is being collected?~~
~~* How is the data being used?~~
  • Loading branch information
tatiana authored Jun 28, 2024
1 parent 32a8959 commit f0d3338
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 5 deletions.
8 changes: 4 additions & 4 deletions anyscale_provider/hooks/anyscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from functools import cached_property
from typing import Any

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook
from anyscale import Anyscale
Expand Down Expand Up @@ -36,16 +37,15 @@ def client(self) -> Anyscale:

# If the token is not found in the connection, try to get it from the environment variable
if not token:
self.log.info(f"Using token from ENV")
token = os.getenv("ANYSCALE_CLI_TOKEN")
self.log.info(f"Using token from config or ENV")
token = conf.get("anyscale", "cli_token", fallback=os.getenv("ANYSCALE_CLI_TOKEN"))

if not token:
raise AirflowException(f"Missing API token for connection id {self.conn_id}")

# 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"]
telemetry_enabled = conf.getboolean("anyscale", "telemetry_enabled", fallback=True)
if telemetry_enabled:
headers["X-Anyscale-Source"] = "airflow"

Expand Down
2 changes: 1 addition & 1 deletion anyscale_provider/operators/anyscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -329,7 +329,7 @@ def on_kill(self) -> None:
self.log.info("Termination request received. Submitted request to terminate the anyscale service rollout.")
return

def execute(self, context: Context) -> str | None:
def execute(self, context: Context) -> None:
"""
Execute the service rollout to Anyscale.
Expand Down

0 comments on commit f0d3338

Please sign in to comment.