Extend module system

Unify loading of flytekit-api and jflyte-api modules. Allows having multiple FileSystem implementation, for instance, one for GCS, and one for S3. Containers can supply implementations of tasks in modules directory that is useful for execute-local mode.
flyteorg · Jun 16, 2020 · 7be868a · 7be868a
1 parent f3d2aac
commit 7be868a
Show file tree

Hide file tree

Showing 13 changed files with 332 additions and 199 deletions.
diff --git a/jflyte-build/Dockerfile → Dockerfile b/jflyte-build/Dockerfile → Dockerfile
@@ -1,12 +1,16 @@
 FROM gcr.io/distroless/java:8
 
-COPY target/lib /jflyte
-COPY target/plugins /jflyte/plugins/
+ARG FLYTE_INTERNAL_IMAGE
 
-ENV FLYTE_INTERNAL_PLUGIN_DIR "/jflyte/plugins"
+COPY jflyte/target/lib /jflyte/
+
+# plugins
+COPY jflyte-google-cloud/target/lib /jflyte/modules/jflyte-google-cloud
+
+ENV FLYTE_INTERNAL_MODULE_DIR "/jflyte/modules"
+ENV FLYTE_INTERNAL_IMAGE=$FLYTE_INTERNAL_IMAGE
 
 ENV FLYTE_PLATFORM_URL "CHANGEME"
-ENV FLYTE_INTERNAL_IMAGE "CHANGEME"
 ENV FLYTE_STAGING_LOCATION "CHANGEME"
 ENV FLYTE_PLATFORM_INSECURE "False"
 

diff --git a/jflyte-api/src/main/java/org/flyte/jflyte/api/FileSystemRegistrar.java b/jflyte-api/src/main/java/org/flyte/jflyte/api/FileSystemRegistrar.java
@@ -16,37 +16,7 @@
  */
 package org.flyte.jflyte.api;
 
-import java.util.ServiceLoader;
-import java.util.logging.Level;
-import java.util.logging.Logger;
-
 /** A registrar that creates {@link FileSystem} instances. */
 public abstract class FileSystemRegistrar {
-  private static final Logger LOG = Logger.getLogger(FileSystemRegistrar.class.getName());
-
-  static {
-    // enable all levels for the actual handler to pick up
-    LOG.setLevel(Level.ALL);
-  }
-
-  public abstract Iterable<FileSystem> load(ClassLoader classLoader);
-
-  public static FileSystem getFileSystem(String scheme, ClassLoader classLoader) {
-    ServiceLoader<FileSystemRegistrar> loader =
-        ServiceLoader.load(FileSystemRegistrar.class, classLoader);
-
-    LOG.fine("Discovering FileSystemRegistrar");
-
-    for (FileSystemRegistrar registrar : loader) {
-      for (FileSystem fileSystem : registrar.load(classLoader)) {
-        LOG.fine(String.format("Discovered FileSystem [%s]", fileSystem.getClass().getName()));
-
-        if (scheme.equals(fileSystem.getScheme())) {
-          return fileSystem;
-        }
-      }
-    }
-
-    return null;
-  }
+  public abstract Iterable<FileSystem> load();
 }
diff --git a/jflyte-build/pom.xml b/jflyte-build/pom.xml
@@ -30,97 +30,36 @@
   <properties>
     <docker.image>${env.FLYTE_INTERNAL_IMAGE}</docker.image>
     <docker.tag>${project.version}</docker.tag>
-    <plugins.dir>${project.build.directory}/plugins</plugins.dir>
 
     <maven.deploy.skip>true</maven.deploy.skip>
+    <!-- don't need enforcer since we only do a docker build -->
+    <enforcer.skip>true</enforcer.skip>
   </properties>
 
   <dependencies>
-    <dependency>
-      <groupId>org.flyte</groupId>
-      <artifactId>flytekit-api</artifactId>
-      <scope>runtime</scope>
-    </dependency>
     <dependency>
       <groupId>org.flyte</groupId>
       <artifactId>jflyte</artifactId>
-      <scope>runtime</scope>
     </dependency>
     <dependency>
       <groupId>org.flyte</groupId>
       <artifactId>jflyte-google-cloud</artifactId>
-      <scope>provided</scope>
     </dependency>
   </dependencies>
 
   <build>
     <finalName>jflyte</finalName>
     <plugins>
-      <!-- if no plugin is enabled (via provided scope), docker copy will fail due to missing directory -->
-      <plugin>
-        <groupId>org.apache.maven.plugins</groupId>
-        <artifactId>maven-antrun-plugin</artifactId>
-        <executions>
-          <execution>
-            <id>create-plugins-dir</id>
-            <phase>prepare-package</phase>
-            <configuration>
-              <target>
-                <delete dir="${plugins.dir}" />
-                <mkdir dir="${plugins.dir}" />
-              </target>
-            </configuration>
-            <goals>
-              <goal>run</goal>
-            </goals>
-          </execution>
-        </executions>
-      </plugin>
-      <plugin>
-        <artifactId>maven-dependency-plugin</artifactId>
-        <executions>
-          <execution>
-            <id>copy-plugins</id>
-            <phase>prepare-package</phase>
-            <goals>
-              <goal>copy-dependencies</goal>
-            </goals>
-            <configuration>
-              <outputDirectory>${plugins.dir}</outputDirectory>
-              <useBaseVersion>false</useBaseVersion>
-              <overWriteReleases>false</overWriteReleases>
-              <overWriteSnapshots>true</overWriteSnapshots>
-              <includeScope>provided</includeScope>
-            </configuration>
-          </execution>
-          <execution>
-            <id>copy-runtime</id>
-            <phase>prepare-package</phase>
-            <goals>
-              <goal>copy-dependencies</goal>
-            </goals>
-            <configuration>
-              <outputDirectory>${project.build.directory}/lib</outputDirectory>
-              <useBaseVersion>false</useBaseVersion>
-              <overWriteReleases>false</overWriteReleases>
-              <overWriteSnapshots>true</overWriteSnapshots>
-              <includeScope>runtime</includeScope>
-              <excludeGroupIds>
-                com.google.android,com.google.auto.value,com.google.code.findbugs,findbugs,com.google.errorprone
-              </excludeGroupIds>
-              <excludeArtifactIds>
-                animal-sniffer-annotations,j2objc-annotations
-              </excludeArtifactIds>
-            </configuration>
-          </execution>
-        </executions>
-      </plugin>
       <plugin>
         <groupId>com.spotify</groupId>
         <artifactId>dockerfile-maven-plugin</artifactId>
         <configuration>
           <repository>${docker.image}</repository>
           <tag>${docker.tag}</tag>
+          <buildArgs>
+            <FLYTE_INTERNAL_IMAGE>${docker.image}:${docker.tag}</FLYTE_INTERNAL_IMAGE>
+          </buildArgs>
+          <contextDirectory>${project.parent.basedir}</contextDirectory>
         </configuration>
       </plugin>
     </plugins>

diff --git a/jflyte-google-cloud/pom.xml b/jflyte-google-cloud/pom.xml
@@ -75,4 +75,15 @@
       <version>2.8.6</version>
     </dependency>
   </dependencies>
+
+  <build>
+    <plugins>
+      <plugin>
+        <artifactId>maven-jar-plugin</artifactId>
+      </plugin>
+      <plugin>
+        <artifactId>maven-dependency-plugin</artifactId>
+      </plugin>
+    </plugins>
+  </build>
 </project>
diff --git a/jflyte-google-cloud/src/main/java/org/flyte/jflyte/gcp/GcsFileSystemRegistrar.java b/jflyte-google-cloud/src/main/java/org/flyte/jflyte/gcp/GcsFileSystemRegistrar.java
@@ -25,7 +25,7 @@
 @AutoService(FileSystemRegistrar.class)
 public class GcsFileSystemRegistrar extends FileSystemRegistrar {
   @Override
-  public Iterable<FileSystem> load(ClassLoader classLoader) {
+  public Iterable<FileSystem> load() {
     return Collections.singletonList(new GcsFileSystem());
   }
 }
diff --git a/jflyte/pom.xml b/jflyte/pom.xml
@@ -147,6 +147,12 @@
           </execution>
         </executions>
       </plugin>
+      <plugin>
+        <artifactId>maven-jar-plugin</artifactId>
+      </plugin>
+      <plugin>
+        <artifactId>maven-dependency-plugin</artifactId>
+      </plugin>
     </plugins>
   </build>
 </project>
diff --git a/jflyte/src/main/java/org/flyte/jflyte/ClassLoaders.java b/jflyte/src/main/java/org/flyte/jflyte/ClassLoaders.java
@@ -17,12 +17,15 @@
 package org.flyte.jflyte;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import java.io.File;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
-import java.util.stream.Stream;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -35,18 +38,39 @@ private ClassLoaders() {
     throw new UnsupportedOperationException();
   }
 
-  static ClassLoader forDirectory(String dir) {
+  static List<ClassLoader> forModuleDir(String dir) {
+    return listDirectory(new File(dir)).stream()
+        .filter(File::isDirectory)
+        .map(ClassLoaders::forDirectory)
+        .collect(Collectors.toList());
+  }
+
+  static ClassLoader forDirectory(File dir) {
+    LOG.debug("Loading jars from [{}]", dir.getAbsolutePath());
+
     return AccessController.doPrivileged(
         (PrivilegedAction<ClassLoader>) () -> new ChildFirstClassLoader(getClassLoaderUrls(dir)));
   }
 
-  static URL[] getClassLoaderUrls(String dir) {
+  static URL[] getClassLoaderUrls(File dir) {
     Preconditions.checkNotNull(dir, "dir is null");
 
-    File file = new File(dir);
+    return listDirectory(dir).stream()
+        .map(
+            file -> {
+              try {
+                URL url = file.toURI().toURL();
+                LOG.debug("Discovered [{}]", url);
 
-    LOG.debug("Loading jars from [{}]", dir);
+                return url;
+              } catch (MalformedURLException e) {
+                throw new RuntimeException(e);
+              }
+            })
+        .toArray(URL[]::new);
+  }
 
+  static List<File> listDirectory(File file) {
     if (!file.exists()) {
       throw new RuntimeException(
           String.format("Directory doesn't exist [%s]", file.getAbsolutePath()));
@@ -58,18 +82,26 @@ static URL[] getClassLoaderUrls(String dir) {
       throw new RuntimeException(String.format("Isn't directory [%s]", file.getAbsolutePath()));
     }
 
-    return Stream.of(files)
-        .map(
-            x -> {
-              try {
-                URL url = x.toURI().toURL();
-                LOG.debug("Discovered [{}]", url);
+    return ImmutableList.copyOf(files);
+  }
 
-                return url;
-              } catch (MalformedURLException e) {
-                throw new RuntimeException(e);
-              }
-            })
-        .toArray(URL[]::new);
+  static <V> V withClassLoader(ClassLoader classLoader, Callable<V> callable) {
+    ClassLoader originalContextClassLoader = Thread.currentThread().getContextClassLoader();
+    Thread.currentThread().setContextClassLoader(classLoader);
+
+    // before we run anything, switch class loader, because we will be touching user classes;
+    // setting it in thread context will give us access to the right class loader
+
+    try {
+      return callable.call();
+    } catch (Exception e) {
+      if (e instanceof RuntimeException) {
+        throw (RuntimeException) e;
+      }
+
+      throw new RuntimeException(e);
+    } finally {
+      Thread.currentThread().setContextClassLoader(originalContextClassLoader);
+    }
   }
 }
diff --git a/jflyte/src/main/java/org/flyte/jflyte/Config.java b/jflyte/src/main/java/org/flyte/jflyte/Config.java
@@ -30,14 +30,14 @@ abstract class Config {
   @Nullable
   abstract String stagingLocation();
 
-  abstract String pluginDir();
+  abstract String moduleDir();
 
   abstract boolean platformInsecure();
 
   static Config load() {
     return Config.builder()
         .platformUrl(getenv("FLYTE_PLATFORM_URL"))
-        .pluginDir(getenv("FLYTE_INTERNAL_PLUGIN_DIR"))
+        .moduleDir(getenv("FLYTE_INTERNAL_MODULE_DIR"))
         .image(getenv("FLYTE_INTERNAL_IMAGE"))
         .stagingLocation(getenvOrNull("FLYTE_STAGING_LOCATION"))
         .platformInsecure(Boolean.parseBoolean(getenv("FLYTE_PLATFORM_INSECURE")))
@@ -71,7 +71,7 @@ abstract static class Builder {
 
     abstract Builder stagingLocation(String stagingLocation);
 
-    abstract Builder pluginDir(String pluginDir);
+    abstract Builder moduleDir(String moduleDir);
 
     abstract Builder platformInsecure(boolean platformInsecure);