Skip to content

Commit

Permalink
Pass environment variables to singularity
Browse files Browse the repository at this point in the history
Previously, environment variables were not visible to singularity if
running with --contain.

Also switch order of setting environment variables and pre_job_script to
allow the script to use the environment variables (useful to write more
generic scripts).
  • Loading branch information
mseitzer committed Apr 26, 2024
1 parent 0568241 commit 1127f54
Showing 1 changed file with 16 additions and 7 deletions.
23 changes: 16 additions & 7 deletions cluster_utils/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,8 @@
import os
import pathlib
import time
import typing
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any, Dict, Optional, Union

import pandas as pd

Expand Down Expand Up @@ -103,15 +102,19 @@ def generate_execution_cmd(self, paths, cmd_prefix: Optional[str] = None):
set_cwd = "cd {}".format(paths["main_path"])

if "variables" in paths:
if not isinstance(paths, dict):
if not isinstance(paths["variables"], dict):
raise ValueError(
'Expected type dict for "variables", but got type'
f' {type(paths["variables"])} instead'
)
env_variables = {
str(name): str(value) for name, value in paths["variables"].items()
}
set_env_variables = "\n".join(
f"export {name}={value}" for name, value in paths["variables"].items()
f"export {name}={value}" for name, value in env_variables.items()
)
else:
env_variables = None
set_env_variables = ""

if "pre_job_script" in paths:
Expand Down Expand Up @@ -181,6 +184,7 @@ def generate_execution_cmd(self, paths, cmd_prefix: Optional[str] = None):
self.singularity_settings,
paths["main_path"],
current_setting["working_dir"],
env_variables,
)

if cmd_prefix:
Expand All @@ -189,10 +193,10 @@ def generate_execution_cmd(self, paths, cmd_prefix: Optional[str] = None):
res = "\n".join(
[
set_cwd,
set_env_variables,
pre_job_script,
virtual_env_activate,
conda_env_activate,
set_env_variables,
exec_cmd,
]
)
Expand All @@ -202,8 +206,9 @@ def singularity_wrap(
self,
exec_cmd: str,
singularity_settings: SingularitySettings,
exec_dir: typing.Union[str, os.PathLike],
working_dir: typing.Union[str, os.PathLike],
exec_dir: Union[str, os.PathLike],
working_dir: Union[str, os.PathLike],
env_variables: Optional[Dict[str, str]],
) -> str:
"""Wrap the given command to execute it in a Singularity container.
Expand All @@ -223,6 +228,9 @@ def singularity_wrap(
# create model directory (so it can be bound into the container)
working_dir.mkdir(exist_ok=True)

if env_variables is None:
env_variables = {}

# construct singularity command
cwd = os.fspath(exec_dir)
bind_dirs = ["/tmp", os.fspath(working_dir), cwd]
Expand All @@ -231,6 +239,7 @@ def singularity_wrap(
"run" if singularity_settings.use_run else "exec",
"--bind=%s" % ",".join(bind_dirs),
"--pwd=%s" % cwd,
", ".join(f"--env {name}={value}" for name, value in env_variables.items()),
*singularity_settings.args,
os.fspath(singularity_image),
]
Expand Down

0 comments on commit 1127f54

Please sign in to comment.