From 96d6b63857eae06329e5ad38c4369f7189076ce7 Mon Sep 17 00:00:00 2001 From: "Ukjae Jeong (Jay)" Date: Sat, 26 Aug 2023 02:59:08 +0900 Subject: [PATCH] support remote storage in modelcardgenerator (#262) --- tfx_addons/model_card_generator/executor.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tfx_addons/model_card_generator/executor.py b/tfx_addons/model_card_generator/executor.py index 1e980a0b..d97a8345 100644 --- a/tfx_addons/model_card_generator/executor.py +++ b/tfx_addons/model_card_generator/executor.py @@ -18,6 +18,7 @@ ModelCardGenerator. """ +import tempfile from typing import Any, Dict, List, Optional from model_card_toolkit import core @@ -25,6 +26,7 @@ from tfx import types from tfx.dsl.components.base.base_executor import BaseExecutor from tfx.types import artifact_utils, standard_component_specs +from tfx.utils import io_utils _DEFAULT_MODEL_CARD_FILE_NAME = 'model_card.html' @@ -128,14 +130,16 @@ def Do(self, input_dict: Dict[str, List[types.Artifact]], metrics_include. """ + # Make local temp directory to support remote storage (gcs, s3, ...) + temp_output_dir = tempfile.mkdtemp() + # Initialize ModelCardToolkit mct = core.ModelCardToolkit(source=src.Source( tfma=self._tfma_source(input_dict, exec_properties), tfdv=self._tfdv_source(input_dict, exec_properties), model=self._model_source(input_dict), ), - output_dir=artifact_utils.get_single_instance( - output_dict['model_card']).uri) + output_dir=temp_output_dir) template_io = exec_properties.get('template_io') or [ (mct.default_template, _DEFAULT_MODEL_CARD_FILE_NAME) ] @@ -144,3 +148,7 @@ def Do(self, input_dict: Dict[str, List[types.Artifact]], mct.scaffold_assets(json=exec_properties.get('json')) for template_path, output_file in template_io: mct.export_format(template_path=template_path, output_file=output_file) + + # Copy all files to output_dir + output_dir = artifact_utils.get_single_uri(output_dict["model_card"]) + io_utils.copy_dir(src=temp_output_dir, dst=output_dir)