added required_disk_space and warning to variations

ctuning · Feb 29, 2024 · b72783f · b72783f
1 parent 974e27c
commit b72783f
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 5 deletions.
diff --git a/cm-mlops/automation/script/module.py b/cm-mlops/automation/script/module.py
@@ -1710,6 +1710,11 @@ def _update_state_from_variations(self, i, meta, variation_tags, variations, env
         import copy
         explicit_variation_tags=copy.deepcopy(variation_tags)
 
+        # Calculate space
+        required_disk_space = {}
+
+        # Check if warning
+        warnings = []
 
         # variation_tags get appended by any aliases
         r = self._get_variations_with_aliases(variation_tags, variations)
@@ -1785,6 +1790,7 @@ def _update_state_from_variations(self, i, meta, variation_tags, variations, env
                 if variation_tag.startswith('~'):
                     # ignore such tag (needed for caching only to differentiate variations)
                     continue
+
                 if variation_tag.startswith('-'):
                     # ignore such tag (needed for caching only to eliminate variations)
                     continue
@@ -1803,14 +1809,19 @@ def _update_state_from_variations(self, i, meta, variation_tags, variations, env
                 if variation_tag_dynamic_suffix:
                     self._update_variation_meta_with_dynamic_suffix(variation_meta, variation_tag_dynamic_suffix)
 
-
-
                 r = update_state_from_meta(variation_meta, env, state, deps, post_deps, prehook_deps, posthook_deps, new_env_keys_from_meta, new_state_keys_from_meta, i)
                 if r['return']>0: return r
 
                 if variation_meta.get('script_name', '')!='':
                     meta['script_name'] = variation_meta['script_name']
 
+                if variation_meta.get('required_disk_space', 0) > 0 and variation_tag not in required_disk_space:
+                    required_disk_space[variation_tag] = variation_meta['required_disk_space']
+
+                if variation_meta.get('warning', '') != '':
+                    x = variation_meta['warning']
+                    if x not in warnings: warnings.append()
+
                 adr=get_adr(variation_meta)
                 if adr:
                     self._merge_dicts_with_tags(add_deps_recursive, adr)
@@ -1839,11 +1850,30 @@ def _update_state_from_variations(self, i, meta, variation_tags, variations, env
                         if combined_variation_meta.get('script_name', '')!='':
                             meta['script_name'] = combined_variation_meta['script_name']
 
+                        if combined_variation_meta.get('required_disk_space', 0) > 0 and combined_variation not in required_disk_space:
+                            required_disk_space[combined_variation] = combined_variation_meta['required_disk_space']
+
+                        if combined_variation_meta.get('warning', '') != '':
+                            x = combined_variation_meta['warning']
+                            if x not in warnings: warnings.append(x)
 
             #Processing them again using updated deps for add_deps_recursive
             r = update_adr_from_meta(deps, post_deps, prehook_deps, posthook_deps, add_deps_recursive)
             if r['return']>0: return r
 
+        if len(required_disk_space)>0:
+            required_disk_space_sum_mb = sum(list(required_disk_space.values()))
+
+            warnings.append('Required disk space: {} MB'.format(required_disk_space_sum_mb))
+
+        if len(warnings)>0:
+            print ('=================================================')
+            print ('WARNINGS:')
+            print ('')
+            for w in warnings:
+                print ('  '+w)
+            print ('=================================================')
+
         return {'return': 0, 'variation_tags_string': variation_tags_string, 'explicit_variation_tags': explicit_variation_tags}
 
     ######################################################################################

diff --git a/cm-mlops/automation/utils/module.py b/cm-mlops/automation/utils/module.py
@@ -206,6 +206,9 @@ def download_file(self, i):
         if os.path.isfile(path_to_file):
             os.remove(path_to_file)
 
+        print ('Downloading to {}'.format(path_to_file))
+        print ('')
+
         # Download
         size = -1
         downloaded = 0

diff --git a/cm-mlops/script/get-ml-model-retinanet/_cm.json b/cm-mlops/script/get-ml-model-retinanet/_cm.json
@@ -76,14 +76,18 @@
     "onnx,fp32": {
       "env": {
         "CM_PACKAGE_URL": "https://zenodo.org/record/6617879/files/resnext50_32x4d_fpn.onnx",
+        "CM_DOWNLOAD_CHECKSUM": "4544f4e56e0a4684215831cc937ea45c",
         "CM_ML_MODEL_ACCURACY": "0.3757"
-      }
+      },
+      "required_disk_space": 150,
+      "warning":"This model is downloaded from Zenodo.org"
     },
     "pytorch,fp32": {
       "env": {
         "CM_PACKAGE_URL": "https://zenodo.org/record/6617981/files/resnext50_32x4d_fpn.pth",
         "CM_ML_MODEL_ACCURACY": "0.3755"
       }
+
     },
     "weights": {
       "env": {
@@ -166,6 +170,7 @@
     }
   },
   "print_env_at_the_end" : {
-    "CM_ML_MODEL_FILE_WITH_PATH": "Path to the ML model"
+    "CM_ML_MODEL_FILE_WITH_PATH": "Path to the ML model",
+    "CM_ML_MODEL_ACCURACY": "Model accuracy"
   }
 }
diff --git a/cm-mlops/script/get-ml-model-tiny-resnet/_cm.json b/cm-mlops/script/get-ml-model-tiny-resnet/_cm.json
@@ -100,7 +100,8 @@
     },
     "tflite,int8": {
       "env": {
-        "CM_PACKAGE_URL": "https://github.com/mlcommons/tiny/raw/master/benchmark/training/image_classification/trained_models/pretrainedResnet_quant.tflite"
+        "CM_PACKAGE_URL": "https://github.com/mlcommons/tiny/raw/master/benchmark/training/image_classification/trained_models/pretrainedResnet_quant.tflite",
+        "CM_DOWNLOAD_CHECKSUM": "2d6dd48722471313e4c4528249205ae3"
       }
     },
     "onnx": {