Skip to content

Commit f155bb6

Browse files
qqndrewqqndrew
committed
Add support for offline mode python envs
1 parent d3f390e commit f155bb6

File tree

12 files changed

+79
-59
lines changed

12 files changed

+79
-59
lines changed

‎API/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
<parent>
3232
<groupId>org.ohnlp.backbone</groupId>
3333
<artifactId>backbone-parent</artifactId>
34-
<version>3.0.31</version>
34+
<version>3.0.32</version>
3535
</parent>
3636

3737
<artifactId>api</artifactId>

‎API/src/main/java/org/ohnlp/backbone/api/components/xlang/python/PythonBridge.java‎

Lines changed: 53 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,14 @@ public class PythonBridge<T> implements Serializable {
4343
private final String entryClass;
4444
private final Class<T> pythonEntryPointClass;
4545
private final File tmpDir;
46+
private final boolean loadEnv;
4647
private transient DefaultExecutor executor;
4748
private transient ClientServer bridgeServer;
4849
private transient File launchFile;
4950
private transient OSType os;
5051

51-
public PythonBridge(File tmpDir, String bundleName, String entryPoint, String entryClass, Class<T> clazz) throws IOException {
52+
public PythonBridge(boolean loadEnv, File tmpDir, String bundleName, String entryPoint, String entryClass, Class<T> clazz) throws IOException {
53+
this.loadEnv = loadEnv;
5254
this.tmpDir = tmpDir;
5355
this.bundleIdentifier = bundleName;
5456
this.entryPoint = entryPoint;
@@ -245,59 +247,60 @@ private void extractPythonResourcesIfNotExists() {
245247
String localEnvName = "env";
246248
File envDir = new File(workDir, localEnvName);
247249
envDir.mkdirs();
248-
if (initEnv) {
249-
// Extract the packaged conda env
250-
String osPath = "linux";
251-
if (this.os.equals(OSType.WINDOWS)) {
252-
osPath = "win32";
253-
} else {
254-
if (this.os.equals(OSType.MAC_DARWIN)) {
255-
osPath = "darwin";
256-
} else if (!this.os.equals(OSType.LINUX)) {
257-
osPath = "unix";
250+
if (loadEnv) {
251+
if (initEnv) {
252+
// Extract the packaged conda env
253+
String osPath = "linux";
254+
if (this.os.equals(OSType.WINDOWS)) {
255+
osPath = "win32";
256+
} else {
257+
if (this.os.equals(OSType.MAC_DARWIN)) {
258+
osPath = "darwin";
259+
} else if (!this.os.equals(OSType.LINUX)) {
260+
osPath = "unix";
261+
}
258262
}
259-
}
260-
File envTar;
261-
try {
262-
InputStream envTarStream = findEnv(osPath);
263-
if (envTarStream == null) {
264-
if (osPath.equals("unix")) {
265-
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS " + osPath + ", attempting online/dynamic environment resolution");
266-
dynamicallyResolveEnvironment(localEnvName, workDir);
267-
} else {
268-
if (!osPath.equals("win32")) {
269-
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS " + osPath + ", attempting to fall back to unix environment, unexpected behaviour may occur");
270-
envTarStream = findEnv("unix");
271-
if (envTarStream == null) {
272-
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS Unix, attempting online/dynamic environment resolution");
263+
File envTar;
264+
try {
265+
InputStream envTarStream = findEnv(osPath);
266+
if (envTarStream == null) {
267+
if (osPath.equals("unix")) {
268+
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS " + osPath + ", attempting online/dynamic environment resolution");
269+
dynamicallyResolveEnvironment(localEnvName, workDir);
270+
} else {
271+
if (!osPath.equals("win32")) {
272+
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS " + osPath + ", attempting to fall back to unix environment, unexpected behaviour may occur");
273+
envTarStream = findEnv("unix");
274+
if (envTarStream == null) {
275+
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS Unix, attempting online/dynamic environment resolution");
276+
dynamicallyResolveEnvironment(localEnvName, workDir);
277+
}
278+
} else {
279+
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS " + osPath + ", win32 cannot fall back to unix, so attempting online/dynamic environment resolution");
273280
dynamicallyResolveEnvironment(localEnvName, workDir);
274281
}
275-
} else {
276-
Logger.getGlobal().log(Level.INFO, "Could not find bundled/offline environment " + this.bundleIdentifier + " for OS " + osPath + ", win32 cannot fall back to unix, so attempting online/dynamic environment resolution");
277-
dynamicallyResolveEnvironment(localEnvName, workDir);
278282
}
279283
}
284+
if (envTarStream != null) {
285+
envTar = new File(workDir, "environment.tar.gz");
286+
Files.copy(envTarStream, envTar.toPath(), StandardCopyOption.REPLACE_EXISTING);
287+
TarGZipUnArchiver unarchiver = new TarGZipUnArchiver();
288+
unarchiver.setSourceFile(envTar);
289+
unarchiver.setDestDirectory(envDir);
290+
unarchiver.extract();
291+
}
292+
} catch (IOException e) {
293+
env.completeExceptionally(e);
294+
return;
280295
}
281-
if (envTarStream != null) {
282-
envTar = new File(workDir, "environment.tar.gz");
283-
Files.copy(envTarStream, envTar.toPath(), StandardCopyOption.REPLACE_EXISTING);
284-
TarGZipUnArchiver unarchiver = new TarGZipUnArchiver();
285-
unarchiver.setSourceFile(envTar);
286-
unarchiver.setDestDirectory(envDir);
287-
unarchiver.extract();
288-
}
289-
} catch (IOException e) {
290-
env.completeExceptionally(e);
291-
return;
292296
}
297+
try (FileWriter fw = new FileWriter(new File(workDir, "modulehash.txt"))) {
298+
fw.write(cachedModuleChecksum);
299+
fw.flush();
300+
} catch (IOException ignored) {
301+
}
302+
// Write new checksums
293303
}
294-
ObjectWriter ow = new ObjectMapper().writer();
295-
try (FileWriter fw = new FileWriter(new File(workDir, "modulehash.txt"))){
296-
fw.write(cachedModuleChecksum);
297-
fw.flush();
298-
} catch (IOException ignored) {
299-
}
300-
// Write new checksums
301304
env.complete(new PythonEnvironment(workDir, envDir));
302305
}
303306

@@ -379,6 +382,10 @@ private void startServer() throws IOException {
379382
// TODO Consider doing the full conda-unpack routine in a future release
380383
// See: https://conda.github.io/conda-pack/ under "Commandline Usage"
381384
String pythonPath = os.equals(OSType.WINDOWS) ? new File(env.envDir, "python.exe").getAbsolutePath() : new File(new File(env.envDir, "bin"), "python").getAbsolutePath();
385+
if (!this.loadEnv) {
386+
// Just use default env without env.envDir , i.e. a blank "python(.exe)"
387+
pythonPath = os.equals(OSType.WINDOWS) ? "python.exe" : "python";
388+
}
382389
String cmd = String.join(" ",
383390
pythonPath,
384391
this.launchFile.getName(),

‎API/src/main/java/org/ohnlp/backbone/api/components/xlang/python/PythonProxyDoFn.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,20 @@ public class PythonProxyDoFn extends DoFn<String, String> implements Serializabl
2121
private final boolean multiInput;
2222
private final boolean multiOutput;
2323
private final File tmpDir;
24+
private final boolean loadEnv;
2425

2526
private transient PythonBridge<? extends PythonProcessingPartitionBasedDoFn<?, ?>> python;
2627
private transient PythonProcessingPartitionBasedDoFn<?, ?> proxiedDoFn;
2728

28-
public PythonProxyDoFn(File tmpDir, String bundleName, String doFnEntryPoint, String doFnEntryClass, String infoFromDriver, boolean multiInput, boolean multiOutput) {
29+
public PythonProxyDoFn(boolean loadEnv, File tmpDir, String bundleName, String doFnEntryPoint, String doFnEntryClass, String infoFromDriver, boolean multiInput, boolean multiOutput) {
2930
this.tmpDir = tmpDir;
3031
this.bundleName = bundleName;
3132
this.doFnEntryPoint = doFnEntryPoint;
3233
this.doFnEntryClass = doFnEntryClass;
3334
this.driverInfo = infoFromDriver;
3435
this.multiInput = multiInput;
3536
this.multiOutput = multiOutput;
37+
this.loadEnv = loadEnv;
3638
}
3739

3840

@@ -49,7 +51,7 @@ public void init() throws IOException {
4951
}
5052
}
5153
// Init python bridge
52-
this.python = new PythonBridge<>(this.tmpDir, this.bundleName, this.doFnEntryPoint, this.doFnEntryClass, implCls);
54+
this.python = new PythonBridge<>(this.loadEnv, this.tmpDir, this.bundleName, this.doFnEntryPoint, this.doFnEntryClass, implCls);
5355
this.python.startBridge();
5456

5557
// Get proxied DoFn

‎API/src/main/java/org/ohnlp/backbone/api/components/xlang/python/PythonProxyTransformComponent.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030

3131
public class PythonProxyTransformComponent extends TransformComponent implements XLangComponent, SingleInputComponent, SchemaInitializable {
3232

33+
private final boolean loadEnv;
3334
private File envTmp;
3435
private String config;
3536

@@ -43,11 +44,12 @@ public class PythonProxyTransformComponent extends TransformComponent implements
4344
private final AtomicBoolean INIT_COMPLETE = new AtomicBoolean(false);
4445
private final AtomicReference<ComponentInitializationException> INIT_ERROR = new AtomicReference(null);
4546

46-
public PythonProxyTransformComponent(File envTmp, String bundleIdentifier, String entryPoint, String entryClass) {
47+
public PythonProxyTransformComponent(boolean loadEnv, File envTmp, String bundleIdentifier, String entryPoint, String entryClass) {
4748
this.envTmp = envTmp;
4849
this.entryPoint = entryPoint;
4950
this.entryClass = entryClass;
5051
this.bundleIdentifier = bundleIdentifier;
52+
this.loadEnv = loadEnv;
5153
}
5254

5355
@Override
@@ -64,7 +66,7 @@ public void init() throws ComponentInitializationException {
6466
if (INIT_LOCK.getAndSet(false)) { // Acquired Lock
6567
if (!INIT_IN_PROGRESS.getAndSet(true)) {
6668
try {
67-
this.python = new PythonBridge<>(this.envTmp, this.bundleIdentifier, this.entryPoint, this.entryClass, PythonBackbonePipelineComponent.class);
69+
this.python = new PythonBridge<>(this.loadEnv, this.envTmp, this.bundleIdentifier, this.entryPoint, this.entryClass, PythonBackbonePipelineComponent.class);
6870
this.python.startBridge();
6971
this.proxiedComponent = python.getPythonEntryPoint();
7072
this.proxiedComponent.init(this.config);
@@ -110,6 +112,7 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
110112
// Get a proxied python DoFn that handles bridge setup on executor nodes, and pass it initialized driver configs
111113
// as well
112114
PythonProxyDoFn proxiedDoFn = new PythonProxyDoFn(
115+
this.loadEnv,
113116
this.envTmp,
114117
this.bundleIdentifier,
115118
this.entryPoint,

‎API/src/main/java/org/ohnlp/backbone/api/config/BackboneConfiguration.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
public class BackboneConfiguration {
1313
private String id;
1414
private String description;
15+
private boolean offlinePythonEnvs = false;
1516
private List<BackbonePipelineComponentConfiguration> pipeline;
1617

1718
public String getId() {
@@ -37,4 +38,12 @@ public List<BackbonePipelineComponentConfiguration> getPipeline() {
3738
public void setPipeline(List<BackbonePipelineComponentConfiguration> pipeline) {
3839
this.pipeline = pipeline;
3940
}
41+
42+
public boolean getOfflinePythonEnvs() {
43+
return offlinePythonEnvs;
44+
}
45+
46+
public void setOfflinePythonEnvs(boolean offlinePythonEnvs) {
47+
this.offlinePythonEnvs = offlinePythonEnvs;
48+
}
4049
}

‎API/src/main/java/org/ohnlp/backbone/api/pipeline/PipelineBuilder.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package org.ohnlp.backbone.api.pipeline;
22

3-
import com.fasterxml.jackson.core.JsonProcessingException;
43
import com.fasterxml.jackson.databind.JsonNode;
54
import com.fasterxml.jackson.databind.ObjectMapper;
65
import org.apache.beam.sdk.schemas.Schema;
@@ -90,7 +89,7 @@ public static ExecutionDAG getPipelineGraph(BackboneConfiguration config) throws
9089
}
9190
} else if (configs[i] instanceof PythonBackbonePipelineComponentConfiguration) {
9291
PythonBackbonePipelineComponentConfiguration pythonconfig = (PythonBackbonePipelineComponentConfiguration) configs[i];
93-
PythonProxyTransformComponent instance = new PythonProxyTransformComponent(new File(System.getProperty("java.io.tmpdir")), pythonconfig.getBundleName(), pythonconfig.getEntryPoint(), pythonconfig.getEntryClass());
92+
PythonProxyTransformComponent instance = new PythonProxyTransformComponent(!config.getOfflinePythonEnvs(), new File(System.getProperty("java.io.tmpdir")), pythonconfig.getBundleName(), pythonconfig.getEntryPoint(), pythonconfig.getEntryClass());
9493
JsonNode configForInstance = pythonconfig.getConfig();
9594
instance.injectConfig(configForInstance);
9695
instance.init();

‎Core/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
<parent>
88
<groupId>org.ohnlp.backbone</groupId>
99
<artifactId>backbone-parent</artifactId>
10-
<version>3.0.31</version>
10+
<version>3.0.32</version>
1111
</parent>
1212

1313
<artifactId>core</artifactId>

‎Example-Backbone-Configs/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
<parent>
88
<groupId>org.ohnlp.backbone</groupId>
99
<artifactId>backbone-parent</artifactId>
10-
<version>3.0.31</version>
10+
<version>3.0.32</version>
1111
</parent>
1212

1313
<artifactId>example-backbone-configs</artifactId>

‎IO/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
<parent>
88
<groupId>org.ohnlp.backbone</groupId>
99
<artifactId>backbone-parent</artifactId>
10-
<version>3.0.31</version>
10+
<version>3.0.32</version>
1111
</parent>
1212

1313
<groupId>org.ohnlp.backbone.io</groupId>

‎Plugin-Manager/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
<parent>
88
<groupId>org.ohnlp.backbone</groupId>
99
<artifactId>backbone-parent</artifactId>
10-
<version>3.0.31</version>
10+
<version>3.0.32</version>
1111
</parent>
1212

1313
<artifactId>plugin-manager</artifactId>

0 commit comments

Comments
 (0)