@@ -347,35 +347,74 @@ def _create_model_variant(
347347 repository and fills directory with config
348348 """
349349
350- variant_name = variant_config .variant_name
351350 if self ._config .triton_launch_mode != "remote" :
352- model_repository = self ._config .model_repository
351+ self ._create_non_remote_mode_model_variant (
352+ original_name , variant_config , ignore_first_config_variant
353+ )
354+ else :
355+ self ._create_remote_mode_model_variant (original_name , variant_config )
353356
354- original_model_dir = os .path .join (model_repository , original_name )
355- new_model_dir = os .path .join (self ._output_model_repo_path , variant_name )
356- try :
357- # Create the directory for the new model
358- os .makedirs (new_model_dir , exist_ok = True )
359- self ._first_config_variant .setdefault (original_name , None )
357+ def _create_non_remote_mode_model_variant (
358+ self ,
359+ original_name : str ,
360+ variant_config : ModelConfigVariant ,
361+ ignore_first_config_variant : bool = False ,
362+ ) -> None :
363+ """
364+ Creates a directory for the model config variant in the output model
365+ repository and fills directory with config
366+ """
367+ variant_name = variant_config .variant_name
368+ model_repository = self ._config .model_repository
369+
370+ original_model_dir = os .path .join (model_repository , original_name )
371+ new_model_dir = os .path .join (self ._output_model_repo_path , variant_name )
372+ try :
373+ # Create the directory for the new model
374+ os .makedirs (new_model_dir , exist_ok = True )
375+ self ._first_config_variant .setdefault (original_name , None )
376+
377+ if ignore_first_config_variant :
378+ variant_config .model_config .write_config_to_file (
379+ new_model_dir , original_model_dir , None
380+ )
381+ else :
382+ variant_config .model_config .write_config_to_file (
383+ new_model_dir ,
384+ original_model_dir ,
385+ self ._first_config_variant [original_name ],
386+ )
360387
361- if ignore_first_config_variant :
362- variant_config .model_config .write_config_to_file (
363- new_model_dir , original_model_dir , None
364- )
365- else :
366- variant_config .model_config .write_config_to_file (
367- new_model_dir ,
368- original_model_dir ,
369- self ._first_config_variant [original_name ],
370- )
388+ if self ._first_config_variant [original_name ] is None :
389+ self ._first_config_variant [original_name ] = os .path .join (
390+ self ._output_model_repo_path , variant_name
391+ )
392+ except FileExistsError :
393+ # Ignore if the file already exists
394+ pass
371395
372- if self ._first_config_variant [original_name ] is None :
373- self ._first_config_variant [original_name ] = os .path .join (
374- self ._output_model_repo_path , variant_name
375- )
376- except FileExistsError :
377- # Ignore if the file already exists
378- pass
396+ def _create_remote_mode_model_variant (
397+ self ,
398+ original_name : str ,
399+ variant_config : ModelConfigVariant ,
400+ ) -> None :
401+ """
402+ Creates a directory for the model config variant in the output model
403+ repository and fills directory with only the config.pbtxt
404+ """
405+ variant_name = variant_config .variant_name
406+ new_model_dir = os .path .join (self ._output_model_repo_path , variant_name )
407+ try :
408+ os .makedirs (new_model_dir , exist_ok = False )
409+ self ._first_config_variant .setdefault (original_name , None )
410+ variant_config .model_config .write_config_to_file (
411+ model_path = new_model_dir ,
412+ src_model_path = new_model_dir ,
413+ first_variant_model_path = None ,
414+ )
415+ except FileExistsError :
416+ # Ignore if the dir already exists
417+ pass
379418
380419 def _load_model_variants (self , run_config ):
381420 """
0 commit comments