diff --git a/aiida_workgraph/utils/__init__.py b/aiida_workgraph/utils/__init__.py index 682b2caa..aec053b8 100644 --- a/aiida_workgraph/utils/__init__.py +++ b/aiida_workgraph/utils/__init__.py @@ -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, diff --git a/docs/source/built-in/pythonjob.ipynb b/docs/source/built-in/pythonjob.ipynb index 07b9c160..88ed8c4d 100644 --- a/docs/source/built-in/pythonjob.ipynb +++ b/docs/source/built-in/pythonjob.ipynb @@ -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",