@@ -32,7 +32,8 @@ def transform_model(model: dict, s3_prefix: str, fixtures_prefix: str) -> dict:
3232 m = dict (model )
3333
3434 if "s3_model" in m :
35- m ["s3_path" ] = f"{ s3_prefix } { m .pop ('s3_model' )}
35+ prefix = m .pop ("s3_prefix" , s3_prefix )
36+ m ["s3_path" ] = f"{ prefix } { m .pop ('s3_model' )}
3637 m ["model_source" ] = "s3"
3738 elif "hf_model" in m :
3839 m ["model_source" ] = "hf"
@@ -62,35 +63,51 @@ def load_yaml(path: str) -> dict:
6263 return json .loads (result .stdout )
6364
6465
65- def matches_customer_type (model : dict , customer_type : str ) -> bool :
66- """A model runs on a config unless it pins a different customer_type.
66+ def _flatten_image_config (image_cfg : dict ) -> dict [str , str ]:
67+ """Flatten nested image config into a single-level dict of string values."""
68+ flat = {}
69+ for section in image_cfg .values ():
70+ if isinstance (section , dict ):
71+ for k , v in section .items ():
72+ flat [k ] = str (v )
73+ return flat
6774
68- Models without a ``customer_type`` field run everywhere (backward
69- compatible). A model that pins e.g. ``customer_type: sagemaker`` only runs
70- when the config's customer type matches — used to gate tests for features
71- that exist on only one container variant (e.g. the SageMaker routing
72- middleware, which adds the JSON->multipart video path absent on EC2).
75+
76+ def _model_matches_image (model : dict , image_fields : dict [str , str ]) -> bool :
77+ """Check if a model's required_image_pattern matches the image config fields.
78+
79+ required_image_pattern can be:
80+ - A dict of {field: value} pairs that must all match in the image config.
81+ - A plain string that must appear in at least one image config value.
82+ - Absent/None → model runs on all images.
7383 """
74- pinned = model .get ("customer_type " )
75- if not pinned or not customer_type :
84+ pattern = model .get ("required_image_pattern " )
85+ if not pattern :
7686 return True
77- return pinned == customer_type
87+ if isinstance (pattern , dict ):
88+ return all (image_fields .get (k ) == str (v ) for k , v in pattern .items ())
89+ return str (pattern ) in image_fields .values ()
7890
7991
8092def parse_config (
81- config_path : str , section : str , runner_type : str , customer_type : str = ""
93+ config_path : str , section : str , runner_type : str , image_config_path : str = ""
8294) -> dict [str , str ]:
8395 cfg = load_yaml (config_path )
8496
8597 s3_prefix = cfg .get ("s3_prefix" , "" )
8698 fixtures_prefix = cfg .get ("test_fixtures_prefix" , "" )
8799
100+ image_fields = {}
101+ if image_config_path and os .path .isfile (image_config_path ):
102+ image_fields = _flatten_image_config (load_yaml (image_config_path ))
103+
88104 results = {}
89105 types = ["codebuild-fleet" , "runner-scale-sets" ] if runner_type == "all" else [runner_type ]
90106
91107 for rt in types :
92108 models = cfg .get (section , {}).get (rt , []) or []
93- models = [m for m in models if matches_customer_type (m , customer_type )]
109+ if image_fields :
110+ models = [m for m in models if _model_matches_image (m , image_fields )]
94111 transformed = [transform_model (m , s3_prefix , fixtures_prefix ) for m in models ]
95112 key = rt if runner_type == "all" else "matrix"
96113 results [key ] = json .dumps (transformed , separators = ("," , ":" ))
@@ -110,18 +127,17 @@ def main():
110127 help = "Runner type: all, codebuild-fleet, or runner-scale-sets" ,
111128 )
112129 parser .add_argument (
113- "--customer-type " ,
130+ "--image-config " ,
114131 default = "" ,
115- help = "Config customer type (e.g. ec2, sagemaker). When set, drops models "
116- "that pin a different customer_type. Empty = include all models." ,
132+ help = "Path to image config YAML; used to filter models by required_image_pattern" ,
117133 )
118134 args = parser .parse_args ()
119135
120136 if not os .path .isfile (args .config ):
121137 print (f"ERROR: Config file not found: { args .config } , file = sys .stderr )
122138 sys .exit (1 )
123139
124- results = parse_config (args .config , args .section , args .runner_type , args .customer_type )
140+ results = parse_config (args .config , args .section , args .runner_type , args .image_config or "" )
125141
126142 output_file = os .environ .get ("GITHUB_OUTPUT" )
127143 if output_file :
0 commit comments