Implement export model and metadata

Author: taewhi

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

@@ -16,7 +16,6 @@
 
 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;
@@ -117,14 +116,15 @@ public MModel trainModel(
  String modeltypeName, String modelName, String schemaName, String tableName,
  List<String> columnNames, RelDataType dataType, @Nullable Long baseTableRows,
  @Nullable Long trainedRows, @Nullable String options) throws CatalogException {
+ MTable mTable;
  try {
  MSchema mSchema = getSchema(schemaName);
  if (mSchema == null) {
  mSchema = new MSchema(schemaName);
  pm.makePersistent(mSchema);
  }
 
- MTable mTable = getTable(schemaName, tableName);
+ mTable = getTable(schemaName, tableName);
  if (mTable == null) {
  mTable = new MTable(tableName, "TABLE", mSchema);
  pm.makePersistent(mTable);
@@ -146,7 +146,7 @@ public MModel trainModel(
  MModeltype mModeltype = getModeltype(modeltypeName);
  MModel mModel = new MModel(
  mModeltype, modelName, schemaName, tableName, columnNames,
- baseTableRows, trainedRows, options == null ? "" : options);
+ baseTableRows, trainedRows, options == null ? "" : options, mTable);
  pm.makePersistent(mModel);
 
  if (mModeltype.getLocation().equals("REMOTE")) {

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

@@ -14,6 +14,7 @@
 
 package traindb.catalog.pm;
 
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import javax.jdo.annotations.Column;
 import javax.jdo.annotations.IdGeneratorStrategy;
 import javax.jdo.annotations.PersistenceCapable;
@@ -22,6 +23,7 @@
 import traindb.catalog.CatalogConstants;
 
 @PersistenceCapable
+@JsonIgnoreProperties({ "table" })
 public final class MColumn {
  @PrimaryKey
  @Persistent(valueStrategy = IdGeneratorStrategy.INCREMENT)

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

@@ -63,10 +63,13 @@ public final class MModel {
  @Persistent(mappedBy = "model", dependentElement = "true")
  private Collection<MTrainingStatus> training_status;
 
+ @Persistent(dependent = "false")
+ private MTable table;
+
  public MModel(
  MModeltype modeltype, String modelName, String schemaName, String tableName,
  List<String> columns, @Nullable Long baseTableRows, @Nullable Long trainedRows,
- String options) {
+ String options, MTable table) {
  this.modeltype = modeltype;
  this.model_name = modelName;
  this.schema_name = schemaName;
@@ -75,6 +78,7 @@ public MModel(
  this.table_rows = (baseTableRows == null) ? 0 : baseTableRows;
  this.trained_rows = (trainedRows == null) ? 0 : trainedRows;
  this.model_options = options.getBytes();
+ this.table = table;
  }
 
  public String getModelName() {
@@ -113,6 +117,10 @@ public Collection<MTrainingStatus> trainingStatus() {
  return training_status;
  }
 
+ public MTable getTable() {
+ return table;
+ }
+
  public boolean isEnabled() {
  if (training_status.isEmpty() || training_status.size() == 0) {
  return true;

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

@@ -14,6 +14,7 @@
 
 package traindb.catalog.pm;
 
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import java.util.ArrayList;
 import java.util.Collection;
 import javax.jdo.annotations.Column;
@@ -24,6 +25,7 @@
 import traindb.catalog.CatalogConstants;
 
 @PersistenceCapable
+@JsonIgnoreProperties({ "schema" })
 public final class MTable {
  @PrimaryKey
  @Persistent(valueStrategy = IdGeneratorStrategy.INCREMENT)

traindb-common/src/main/java/traindb/util/ZipUtils.java

@@ -0,0 +1,80 @@
+/*
+ * 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.util;
+
+import java.io.IOException;
+import java.io.Writer;
+import java.net.URI;
+import java.nio.file.FileSystem;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.StandardCopyOption;
+import java.nio.file.StandardOpenOption;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
+
+public final class ZipUtils {
+
+ private ZipUtils() {
+ }
+
+ public static void pack(String sourceDirPath, String zipFilePath) throws IOException {
+ Path zp = Files.createFile(Paths.get(zipFilePath));
+ try (ZipOutputStream zs = new ZipOutputStream(Files.newOutputStream(zp))) {
+ Path sp = Paths.get(sourceDirPath);
+ Files.walk(sp)
+ .filter(path -> !Files.isDirectory(path))
+ .forEach(path -> {
+ ZipEntry zipEntry = new ZipEntry(sp.relativize(path).toString());
+ try {
+ zs.putNextEntry(zipEntry);
+ Files.copy(path, zs);
+ zs.closeEntry();
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ });
+ }
+ }
+
+ public static void addFileToZip(Path file, Path zip) throws IOException {
+ Map<String, String> env = new HashMap<>();
+ env.put("create", "false");
+
+ URI uri = URI.create("jar:file:" + zip.toString());
+ try (FileSystem fs = FileSystems.newFileSystem(uri, env)) {
+ Path p = fs.getPath(file.getFileName().toString());
+ Files.copy(file, p, StandardCopyOption.REPLACE_EXISTING);
+ }
+ }
+
+ public static void addNewFileFromStringToZip(String newFilename, String contents, Path zip)
+ throws IOException {
+ Map<String, String> env = new HashMap<>();
+ env.put("create", "false");
+
+ URI uri = URI.create("jar:file:" + zip.toString());
+ try (FileSystem fs = FileSystems.newFileSystem(uri, env)) {
+ Path p = fs.getPath(newFilename);
+ try (Writer writer = Files.newBufferedWriter(p, StandardOpenOption.CREATE)) {
+ writer.write(contents);
+ }
+ }
+ }
+}

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

@@ -28,6 +28,7 @@
 import org.json.simple.JSONObject;
 import traindb.catalog.CatalogContext;
 import traindb.common.TrainDBConfiguration;
+import traindb.common.TrainDBException;
 import traindb.jdbc.TrainDBConnectionImpl;
 import traindb.schema.TrainDBTable;
 
@@ -57,6 +58,8 @@ public abstract String infer(String aggregateExpression, String groupByColumn,
 
  public abstract String listHyperparameters(String className, String uri) throws Exception;
 
+ public abstract void exportModel(String outputPath) throws Exception;
+
  public boolean checkAvailable(String modelName) throws Exception {
  return true;
  }

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

@@ -231,6 +231,11 @@ public String listHyperparameters(String className, String uri) throws Exception
  return unescapeString(response.toString());
  }
 
+ @Override
+ public void exportModel(String outputPath) throws Exception {
+ throw new TrainDBException("Not supported yet");
+ }
+
  @Override
  public boolean checkAvailable(String modelName) throws Exception {
  MModeltype mModeltype = catalogContext.getModel(modelName).getModeltype();

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

@@ -32,6 +32,7 @@
 import traindb.common.TrainDBException;
 import traindb.jdbc.TrainDBConnectionImpl;
 import traindb.schema.TrainDBTable;
+import traindb.util.ZipUtils;
 
 public class TrainDBFileModelRunner extends AbstractTrainDBModelRunner {
 
@@ -151,4 +152,9 @@ public String listHyperparameters(String className, String uri) throws Exception
  return hyperparamsInfo;
  }
 
+ @Override
+ public void exportModel(String outputPath) throws Exception {
+ String modelPath = getModelPath().toString();
+ ZipUtils.pack(modelPath, outputPath);
+ }
 }

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

@@ -37,6 +37,7 @@
 import traindb.common.TrainDBException;
 import traindb.jdbc.TrainDBConnectionImpl;
 import traindb.schema.TrainDBTable;
+import traindb.util.ZipUtils;
 
 public class TrainDBPy4JModelRunner extends AbstractTrainDBModelRunner {
 
@@ -109,6 +110,7 @@ public void trainModel(
  server.shutdown();
  }
 
+ @Override
  public void generateSynopsis(String outputPath, int rows) throws Exception {
  String modelPath = getModelPath().toString();
  MModeltype mModeltype = catalogContext.getModel(modelName).getModeltype();
@@ -153,6 +155,12 @@ public String listHyperparameters(String className, String uri) throws Exception
  return hyperparamsInfo;
  }
 
+ @Override
+ public void exportModel(String outputPath) throws Exception {
+ String modelPath = getModelPath().toString();
+ ZipUtils.pack(modelPath, outputPath);
+ }
+
  private int getAvailablePort() throws Exception {
  ServerSocket s;
  try {

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

@@ -14,10 +14,16 @@
 
 package traindb.engine;
 
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableList;
 import com.opencsv.CSVReader;
 import com.opencsv.CSVReaderBuilder;
+import java.io.File;
+import java.io.FileInputStream;
 import java.io.FileReader;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -27,8 +33,6 @@
 import java.util.List;
 import java.util.Map;
 import org.apache.calcite.sql.SqlDialect;
-import org.codehaus.jackson.annotate.JsonIgnoreProperties;
-import org.codehaus.jackson.map.ObjectMapper;
 import org.json.simple.JSONObject;
 import traindb.adapter.TrainDBSqlDialect;
 import traindb.catalog.CatalogContext;
@@ -39,14 +43,15 @@
 import traindb.catalog.pm.MSynopsis;
 import traindb.catalog.pm.MTask;
 import traindb.catalog.pm.MTrainingStatus;
-import traindb.common.TrainDBConfiguration;
 import traindb.common.TrainDBException;
 import traindb.common.TrainDBLogger;
+import traindb.engine.nio.ByteArray;
 import traindb.jdbc.TrainDBConnectionImpl;
 import traindb.schema.SchemaManager;
 import traindb.schema.TrainDBTable;
 import traindb.sql.TrainDBSqlRunner;
 import traindb.task.TaskTracer;
+import traindb.util.ZipUtils;
 
 
 public class TrainDBQueryEngine implements TrainDBSqlRunner {
@@ -683,7 +688,68 @@ public void deleteTasks(Integer cnt) throws Exception {
 
  @Override
  public TrainDBListResultSet exportModel(String modelName) throws Exception {
- throw new TrainDBException("Not supported yet");
+ T_tracer.startTaskTracer("export model " + modelName);
+
+ T_tracer.openTaskTime("find : model");
+ if (!catalogContext.modelExists(modelName)) {
+ String msg = "model '" + modelName + "' does not exist";
+
+ T_tracer.closeTaskTime(msg);
+ T_tracer.endTaskTracer();
+
+ throw new CatalogException(msg);
+ }
+ T_tracer.closeTaskTime("SUCCESS");
+
+ T_tracer.openTaskTime("export model");
+ MModel mModel = catalogContext.getModel(modelName);
+ MModeltype mModeltype = mModel.getModeltype();
+
+ AbstractTrainDBModelRunner runner = createModelRunner(
+ mModeltype.getModeltypeName(), modelName, mModeltype.getLocation());
+
+ if (!mModel.isEnabled()) { // remote model
+ if (!runner.checkAvailable(modelName)) {
+ throw new TrainDBException(
+ "model '" + modelName + "' is not available (training is not finished)");
+ }
+ catalogContext.updateTrainingStatus(modelName, "FINISHED");
+ }
+
+ Path outputPath = Paths.get(runner.getModelPath().getParent().toString(), modelName + ".zip");
+ runner.exportModel(outputPath.toString());
+
+ ObjectMapper mapper = new ObjectMapper();
+ ZipUtils.addNewFileFromStringToZip("export_metadata.json",
+ mapper.writeValueAsString(mModel), outputPath);
+
+ List<String> header = Arrays.asList("export_model");
+ List<List<Object>> exportModelInfo = new ArrayList<>();
+ ByteArray byteArray = convertFileToByteArray(new File(outputPath.toString()));
+ exportModelInfo.add(Arrays.asList(byteArray));
+
+ T_tracer.closeTaskTime("SUCCESS");
+ T_tracer.endTaskTracer();
+
+ return new TrainDBListResultSet(header, exportModelInfo);
+ }
+
+ private ByteArray convertFileToByteArray(File file) throws Exception {
+ byte[] bytes;
+ FileInputStream inputStream = null;
+
+ try {
+ bytes = new byte[(int) file.length()];
+ inputStream = new FileInputStream(file);
+ inputStream.read(bytes);
+ } catch (Exception e) {
+ throw e;
+ } finally {
+ if (inputStream != null) {
+ inputStream.close();
+ }
+ }
+ return new ByteArray(bytes);
  }
 
  private void checkShowWhereColumns(Map<String, Object> patterns, List<String> columns)