File size: 10,109 Bytes
1499d83 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 | # Imports the function named by the config file, and runs it with the arguments inside the config specified
import importlib
import asyncio
import sys
import traceback
import yaml
import argparse
from pathlib import Path
import json
import io
# ==============================================================================
# FORCE UTF-8 FOR ALL I/O
# where the default encoding is not UTF-8. It prevents UnicodeEncodeError
# when printing characters that are not in the default system codepage.
# ==============================================================================
if sys.stdout.encoding != 'utf-8':
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
if sys.stderr.encoding != 'utf-8':
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
# ==============================================================================
from resolve_path import resolve_path
def load_function_from_path(function_path):
"""Dynamically import a function from a module path string"""
module_path, function_name = function_path.rsplit(".", 1)
# Ensure module path uses dots for importlib
module_import_path = module_path.replace("/", ".")
module = importlib.import_module(module_import_path)
return getattr(module, function_name)
def flatten_config(config, no_flatten_keys=None):
"""Flatten config while preserving specified nested structures"""
flattened = {}
no_flatten = set(no_flatten_keys or [])
for key, value in config.items():
if key in no_flatten:
flattened[key] = value
continue
if isinstance(value, dict):
nested_flat = flatten_config(value, no_flatten_keys)
for nested_key, nested_value in nested_flat.items():
if nested_key in flattened:
raise ValueError(f"Key conflict: '{nested_key}'")
flattened[nested_key] = nested_value
else:
if key in flattened:
raise ValueError(f"Key conflict: '{key}'")
flattened[key] = value
print("FLATTENED")
print(flattened)
return flattened
try:
with open("super_config.yaml", "r", encoding="utf-8") as f:
super_config = yaml.safe_load(f)
path_aliases = super_config.get(
"path_aliases", {}
) # Get aliases, default to empty dict if not present
except FileNotFoundError:
print(f"Error: Super config file not found at super_config.yaml")
sys.exit(1)
# Load super config
except yaml.YAMLError as e:
print(f"Error parsing super config file super_config.yaml: {e}")
sys.exit(1)
def run_pipeline(node, config, override_fields={}):
# Resolve node and config paths using aliases
resolved_node_path = resolve_path(node, path_aliases)
# print(f"DEBUG: Resolved node path: {resolved_node_path}")
# Handle optional config key and resolve its path if present
config_path_str = config
resolved_config_path_str = (
resolve_path(config_path_str, path_aliases) if config_path_str else None
)
# print(f"DEBUG: Resolved config path: {resolved_config_path_str}")
resolved_config_path_str = (
resolved_config_path_str
if resolved_config_path_str and resolved_config_path_str.endswith(".yaml")
else (
resolved_config_path_str + ".yaml"
if resolved_config_path_str
else resolved_config_path_str
)
)
# Load pipeline-specific config if a path is provided
config = {}
if resolved_config_path_str:
# Resolve config path relative to the script's directory
script_dir = Path(__file__).parent
config_path = (script_dir / resolved_config_path_str).resolve()
# print(f"DEBUG: Final resolved config path: {config_path}")
try:
with open(config_path, "r", encoding="utf-8") as f:
config = (
yaml.safe_load(f) or {}
) # Ensure config is at least an empty dict
except FileNotFoundError:
print(
f"Warning: Config file not found at {config_path}. Proceeding without it."
)
except Exception as e:
print(f"Error loading config file {config_path}: {e}")
# Decide if you want to raise the error or continue
# raise e
print("Proceeding with empty configuration for this pipeline.")
# print("DEBUG: CONFIG")
# print(config)
run_pipeline_config(
config=config,
resolved_node_path=resolved_node_path,
override_fields=override_fields,
)
def run_pipeline_config(
config, resolved_node_path, override_fields={}
): # the second half of run_pipeline, extracted so that it is easier to use in isolation as an api.
# Merge pipeline config from super_config with loaded config file parameters
# Parameters defined directly in the pipeline entry override those in the loaded config file.
# Parameters passed in as overrides override those in either.
# Flatten the nested configuration
no_flatten_keys = config.get(
"no_flatten", []
) # Get no_flatten from loaded/merged config
flattened_config = flatten_config(config, no_flatten_keys=no_flatten_keys)
flattened_config.update(override_fields)
# Import the target function using the resolved node path
try:
function = load_function_from_path(resolved_node_path)
except (ImportError, AttributeError, ValueError) as e:
print(f"Error loading function from node path '{resolved_node_path}': {e}")
print(f"Skipping pipeline: resolved_node_path") # Use name if available
return # Skip this pipeline if function cannot be loaded
if asyncio.iscoroutinefunction(function):
print(f"Running async pipeline: {resolved_node_path}")
# print("DEBUG: Flattened config")
# print(flattened_config)
try:
asyncio.run(function(**flattened_config))
except Exception as e:
print(f"Error running async pipeline {resolved_node_path}: {e}")
traceback.print_exc()
raise
# Optionally re-raise or handle error reporting
else:
print(f"Running sync pipeline: {resolved_node_path}")
try:
function(**flattened_config)
except Exception as e:
print(f"Error running sync pipeline {resolved_node_path}: {e}")
# Optionally re-raise or handle error reporting
raise
print(f"Completed pipeline: {resolved_node_path}")
def main():
parser = argparse.ArgumentParser(description="Run Augmentoolkit pipelines.")
parser.add_argument(
"--node",
type=str,
help="Path (potentially aliased) to the pipeline node function (e.g., 'pipelines/my_pipeline.run').",
)
parser.add_argument(
"--config",
type=str,
help="Path (potentially aliased) to the pipeline-specific configuration YAML file.",
)
parser.add_argument(
"--override-json",
type=str,
help="JSON string of parameters to override pipeline config.",
)
args = parser.parse_args()
# Always load path aliases
# path_aliases are already loaded globally before main()
override_params = {}
if args.override_json:
try:
override_params = json.loads(args.override_json)
if not isinstance(override_params, dict):
print(
"Error: --override-json must be a valid JSON object (dictionary)."
)
sys.exit(1)
except json.JSONDecodeError as e:
print(f"Error parsing --override-json: {e}")
sys.exit(1)
# Decide execution path based on provided arguments
if args.node:
print(f"Running single pipeline specified via command line arguments:")
print(f" Node: {args.node}")
print(f" Config: {args.config}")
print(f" Overrides: {override_params}")
run_pipeline(args.node, args.config, override_params)
else:
print("Running pipelines defined in super_config.yaml 'pipeline_order'.")
# Run pipelines in order specified in super_config
pipelines_to_run = super_config.get("pipeline_order", [])
if not pipelines_to_run:
print(
"No pipelines specified in 'pipeline_order' and no specific pipeline provided via args. Exiting."
)
return
for pipeline in pipelines_to_run:
if not isinstance(pipeline, dict) or "node" not in pipeline:
print(
f"Warning: Skipping invalid pipeline entry in super_config: {pipeline}. Must be a dictionary with a 'node' key."
)
continue
# Merge super_config parameters with any CLI overrides (though CLI overrides usually imply single pipeline run)
pipeline_params = pipeline.get("parameters", {})
# Note: If running from super_config, CLI overrides are *not* typically used per-pipeline.
# The logic here prioritizes the CLI override if BOTH --override-json and super_config parameters exist,
# but this scenario is less common when running the whole sequence.
# If you intended CLI overrides to *only* apply when --node is used, keep override_params empty here.
# If CLI overrides should *globally* apply even to super_config runs, update pipeline_params:
# pipeline_params.update(override_params) # Uncomment this if CLI overrides should apply globally
print(f"Running pipeline from super_config:")
print(f" Node: {pipeline['node']}")
print(f" Config: {pipeline.get('config')}") # Use .get for optional config
print(f" Parameters: {pipeline_params}")
run_pipeline(
pipeline["node"], pipeline.get("config"), pipeline_params
) # Use .get for config key
if __name__ == "__main__":
main()
|