Add alpha preview feature into compiler (kubeflow#991)

* add alpha preview feature into compiler * address comments * fix typo
red-hat-data-services · Jul 18, 2022 · 7d9f8c9 · 7d9f8c9
1 parent 852b9ac
commit 7d9f8c9
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/sdk/python/kfp_tekton/compiler/_data_passing_rewriter.py b/sdk/python/kfp_tekton/compiler/_data_passing_rewriter.py
@@ -634,13 +634,20 @@ def append_taskrun_params(task_name_append: str):
                         dst = '$(results.%s.path)' % sanitize_k8s_name(result['name'])
                         if artifact_name == result['name'] and src != dst:
                             add_copy_results_artifacts_step = True
+                            total_size_command = 'ARTIFACT_SIZE=`wc -c %s | awk \'{print $1}\'`\n' % src + \
+                                                 'TOTAL_SIZE=$( expr $TOTAL_SIZE + $ARTIFACT_SIZE)\n'
+                            copy_command = '    cp ' + src + ' ' + dst + '\n'
+                            if env.get('OUTPUT_PREVIEW', 'false').lower() == 'true':
+                                preview_size = env.get('OUTPUT_PREVIEW_SIZE', '100')
+                                total_size_command = 'TOTAL_SIZE=$( expr $TOTAL_SIZE + %s)\n' % preview_size
+                                copy_command = '    dd if=' + src + ' of=' + \
+                                               dst + ' bs=' + preview_size + ' count=1\n'
                             script += (
-                                    'ARTIFACT_SIZE=`wc -c %s | awk \'{print $1}\'`\n' % src +
-                                    'TOTAL_SIZE=$( expr $TOTAL_SIZE + $ARTIFACT_SIZE)\n' +
+                                    total_size_command +
                                     'touch ' + dst + '\n' +  # create an empty file by default.
                                     'if [[ $TOTAL_SIZE -lt 3072 ]]; then\n' +
                                     '  if ! awk "/[^[:print:]]/{f=1} END{exit !f}" %s; then\n' % src +
-                                    '    cp ' + src + ' ' + dst + '\n' +
+                                    copy_command +
                                     '  fi\n'
                                     'fi\n'
                             )