From 68333337b8e198e9c720b626c890a27f62c60645 Mon Sep 17 00:00:00 2001 From: Felix Abecassis Date: Fri, 6 Aug 2021 15:22:40 -0700 Subject: [PATCH] Add support for containerized sbatch/salloc Signed-off-by: Felix Abecassis --- Makefile | 2 +- README.md | 17 ++++++++++++ args.c | 2 +- args.h | 2 ++ common.h | 5 ++++ config.c | 9 ++++++ config.h | 1 + pyxis_alloc.c | 59 ++++++++++++++++++++++++++++++++++++++++ pyxis_alloc.h | 16 +++++++++++ pyxis_dispatch.c | 7 +++++ pyxis_slurmstepd.c | 24 ++++++++++++++++ tests/common.bash | 18 ++++++++++++ tests/sbatch/sbatch.bats | 29 ++++++++++++++++++++ 13 files changed, 189 insertions(+), 2 deletions(-) create mode 100644 pyxis_alloc.c create mode 100644 pyxis_alloc.h create mode 100644 tests/sbatch/sbatch.bats diff --git a/Makefile b/Makefile index bdfe2a6..764624c 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,7 @@ CPPFLAGS := -D_GNU_SOURCE -D_FORTIFY_SOURCE=2 -DPYXIS_VERSION=\"$(VERSION)\" $(C CFLAGS := -std=gnu11 -O2 -g -Wall -Wunused-variable -fstack-protector-strong -fpic $(CFLAGS) LDFLAGS := -Wl,-znoexecstack -Wl,-zrelro -Wl,-znow $(LDFLAGS) -C_SRCS := common.c args.c pyxis_slurmstepd.c pyxis_slurmd.c pyxis_srun.c pyxis_dispatch.c config.c enroot.c +C_SRCS := common.c args.c pyxis_slurmstepd.c pyxis_slurmd.c pyxis_srun.c pyxis_alloc.c pyxis_dispatch.c config.c enroot.c C_OBJS := $(C_SRCS:.c=.o) DEPS := $(C_OBJS:%.o=%.d) diff --git a/README.md b/README.md index 7f03528..74872cc 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,7 @@ $ srun --help ## Examples +### `srun` ```console $ # Run a command on a worker node $ srun grep PRETTY /etc/os-release @@ -105,6 +106,21 @@ $ srun --container-image=centos --container-mounts=/etc/os-release:/host/os-rele PRETTY_NAME="Ubuntu 18.04.2 LTS" ``` +### `sbatch` +```console +$ # execute the sbatch script inside a container image +$ sbatch --wait -o slurm.out <remap_root = true; config->execute_entrypoint = false; config->container_scope = SCOPE_GLOBAL; + config->sbatch_support = true; for (int i = 0; i < ac; ++i) { if (strncmp("runtime_path=", av[i], 13) == 0) { @@ -69,6 +70,14 @@ int pyxis_config_parse(struct plugin_config *config, int ac, char **av) slurm_error("pyxis: container_scope: invalid value: %s", optarg); return (-1); } + } else if (strncmp("sbatch_support=", av[i], 15) == 0) { + optarg = av[i] + 15; + ret = parse_bool(optarg); + if (ret < 0) { + slurm_error("pyxis: sbatch_support: invalid value: %s", optarg); + return (-1); + } + config->sbatch_support = ret; } else { slurm_error("pyxis: unknown configuration option: %s", av[i]); return (-1); diff --git a/config.h b/config.h index e0486dc..91fb230 100644 --- a/config.h +++ b/config.h @@ -18,6 +18,7 @@ struct plugin_config { bool remap_root; bool execute_entrypoint; enum container_scope container_scope; + bool sbatch_support; }; int pyxis_config_parse(struct plugin_config *config, int ac, char **av); diff --git a/pyxis_alloc.c b/pyxis_alloc.c new file mode 100644 index 0000000..de68844 --- /dev/null +++ b/pyxis_alloc.c @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + */ + +#include +#include +#include + +#include "pyxis_alloc.h" +#include "args.h" +#include "config.h" + +struct plugin_context { + struct plugin_args *args; +}; + +static struct plugin_context context = { + .args = NULL, +}; + +int pyxis_alloc_init(spank_t sp, int ac, char **av) +{ + int ret; + struct plugin_config config; + + ret = pyxis_config_parse(&config, ac, av); + if (ret < 0) { + slurm_error("pyxis: failed to parse configuration"); + return (-1); + } + + if (!config.sbatch_support) + return (0); + + context.args = pyxis_args_register(sp); + if (context.args == NULL) { + slurm_error("pyxis: failed to register arguments"); + return (-1); + } + + return (0); +} + +int pyxis_alloc_post_opt(spank_t sp, int ac, char **av) +{ + /* Calling pyxis_args_enabled() for arguments validation */ + pyxis_args_enabled(); + + return (0); +} + +int pyxis_alloc_exit(spank_t sp, int ac, char **av) +{ + pyxis_args_free(); + + memset(&context, 0, sizeof(context)); + + return (0); +} diff --git a/pyxis_alloc.h b/pyxis_alloc.h new file mode 100644 index 0000000..576beb8 --- /dev/null +++ b/pyxis_alloc.h @@ -0,0 +1,16 @@ +/* + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + */ + +#ifndef PYXIS_ALLOC_H_ +#define PYXIS_ALLOC_H_ + +#include + +int pyxis_alloc_init(spank_t sp, int ac, char **av); + +int pyxis_alloc_post_opt(spank_t sp, int ac, char **av); + +int pyxis_alloc_exit(spank_t sp, int ac, char **av); + +#endif /* PYXIS_ALLOC_H_ */ diff --git a/pyxis_dispatch.c b/pyxis_dispatch.c index 4c05dc6..d73004e 100644 --- a/pyxis_dispatch.c +++ b/pyxis_dispatch.c @@ -6,6 +6,7 @@ #include "pyxis_slurmd.h" #include "pyxis_srun.h" +#include "pyxis_alloc.h" #include "pyxis_slurmstepd.h" SPANK_PLUGIN(pyxis, 1) @@ -17,6 +18,8 @@ int slurm_spank_init(spank_t sp, int ac, char **av) return pyxis_slurmd_init(sp, ac, av); case S_CTX_LOCAL: return pyxis_srun_init(sp, ac, av); + case S_CTX_ALLOCATOR: + return pyxis_alloc_init(sp, ac, av); case S_CTX_REMOTE: return pyxis_slurmstepd_init(sp, ac, av); default: @@ -29,6 +32,8 @@ int slurm_spank_init_post_opt(spank_t sp, int ac, char **av) switch (spank_context()) { case S_CTX_LOCAL: return pyxis_srun_post_opt(sp, ac, av); + case S_CTX_ALLOCATOR: + return pyxis_alloc_post_opt(sp, ac, av); case S_CTX_REMOTE: return pyxis_slurmstepd_post_opt(sp, ac, av); default: @@ -43,6 +48,8 @@ int slurm_spank_exit(spank_t sp, int ac, char **av) return pyxis_slurmd_exit(sp, ac, av); case S_CTX_LOCAL: return pyxis_srun_exit(sp, ac, av); + case S_CTX_ALLOCATOR: + return pyxis_alloc_exit(sp, ac, av); case S_CTX_REMOTE: return pyxis_slurmstepd_exit(sp, ac, av); default: diff --git a/pyxis_slurmstepd.c b/pyxis_slurmstepd.c index 49dd934..54e6ce3 100644 --- a/pyxis_slurmstepd.c +++ b/pyxis_slurmstepd.c @@ -866,6 +866,9 @@ static int shm_destroy(struct shared_memory *shm) int slurm_spank_user_init(spank_t sp, int ac, char **av) { int ret; + spank_err_t rc; + int spank_argc = 0; + char **spank_argv = NULL; char *container_name = NULL; pid_t pid; int rv = -1; @@ -881,6 +884,27 @@ int slurm_spank_user_init(spank_t sp, int ac, char **av) if (ret < 0) goto fail; + if (context.job.stepid == SLURM_BATCH_SCRIPT) { + rc = spank_get_item(sp, S_JOB_ARGV, &spank_argc, &spank_argv); + if (rc != ESPANK_SUCCESS) { + slurm_error("pyxis: couldn't get job argv: %s", spank_strerror(rc)); + goto fail; + } + + if (spank_argc == 0) { + slurm_error("pyxis: couldn't get sbatch script: argc == 0"); + goto fail; + } + + /* Mount the sbatch script (from the Slurmd spool dir) inside the container */ + ret = add_mount(spank_argv[0], spank_argv[0], + "x-create=file,bind,ro,nosuid,nodev,private"); + if (ret < 0) { + slurm_error("pyxis: couldn't add bind mount for sbatch script"); + goto fail; + } + } + if (context.args->container_name != NULL) { if (context.config.container_scope == SCOPE_JOB) ret = xasprintf(&container_name, "pyxis_%u_%s", context.job.jobid, context.args->container_name); diff --git a/tests/common.bash b/tests/common.bash index 50e0da8..778b29c 100644 --- a/tests/common.bash +++ b/tests/common.bash @@ -32,6 +32,24 @@ function run_srun() { [ "${status}" -eq 0 ] } +function run_sbatch_unchecked() { + log "+ sbatch --wait $@" + + slurm_log=$(mktemp) + run sbatch --wait -o "${slurm_log}" "$@" + + log "${output}" + logf "+ job log (${slurm_log}):\n$(cat ${slurm_log})" + rm -f "${slurm_log}" + + echo "+ exit status: ${status}" +} + +function run_sbatch() { + run_sbatch_unchecked "$@" + [ "${status}" -eq 0 ] +} + function enroot_cleanup() { for name in "$@"; do { srun enroot remove -f pyxis_${name} || srun enroot remove -f pyxis_${SLURM_JOB_ID}_${name}; } >/dev/null 2>&1 diff --git a/tests/sbatch/sbatch.bats b/tests/sbatch/sbatch.bats new file mode 100644 index 0000000..0cad350 --- /dev/null +++ b/tests/sbatch/sbatch.bats @@ -0,0 +1,29 @@ +#!/usr/bin/env bats + +load ../common + +@test "sbatch ubuntu:16.04" { + run_sbatch --container-image=ubuntu:16.04 <