mirror of
https://github.com/vale981/ray
synced 2025-03-06 02:21:39 -05:00
b0fa9d6766
After #24459, this PR will add similar support for model artifact saving and an example for experiment tracking with Ray AIR for Comet ML.
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""
|
|
This script is used to set up credentials for some services in the
|
|
CI environment. For instance, it can fetch WandB API tokens and write
|
|
the WandB configuration file so test scripts can use the service.
|
|
"""
|
|
import os
|
|
import sys
|
|
|
|
import boto3
|
|
|
|
AWS_WANDB_SECRET_ARN = (
|
|
"arn:aws:secretsmanager:us-west-2:029272617770:secret:oss-ci/wandb-key-V8UeE5"
|
|
)
|
|
AWS_COMET_SECRET_ARN = (
|
|
"arn:aws:secretsmanager:us-west-2:029272617770:secret:oss-ci/comet-ml-token-vw81C3"
|
|
)
|
|
|
|
|
|
def get_and_write_wandb_api_key(client):
|
|
api_key = client.get_secret_value(SecretId=AWS_WANDB_SECRET_ARN)["SecretString"]
|
|
with open(os.path.expanduser("~/.netrc"), "w") as fp:
|
|
fp.write(f"machine api.wandb.ai\n" f" login user\n" f" password {api_key}\n")
|
|
|
|
|
|
def get_and_write_comet_ml_api_key(client):
|
|
api_key = client.get_secret_value(SecretId=AWS_COMET_SECRET_ARN)["SecretString"]
|
|
with open(os.path.expanduser("~/.comet.config"), "w") as fp:
|
|
fp.write(f"[comet]\napi_key={api_key}\n")
|
|
|
|
|
|
SERVICES = {
|
|
"wandb": get_and_write_wandb_api_key,
|
|
"comet_ml": get_and_write_comet_ml_api_key,
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print(f"Usage: python {sys.argv[0]} <service1> [service2] ...")
|
|
sys.exit(0)
|
|
|
|
services = sys.argv[1:]
|
|
|
|
if any(service not in SERVICES for service in services):
|
|
raise RuntimeError(
|
|
f"All services must be included in {list(SERVICES.keys())}. "
|
|
f"Got: {services}"
|
|
)
|
|
|
|
client = boto3.client("secretsmanager", region_name="us-west-2")
|
|
for service in services:
|
|
SERVICES[service](client)
|