Feat: Train models on remote model servers

traindb-catalog/src/main/java/traindb/catalog/CatalogContext.java

@@ -26,6 +26,7 @@
 import traindb.catalog.pm.MSynopsis;
 import traindb.catalog.pm.MTable;
 import traindb.catalog.pm.MTask;
+import traindb.catalog.pm.MTrainingStatus;
 
 public interface CatalogContext {
 
@@ -62,6 +63,11 @@ Collection<MModel> getInferenceModels(String baseSchema, String baseTable)
 
  MModel getModel(String name);
 
+ Collection<MTrainingStatus> getTrainingStatus(Map<String, Object> filterPatterns)
+ throws CatalogException;
+
+ void updateTrainingStatus(String modelName, String status) throws CatalogException;
+
  /* Synopsis */
  MSynopsis createSynopsis(String synopsisName, String modelName, Integer rows, Double ratio)
  throws CatalogException;

traindb-catalog/src/main/java/traindb/catalog/JDOCatalogContext.java

@@ -15,7 +15,10 @@
 package traindb.catalog;
 
 import com.google.common.collect.ImmutableMap;
+import java.sql.Timestamp;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import javax.jdo.PersistenceManager;
@@ -32,6 +35,7 @@
 import traindb.catalog.pm.MSynopsis;
 import traindb.catalog.pm.MTable;
 import traindb.catalog.pm.MTask;
+import traindb.catalog.pm.MTrainingStatus;
 import traindb.common.TrainDBLogger;
 
 public final class JDOCatalogContext implements CatalogContext {
@@ -139,10 +143,18 @@ public MModel trainModel(
  }
  }
 
+ MModeltype mModeltype = getModeltype(modeltypeName);
  MModel mModel = new MModel(
- getModeltype(modeltypeName), modelName, schemaName, tableName, columnNames,
+ mModeltype, modelName, schemaName, tableName, columnNames,
  baseTableRows, trainedRows, options == null ? "" : options);
  pm.makePersistent(mModel);
+
+ if (mModeltype.getLocation().equals("REMOTE")) {
+ MTrainingStatus mTrainingStatus = new MTrainingStatus(modelName, "TRAINING",
+ new Timestamp(System.currentTimeMillis()), mModel);
+ pm.makePersistent(mTrainingStatus);
+ }
+
  return mModel;
  } catch (RuntimeException e) {
  throw new CatalogException("failed to train model '" + modelName + "'", e);
@@ -178,6 +190,13 @@ public void dropModel(String name) throws CatalogException {
  tx.commit();
  }
 
+ Collection<MTrainingStatus> trainingStatus =
+ getTrainingStatus(ImmutableMap.of("model_name", name));
+ if (trainingStatus != null && trainingStatus.size() > 0) {
+ tx.begin();
+ pm.deletePersistentAll(trainingStatus);
+ tx.commit();
+ }
  } catch (RuntimeException e) {
  throw new CatalogException("failed to drop model '" + name + "'", e);
  } finally {
@@ -228,6 +247,33 @@ public boolean modelExists(String name) {
  return null;
  }
 
+ @Override
+ public Collection<MTrainingStatus> getTrainingStatus(Map<String, Object> filterPatterns)
+ throws CatalogException {
+ try {
+ Query query = pm.newQuery(MTrainingStatus.class);
+ setFilterPatterns(query, filterPatterns);
+ return (List<MTrainingStatus>) query.execute();
+ } catch (RuntimeException e) {
+ throw new CatalogException("failed to get training status", e);
+ }
+ }
+
+ @Override
+ public void updateTrainingStatus(String modelName, String status) throws CatalogException {
+ try {
+ Query query = pm.newQuery(MTrainingStatus.class);
+ setFilterPatterns(query, ImmutableMap.of("model_name", modelName));
+ List<MTrainingStatus> trainingStatus = (List<MTrainingStatus>) query.execute();
+ Comparator<MTrainingStatus> comparator = Comparator.comparing(MTrainingStatus::getStartTime);
+ MTrainingStatus latestStatus = trainingStatus.stream().max(comparator).get();
+ latestStatus.setTrainingStatus(status);
+ pm.makePersistent(latestStatus);
+ } catch (RuntimeException e) {
+ throw new CatalogException("failed to get training status", e);
+ }
+ }
+
  @Override
  public MSynopsis createSynopsis(String synopsisName, String modelName, Integer rows,
  @Nullable Double ratio) throws CatalogException {

traindb-catalog/src/main/java/traindb/catalog/pm/MModel.java

@@ -14,6 +14,8 @@
 
 package traindb.catalog.pm;
 
+import java.util.Collection;
+import java.util.Comparator;
 import java.util.List;
 import javax.jdo.annotations.Column;
 import javax.jdo.annotations.IdGeneratorStrategy;
@@ -58,6 +60,9 @@ public final class MModel {
  @Persistent
  private byte[] model_options;
 
+ @Persistent(mappedBy = "model", dependentElement = "true")
+ private Collection<MTrainingStatus> training_status;
+
  public MModel(
  MModeltype modeltype, String modelName, String schemaName, String tableName,
  List<String> columns, @Nullable Long baseTableRows, @Nullable Long trainedRows,
@@ -103,4 +108,17 @@ public long getTrainedRows() {
  public String getModelOptions() {
  return new String(model_options);
  }
+
+ public Collection<MTrainingStatus> trainingStatus() {
+ return training_status;
+ }
+
+ public boolean isEnabled() {
+ if (training_status.isEmpty() || training_status.size() == 0) {
+ return true;
+ }
+ Comparator<MTrainingStatus> comparator = Comparator.comparing(MTrainingStatus::getStartTime);
+ MTrainingStatus latestStatus = training_status.stream().max(comparator).get();
+ return latestStatus.getTrainingStatus().equals("FINISHED");
+ }
 }

traindb-catalog/src/main/java/traindb/catalog/pm/MTrainingStatus.java

@@ -0,0 +1,74 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package traindb.catalog.pm;
+
+import java.sql.Timestamp;
+import javax.jdo.annotations.Column;
+import javax.jdo.annotations.IdGeneratorStrategy;
+import javax.jdo.annotations.Index;
+import javax.jdo.annotations.PersistenceCapable;
+import javax.jdo.annotations.Persistent;
+import javax.jdo.annotations.PrimaryKey;
+import traindb.catalog.CatalogConstants;
+
+@PersistenceCapable
+@Index(name="TRAINING_STATUS_IDX", members={"model_name", "start_time"})
+public final class MTrainingStatus {
+ @PrimaryKey
+ @Persistent(valueStrategy = IdGeneratorStrategy.INCREMENT)
+ private long id;
+
+ @Persistent
+ @Column(length = CatalogConstants.IDENTIFIER_MAX_LENGTH)
+ private String model_name;
+
+ @Persistent
+ private Timestamp start_time;
+
+ @Persistent
+ @Column(length = 9)
+ // Status: TRAINING, FINISHED
+ private String training_status;
+
+ @Persistent(dependent = "false")
+ private MModel model;
+
+ public MTrainingStatus(String modelName, String status, Timestamp startTime, MModel model) {
+ this.model_name = modelName;
+ this.training_status = status;
+ this.start_time = startTime;
+ this.model = model;
+ }
+
+ public String getModelName() {
+ return model_name;
+ }
+
+ public Timestamp getStartTime() {
+ return start_time;
+ }
+
+ public String getTrainingStatus() {
+ return training_status;
+ }
+
+ public MModel getModel() {
+ return model;
+ }
+
+ public void setTrainingStatus(String status) {
+ this.training_status = status;
+ }
+}

traindb-core/src/main/antlr4/traindb/sql/TrainDBSql.g4

@@ -125,6 +125,7 @@ showTargets
  | K_SCHEMAS
  | K_TABLES
  | K_HYPERPARAMETERS
+ | K_TRAININGS
  | K_QUERYLOGS
  | K_TASKS
  ;
@@ -241,6 +242,7 @@ K_SYNOPSIS : S Y N O P S I S ;
 K_TABLES : T A B L E S ;
 K_TASKS : T A S K S ;
 K_TRAIN : T R A I N ;
+K_TRAININGS : T R A I N I N G S;
 K_USE : U S E ;
 K_WHERE : W H E R E ;
 

traindb-core/src/main/java/traindb/engine/AbstractTrainDBModelRunner.java

@@ -53,11 +53,30 @@ public abstract String infer(String aggregateExpression, String groupByColumn,
 
  public abstract String listHyperparameters(String className, String uri) throws Exception;
 
+ public boolean checkAvailable(String modelName) throws Exception {
+ return true;
+ }
+
  public Path getModelPath() {
  return Paths.get(TrainDBConfiguration.getTrainDBPrefixPath(), "models",
  modeltypeName, modelName);
  }
 
+ public static AbstractTrainDBModelRunner createModelRunner(
+ TrainDBConnectionImpl conn, CatalogContext catalogContext, TrainDBConfiguration config,
+ String modeltypeName, String modelName, String location) {
+ if (location.equals("REMOTE")) {
+ return new TrainDBFastApiModelRunner(conn, catalogContext, modeltypeName, modelName);
+ }
+ // location.equals("LOCAL")
+ if (config.getModelRunner().equals("py4j")) {
+ return new TrainDBPy4JModelRunner(conn, catalogContext, modeltypeName, modelName);
+ }
+
+ return new TrainDBFileModelRunner(conn, catalogContext, modeltypeName, modelName);
+ }
+
+
  protected String buildSelectTrainingDataQuery(String schemaName, String tableName,
  List<String> columnNames) {
  StringBuilder sb = new StringBuilder();