yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-12-29 08:06:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Secrets management

Browse Source
This commit is contained in:
Jake Poznanski 2024-11-13 11:26:46 -08:00
parent 39256c19bb
commit 9610eac4f0
1 changed files with 27 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
pdelfin/beakerpipeline.py
View File

@ -536,6 +536,7 @@ def submit_beaker_job(args):
b = Beaker.from_env(default_workspace=args.beaker_workspace)
account = b.account.whoami()
owner = account.name
beaker_image = "jakep/pdelfin-inference"
task_name = f"pdelfin-{os.path.basename(args.workspace.rstrip('/'))}"
@ -543,6 +544,22 @@ def submit_beaker_job(args):
args_list = [arg for arg in sys.argv[1:] if arg != "--beaker"]
try:
b.secret.get(f"{owner}-WEKA_ACCESS_KEY_ID", args.beaker_workspace)
b.secret.get(f"{owner}-WEKA_SECRET_ACCESS_KEY", args.beaker_workspace)
b.secret.get(f"{owner}-AWS_CREDENTIALS_FILE", args.beaker_workspace)
except SecretNotFound:
print(f"Expected beaker secrets for accessing Weka and S3 are not found. Are you okay to write those to your beaker workspace {args.beaker_workspace}? [y/n]")
if input().strip().lower() != "y":
print("Exiting...")
sys.exit(1)
b.secret.write(f"{owner}-WEKA_ACCESS_KEY_ID", os.environ.get("WEKA_ACCESS_KEY_ID", ""), args.beaker_workspace)
b.secret.write(f"{owner}-WEKA_SECRET_ACCESS_KEY", os.environ.get("WEKA_SECRET_ACCESS_KEY", ""), args.beaker_workspace)
b.secret.write(f"{owner}-AWS_CREDENTIALS_FILE", open(os.path.join(os.path.expanduser('~'), '.aws', 'credentials')).read(), args.beaker_workspace)
# Create the experiment spec
experiment_spec = ExperimentSpec(
budget="ai2/oe-data",
@ -561,9 +578,10 @@ def submit_beaker_job(args):
command=["python", "-m", "pdelfin.beakerpipeline"] + args_list,
env_vars=[
EnvVar(name="BEAKER_JOB_NAME", value=task_name),
EnvVar(name="OWNER", value=account.name),
#EnvVar(name="AWS_ACCESS_KEY_ID", secret=f"{account.name}-{S2_AWS_ACCESS_KEY_ID_SECRET_NAME}"),
#EnvVar(name="AWS_SECRET_ACCESS_KEY", secret=f"{account.name}-{S2_AWS_SECRET_ACCESS_KEY_SECRET_NAME}"),
EnvVar(name="OWNER", value=owner),
EnvVar(name="WEKA_ACCESS_KEY_ID", secret=f"{owner}-WEKA_ACCESS_KEY_ID"),
EnvVar(name="WEKA_SECRET_ACCESS_KEY", secret=f"{owner}-WEKA_SECRET_ACCESS_KEY"),
EnvVar(name="AWS_CREDENTIALS_FILE", secret=f"{owner}-AWS_CREDENTIALS_FILE"),
],
resources=TaskResources(gpu_count=1),
constraints=Constraints(cluster=args.beaker_cluster),
@ -602,6 +620,12 @@ async def main():
parser.add_argument('--beaker_cluster', help='Beaker clusters you want to run on', default=["ai2/jupiter-cirrascale-2", "ai2/pluto-cirrascale", "ai2/saturn-cirrascale"])
args = parser.parse_args()
if "AWS_CREDENTIALS_FILE" in os.environ:
cred_path = os.path.join(os.path.expanduser('~'), '.aws', 'credentials')
os.path.makedirs(os.path.dirname(cred_path), exist_ok=True)
with open(cred_path, "w") as f:
f.write(os.environ.get("AWS_CREDENTIALS_FILE"))
if args.workspace_profile:
global workspace_s3
workspace_session = boto3.Session(profile_name=args.workspace_profile)