Add support for containerized sbatch/salloc

Signed-off-by: Felix Abecassis <[email protected]>

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)

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 <<EOF
+#!/bin/bash -eux
+#SBATCH --container-image centos:8
+grep PRETTY /etc/os-release
+EOF
+
+$ cat slurm.out
+pyxis: importing docker image ...
++ grep PRETTY /etc/os-release
+PRETTY_NAME="CentOS Linux 8"
+```
+
 ## Advanced Documentation (wiki)
 1. [Home](https://github.com/NVIDIA/pyxis/wiki/Home)
 1. [Installation](https://github.com/NVIDIA/pyxis/wiki/Installation)
@@ -124,6 +140,7 @@ This project is released under the [Apache License Version 2.0](https://github.c
 Integration tests can be ran with [bats](https://github.com/bats-core/bats-core) from within a Slurm job allocation:
 ```console
 $ salloc --overcommit bats tests
+$ bats tests/sbatch
 ```
 Some tests assume a specific enroot configuration (such as PMIx/PyTorch hooks), so they might not pass on all systems.
 

args.c

@@ -183,7 +183,7 @@ static int add_mount_entry(const char *entry)
 	return (rv);
 }
 
-static int add_mount(const char *source, const char *target, const char *flags)
+int add_mount(const char *source, const char *target, const char *flags)
 {
 	int ret;
 	char *entry = NULL;

args.h

@@ -27,6 +27,8 @@ struct plugin_args *pyxis_args_register(spank_t sp);
 
 bool pyxis_args_enabled(void);
 
+int add_mount(const char *source, const char *target, const char *flags);
+
 void pyxis_args_free(void);
 
 #endif /* ARGS_H_ */

common.h

@@ -16,6 +16,11 @@ static inline void xclose(int fd)
 
 int xasprintf(char **strp, const char *fmt, ...);
 
+/* https://github.com/SchedMD/slurm/blob/slurm-20-11-8-1/slurm/slurm.h.in#L161-L162 */
+#if !defined(SLURM_BATCH_SCRIPT)
+# define SLURM_BATCH_SCRIPT (0xfffffffb)
+#endif
+
 #if !defined(MFD_CLOEXEC)
 # define MFD_CLOEXEC 0x0001U
 #endif /* !defined(MFD_CLOEXEC) */

config.c

@@ -35,6 +35,7 @@ int pyxis_config_parse(struct plugin_config *config, int ac, char **av)
 	config->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);

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);

pyxis_alloc.c

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
+ */
+
+#include <errno.h>
+#include <stdlib.h>
+#include <string.h>
+
+#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);
+}

pyxis_alloc.h

@@ -0,0 +1,16 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
+ */
+
+#ifndef PYXIS_ALLOC_H_
+#define PYXIS_ALLOC_H_
+
+#include <slurm/spank.h>
+
+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_ */

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:

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);

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

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 <<EOF
+#!/bin/bash
+grep 'Ubuntu 16.04' /etc/os-release
+EOF
+
+    run_sbatch <<EOF
+#!/bin/bash
+#SBATCH --container-image=ubuntu:16.04
+grep 'Ubuntu 16.04' /etc/os-release
+EOF
+}
+
+@test "sbatch centos:8" {
+    run_sbatch --container-image=centos:8 <<EOF
+#!/bin/bash
+grep 'CentOS Linux release 8.' /etc/redhat-release
+EOF
+
+    run_sbatch <<EOF
+#!/bin/bash
+#SBATCH --container-image=centos:8
+grep 'CentOS Linux release 8.' /etc/redhat-release
+EOF
+}