Skip to content

Commit 32389de

Browse files
authored
Output config.pbtxt when in remote mode (#784)
1 parent b1431e0 commit 32389de

File tree

1 file changed

+64
-25
lines changed

1 file changed

+64
-25
lines changed

model_analyzer/record/metrics_manager.py

Lines changed: 64 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)