Update backend.py

backend.py

@@ -1,4 +1,4 @@
-# backend.py — USERNAME/PASSWORD & PARALLEL PROCESSING VERSION
+# backend.py — FINAL VERSION
 import sqlite3
 import threading
 import time
@@ -21,7 +21,7 @@ scheduler_lock = threading.Lock()
 RUN_TIMEOUT = 48 * 3600
 MAX_RAM_PER_RUN_GB = 1.5
 
-# ------------------------------ DATABASE (NEW SCHEMA) ------------------------------
+# ------------------------------ DATABASE ------------------------------
 
 def init_db():
     conn = sqlite3.connect(DB_PATH, check_same_thread=False)
@@ -66,7 +66,7 @@ def get_user_by_username(username):
     rows, _ = db_query("SELECT id, password_hash FROM users WHERE username = ?", (username,))
     return rows[0] if rows else None
 
-# ------------------------------ NEW AUTHENTICATION ------------------------------
+# ------------------------------ AUTHENTICATION ------------------------------
 
 def signup_user(username, password):
     if not username or not password:
@@ -81,11 +81,10 @@ def signup_user(username, password):
 def login_user(username, password):
     user = get_user_by_username(username)
     if user and check_password_hash(user[1], password):
-        user_id = user[0]
-        return user_id, f"Welcome back, {username}!"
+        return user[0], f"Welcome back, {username}!"
     return None, "Invalid username or password."
 
-# ------------------------------ PARALLEL TRAINING QUEUE ------------------------------
+# ------------------------------ TRAINING QUEUE & SCHEDULER ------------------------------
 
 def ram_available():
     return (psutil.virtual_memory().available / (1024**3)) >= MAX_RAM_PER_RUN_GB
@@ -93,7 +92,6 @@ def ram_available():
 def queue_training_run(user_id, config):
     _, run_id = db_query("INSERT INTO training_runs (user_id, arch_type, num_layers, learning_rate, epochs, batch_size) VALUES (?, ?, ?, ?, ?, ?)", (user_id, config['arch_type'], config['num_layers'], config['learning_rate'], config['epochs'], config['batch_size']))
     training_queue.append({"run_id": run_id, "user_id": user_id, **config})
-    # Trigger the scheduler every time a new job is added
     start_training_if_free()
     return run_id
 
@@ -131,8 +129,9 @@ def get_user_runs(user_id):
     rows, _ = db_query("SELECT id, arch_type, num_layers, status, started_at FROM training_runs WHERE user_id = ? ORDER BY id DESC", (user_id,))
     return rows
 
-def get_run_logs(run_id):
-    rows, _ = db_query("SELECT logs, status FROM training_runs WHERE id = ?", (run_id,))
+def get_run_logs(user_id, run_id):
+    """Securely fetches logs by checking ownership (user_id)."""
+    rows, _ = db_query("SELECT logs, status FROM training_runs WHERE id = ? AND user_id = ?", (run_id, user_id))
     return rows[0] if rows else ("", "unknown")
 
 def update_run_status(run_id, status):
@@ -273,8 +272,6 @@ def run_inference(run_id, prompt):
         generated_ids = torch.argmax(logits, dim=-1)
         return f"🧑‍🍳 Model says:\n{tokenizer.decode(generated_ids[0], skip_special_tokens=True)}"
 
-# ------------------------------ PUBLISH (HF Token passed as argument) ------------------------------
-
 def publish_run_to_hub(run_id, hf_token, repo_name, user_description=""):
     try:
         user_info = whoami(token=hf_token)