Skip to content

Commit e64446c

Browse files
tudormdavorbonaci
authored andcommitted
Add support to configure the shuffle client library with a property.
----Release Notes---- [] ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=113307012
1 parent 76b8689 commit e64446c

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/ShuffleLibrary.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,13 @@ class ShuffleLibrary {
3232
*/
3333
static void load() {
3434
try {
35+
final String shuffleClientLibraryPropertyKey = "batch.shuffle_client_library";
36+
String shuffleClientLibrary = "libshuffle_client_jni.so.stripped";
37+
if (System.getProperties().containsKey(shuffleClientLibraryPropertyKey)) {
38+
shuffleClientLibrary = System.getProperty(shuffleClientLibraryPropertyKey);
39+
}
3540
File tempfile = File.createTempFile("libshuffle_client_jni", ".so");
36-
InputStream input = ClassLoader.getSystemResourceAsStream(
37-
"libshuffle_client_jni.so.stripped");
41+
InputStream input = ClassLoader.getSystemResourceAsStream(shuffleClientLibrary);
3842
Files.copy(input, tempfile.toPath(), StandardCopyOption.REPLACE_EXISTING);
3943
System.load(tempfile.getAbsolutePath());
4044
} catch (IOException e) {

0 commit comments

Comments
 (0)