refactor(slurm): simplify preflight and run preparation

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,
@@ -1020,17 +1004,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 +1035,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"] = {

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,70 +327,58 @@ 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)
 
 
-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,40 +388,38 @@ 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
     assert "readiness timed out" in completed.stderr
 
 
 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

packages/data-designer-slurm/tests/services/test_wiring.py

@@ -491,14 +491,15 @@ 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")
     payload["client"]["dependencies"]["requirements"] = ["example-plugin==1.0.0"]
     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