update tank and add validation distinction for uploaded names (#548)

This commit makes minor changes to the model testing and base models, so
that there is a difference between validated and failing mlir names in
the turbinetank.

Author: saienduri

models/turbine_models/custom_models/sd_inference/clip.py

@@ -138,14 +138,16 @@ def main(self, inp=AbstractTensor(1, input_len, dtype=torch.int64)):
             f.write(module_str)
         model_name_upload = hf_model_name.replace("/", "_")
         model_name_upload += "-clip"
-        turbine_tank.uploadToBlobStorage(
+        blob_name = turbine_tank.uploadToBlobStorage(
             str(os.path.abspath(f"{safe_name}.mlir")),
             f"{model_name_upload}/{model_name_upload}.mlir",
         )
     if compile_to != "vmfb":
         return module_str, tokenizer
     else:
         utils.compile_to_vmfb(module_str, device, target_triple, max_alloc, safe_name)
+        if upload_ir:
+            return blob_name
 
 
 if __name__ == "__main__":

models/turbine_models/custom_models/sd_inference/schedulers.py

@@ -153,14 +153,16 @@ def main(
             f.write(module_str)
         model_name_upload = hf_model_name.replace("/", "-")
         model_name_upload = model_name_upload + "_scheduler"
-        turbine_tank.uploadToBlobStorage(
+        blob_name = turbine_tank.uploadToBlobStorage(
             str(os.path.abspath(f"{safe_name}.mlir")),
             f"{model_name_upload}/{model_name_upload}.mlir",
         )
     if compile_to != "vmfb":
         return module_str
     else:
         utils.compile_to_vmfb(module_str, device, target_triple, max_alloc, safe_name)
+        if upload_ir:
+            return blob_name
 
 
 if __name__ == "__main__":

models/turbine_models/custom_models/sd_inference/unet.py

@@ -132,14 +132,16 @@ def main(
             f.write(module_str)
         model_name_upload = hf_model_name.replace("/", "-")
         model_name_upload += "_unet"
-        turbine_tank.uploadToBlobStorage(
+        blob_name = turbine_tank.uploadToBlobStorage(
             str(os.path.abspath(f"{safe_name}.mlir")),
             f"{model_name_upload}/{model_name_upload}.mlir",
         )
     if compile_to != "vmfb":
         return module_str
     else:
         utils.compile_to_vmfb(module_str, device, target_triple, max_alloc, safe_name)
+        if upload_ir:
+            return blob_name
 
 
 if __name__ == "__main__":

models/turbine_models/custom_models/sd_inference/utils.py

@@ -83,7 +83,6 @@ def compile_to_vmfb(module_str, device, target_triple, max_alloc, safe_name):
     with open(f"{safe_name}.vmfb", "wb+") as f:
         f.write(flatbuffer_blob)
     print("Saved to", safe_name + ".vmfb")
-    exit()
 
 
 def create_safe_name(hf_model_name, model_name_str):

models/turbine_models/custom_models/sd_inference/vae.py

@@ -120,14 +120,16 @@ def main(self, inp=AbstractTensor(*sample, dtype=torch.float32)):
             f.write(module_str)
         model_name_upload = hf_model_name.replace("/", "_")
         model_name_upload = model_name_upload + "-vae-" + variant
-        turbine_tank.uploadToBlobStorage(
+        blob_name = turbine_tank.uploadToBlobStorage(
             str(os.path.abspath(f"{safe_name}.mlir")),
             f"{model_name_upload}/{model_name_upload}.mlir",
         )
     if compile_to != "vmfb":
         return module_str
     else:
         utils.compile_to_vmfb(module_str, device, target_triple, max_alloc, safe_name)
+        if upload_ir:
+            return blob_name
 
 
 if __name__ == "__main__":

models/turbine_models/custom_models/stateless_llama.py

@@ -416,7 +416,7 @@ def evict_kvcache_space(self):
         with open(f"{safe_name}.mlir", "w+") as f:
             f.write(module_str)
         model_name_upload = hf_model_name.replace("/", "_")
-        turbine_tank.uploadToBlobStorage(
+        blob_name = turbine_tank.uploadToBlobStorage(
             str(os.path.abspath(f"{safe_name}.mlir")),
             f"{model_name_upload}/{model_name_upload}.mlir",
         )
@@ -478,6 +478,8 @@ def evict_kvcache_space(self):
         with open(vmfb_path, "wb+") as f:
             f.write(flatbuffer_blob)
         print("saved to ", safe_name + ".vmfb")
+        if upload_ir:
+            return blob_name
         return module_str, tokenizer