Skip to content

Commit 38d8d36

Browse files
committed
Allow setting pyxis arguments through environment variables
Closes #153 Signed-off-by: Felix Abecassis <fabecassis@nvidia.com>
1 parent 521603b commit 38d8d36

File tree

8 files changed

+293
-3
lines changed

8 files changed

+293
-3
lines changed

args.c

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
#include "args.h"
1010
#include "common.h"
11+
#include "config.h"
1112

1213
static struct plugin_args pyxis_args = {
1314
.image = NULL,
@@ -149,6 +150,92 @@ struct spank_option spank_opts[] =
149150
SPANK_OPTIONS_TABLE_END
150151
};
151152

153+
static const char *get_env_var(spank_t sp, const char *var_name, char *buf, size_t buf_size)
154+
{
155+
spank_context_t ctx = spank_context();
156+
spank_err_t rc;
157+
158+
if (ctx == S_CTX_REMOTE) {
159+
rc = spank_getenv(sp, var_name, buf, buf_size);
160+
return (rc == ESPANK_SUCCESS) ? buf : NULL;
161+
} else {
162+
return getenv(var_name);
163+
}
164+
}
165+
166+
void pyxis_args_check_environment_variables(spank_t sp)
167+
{
168+
const char *env_val;
169+
char buf[PATH_MAX];
170+
int ret;
171+
172+
env_val = get_env_var(sp, "PYXIS_CONTAINER_IMAGE", buf, sizeof(buf));
173+
if (env_val != NULL && pyxis_args.image == NULL)
174+
spank_option_image(0, env_val, 0);
175+
176+
env_val = get_env_var(sp, "PYXIS_CONTAINER_MOUNTS", buf, sizeof(buf));
177+
if (env_val != NULL && pyxis_args.mounts_len == 0)
178+
spank_option_mount(0, env_val, 0);
179+
180+
env_val = get_env_var(sp, "PYXIS_CONTAINER_WORKDIR", buf, sizeof(buf));
181+
if (env_val != NULL && pyxis_args.workdir == NULL)
182+
spank_option_workdir(0, env_val, 0);
183+
184+
env_val = get_env_var(sp, "PYXIS_CONTAINER_NAME", buf, sizeof(buf));
185+
if (env_val != NULL && pyxis_args.container_name == NULL)
186+
spank_option_container_name(0, env_val, 0);
187+
188+
env_val = get_env_var(sp, "PYXIS_CONTAINER_SAVE", buf, sizeof(buf));
189+
if (env_val != NULL && pyxis_args.container_save == NULL)
190+
spank_option_container_save(0, env_val, 0);
191+
192+
env_val = get_env_var(sp, "PYXIS_CONTAINER_MOUNT_HOME", buf, sizeof(buf));
193+
if (env_val != NULL && pyxis_args.mount_home == -1) {
194+
ret = parse_bool(env_val);
195+
if (ret >= 0)
196+
spank_option_container_mount_home(ret, NULL, 0);
197+
}
198+
199+
env_val = get_env_var(sp, "PYXIS_CONTAINER_REMAP_ROOT", buf, sizeof(buf));
200+
if (env_val != NULL && pyxis_args.remap_root == -1) {
201+
ret = parse_bool(env_val);
202+
if (ret >= 0)
203+
spank_option_container_remap_root(ret, NULL, 0);
204+
}
205+
206+
env_val = get_env_var(sp, "PYXIS_CONTAINER_ENTRYPOINT", buf, sizeof(buf));
207+
if (env_val != NULL && pyxis_args.entrypoint == -1) {
208+
ret = parse_bool(env_val);
209+
if (ret >= 0)
210+
spank_option_container_entrypoint(ret, NULL, 0);
211+
}
212+
213+
env_val = get_env_var(sp, "PYXIS_CONTAINER_ENTRYPOINT_LOG", buf, sizeof(buf));
214+
if (env_val != NULL && pyxis_args.entrypoint_log == -1) {
215+
ret = parse_bool(env_val);
216+
if (ret >= 0)
217+
spank_option_container_entrypoint_log(ret, NULL, 0);
218+
}
219+
220+
env_val = get_env_var(sp, "PYXIS_CONTAINER_WRITABLE", buf, sizeof(buf));
221+
if (env_val != NULL && pyxis_args.writable == -1) {
222+
ret = parse_bool(env_val);
223+
if (ret >= 0)
224+
spank_option_container_writable(ret, NULL, 0);
225+
}
226+
227+
env_val = get_env_var(sp, "PYXIS_CONTAINER_READONLY", buf, sizeof(buf));
228+
if (env_val != NULL && pyxis_args.writable == -1) {
229+
ret = parse_bool(env_val);
230+
if (ret >= 0)
231+
spank_option_container_writable(!ret, NULL, 0);
232+
}
233+
234+
env_val = get_env_var(sp, "PYXIS_CONTAINER_ENV", buf, sizeof(buf));
235+
if (env_val != NULL && pyxis_args.env_vars_len == 0)
236+
spank_option_container_env(0, env_val, 0);
237+
}
238+
152239
static int spank_option_image(int val, const char *optarg, int remote)
153240
{
154241
if (optarg == NULL || *optarg == '\0') {

args.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,4 +37,6 @@ void remove_all_mounts(void);
3737

3838
void pyxis_args_free(void);
3939

40+
void pyxis_args_check_environment_variables(spank_t sp);
41+
4042
#endif /* ARGS_H_ */

config.c

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,12 @@
99

1010
#include "config.h"
1111

12-
static int parse_bool(const char *s)
12+
int parse_bool(const char *s)
1313
{
14-
if (strcmp(s, "1") == 0 || strcmp(s, "true") == 0)
14+
if (strcmp(s, "1") == 0 || strcmp(s, "true") == 0 || strcmp(s, "yes") == 0 || strcmp(s, "y") == 0)
1515
return (true);
1616

17-
if (strcmp(s, "0") == 0 || strcmp(s, "false") == 0)
17+
if (strcmp(s, "0") == 0 || strcmp(s, "false") == 0 || strcmp(s, "no") == 0 || strcmp(s, "n") == 0)
1818
return (false);
1919

2020
return (-1);

config.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,6 @@ struct plugin_config {
2222

2323
int pyxis_config_parse(struct plugin_config *config, int ac, char **av);
2424

25+
int parse_bool(const char *s);
26+
2527
#endif /* CONFIG_H_ */

pyxis_alloc.c

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,9 @@ int pyxis_alloc_init(spank_t sp, int ac, char **av)
4343

4444
int pyxis_alloc_post_opt(spank_t sp, int ac, char **av)
4545
{
46+
/* Check environment variables for default values after command-line processing */
47+
pyxis_args_check_environment_variables(sp);
48+
4649
/* Calling pyxis_args_enabled() for arguments validation */
4750
pyxis_args_enabled();
4851

pyxis_slurmstepd.c

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,9 @@ int pyxis_slurmstepd_post_opt(spank_t sp, int ac, char **av)
255255
{
256256
int ret;
257257

258+
/* Check environment variables for default values after command-line processing */
259+
pyxis_args_check_environment_variables(sp);
260+
258261
if (!pyxis_args_enabled())
259262
return (0);
260263

pyxis_srun.c

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@ int pyxis_srun_init(spank_t sp, int ac, char **av)
3030

3131
int pyxis_srun_post_opt(spank_t sp, int ac, char **av)
3232
{
33+
/* Check environment variables for default values after command-line processing */
34+
pyxis_args_check_environment_variables(sp);
35+
3336
/* Calling pyxis_args_enabled() for arguments validation */
3437
pyxis_args_enabled();
3538

tests/env_vars.bats

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
#!/usr/bin/env bats
2+
3+
load ./common
4+
5+
function setup() {
6+
enroot_cleanup env-test env-test2 || true
7+
}
8+
9+
function teardown() {
10+
enroot_cleanup env-test env-test2 || true
11+
}
12+
13+
@test "PYXIS_CONTAINER_IMAGE environment variable" {
14+
PYXIS_CONTAINER_IMAGE=ubuntu:22.04 run_srun grep 'Ubuntu 22.04' /etc/os-release
15+
}
16+
17+
@test "PYXIS_CONTAINER_IMAGE: command line precedence" {
18+
PYXIS_CONTAINER_IMAGE=ubuntu:22.04 run_srun --container-image=ubuntu:24.04 grep 'Ubuntu 24.04' /etc/os-release
19+
}
20+
21+
@test "PYXIS_CONTAINER_MOUNTS environment variable" {
22+
PYXIS_CONTAINER_MOUNTS=/home:/test-mnt run_srun --container-image=ubuntu:24.04 findmnt /test-mnt
23+
}
24+
25+
@test "PYXIS_CONTAINER_MOUNTS: command line precedence" {
26+
PYXIS_CONTAINER_MOUNTS=/home:/mnt-env run_srun --container-image=ubuntu:24.04 --container-mounts=/tmp:/mnt-cli bash -c 'findmnt /mnt-cli && ! findmnt /mnt-env'
27+
}
28+
29+
@test "PYXIS_CONTAINER_WORKDIR environment variable" {
30+
PYXIS_CONTAINER_WORKDIR=/usr/local/bin run_srun --container-image=ubuntu:24.04 pwd
31+
[ "${lines[-1]}" == "/usr/local/bin" ]
32+
}
33+
34+
@test "PYXIS_CONTAINER_WORKDIR: command line precedence" {
35+
PYXIS_CONTAINER_WORKDIR=/usr/local/bin run_srun --container-image=ubuntu:24.04 --container-workdir=/tmp pwd
36+
[ "${lines[-1]}" == "/tmp" ]
37+
}
38+
39+
@test "PYXIS_CONTAINER_NAME environment variable" {
40+
PYXIS_CONTAINER_NAME=env-test run_srun --container-image=ubuntu:24.04 --container-remap-root bash -c 'apt-get update && apt-get install -y --no-install-recommends file'
41+
run_srun --container-name=env-test which file
42+
}
43+
44+
@test "PYXIS_CONTAINER_NAME: command line precedence" {
45+
PYXIS_CONTAINER_NAME=env-test run_srun --container-image=ubuntu:24.04 --container-name=env-test2 --container-remap-root bash -c 'apt-get update && apt-get install -y --no-install-recommends make'
46+
run_srun --container-name=env-test2 which make
47+
run_srun_unchecked --container-name=env-test which make
48+
[ "${status}" -ne 0 ]
49+
}
50+
51+
@test "PYXIS_CONTAINER_SAVE environment variable" {
52+
readonly image="./env-test.sqsh"
53+
rm -f "${image}"
54+
PYXIS_CONTAINER_SAVE=${image} run_srun --container-image=ubuntu:24.04 sh -c 'echo pyxis > /test'
55+
sleep 1s
56+
run_srun --container-image=${image} cat /test
57+
[ "${lines[-1]}" == "pyxis" ]
58+
rm -f "${image}"
59+
}
60+
61+
@test "PYXIS_CONTAINER_SAVE: command line precedence" {
62+
readonly env_image="./env-test.sqsh"
63+
readonly cli_image="./cli-test.sqsh"
64+
rm -f "${env_image}" "${cli_image}"
65+
PYXIS_CONTAINER_SAVE=${env_image} run_srun --container-image=ubuntu:24.04 --container-save=${cli_image} sh -c 'echo pyxis > /test'
66+
sleep 1s
67+
run_srun --container-image=${cli_image} cat /test
68+
[ "${lines[-1]}" == "pyxis" ]
69+
[ ! -f "${env_image}" ]
70+
rm -f "${cli_image}"
71+
}
72+
73+
@test "PYXIS_CONTAINER_MOUNT_HOME=1 environment variable" {
74+
PYXIS_CONTAINER_MOUNT_HOME=1 run_srun --container-remap-root --container-image=ubuntu:24.04 findmnt /root
75+
}
76+
77+
@test "PYXIS_CONTAINER_MOUNT_HOME=0 environment variable" {
78+
PYXIS_CONTAINER_MOUNT_HOME=0 run_srun --container-remap-root --container-image=ubuntu:24.04 bash -c '! findmnt /root'
79+
}
80+
81+
@test "PYXIS_CONTAINER_MOUNT_HOME: command line precedence" {
82+
PYXIS_CONTAINER_MOUNT_HOME=1 run_srun --no-container-mount-home --container-remap-root --container-image=ubuntu:24.04 bash -c '! findmnt /root'
83+
}
84+
85+
@test "PYXIS_CONTAINER_REMAP_ROOT=1 environment variable" {
86+
PYXIS_CONTAINER_REMAP_ROOT=1 run_srun --container-image=ubuntu:24.04 bash -c 'cat /proc/self/uid_map'
87+
uidmap=(${lines[-1]})
88+
[ "${uidmap[0]}" -eq 0 ]
89+
[ "${uidmap[1]}" -eq $(id -u) ]
90+
[ "${uidmap[2]}" -eq 1 ]
91+
}
92+
93+
@test "PYXIS_CONTAINER_REMAP_ROOT=0 environment variable" {
94+
PYXIS_CONTAINER_REMAP_ROOT=0 run_srun --container-image=ubuntu:24.04 bash -c 'cat /proc/self/uid_map'
95+
uidmap=(${lines[-1]})
96+
[ "${uidmap[0]}" -eq $(id -u) ]
97+
[ "${uidmap[1]}" -eq $(id -u) ]
98+
[ "${uidmap[2]}" -eq 1 ]
99+
}
100+
101+
@test "PYXIS_CONTAINER_REMAP_ROOT: command line precedence" {
102+
PYXIS_CONTAINER_REMAP_ROOT=1 run_srun --no-container-remap-root --container-image=ubuntu:24.04 bash -c 'cat /proc/self/uid_map'
103+
uidmap=(${lines[-1]})
104+
[ "${uidmap[0]}" -eq $(id -u) ]
105+
[ "${uidmap[1]}" -eq $(id -u) ]
106+
[ "${uidmap[2]}" -eq 1 ]
107+
}
108+
109+
@test "PYXIS_CONTAINER_ENTRYPOINT=1 environment variable" {
110+
if srun bash -c '[ -f /etc/enroot/entrypoint ]'; then
111+
skip "entrypoint disabled by enroot"
112+
fi
113+
114+
PYXIS_CONTAINER_ENTRYPOINT=1 run_srun --container-image=docker:26.1.0-dind-rootless sh -c '[ -n "${DOCKER_HOST}" ]'
115+
}
116+
117+
@test "PYXIS_CONTAINER_ENTRYPOINT=0 environment variable" {
118+
PYXIS_CONTAINER_ENTRYPOINT=0 run_srun --container-image=docker:26.1.0-dind-rootless sh -c '[ -z "${DOCKER_HOST}" ]'
119+
}
120+
121+
@test "PYXIS_CONTAINER_ENTRYPOINT: command line precedence" {
122+
if srun bash -c '[ -f /etc/enroot/entrypoint ]'; then
123+
skip "entrypoint disabled by enroot"
124+
fi
125+
126+
PYXIS_CONTAINER_ENTRYPOINT=0 run_srun --container-entrypoint --container-image=docker:26.1.0-dind-rootless sh -c '[ -n "${DOCKER_HOST}" ]'
127+
}
128+
129+
@test "PYXIS_CONTAINER_ENTRYPOINT_LOG=1 environment variable" {
130+
if srun bash -c '[ -f /etc/enroot/entrypoint ]'; then
131+
skip "entrypoint disabled by enroot"
132+
fi
133+
134+
PYXIS_CONTAINER_ENTRYPOINT_LOG=1 run_srun --container-entrypoint --container-image=nvidia/cuda:12.9.1-runtime-ubuntu24.04 true
135+
grep -q "== CUDA ==" <<< "${output}"
136+
}
137+
138+
@test "PYXIS_CONTAINER_ENTRYPOINT_LOG=0 environment variable" {
139+
if srun bash -c '[ -f /etc/enroot/entrypoint ]'; then
140+
skip "entrypoint disabled by enroot"
141+
fi
142+
143+
PYXIS_CONTAINER_ENTRYPOINT_LOG=0 run_srun --container-entrypoint --container-image=nvidia/cuda:12.9.1-runtime-ubuntu24.04 true
144+
! grep -q "== CUDA ==" <<< "${output}"
145+
}
146+
147+
@test "PYXIS_CONTAINER_WRITABLE=1 environment variable" {
148+
PYXIS_CONTAINER_WRITABLE=1 run_srun --container-image=ubuntu:24.04 touch /newfile
149+
}
150+
151+
@test "PYXIS_CONTAINER_WRITABLE=0 environment variable" {
152+
PYXIS_CONTAINER_WRITABLE=0 run_srun_unchecked --container-image=ubuntu:24.04 touch /newfile
153+
[ "${status}" -ne 0 ]
154+
grep -q 'Read-only file system' <<< "${output}"
155+
}
156+
157+
@test "PYXIS_CONTAINER_WRITABLE: command line precedence" {
158+
PYXIS_CONTAINER_WRITABLE=1 run_srun_unchecked --container-readonly --container-image=ubuntu:24.04 touch /newfile
159+
[ "${status}" -ne 0 ]
160+
grep -q 'Read-only file system' <<< "${output}"
161+
}
162+
163+
@test "PYXIS_CONTAINER_READONLY=1 environment variable" {
164+
PYXIS_CONTAINER_READONLY=1 run_srun_unchecked --container-image=ubuntu:24.04 touch /newfile
165+
[ "${status}" -ne 0 ]
166+
grep -q 'Read-only file system' <<< "${output}"
167+
}
168+
169+
@test "PYXIS_CONTAINER_READONLY=0 environment variable" {
170+
PYXIS_CONTAINER_READONLY=0 run_srun_unchecked --container-image=ubuntu:24.04 touch /newfile
171+
}
172+
173+
@test "PYXIS_CONTAINER_READONLY: command line precedence" {
174+
PYXIS_CONTAINER_READONLY=0 run_srun_unchecked --container-readonly --container-image=ubuntu:24.04 touch /newfile
175+
[ "${status}" -ne 0 ]
176+
grep -q 'Read-only file system' <<< "${output}"
177+
}
178+
179+
@test "PYXIS_CONTAINER_ENV environment variable" {
180+
export CUDA_VERSION=11.0.0
181+
PYXIS_CONTAINER_ENV=CUDA_VERSION run_srun --no-container-mount-home --container-image=nvidia/cuda:12.9.1-base-ubuntu24.04 sh -c 'echo $CUDA_VERSION'
182+
[ "${lines[-1]}" == "11.0.0" ]
183+
}
184+
185+
@test "PYXIS_CONTAINER_ENV: command line precedence" {
186+
export CUDA_VERSION=12.0.0
187+
export LD_LIBRARY_PATH=/usr/local/cuda/lib64
188+
PYXIS_CONTAINER_ENV=CUDA_VERSION run_srun --no-container-mount-home --container-image=nvidia/cuda:12.9.1-base-ubuntu24.04 --container-env=LD_LIBRARY_PATH sh -c 'echo $CUDA_VERSION $LD_LIBRARY_PATH'
189+
[[ "${lines[-1]}" == "12.9.1 /usr/local/cuda/lib64" ]]
190+
}

0 commit comments

Comments
 (0)