Created
January 4, 2025 23:36
-
-
Save drscotthawley/8e1dfc7b7bda6943a8829a97e50e5454 to your computer and use it in GitHub Desktop.
Wrapper to run another script with CLI args read from a JSON config file stored on WandB; allows override via new CLI args. Gen'd by Claude 3.5 Sonnet.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
import json | |
import subprocess | |
import argparse | |
import sys | |
import os | |
def parse_override_args(args): | |
"""Parse CLI arguments that start with '--' into a dict.""" | |
overrides = {} | |
i = 0 | |
while i < len(args): | |
arg = args[i] | |
if arg.startswith('--'): | |
key = arg[2:] # Remove '--' | |
if '=' in key: # Handle --key=value format | |
key, value = key.split('=', 1) | |
else: # Handle --key value format | |
if i + 1 >= len(args) or args[i + 1].startswith('--'): | |
value = 'true' # Treat as boolean flag | |
else: | |
value = args[i + 1] | |
i += 1 | |
overrides[key.replace('-', '_')] = value | |
i += 1 | |
return overrides | |
def wandb_config_to_args(config, overrides=None): | |
"""Convert WandB config JSON to command line arguments.""" | |
args = [] | |
overrides = overrides or {} | |
# Skip these keys | |
skip_keys = {"_wandb", "learning_rate"} | |
for key, value_dict in config.items(): | |
if key in skip_keys: # Skip wandb internal config and derived values | |
continue | |
# Get value from config, but allow override | |
cli_key = key.replace('_', '-') | |
if key in overrides: | |
value = overrides.pop(key) # Remove from overrides to track usage | |
else: | |
value = value_dict.get("value") | |
if value is not None: # Only add if value exists | |
# Handle boolean values | |
if isinstance(value, bool): | |
if value: # Only add flag if True | |
args.append(f"--{cli_key}") | |
else: | |
args.extend([f"--{cli_key}", str(value)]) | |
# Add any remaining overrides (new parameters not in config) | |
for key, value in overrides.items(): | |
cli_key = key.replace('_', '-') | |
# Handle boolean values in overrides | |
if value.lower() in ('true', 'yes', '1', 'on'): | |
args.append(f"--{cli_key}") | |
elif value.lower() in ('false', 'no', '0', 'off'): | |
continue # Skip false boolean flags | |
else: | |
args.extend([f"--{cli_key}", value]) | |
return args | |
def main(): | |
parser = argparse.ArgumentParser(description='Run a Python script with WandB config JSON') | |
parser.add_argument('script', help='Python script to run') | |
parser.add_argument('config', help='WandB config JSON file') | |
args, unknown = parser.parse_known_args() # Capture unknown args for overrides | |
# Read config file | |
with open(args.config) as f: | |
config = json.load(f) | |
# Parse override arguments | |
overrides = parse_override_args(unknown) | |
# Convert config to command line arguments, applying overrides | |
cmd_args = wandb_config_to_args(config, overrides) | |
# Print the final command for debugging | |
cmd = [sys.executable, args.script] + cmd_args | |
print("Running command:", ' '.join(cmd), "\n", flush=True) | |
# Run the command and pass through all output directly | |
process = subprocess.Popen( | |
cmd, | |
stdout=None, # Use parent process's stdout directly | |
stderr=None, # Use parent process's stderr directly | |
bufsize=1, | |
universal_newlines=True, | |
env={**os.environ, 'PYTHONUNBUFFERED': '1'} # Ensure Python output is unbuffered | |
) | |
# Wait for completion and get return code | |
return process.wait() | |
if __name__ == '__main__': | |
sys.exit(main()) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment