Fix artifact upload using wrong slurm job in SharedSlurmJob mode

core/modes/slurm/shared.py

@@ -4,6 +4,8 @@ from .srun import Slurm
 
 
 class SingleSlurmJobAcrossStages(Slurm, ABC):
+    substeps_in_shared_job = ('build_script', 'step_script', 'get_sources',
+                            'upload_artifacts_on_success', 'upload_artifacts_on_failure')
 
     def get_jobid_from_file(self, path):
         with open(path, 'r') as node_index_fp:
@@ -66,12 +68,12 @@ class SingleSlurmJobAcrossStages(Slurm, ABC):
         return parameters
 
     def get_simple_run_wrapper(self):
-        if self.job.args[1] == 'run' and self.job.args[3] in ['build_script', 'step_script', 'get_sources']:
+        if self.job.args[1] == 'run' and self.job.args[3] in self.substeps_in_shared_job:
             return self.get_run_wrapper()
         return Slurm.get_run_wrapper(self)
 
     def get_simple_run_parameters(self):
-        if self.job.args[1] == 'run' and self.job.args[3] in ['build_script', 'step_script', 'get_sources']:
+        if self.job.args[1] == 'run' and self.job.args[3] in self.substeps_in_shared_job:
             return self.get_run_parameters()
         return Slurm.get_simple_run_parameters(self)