Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add python path to conda env #1165

Merged
merged 2 commits into from
Mar 12, 2024
Merged

Add python path to conda env #1165

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e658f16
Add python path to conda env
ArtoLord Mar 11, 2024
9c8c45b
Fix
ArtoLord Mar 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 37 additions & 3 deletions lzy/execution-env/src/main/java/ai/lzy/env/aux/CondaEnvironment.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}

Expand All @@ -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;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why don't you add this env var to the envList?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because conda can change this variable after activation

}

return execInEnv(cmd, envList.toArray(String[]::new), workingDir);
}

@Override
Expand Down
Loading