Skip to content

Instantly share code, notes, and snippets.

@drscotthawley
Created January 4, 2025 23:36
Show Gist options
  • Save drscotthawley/8e1dfc7b7bda6943a8829a97e50e5454 to your computer and use it in GitHub Desktop.
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.
#!/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