Skip to content

Commit

Permalink
Remove create_conda_env (#366)
Browse files Browse the repository at this point in the history
Move `create_conda_env`  to aiida-pythonjob
  • Loading branch information
superstar54 authored Nov 28, 2024
1 parent 8d68a08 commit fe3ede7
Show file tree
Hide file tree
Showing 2 changed files with 1 addition and 140 deletions.
139 changes: 0 additions & 139 deletions aiida_workgraph/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -497,145 +497,6 @@ def serialize_properties(wgdata):
prop["value"] = PickledLocalFunction(prop["value"]).store()


def generate_bash_to_create_python_env(
name: str,
pip: list = None,
conda: dict = None,
modules: list = None,
python_version: str = None,
variables: dict = None,
shell: str = "posix",
):
"""
Generates a bash script for creating or updating a Python environment on a remote computer.
If python_version is None, it uses the Python version from the local environment.
Conda is a dictionary that can include 'channels' and 'dependencies'.
"""
import sys

pip = pip or []
conda_channels = conda.get("channels", []) if conda else []
conda_dependencies = conda.get("dependencies", []) if conda else []
# Determine the Python version from the local environment if not provided
local_python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
desired_python_version = (
python_version if python_version is not None else local_python_version
)

# Start of the script
script = "#!/bin/bash\n\n"

# Load modules if provided
if modules:
script += "# Load specified system modules\n"
for module in modules:
script += f"module load {module}\n"

# Conda shell hook initialization for proper conda activation
script += "# Initialize Conda for this shell\n"
script += f'eval "$(conda shell.{shell} hook)"\n'

script += "# Setup the Python environment\n"
script += "if ! conda info --envs | grep -q ^{name}$; then\n"
script += " # Environment does not exist, create it\n"
if conda_dependencies:
dependencies_string = " ".join(conda_dependencies)
script += f" conda create -y -n {name} python={desired_python_version} {dependencies_string}\n"
else:
script += f" conda create -y -n {name} python={desired_python_version}\n"
script += "fi\n"
if conda_channels:
script += "EXISTING_CHANNELS=$(conda config --show channels)\n"
script += "for CHANNEL in " + " ".join(conda_channels) + ";\n"
script += "do\n"
script += ' if ! echo "$EXISTING_CHANNELS" | grep -q $CHANNEL; then\n'
script += " conda config --prepend channels $CHANNEL\n"
script += " fi\n"
script += "done\n"
script += f"conda activate {name}\n"

# Install pip packages
if pip:
script += f"pip install {' '.join(pip)}\n"

# Set environment variables
if variables:
for var, value in variables.items():
script += f"export {var}='{value}'\n"

# End of the script
script += "echo 'Environment setup is complete.'\n"

return script


def create_conda_env(
computer: Union[str, orm.Computer],
name: str,
pip: list = None,
conda: list = None,
modules: list = None,
python_version: str = None,
variables: dict = None,
shell: str = "posix",
) -> tuple:
"""Test that there is no unexpected output from the connection."""
# Execute a command that should not return any error, except ``NotImplementedError``
# since not all transport plugins implement remote command execution.
from aiida.common.exceptions import NotExistent
from aiida import orm

user = orm.User.collection.get_default()
if isinstance(computer, str):
computer = orm.load_computer(computer)
try:
authinfo = computer.get_authinfo(user)
except NotExistent:
raise f"Computer<{computer.label}> is not yet configured for user<{user.email}>"

scheduler = authinfo.computer.get_scheduler()
transport = authinfo.get_transport()

script = generate_bash_to_create_python_env(
name, pip, conda, modules, python_version, variables, shell
)
with transport:
scheduler.set_transport(transport)
try:
retval, stdout, stderr = transport.exec_command_wait(script)
except NotImplementedError:
return (
True,
f"Skipped, remote command execution is not implemented for the "
f"`{computer.transport_type}` transport plugin",
)

if retval != 0:
return (
False,
f"The command `echo -n` returned a non-zero return code ({retval})",
)

template = """
We detected an error while creating the environemnt on the remote computer, as shown between the bars
=============================================================================================
{}
=============================================================================================
Please check!
"""
if stderr:
return False, template.format(stderr)

if stdout:
# the last line is the echo 'Environment setup is complete.'
if not stdout.strip().endswith("Environment setup is complete."):
return False, template.format(stdout)
else:
return True, "Environment setup is complete."

return True, None


def create_and_pause_process(
runner: Runner = None,
process_class: Callable = None,
Expand Down
2 changes: 1 addition & 1 deletion docs/source/built-in/pythonjob.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@
"One can use the `create_conda_env` function to create a conda environment on the remote computer. The function will create a conda environment with the specified packages and modules. The function will update the packages if the environment already exists.\n",
"\n",
"```python\n",
"from aiida_workgraph.utils import create_conda_env\n",
"from aiida_pythonjob.utils import create_conda_env\n",
"# create a conda environment on remote computer\n",
"create_conda_env(\"merlin6\", \"test_pythonjob\", modules=[\"anaconda\"],\n",
" pip=[\"numpy\", \"matplotlib\"],\n",
Expand Down

0 comments on commit fe3ede7

Please sign in to comment.