diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/aux/CondaEnvironment.java b/lzy/execution-env/src/main/java/ai/lzy/env/aux/CondaEnvironment.java index 9fc88b82e..f8b4ee76a 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/aux/CondaEnvironment.java +++ b/lzy/execution-env/src/main/java/ai/lzy/env/aux/CondaEnvironment.java @@ -5,12 +5,14 @@ import ai.lzy.env.logs.LogStream; import com.google.common.annotations.VisibleForTesting; import jakarta.annotation.Nullable; +import org.apache.commons.io.IOUtils; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import java.io.File; import java.io.FileWriter; import java.io.IOException; +import java.nio.charset.Charset; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; @@ -37,6 +39,9 @@ public class CondaEnvironment implements AuxEnvironment { private final Path hostWorkingDir; private final Path baseEnvWorkingDir; + @Nullable + private String pythonPath; + @VisibleForTesting public static void reconfigureConda(boolean reconfigure) { RECONFIGURE_CONDA = reconfigure; @@ -114,6 +119,25 @@ public void install(LogStream outStream, LogStream errStream) throws Environment condaFileOnHost.delete(); } } + + try { + var proc = execInEnv("echo $PYTHONPATH", null); + var rc = proc.waitFor(); + + if (rc != 0) { + LOG.error("Cannot get PYTHONPATH from env"); + return; + } + + var path = IOUtils.toString(proc.out(), Charset.defaultCharset()).strip(); + if (!path.isEmpty()) { + path += ":"; + } + path += baseEnvWorkingDir.toAbsolutePath(); + this.pythonPath = path; + } catch (Exception e) { + LOG.error("Cannot get PYTHONPATH from env: ", e); + } } catch (IOException | InterruptedException | ExecutionException e) { LOG.error("CondaEnvironment setup failed", e); throw new RuntimeException(e); @@ -155,8 +179,10 @@ private LzyProcess execInEnv(String command, @Nullable String[] envp, @Nullable return baseEnv.runProcess(bashCmd, envp, workingDir == null ? baseEnvWorkingDir.toString() : workingDir); } - private LzyProcess execInEnv(String command, LogStream out) { - out.log("RunCmd: %s".formatted(command)); + private LzyProcess execInEnv(String command, @Nullable LogStream out) { + if (out != null) { + out.log("RunCmd: %s".formatted(command)); + } return execInEnv(command, null, null); } @@ -167,7 +193,15 @@ public LzyProcess runProcess(String[] command, @Nullable String[] envp, @Nullabl if (envp != null) { envList.addAll(Arrays.asList(envp)); } - return execInEnv(String.join(" ", command), envList.toArray(String[]::new), workingDir); + + var cmd = String.join(" ", command); + + if (pythonPath != null) { + // Adding export here to prevent conda from updating PYTHONPATH + cmd = "export PYTHONPATH=" + pythonPath + " && " + cmd; + } + + return execInEnv(cmd, envList.toArray(String[]::new), workingDir); } @Override