move models dir inside pg\

levkk · levkk · commit 74bfa5e64328 · 2022-04-08T13:27:48.000-07:00
diff --git a/README.md b/README.md
@@ -71,13 +71,6 @@ sudo python3 setup.py install
 cd ../
 ```
 
-Create a directory for PostgresML to store its models (this is currently hardcoded):
-
-```
-mkdir -p /app/models
-chown postgres:postgres /app/models
-```
-
 Run the test:
 
 ```
diff --git a/pgml/pgml/score.py b/pgml/pgml/score.py
@@ -0,0 +1,17 @@
+"""Score"""
+
+import os
+import pickle
+
+from pgml.exceptions import PgMLException
+
+
+def load(name, source):
+    """Load a model from file."""
+    path = os.path.join(source, name)
+
+    if not os.path.exists(path):
+        raise PgMLException(f"Model source directory `{path}` does not exist.")
+
+    with open(path, "rb") as f:
+        return pickle.load(f)
diff --git a/pgml/pgml/sql.py b/pgml/pgml/sql.py
@@ -1,5 +1,6 @@
 """Tools to run SQL.
 """
+import os
 
 
 def all_rows(cursor):
@@ -11,3 +12,21 @@ def all_rows(cursor):
 
         for row in rows:
             yield row
+
+
+def models_directory(plpy):
+    """Get the directory where we store our models."""
+    data_directory = plpy.execute(
+        """
+        SELECT setting FROM pg_settings WHERE name = 'data_directory'
+    """,
+        1,
+    )[0]["setting"]
+
+    models_dir = os.path.join(data_directory, "pgml_models")
+
+    # TODO: Ideally this happens during extension installation.
+    if not os.path.exists(models_dir):
+        os.mkdir(models_dir, 0o770)
+
+    return models_dir
diff --git a/pgml/pgml/train.py b/pgml/pgml/train.py
@@ -8,13 +8,14 @@
 from sklearn.metrics import mean_squared_error, r2_score
 
 import pickle
+import os
 
 from pgml.sql import all_rows
 from pgml.exceptions import PgMLException
 from pgml.validate import check_type
 
 
-def train(cursor, y_column, name, save=True):
+def train(cursor, y_column, name, save=True, destination="/tmp/pgml_models"):
     """Train the model on data on some rows.
 
     Arguments:
@@ -65,7 +66,8 @@ def train(cursor, y_column, name, save=True):
     msq = mean_squared_error(y_test, y_pred)
     r2 = r2_score(y_test, y_pred)
 
-    path = f"/app/models/{name}"
+    path = os.path.join(destination, name)
+
     if save:
         with open(path, "wb") as f:
             pickle.dump(lr, f)
diff --git a/sql/install.sql b/sql/install.sql
@@ -21,6 +21,7 @@ CREATE TABLE pgml.model_versions(
 	name VARCHAR,
 	location VARCHAR NULL,
 	data_source TEXT,
+	y_column VARCHAR,
 	started_at TIMESTAMP WITHOUT TIME ZONE DEFAULT CURRENT_TIMESTAMP,
 	ended_at TIMESTAMP WITHOUT TIME ZONE NULL,
 	mean_squared_error DOUBLE PRECISION,
@@ -50,29 +51,34 @@ CREATE OR REPLACE FUNCTION pgml_train(table_name TEXT, y TEXT)
 RETURNS TEXT
 AS $$
 	from pgml.train import train
+	from pgml.sql import models_directory
+	import os
 
 	data_source = f"SELECT * FROM {table_name}"
 
 	# Start training.
 	start = plpy.execute(f"""
 		INSERT INTO pgml.model_versions
-			(name, data_source)
-		VALUES ('{table_name}', '{data_source}')
+			(name, data_source, y_column)
+		VALUES
+			('{table_name}', '{data_source}', '{y}')
 		RETURNING *""", 1)
 
 	id_ = start[0]["id"]
 	name = f"{table_name}_{id_}"
 
+	destination = models_directory(plpy)
+
 	# Train!
-	location, msq, r2 = train(plpy.cursor(data_source), y_column=y, name=name)
+	location, msq, r2 = train(plpy.cursor(data_source), y_column=y, name=name, destination=destination)
 
 	plpy.execute(f"""
 		UPDATE pgml.model_versions
 		SET location = '{location}',
 			successful = true,
-			ended_at = NOW(),
 			mean_squared_error = '{msq}',
-			r2_score = '{r2}'
+			r2_score = '{r2}',
+			ended_at = clock_timestamp()
 		WHERE id = {id_}""")
 
 	return name
@@ -86,14 +92,15 @@ DROP FUNCTION pgml_score(model_name TEXT, VARIADIC features DOUBLE PRECISION[]);
 CREATE OR REPLACE FUNCTION pgml_score(model_name TEXT, VARIADIC features DOUBLE PRECISION[])
 RETURNS DOUBLE PRECISION
 AS $$
+	from pgml.sql import models_directory
+	from pgml.score import load
 	import pickle
 
 	if model_name in SD:
 		model = SD[model_name]
 	else:
-		with open(f"/app/models/{model_name}", "rb") as f:
-		    model = pickle.load(f)
-		    SD[model_name] = model
+		SD[model_name] = load(model_name, models_directory(plpy))
+		model = SD[model_name]
 
 	scores = model.predict([features,])
 	return scores[0]
