Add upload progress and BPS tracking to upload.py

emblem5/ai/upload.py

@@ -2,6 +2,7 @@
 import oss2
 import sys
 import os
+import time
 
 oss_ak = 'LTAI5tC2qXGxwHZUZP7DoD1A'
 oss_sk = 'qPo9O6ZvEfqo4t8oflGEm0DoxLHJhm'
@@ -10,6 +11,25 @@ auth = oss2.Auth(oss_ak, oss_sk)
 endpoint = 'https://oss-rg-china-mainland.aliyuncs.com'
 bucket = oss2.Bucket(auth, endpoint, 'emblem-models')
 
-for x in sys.argv[1:]:
+class ProgressCallback:
-    bucket.put_object_from_file(os.path.basename(x), x)
+    def __init__(self, filename, total_size):
+        self.filename = filename
+        self.total_size = total_size
+        self.consumed = 0
+        self.start_time = time.time()
+        
+    def __call__(self, consumed, total):
+        self.consumed = consumed
+        elapsed = time.time() - self.start_time
+        if elapsed > 0:
+            bps = consumed / elapsed
+            percent = 100 * consumed / total if total else 0
+            print(f"\r{self.filename}: {percent:.1f}% ({consumed}/{total} bytes, {bps:.0f} B/s)", end='', flush=True)
+
+for x in sys.argv[1:]:
+    file_size = os.path.getsize(x)
+    filename = os.path.basename(x)
+    print(f"Uploading {filename} ({file_size} bytes)...")
+    bucket.put_object_from_file(filename, x, progress_callback=ProgressCallback(filename, file_size))
+    print()  # New line after progress completes