rollbacks

levkk · levkk · commit c24419eb3ab8 · 2022-04-16T11:48:51.000-07:00
diff --git a/pgml/pgml/model.py b/pgml/pgml/model.py
@@ -440,13 +440,23 @@ def train(
         )
 
     snapshot = Snapshot.create(relation_name, y_column_name, test_size, test_sampling)
-    best_model = None
-    best_error = None
+    deployed = Model.find_deployed(project.id)
+
+    # Let's assume that the deployed model is better for now.
+    best_model = deployed
+    best_error = best_model.mean_squared_error if best_model else None
 
     for algorithm_name in algorithms:
         model = Model.create(project, snapshot, algorithm_name)
         model.fit(snapshot)
+
+        # Find the better model and deploy that.
         if best_error is None or model.mean_squared_error < best_error:
             best_error = model.mean_squared_error
             best_model = model
-    best_model.deploy()
+
+    if deployed and deployed.id == best_model.id:
+        return "rolled back"
+    else:
+        best_model.deploy()
+        return "deployed"
diff --git a/sql/install.sql b/sql/install.sql
@@ -109,9 +109,9 @@ RETURNS TABLE(project_name TEXT, objective TEXT, status TEXT)
 AS $$
 	from pgml.model import train
 
-	train(project_name, objective, relation_name, y_column_name)
+	status = train(project_name, objective, relation_name, y_column_name)
 
-	return [(project_name, objective, "deployed")]
+	return [(project_name, objective, status)]
 $$ LANGUAGE plpython3u;
 
 ---
