From d1f14245d1b8bd8f8072de1fcb84b4747bade5cf Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 16 Oct 2026 15:09:12 -0600 Subject: [PATCH 1/3] refactor(slurm): simplify preflight and readiness tests Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/services/wiring.py | 31 +++---- .../tests/runtime/test_bootstrap.py | 86 ++++++++----------- .../tests/services/test_wiring.py | 3 +- 3 files changed, 49 insertions(+), 71 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index 2d90c069b..fb31c264f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -1020,17 +1020,18 @@ def _validate_builder_for_submission( "seed-dataset", "validation", } + builtin_processors = {"drop_columns", "schema_transform"} + + def is_unrecognized_type(value: object, discriminator: str, known: set[str]) -> bool: + return ( + isinstance(value, dict) and isinstance(value.get(discriminator), str) and value[discriminator] not in known + ) + unknown_column = isinstance(columns, list) and any( - isinstance(column, dict) - and isinstance(column.get("column_type"), str) - and column["column_type"] not in builtin_columns - for column in columns + is_unrecognized_type(column, "column_type", builtin_columns) for column in columns ) unknown_processor = isinstance(processors, list) and any( - isinstance(processor, dict) - and isinstance(processor.get("processor_type"), str) - and processor["processor_type"] not in {"drop_columns", "schema_transform"} - for processor in processors + is_unrecognized_type(processor, "processor_type", builtin_processors) for processor in processors ) unknown_seed_source = ( isinstance(seed_source, dict) @@ -1050,23 +1051,13 @@ def _validate_builder_for_submission( data_designer = dict(data_designer) if plugin_dependency and unknown_column: data_designer["columns"] = [ - column - for column in columns - if not ( - isinstance(column, dict) - and isinstance(column.get("column_type"), str) - and column["column_type"] not in builtin_columns - ) + column for column in columns if not is_unrecognized_type(column, "column_type", builtin_columns) ] if plugin_dependency and unknown_processor: data_designer["processors"] = [ processor for processor in processors - if not ( - isinstance(processor, dict) - and isinstance(processor.get("processor_type"), str) - and processor["processor_type"] not in {"drop_columns", "schema_transform"} - ) + if not is_unrecognized_type(processor, "processor_type", builtin_processors) ] if (plugin_dependency and unknown_seed_source) or deferred_local_seed_path: data_designer["seed_config"] = { diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index a4ea426a6..6966f3ee0 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -334,65 +334,53 @@ def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path assert all(actual == expected for actual, expected in budgets) -def test_server_readiness_uses_one_deadline_for_all_probes() -> None: +def _run_readiness_command(body: str) -> subprocess.CompletedProcess[str]: entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" - command = f""" -source {shlex.quote(str(entrypoint))} -DD_RUNTIME_MANIFEST=unused -SECONDS=0 -DD_STEP_STARTED_AT[server]=0 -dd_read_step_ids() {{ DD_STEP_IDS=(server); }} -dd_read_step() {{ DD_STEP_PROBE_FIELDS=(127.0.0.1 10001 /health 4 127.0.0.1 10002 /health 4); }} -dd_require_running() {{ :; }} -dd_sleep() {{ SECONDS=$((SECONDS + 1)); }} -curl() {{ [[ "${{@: -1}}" == *:10001/* && $SECONDS -ge 3 ]]; }} -if dd_wait_for_role_readiness server; then - echo unexpected-success -else - echo "timeout-at:$SECONDS" -fi -""" - completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) - - assert completed.returncode == 0 - assert "timeout-at:4" in completed.stdout - assert "readiness timed out" in completed.stderr + return subprocess.run( + ("bash", "-c", f"source {shlex.quote(str(entrypoint))}\n{body}"), + capture_output=True, + text=True, + timeout=10, + ) -def test_server_readiness_rejects_late_success_after_earlier_probe() -> None: - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" - command = f""" -source {shlex.quote(str(entrypoint))} +def test_server_readiness_uses_one_deadline_for_all_probes() -> None: + command = """ DD_RUNTIME_MANIFEST=unused SECONDS=0 DD_STEP_STARTED_AT[server]=0 -dd_read_step_ids() {{ DD_STEP_IDS=(server); }} -dd_read_step() {{ DD_STEP_PROBE_FIELDS=(127.0.0.1 10001 /health 4 127.0.0.1 10002 /health 4); }} -dd_require_running() {{ :; }} -dd_sleep() {{ :; }} -curl() {{ SECONDS=5; return 0; }} +dd_read_step_ids() { DD_STEP_IDS=(server); } +dd_read_step() { DD_STEP_PROBE_FIELDS=(127.0.0.1 10001 /health 4 127.0.0.1 10002 /health 4); } +dd_require_running() { :; } +dd_sleep() { SECONDS=$((SECONDS + 1)); } +curl() { [[ "${@: -1}" == *:10001/* && $SECONDS -ge 3 ]]; } if dd_wait_for_role_readiness server; then echo unexpected-success else echo "timeout-at:$SECONDS" fi """ - completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0 - assert "timeout-at:5" in completed.stdout + assert "timeout-at:4" in completed.stdout assert "readiness timed out" in completed.stderr -def test_server_readiness_rejects_late_success_for_single_probe() -> None: - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" +@pytest.mark.parametrize( + "probe_fields", + ( + pytest.param("127.0.0.1 10001 /health 4", id="single-probe"), + pytest.param("127.0.0.1 10001 /health 4 127.0.0.1 10002 /health 4", id="second-probe"), + ), +) +def test_server_readiness_rejects_late_success(probe_fields: str) -> None: command = f""" -source {shlex.quote(str(entrypoint))} DD_RUNTIME_MANIFEST=unused SECONDS=0 DD_STEP_STARTED_AT[server]=0 dd_read_step_ids() {{ DD_STEP_IDS=(server); }} -dd_read_step() {{ DD_STEP_PROBE_FIELDS=(127.0.0.1 10001 /health 4); }} +dd_read_step() {{ DD_STEP_PROBE_FIELDS=({probe_fields}); }} dd_require_running() {{ :; }} dd_sleep() {{ :; }} curl() {{ SECONDS=5; return 0; }} @@ -402,7 +390,7 @@ def test_server_readiness_rejects_late_success_for_single_probe() -> None: echo "timeout-at:$SECONDS" fi """ - completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0 assert "timeout-at:5" in completed.stdout @@ -410,32 +398,30 @@ def test_server_readiness_rejects_late_success_for_single_probe() -> None: def test_server_readiness_checks_later_steps_before_earlier_steps_finish() -> None: - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" - command = f""" -source {shlex.quote(str(entrypoint))} + command = """ DD_RUNTIME_MANIFEST=unused SECONDS=0 DD_STEP_STARTED_AT[first]=0 DD_STEP_STARTED_AT[second]=0 -dd_read_step_ids() {{ DD_STEP_IDS=(first second); }} -dd_read_step() {{ +dd_read_step_ids() { DD_STEP_IDS=(first second); } +dd_read_step() { case "$2" in first) DD_STEP_PROBE_FIELDS=(127.0.0.1 10001 /health 120);; second) DD_STEP_PROBE_FIELDS=(127.0.0.1 10002 /health 60);; esac -}} -dd_require_running() {{ :; }} -dd_sleep() {{ SECONDS=$((SECONDS + 1)); }} -curl() {{ - case "${{@: -1}}" in +} +dd_require_running() { :; } +dd_sleep() { SECONDS=$((SECONDS + 1)); } +curl() { + case "${@: -1}" in *:10001/*) ((SECONDS >= 110));; *:10002/*) ((SECONDS >= 10));; esac -}} +} dd_wait_for_role_readiness server echo "ready-at:$SECONDS" """ - completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0, completed.stderr assert "ready-at:110" in completed.stdout diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index d4a5a02d8..042152a50 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -491,7 +491,7 @@ def test_execute_rejects_invalid_builtin_builder_before_submission( assert not launcher.submissions -def test_submit_validation_defers_overlay_plugin_columns( +def test_submit_validation_defers_overlay_plugin_components( authored_run_single: DataDesignerSlurmConfig, ) -> None: payload = authored_run_single.model_dump(mode="json") @@ -499,6 +499,7 @@ def test_submit_validation_defers_overlay_plugin_columns( payload["builder"]["inline"]["data_designer"]["columns"] = [ {"name": "plugin", "column_type": "example-plugin-column"} ] + payload["builder"]["inline"]["data_designer"]["processors"] = [{"processor_type": "example-plugin-processor"}] authored = DataDesignerSlurmConfig.model_validate(payload) wiring_module._validate_builder_for_submission( # noqa: SLF001 From 3da428a06738746672d1ae70f128b5a1b09471ce Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 16 Oct 2026 15:15:29 -0600 Subject: [PATCH 2/3] test(slurm): reuse readiness command helper Signed-off-by: Nabin Mulepati --- packages/data-designer-slurm/tests/runtime/test_bootstrap.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 6966f3ee0..5c92d5a38 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -316,9 +316,7 @@ def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path ) budgets.append((probe.deadline_seconds, process.launch_delay_seconds + 10)) cases.append(f"*://{probe.host}:{probe.port}/*) ((SECONDS >= {process.launch_delay_seconds + 1}));;") - entrypoint = Path(__file__).parents[2] / "src/data_designer/slurm/runtime/entrypoint.sh" command = f""" -source {shlex.quote(str(entrypoint))} DD_RUNTIME_MANIFEST={shlex.quote(str(manifest_path))} SECONDS=0 dd_start_step() {{ DD_LAST_PID=1; }} @@ -329,7 +327,7 @@ def _assert_launch_readiness(manifest: RuntimeBootstrapManifest, workspace: Path dd_start_servers dd_wait_for_role_readiness server """ - completed = subprocess.run(("bash", "-c", command), capture_output=True, text=True, timeout=10) + completed = _run_readiness_command(command) assert completed.returncode == 0, completed.stderr assert all(actual == expected for actual, expected in budgets) From 77e45e8cbfe0e13bc39e60919b2f7f2be4687977 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Fri, 16 Oct 2026 15:18:55 -0600 Subject: [PATCH 3/3] refactor(slurm): simplify run preparation ownership Signed-off-by: Nabin Mulepati --- .../data_designer/slurm/services/wiring.py | 36 ++++++------------- 1 file changed, 10 insertions(+), 26 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py index fb31c264f..06f421c52 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/wiring.py @@ -175,12 +175,7 @@ def prepare( ) -> Iterator[_PreparedRun]: stack = ExitStack() try: - prepared = self._prepare(authored, source_root=source_root, operation=operation, stack=stack) - except BaseException: - stack.close() - raise - try: - yield prepared + yield self._prepare(authored, source_root=source_root, operation=operation, stack=stack) finally: stack.close() @@ -352,13 +347,7 @@ def execute( ) self._materialize_default_managed_assets(config, plan) export_environment = self._build_export_environment(config) - publisher = self._publisher - self._initialize_run( - publisher, - config, - prepared, - force=False, - ) + self._initialize_run(config, prepared) self._report_progress("Rendered batch script; invoking sbatch --parsable --hold (script via stdin).") try: receipt = self._launcher.submit_script( @@ -379,7 +368,7 @@ def execute( "Slurm submission is unavailable", ) from None try: - self._record_submission(publisher, plan, receipt.job_id) + self._record_submission(plan, receipt.job_id) except Exception as error: try: self._launcher.cancel(receipt.job_id) @@ -390,7 +379,7 @@ def execute( f"Slurm job {receipt.job_id} was submitted but could not be recorded or cancelled", ) from error try: - self._record_submission_failure(publisher, plan, receipt.job_id) + self._record_submission_failure(plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -405,7 +394,7 @@ def execute( pass else: try: - self._record_submission_failure(publisher, plan, receipt.job_id) + self._record_submission_failure(plan, receipt.job_id) except BaseException: pass raise @@ -421,7 +410,7 @@ def execute( f"held Slurm job {receipt.job_id} could not be released or cancelled", ) from error try: - self._record_submission_failure(publisher, plan, receipt.job_id) + self._record_submission_failure(plan, receipt.job_id) except SlurmServiceError as state_error: raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL, @@ -532,19 +521,16 @@ def _build_export_environment(self, config: DataDesignerSlurmConfig) -> dict[str def _initialize_run( self, - publisher: SlurmRunArtifactPublisher, config: DataDesignerSlurmConfig, prepared: _PreparedRun, - *, - force: bool, ) -> None: try: - publisher.initialize( + self._publisher.initialize( config, prepared.plan, prepared.dependencies, prepared.builder_payload, - force=force, + force=False, ) except StateNotFoundError: raise SlurmServiceError( @@ -567,12 +553,11 @@ def _initialize_run( def _record_submission( self, - publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan, job_id: int, ) -> None: try: - publisher.record_submission(plan, job_id, submitted_at=self._clock()) + self._publisher.record_submission(plan, job_id, submitted_at=self._clock()) except StateNotFoundError: raise SlurmServiceError( SlurmServiceErrorCode.NOT_FOUND, @@ -594,12 +579,11 @@ def _record_submission( def _record_submission_failure( self, - publisher: SlurmRunArtifactPublisher, plan: ResolvedSlurmRunPlan, job_id: int, ) -> None: try: - publisher.record_submission_failure(plan, job_id, failed_at=self._clock()) + self._publisher.record_submission_failure(plan, job_id, failed_at=self._clock()) except (StateNotFoundError, StateConflictError, SlurmStateError): raise SlurmServiceError( SlurmServiceErrorCode.INTERNAL,