Set all TensorFlow version directories to "1" (#560)

(cherry picked from commit 0147e52)

Author: deliahu

pkg/operator/workloads/api_workload.go

@@ -234,10 +234,11 @@ func (aw *APIWorkload) IsFailed(ctx *context.Context) (bool, error) {
 }
 
 type downloadContainerArg struct {
-	From     string `json:"from"`
-	To       string `json:"to"`
-	Unzip    bool   `json:"unzip"`
-	ItemName string `json:"item_name"` // name of the item being downloaded, just for logging (if "" nothing will be logged)
+	From                 string `json:"from"`
+	To                   string `json:"to"`
+	Unzip                bool   `json:"unzip"`
+	ItemName             string `json:"item_name"`               // name of the item being downloaded, just for logging (if "" nothing will be logged)
+	TFModelVersionRename string `json:"tf_model_version_rename"` // e.g. passing in /mnt/model/1 will rename /mnt/model/* to /mnt/model/1 only if there is one item in /mnt/model/
 }
 
 func tfAPISpec(
@@ -269,10 +270,11 @@ func tfAPISpec(
 
 	downloadArgs := []downloadContainerArg{
 		{
-			From:     ctx.APIs[api.Name].TensorFlow.Model,
-			To:       path.Join(consts.EmptyDirMountPath, "model"),
-			Unzip:    strings.HasSuffix(ctx.APIs[api.Name].TensorFlow.Model, ".zip"),
-			ItemName: "model",
+			From:                 ctx.APIs[api.Name].TensorFlow.Model,
+			To:                   path.Join(consts.EmptyDirMountPath, "model"),
+			Unzip:                strings.HasSuffix(ctx.APIs[api.Name].TensorFlow.Model, ".zip"),
+			ItemName:             "model",
+			TFModelVersionRename: path.Join(consts.EmptyDirMountPath, "model", "1"),
 		},
 	}
 

pkg/workloads/cortex/downloader/download.py

@@ -37,13 +37,21 @@ def start(args):
             logger.info("downloading {} from {}".format(item_name, from_path))
         s3_client.download(prefix, to_path)
 
-        if download_arg["unzip"]:
+        if download_arg.get("unzip", False):
             if item_name != "":
                 logger.info("unzipping {}".format(item_name))
             util.extract_zip(
                 os.path.join(to_path, os.path.basename(from_path)), delete_zip_file=True
             )
 
+        if download_arg.get("tf_model_version_rename", "") != "":
+            dest = util.trim_suffix(download_arg["tf_model_version_rename"], "/")
+            dir_path = os.path.dirname(dest)
+            entries = os.listdir(dir_path)
+            if len(entries) == 1:
+                src = os.path.join(dir_path, entries[0])
+                os.rename(src, dest)
+
 
 def main():
     parser = argparse.ArgumentParser()