diff --git a/easybuild/easyblocks/d/deepspeed.py b/easybuild/easyblocks/d/deepspeed.py
new file mode 100644
index 0000000000..55e9db3f97
--- /dev/null
+++ b/easybuild/easyblocks/d/deepspeed.py
@@ -0,0 +1,105 @@
+##
+# Copyright 2009-2024 Ghent University
+#
+# This file is part of EasyBuild,
+# originally created by the HPC team of Ghent University (http://ugent.be/hpc/en),
+# with support of Ghent University (http://ugent.be/hpc),
+# the Flemish Supercomputer Centre (VSC) (https://www.vscentrum.be),
+# Flemish Research Foundation (FWO) (http://www.fwo.be/en)
+# and the Department of Economy, Science and Innovation (EWI) (http://www.ewi-vlaanderen.be/en).
+#
+# https://github.com/easybuilders/easybuild
+#
+# EasyBuild is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation v2.
+#
+# EasyBuild is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with EasyBuild. If not, see .
+##
+"""
+EasyBuild support for building and installing DeepSpeed, implemented as an easyblock
+
+author: Viktor Rehnberg (Chalmers University of Technology)
+"""
+from easybuild.easyblocks.generic.pythonpackage import PythonPackage
+from easybuild.framework.easyconfig import CUSTOM
+from easybuild.tools.build_log import EasyBuildError
+from easybuild.tools.config import build_option
+import easybuild.tools.environment as env
+
+
+class EB_DeepSpeed(PythonPackage):
+ """Custom easyblock for DeepSpeed"""
+
+ @staticmethod
+ def extra_options():
+ """Change some defaults for easyconfig parameters."""
+ extra_vars = PythonPackage.extra_options()
+ extra_vars['use_pip'][0] = True
+ extra_vars['download_dep_fail'][0] = True
+ extra_vars['sanity_pip_check'][0] = True
+
+ # Add DeepSpeed specific vars
+ extra_vars['ds_build_ops_to_skip'] = [[], "For in list will set DS_BUILD_=0 (default: [])", CUSTOM]
+ return extra_vars
+
+ def __init__(self, *args, **kwargs):
+ """Initialize DeepSpeed easyblock."""
+ super().__init__(*args, **kwargs)
+
+ dep_names = set(dep['name'] for dep in self.cfg.dependencies())
+
+ # enable building with GPU support if CUDA is included as dependency
+ if 'CUDA' in dep_names:
+ self.with_cuda = True
+ else:
+ self.with_cuda = False
+
+ @property
+ def cuda_compute_capabilities(self):
+ return self.cfg['cuda_compute_capabilities'] or build_option('cuda_compute_capabilities')
+
+ def configure_step(self):
+ """Set up DeepSpeed config"""
+ # require that PyTorch is listed as dependency
+ dep_names = set(dep['name'] for dep in self.cfg.dependencies())
+ if 'PyTorch' not in dep_names:
+ raise EasyBuildError('PyTorch not found as a dependency')
+
+ if self.with_cuda:
+ # https://github.com/microsoft/DeepSpeed/issues/3358
+ env.setvar('NVCC_PREPEND_FLAGS', '--forward-unknown-opts')
+
+ if self.cuda_compute_capabilities:
+ # specify CUDA compute capabilities via $TORCH_CUDA_ARCH_LIST
+ env.setvar('TORCH_CUDA_ARCH_LIST', ';'.join(self.cuda_compute_capabilities))
+
+ # By default prebuild all ops with a few exceptions
+ # http://www.deepspeed.ai/tutorials/advanced-install/#pre-install-deepspeed-ops
+ # > DeepSpeed will only install any ops that are compatible with your machine
+ env.setvar('DS_BUILD_OPS', '1')
+
+ # Some may be problematic for different reasons, these are specified in the easyconfig
+ for opt in self.cfg['ds_build_ops_to_skip']:
+ env.setvar('DS_BUILD_{}'.format(opt), '0')
+
+ super().configure_step()
+
+ def sanity_check_step(self):
+ '''Custom sanity check for DeepSpeed.'''
+ custom_paths = {
+ 'files': ['bin/deepspeed'],
+ 'dirs': [],
+ }
+ custom_commands = [
+ 'deepspeed --help',
+ 'python -m deepspeed.env_report',
+ '[ "$(ds_report | grep -c "\\[NO\\]")" -eq "{:d}" ]'.format(len(self.cfg['ds_build_ops_to_skip']))
+ ]
+ return super().sanity_check_step(custom_paths=custom_paths, custom_commands=custom_commands)
diff --git a/easybuild/easyblocks/generic/pythonpackage.py b/easybuild/easyblocks/generic/pythonpackage.py
index eae50cae7c..93142dbb9e 100644
--- a/easybuild/easyblocks/generic/pythonpackage.py
+++ b/easybuild/easyblocks/generic/pythonpackage.py
@@ -861,8 +861,13 @@ def test_step(self, return_output_ec=False):
# print Python search path (just debugging purposes)
run_cmd("%s -c 'import sys; print(sys.path)'" % self.python_cmd, verbose=False, trace=False)
+ abs_bindirs = [os.path.join(actual_installdir, 'bin')]
abs_pylibdirs = [os.path.join(actual_installdir, pylibdir) for pylibdir in self.all_pylibdirs]
- extrapath = "export PYTHONPATH=%s &&" % os.pathsep.join(abs_pylibdirs + ['$PYTHONPATH'])
+ extrapath = "export PATH={path} PYTHONPATH={pythonpath} LD_LIBRARY_PATH={ld_library_path} &&".format(
+ path=os.pathsep.join(abs_bindirs + ['$PATH']),
+ pythonpath=os.pathsep.join(abs_pylibdirs + ['$PYTHONPATH']),
+ ld_library_path=os.pathsep.join(abs_pylibdirs + ['$LD_LIBRARY_PATH']),
+ )
cmd = self.compose_install_command(test_installdir, extrapath=extrapath)
run_cmd(cmd, log_all=True, simple=True, verbose=False)