progress

pkg/operator/context/models.go

@@ -85,7 +85,6 @@ func getModels(
 		datasetIDWithTags := hash.Bytes(buf.Bytes())
 
 		datasetRoot := filepath.Join(root, consts.TrainingDataDir, datasetID)
-
 		trainingDatasetName := strings.Join([]string{
 			modelConfig.Name,
 			resource.TrainingDatasetType.String(),
@@ -97,7 +96,7 @@ func getModels(
 					ID:           modelID,
 					IDWithTags:   modelID,
 					ResourceType: resource.ModelType,
-					MetadataKey:  filepath.Join(datasetRoot, "metadata.json"),
+					MetadataKey:  filepath.Join(root, consts.ModelsDir, modelID+"_metadata.json"),
 				},
 			},
 			Model:   modelConfig,
@@ -115,6 +114,7 @@ func getModels(
 						ID:           datasetID,
 						IDWithTags:   datasetIDWithTags,
 						ResourceType: resource.TrainingDatasetType,
+						MetadataKey:  filepath.Join(datasetRoot, "metadata.json"),
 					},
 				},
 				ModelName: modelConfig.Name,

pkg/workloads/lib/context.py

@@ -479,6 +479,17 @@ def update_metadata(self, metadata, context_key, context_item=""):
         self.ctx[context_key][context_item]["metadata"] = metadata
         self.storage.put_json(metadata, self.ctx[context_key][context_item]["metadata_key"])
 
+    def get_metadata(self, context_key, context_item, use_cache=True):
+        if use_cache and self.ctx[context_key][context_item]["metadata"]:
+            return self.ctx[context_key][context_item]["metadata"]
+
+        metadata_uri = self.ctx[context_key][context_item]["metadata_key"]
+        metadata = self.storage.get_json(metadata_uri, allow_missing=True)
+        self.ctx[context_key][context_item]["metadata"] = metadata
+        return metadata
+
+
+
     def fetch_metadata(self):
         resources = [
             "python_packages",
@@ -499,6 +510,13 @@ def fetch_metadata(self):
                     metadata = {}
                 self.ctx[resource][k]["metadata"] = metadata
 
+        # fetch dataset metadata for models
+        for k, v in self.ctx["models"].items():
+            metadata = self.storage.get_json(v["dataset"]["metadata_key"], allow_missing=True)
+            if not metadata:
+                metadata = {}
+            self.ctx["models"][k]["dataset"]["metadata"] = metadata
+
         metadata = self.storage.get_json(self.raw_dataset["metadata_key"], allow_missing=True)
         if not metadata:
             metadata = {}

pkg/workloads/tf_train/train_util.py

@@ -149,9 +149,10 @@ def train(model_name, model_impl, ctx, model_dir):
     exporter = tf.estimator.FinalExporter("estimator", serving_input_fn, as_text=False)
 
     train_num_steps = model["training"]["num_steps"]
+    dataset_metadata = model["dataset"]["metadata"]
     if model["training"]["num_epochs"]:
         train_num_steps = (
-            math.ceil(model["metadata"]["training_size"] / float(model["training"]["batch_size"]))
+            math.ceil(dataset_metadata["training_size"] / float(model["training"]["batch_size"]))
             * model["training"]["num_epochs"]
         )
 
@@ -160,7 +161,7 @@ def train(model_name, model_impl, ctx, model_dir):
     eval_num_steps = model["evaluation"]["num_steps"]
     if model["evaluation"]["num_epochs"]:
         eval_num_steps = (
-            math.ceil(model["metadata"]["training_size"] / float(model["evaluation"]["batch_size"]))
+            math.ceil(model["metadata"]["eval_size"] / float(model["evaluation"]["batch_size"]))
             * model["evaluation"]["num_epochs"]
         )