Skip to content

Commit 146e775

Browse files
jeremymanningclaude
andcommitted
Issue #120: Make PBS and SGE run the same job body as SLURM
PBS did not work at all. Its script ended with cd {job_dir} source venv/bin/activate python execute_function.py and `execute_function.py` appears exactly once in the whole codebase -- on that line. Nothing in clustrix has ever created it, so every PBS job died immediately with "can't open file". SGE carried its own copy of the older single-venv script, which meant it silently missed everything fixed on the SLURM one: the two-venv path, the dill handoff, conda sourcing, the per-stage error files and result signing. Three copies of the same logic drifting apart is how PBS came to run a nonexistent file without anyone noticing. The execution body is now one function, `job_execution_lines()`, and all three schedulers call it. Only the directive header differs, which a test asserts directly by comparing the bodies. slurm two_venv=True conda_sourced=True signs_result=True pbs two_venv=True conda_sourced=True signs_result=True sge two_venv=True conda_sourced=True signs_result=True Neither PBS nor SGE has been run against real hardware -- there is none to hand -- so this fixes a script that could not possibly have worked rather than claiming the backends are verified. That distinction stays in the release notes. 16 new tests in tests/unit/test_two_venv_execution.py. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012gTBDPK16HUZ3kHQ2QyjuU
1 parent d132b25 commit 146e775

2 files changed

Lines changed: 115 additions & 104 deletions

File tree

‎clustrix/utils.py‎

Lines changed: 48 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -1425,37 +1425,16 @@ def create_job_script(
14251425
raise ValueError(f"Unsupported cluster type: {cluster_type}")
14261426

14271427

1428-
def _create_slurm_script(
1429-
job_config: Dict[str, Any], remote_job_dir: str, config: ClusterConfig
1430-
) -> str:
1431-
"""Create SLURM job script."""
1432-
1433-
script_lines = [
1434-
"#!/bin/bash",
1435-
"#SBATCH --job-name=clustrix",
1436-
f"#SBATCH --output={remote_job_dir}/slurm-%j.out",
1437-
f"#SBATCH --error={remote_job_dir}/slurm-%j.err",
1438-
f"#SBATCH --cpus-per-task={job_config['cores']}",
1439-
f"#SBATCH --mem={normalize_memory(job_config['memory'], 'slurm')}",
1440-
f"#SBATCH --time={job_config['time']}",
1441-
]
1442-
1443-
if job_config.get("partition"):
1444-
script_lines.append(f"#SBATCH --partition={job_config['partition']}")
1445-
1446-
# Add environment setup
1447-
if config.module_loads:
1448-
for module in config.module_loads:
1449-
script_lines.append(f"module load {module}")
1450-
1451-
if config.environment_variables:
1452-
for var, value in config.environment_variables.items():
1453-
script_lines.append(f"export {var}={value}")
1454-
1455-
if config.pre_execution_commands:
1456-
for cmd in config.pre_execution_commands:
1457-
script_lines.append(cmd)
1458-
1428+
def job_execution_lines(remote_job_dir: str, config: ClusterConfig) -> list:
1429+
"""The lines that actually run the user's function in a job script.
1430+
1431+
Shared by every scheduler. It used to live only inside the SLURM
1432+
generator: PBS ran `python execute_function.py`, a file nothing in
1433+
clustrix has ever created, and SGE carried its own divergent copy of
1434+
the single-venv script. Both therefore missed the two-venv path, the
1435+
result signing and every fix made to the SLURM one.
1436+
"""
1437+
script_lines: list = []
14591438
# Add execution commands
14601439
python_cmd = config.python_executable if config.python_executable else "python"
14611440

@@ -1518,6 +1497,42 @@ def _create_slurm_script(
15181497
]
15191498
)
15201499

1500+
return script_lines
1501+
1502+
1503+
def _create_slurm_script(
1504+
job_config: Dict[str, Any], remote_job_dir: str, config: ClusterConfig
1505+
) -> str:
1506+
"""Create SLURM job script."""
1507+
1508+
script_lines = [
1509+
"#!/bin/bash",
1510+
"#SBATCH --job-name=clustrix",
1511+
f"#SBATCH --output={remote_job_dir}/slurm-%j.out",
1512+
f"#SBATCH --error={remote_job_dir}/slurm-%j.err",
1513+
f"#SBATCH --cpus-per-task={job_config['cores']}",
1514+
f"#SBATCH --mem={normalize_memory(job_config['memory'], 'slurm')}",
1515+
f"#SBATCH --time={job_config['time']}",
1516+
]
1517+
1518+
if job_config.get("partition"):
1519+
script_lines.append(f"#SBATCH --partition={job_config['partition']}")
1520+
1521+
# Add environment setup
1522+
if config.module_loads:
1523+
for module in config.module_loads:
1524+
script_lines.append(f"module load {module}")
1525+
1526+
if config.environment_variables:
1527+
for var, value in config.environment_variables.items():
1528+
script_lines.append(f"export {var}={value}")
1529+
1530+
if config.pre_execution_commands:
1531+
for cmd in config.pre_execution_commands:
1532+
script_lines.append(cmd)
1533+
1534+
script_lines.extend(job_execution_lines(remote_job_dir, config))
1535+
15211536
return "\n".join(script_lines)
15221537

15231538

@@ -1550,14 +1565,7 @@ def _create_pbs_script(
15501565
for cmd in config.pre_execution_commands:
15511566
script_lines.append(cmd)
15521567

1553-
# Add similar execution logic as SLURM
1554-
script_lines.extend(
1555-
[
1556-
f"cd {remote_job_dir}",
1557-
"source venv/bin/activate",
1558-
"python execute_function.py",
1559-
]
1560-
)
1568+
script_lines.extend(job_execution_lines(remote_job_dir, config))
15611569

15621570
return "\n".join(script_lines)
15631571

@@ -1590,71 +1598,7 @@ def _create_sge_script(
15901598
for cmd in config.pre_execution_commands:
15911599
script_lines.append(cmd)
15921600

1593-
script_lines.extend(
1594-
[
1595-
f"cd {remote_job_dir}",
1596-
"source venv/bin/activate",
1597-
f'{config.python_executable if config.python_executable else "python"} -c "',
1598-
"import pickle",
1599-
"import sys",
1600-
"import traceback",
1601-
"",
1602-
"try:",
1603-
" import dill",
1604-
"except ImportError:",
1605-
" dill = None",
1606-
"try:",
1607-
" import cloudpickle",
1608-
"except ImportError:",
1609-
" cloudpickle = None",
1610-
"",
1611-
"try:",
1612-
" with open('function_data.pkl', 'rb') as f:",
1613-
" data = pickle.load(f)",
1614-
" ",
1615-
" # Try dill first, then cloudpickle, then source code",
1616-
" func = None",
1617-
" if dill:",
1618-
" try:",
1619-
" func = dill.loads(data['function'])",
1620-
" except Exception as e:",
1621-
" pass",
1622-
" if func is None and cloudpickle:",
1623-
" try:",
1624-
" func = cloudpickle.loads(data['function'])",
1625-
" except Exception as e:",
1626-
" pass",
1627-
" if func is None and data.get('function_source'):",
1628-
" try:",
1629-
" import textwrap",
1630-
" source = data['function_source']",
1631-
" # Execute the source code to create the function",
1632-
" namespace = {}",
1633-
" exec(source, namespace)",
1634-
" # Get the function name from func_info",
1635-
" func_name = data['func_info']['name']",
1636-
" func = namespace[func_name]",
1637-
" except Exception as e:",
1638-
" pass",
1639-
" if func is None:",
1640-
" error_msg = 'Could not deserialize function with dill, cloudpickle, or source code. '",
1641-
" error_msg += f'dill available: {dill is not None}, cloudpickle available: {cloudpickle is not None}, '",
1642-
" raise RuntimeError('Could not deserialize function with dill, cloudpickle, or source code')",
1643-
" ",
1644-
" args = pickle.loads(data['args'])",
1645-
" kwargs = pickle.loads(data['kwargs'])",
1646-
" ",
1647-
" result = func(*args, **kwargs)",
1648-
" ",
1649-
" with open('result.pkl', 'wb') as f:",
1650-
" pickle.dump(result, f, protocol=4)",
1651-
"except Exception as e:",
1652-
" with open('error.pkl', 'wb') as f:",
1653-
" pickle.dump({'error': str(e), 'traceback': traceback.format_exc()}, f, protocol=4)",
1654-
" raise",
1655-
'"',
1656-
]
1657-
)
1601+
script_lines.extend(job_execution_lines(remote_job_dir, config))
16581602

16591603
return "\n".join(script_lines)
16601604

‎tests/unit/test_two_venv_execution.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,3 +126,70 @@ def test_virtualenv_mode_deactivates_between_stages(self):
126126
def test_conda_mode_does_not_emit_deactivate(self):
127127
"""`conda run` does not leave an environment to deactivate."""
128128
assert "deactivate" not in generate_two_venv_execution_commands(*CONDA)
129+
130+
131+
class TestEverySchedulerRunsTheSameBody:
132+
"""PBS and SGE used to diverge from SLURM, and PBS did not work at all.
133+
134+
PBS ended its script with ``python execute_function.py`` -- a file nothing
135+
in clustrix has ever created, so every PBS job died immediately. SGE
136+
carried its own copy of the single-venv script, which meant it silently
137+
missed the two-venv path, the conda sourcing and the result signing as
138+
those were fixed on the SLURM one. All three now share one body.
139+
"""
140+
141+
SCHEDULERS = ("slurm", "pbs", "sge")
142+
143+
def _script(self, scheduler, *, two_venv):
144+
from clustrix.config import ClusterConfig
145+
from clustrix.utils import create_job_script
146+
147+
config = ClusterConfig(cluster_type=scheduler)
148+
if two_venv:
149+
config.venv_info = {
150+
"conda_env1_name": "e1",
151+
"conda_env2_name": "e2",
152+
"conda_setup_prefix": "source /opt/conda/etc/profile.d/conda.sh",
153+
}
154+
return create_job_script(
155+
scheduler,
156+
{"cores": 2, "memory": "8GB", "time": "01:00:00"},
157+
"/remote/job",
158+
config,
159+
)
160+
161+
@pytest.mark.parametrize("scheduler", SCHEDULERS)
162+
def test_no_scheduler_runs_a_file_that_is_never_created(self, scheduler):
163+
assert "execute_function.py" not in self._script(scheduler, two_venv=False)
164+
165+
@pytest.mark.parametrize("scheduler", SCHEDULERS)
166+
def test_every_scheduler_reads_the_function_payload(self, scheduler):
167+
assert "function_data.pkl" in self._script(scheduler, two_venv=False)
168+
169+
@pytest.mark.parametrize("scheduler", SCHEDULERS)
170+
def test_every_scheduler_uses_the_two_venv_path_when_available(self, scheduler):
171+
script = self._script(scheduler, two_venv=True)
172+
assert "VENV1" in script
173+
assert "conda run -n e1" in script
174+
assert "conda run -n e2" in script
175+
176+
@pytest.mark.parametrize("scheduler", SCHEDULERS)
177+
def test_every_scheduler_sources_conda(self, scheduler):
178+
script = self._script(scheduler, two_venv=True)
179+
assert "source /opt/conda/etc/profile.d/conda.sh" in script
180+
181+
@pytest.mark.parametrize("scheduler", SCHEDULERS)
182+
def test_every_scheduler_signs_its_result(self, scheduler):
183+
script = self._script(scheduler, two_venv=True)
184+
assert "result.pkl.hmac" in script
185+
assert "CLUSTRIX_RESULT_KEY" in script
186+
187+
def test_the_bodies_are_identical_across_schedulers(self):
188+
"""Only the directive header should differ between schedulers."""
189+
bodies = {}
190+
for scheduler in self.SCHEDULERS:
191+
script = self._script(scheduler, two_venv=True)
192+
# Drop the scheduler directives; keep everything from `cd` onward.
193+
body = script[script.index("cd /remote/job") :]
194+
bodies[scheduler] = body
195+
assert bodies["slurm"] == bodies["pbs"] == bodies["sge"]

0 commit comments

Comments
 (0)