diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py index 6e96d15ea..2d7780948 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/environment.py @@ -272,6 +272,8 @@ def _install_overlay( def inspect_distributions(path: Path | None) -> tuple[InstalledDistribution, ...]: """Return an exact immutable distribution inventory.""" + if path is not None: + importlib.invalidate_caches() distributions = ( importlib.metadata.distributions() if path is None else importlib.metadata.distributions(path=[path.as_posix()]) ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/process.py b/packages/data-designer-slurm/src/data_designer/slurm/client/process.py new file mode 100644 index 000000000..f66895e42 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/process.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bootstrap-safe process boundary for the plugin-aware client worker. + +This module must remain free of Data Designer configuration, interface, plugin, +and worker imports. The child process activates its verified dependency overlay +before importing any of those modules. +""" + +from __future__ import annotations + +import subprocess +import sys +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field + +ProcessExecutor = Callable[[tuple[str, ...]], int] + + +def execute_process(command: tuple[str, ...]) -> int: + """Execute a child process and return its exit status.""" + return subprocess.run(command, check=False).returncode + + +@dataclass(frozen=True) +class ClientWorkerProcess: + """Launch the plugin-aware client worker in a fresh Python interpreter.""" + + executable: str = field(default_factory=lambda: sys.executable) + executor: ProcessExecutor = execute_process + + def run(self, arguments: Sequence[str]) -> int: + """Run one client-worker operation without inheriting imported modules.""" + command = (self.executable, "-m", "data_designer.slurm.client.worker", *arguments) + return self.executor(command) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py index db0462981..dd43a06f2 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.py @@ -13,7 +13,7 @@ from pathlib import Path from data_designer.slurm.client.filesystem import ensure_private_directory, replace_private_text -from data_designer.slurm.client.worker import main as client_worker_main +from data_designer.slurm.client.process import ClientWorkerProcess from data_designer.slurm.runtime.bootstrap import build_runtime_manifest from data_designer.slurm.runtime.context import load_allocation_context from data_designer.slurm.runtime.errors import SlurmRuntimeError, SlurmRuntimeErrorCode @@ -145,8 +145,14 @@ def _ready(arguments: argparse.Namespace, environment: Mapping[str, str]) -> Non ) -def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> None: +def _client( + arguments: argparse.Namespace, + environment: Mapping[str, str], + *, + client_worker: ClientWorkerProcess | None = None, +) -> None: context, writer = _load_context(arguments, environment) + worker_process = client_worker or ClientWorkerProcess() generation_started_at = _now(context.attempt, _load_optional_readiness(context, writer)) resume_mode = ( context.plan.invocation.authored.resume @@ -158,7 +164,7 @@ def _client(arguments: argparse.Namespace, environment: Mapping[str, str]) -> No context.attempt.attempt_id, resume_mode, ): - return_code = client_worker_main( + return_code = worker_process.run( ( "run", "--plan", diff --git a/packages/data-designer-slurm/tests/client/test_environment.py b/packages/data-designer-slurm/tests/client/test_environment.py index c42e284fd..51c920589 100644 --- a/packages/data-designer-slurm/tests/client/test_environment.py +++ b/packages/data-designer-slurm/tests/client/test_environment.py @@ -6,6 +6,7 @@ import hashlib import importlib import json +import os import subprocess import sys from pathlib import Path @@ -144,6 +145,20 @@ def distributions(**kwargs: object) -> tuple[()]: assert calls == [{}] +def test_inspect_distributions_refreshes_overlay_with_unchanged_directory_mtime(tmp_path: Path) -> None: + original_stat = tmp_path.stat() + assert inspect_distributions(tmp_path) == () + distribution = tmp_path / "cache_probe-1.0.0.dist-info" + distribution.mkdir() + (distribution / "METADATA").write_text( + "Metadata-Version: 2.1\nName: cache-probe\nVersion: 1.0.0\n", + encoding="utf-8", + ) + os.utime(tmp_path, ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns)) + + assert inspect_distributions(tmp_path) == (InstalledDistribution(name="cache-probe", version="1.0.0"),) + + def test_inspect_distributions_rejects_unhashed_direct_url(monkeypatch: pytest.MonkeyPatch) -> None: distribution = Mock( metadata={"Name": "example"}, diff --git a/packages/data-designer-slurm/tests/client/test_process.py b/packages/data-designer-slurm/tests/client/test_process.py new file mode 100644 index 000000000..98450e5fa --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_process.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.slurm.client.process import ClientWorkerProcess + + +def test_client_worker_process_uses_a_fresh_python_interpreter() -> None: + commands: list[tuple[str, ...]] = [] + process = ClientWorkerProcess( + executable="/client/python", + executor=lambda command: commands.append(command) or 17, + ) + + assert process.run(("run", "--plan", "/workspace/resolved-plan.json")) == 17 + assert commands == [ + ( + "/client/python", + "-m", + "data_designer.slurm.client.worker", + "run", + "--plan", + "/workspace/resolved-plan.json", + ) + ] diff --git a/packages/data-designer-slurm/tests/client/test_worker.py b/packages/data-designer-slurm/tests/client/test_worker.py index 7da326e0c..bb5805573 100644 --- a/packages/data-designer-slurm/tests/client/test_worker.py +++ b/packages/data-designer-slurm/tests/client/test_worker.py @@ -500,7 +500,14 @@ def test_preflight_rejects_plugin_secondary_model_alias( ) -> None: payload = client_worker_case.plan.model_dump(mode="json") builder = payload["builder"]["inline"] - builder["data_designer"]["columns"] = [{"name": "custom", "column_type": "fake-slurm-column"}] + builder["data_designer"]["columns"] = [ + { + "name": "custom", + "column_type": "fake-slurm-column", + "model_alias": "generator", + "judge_model_alias": "missing", + } + ] payload["builder"]["content_sha256"] = compute_serialized_json_sha256(builder) lock_payload = client_worker_case.lock.model_dump(mode="json") wheel_path = ( diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/implementation.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/implementation.py new file mode 100644 index 000000000..ea90846c8 --- /dev/null +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/implementation.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +from data_designer.engine.column_generators.generators.base import ColumnGeneratorCellByCell +from fake_data_designer_plugin.plugin import FakePluginConfig + + +class FakePluginImplementation(ColumnGeneratorCellByCell[FakePluginConfig]): + """Add a deterministic marker through the real generator contract.""" + + def generate(self, data: dict[str, Any]) -> dict[str, Any]: + """Return the complete record with the plugin-owned column.""" + return {**data, self.config.name: "plugin-marker"} diff --git a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py index 67ce6aae3..dea69fc26 100644 --- a/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py +++ b/packages/data-designer-slurm/tests/fixtures/fake_plugin_overlay/fake_data_designer_plugin/plugin.py @@ -13,8 +13,8 @@ class FakePluginConfig(SingleColumnConfig): """Minimal custom column configuration.""" column_type: Literal["fake-slurm-column"] = "fake-slurm-column" - model_alias: str = "generator" - judge_model_alias: str = "judge" + model_alias: str | None = None + judge_model_alias: str | None = None @property def required_columns(self) -> list[str]: @@ -25,19 +25,11 @@ def side_effect_columns(self) -> list[str]: return [] def get_model_aliases(self) -> list[str]: - return [self.model_alias, self.judge_model_alias] - - -class FakePluginImplementation: - """Minimal loadable implementation for entry-point verification.""" - - def generate(self, data: dict[str, object]) -> dict[str, object]: - """Return the provided record unchanged.""" - return data + return [alias for alias in (self.model_alias, self.judge_model_alias) if alias is not None] plugin = Plugin( config_qualified_name="fake_data_designer_plugin.plugin.FakePluginConfig", - impl_qualified_name="fake_data_designer_plugin.plugin.FakePluginImplementation", + impl_qualified_name="fake_data_designer_plugin.implementation.FakePluginImplementation", plugin_type=PluginType.COLUMN_GENERATOR, ) diff --git a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py index 17c12005d..561588810 100644 --- a/packages/data-designer-slurm/tests/runtime/test_entrypoint.py +++ b/packages/data-designer-slurm/tests/runtime/test_entrypoint.py @@ -3,21 +3,50 @@ from __future__ import annotations +import csv +import hashlib +import io +import json +import os +import sys +import venv +import zipfile +from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace +from typing import cast +import pyarrow.parquet as pq import pytest from conftest import FakeStateStore, RuntimeCase +from packaging.tags import interpreter_name, interpreter_version import data_designer.slurm.runtime.entrypoint as entrypoint -from data_designer.slurm.contracts import ArtifactReference -from data_designer.slurm.state import AttemptLifecycleState, ReadinessState +from data_designer.slurm.client.process import ClientWorkerProcess +from data_designer.slurm.client.records import ClientEnvironmentManifest, ClientInstallerOutcome, ClientResult +from data_designer.slurm.contracts import ( + ArtifactReference, + InstalledDistribution, + compute_canonical_json_sha256, + compute_serialized_json_sha256, +) +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.runtime.models import AllocationContext +from data_designer.slurm.runtime.ports import resolve_allocation_plan +from data_designer.slurm.state import AttemptLifecycleState, CandidateOutputManifest, ReadinessState class _InjectedFailure(Exception): pass +@dataclass(frozen=True) +class _PluginRuntimeCase: + context: AllocationContext + plan_path: Path + image_metadata_directory: Path + + def test_entrypoint_rejects_relative_paths_without_traceback(capsys: pytest.CaptureFixture[str]) -> None: assert ( entrypoint.main( @@ -65,8 +94,6 @@ def verify_attempt_directory(path: Path) -> None: with pytest.raises(_InjectedFailure): entrypoint._prepare(prepare, {}) - monkeypatch.setattr(entrypoint, "client_worker_main", lambda arguments: 0) - def load_candidate(*args: object, attempt_directory: Path | None = None) -> None: assert attempt_directory == container_attempt_directory raise _InjectedFailure @@ -76,7 +103,58 @@ def load_candidate(*args: object, attempt_directory: Path | None = None) -> None _phase_arguments("client", runtime_case, attempt_directory=container_attempt_directory) ) with pytest.raises(_InjectedFailure): - entrypoint._client(client, {}) + entrypoint._client(client, {}, client_worker=ClientWorkerProcess(executor=lambda command: 0)) + + +def test_client_phase_installs_and_runs_plugin_in_default_worker_process( + monkeypatch: pytest.MonkeyPatch, + runtime_case: RuntimeCase, + fake_plugin_overlay: Path, + tmp_path: Path, +) -> None: + case = _prepare_plugin_runtime_case(runtime_case, fake_plugin_overlay, tmp_path) + state = FakeStateStore(case.context.attempt) + runtime_case.context = case.context + _patch_runtime_context(monkeypatch, runtime_case, state) + assert "data_designer.config.column_types" in sys.modules + _use_controlled_image_inventory(monkeypatch, case.image_metadata_directory) + + allocation = resolve_allocation_plan(case.context.plan, {"SLURM_JOB_GPUS": "0"}) + endpoint = f"generator=http://127.0.0.1:{allocation.client.ports[0].port}/v1" + worker_arguments = ( + "preflight", + "--plan", + case.plan_path.as_posix(), + "--shard-id", + case.context.shard.shard_id, + "--attempt-id", + case.context.attempt.attempt_id, + "--attempt-dir", + case.context.attempt_directory.as_posix(), + "--endpoint", + endpoint, + ) + assert ClientWorkerProcess().run(worker_arguments) == 0 + + client = entrypoint._parse_arguments((*_phase_arguments("client", runtime_case), "--endpoint", endpoint)) + entrypoint._client(client, {}) + + environment = ClientEnvironmentManifest.model_validate_json( + (case.context.attempt_directory / "client-environment.json").read_text() + ) + result = ClientResult.model_validate_json((case.context.attempt_directory / "client-result.json").read_text()) + candidate = CandidateOutputManifest.model_validate_json( + (case.context.attempt_directory / "output-manifest.json").read_text() + ) + dataset = pq.read_table(Path(candidate.dataset_path) / candidate.files[0].relative_path) + + assert environment.installer_outcome is ClientInstallerOutcome.INSTALLED + assert [plugin.plugin_name for plugin in environment.plugins] == ["fake-slurm-column"] + assert result.candidate_output_manifest is not None + assert result.candidate_output_manifest.sha256 == candidate.compute_sha256() + assert state.attempt.candidate_output == result.candidate_output_manifest + assert set(dataset.column_names) == {"record_id", "custom"} + assert dataset.column("custom").to_pylist() == ["plugin-marker"] * case.context.shard.requested_records def test_control_phases_record_running_ready_and_failed( @@ -184,6 +262,192 @@ def _patch_runtime_context( monkeypatch.setenv("SLURM_JOB_GPUS", "0") +def _prepare_plugin_runtime_case( + runtime_case: RuntimeCase, + fake_plugin_overlay: Path, + tmp_path: Path, +) -> _PluginRuntimeCase: + image_distributions = runtime_case.context.plan.client.image.inspection_facts.distributions + wheel_path = Path(runtime_case.context.plan.authored_config.path).parent / ( + "dependencies/fake_data_designer_plugin-1.0.0-py3-none-any.whl" + ) + _build_plugin_wheel(fake_plugin_overlay, wheel_path) + payload = runtime_case.context.plan.model_dump(mode="json") + _configure_identity_mount(payload, runtime_case.workspace) + _configure_plugin_builder(payload) + client = cast(dict[str, object], payload["client"]) + _configure_client_inspection(client, image_distributions, _write_test_installer(tmp_path)) + lock = _plugin_dependency_lock(client, image_distributions, wheel_path) + _write_dependency_lock(client, lock) + + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + plan_path = Path(plan.authored_config.path).with_name("resolved-plan.json") + plan_path.write_text(plan.serialize_json()) + Path(plan.invocation.effective_input_bindings.managed_assets_path).mkdir(parents=True, exist_ok=True) + attempt = runtime_case.context.attempt.model_copy( + update={ + "resolved_plan": ArtifactReference(path=plan_path.as_posix(), sha256=plan.compute_sha256()), + "state": AttemptLifecycleState.RUNNING, + } + ) + context = AllocationContext(plan, plan.shards[0], attempt, runtime_case.context.attempt_directory) + image_metadata_directory = _write_inventory_bootstrap(image_distributions, tmp_path) + return _PluginRuntimeCase(context, plan_path, image_metadata_directory) + + +def _configure_identity_mount(payload: dict[str, object], workspace: Path) -> None: + mount = {"source": workspace.as_posix(), "target": workspace.as_posix(), "read_only": False} + payload["container_mounts"] = [mount] + selected_profile = cast(dict[str, object], payload["selected_profile"]) + profile = cast(dict[str, object], selected_profile["profile"]) + profile["container_mounts"] = [mount] + selected_profile["profile_sha256"] = compute_canonical_json_sha256(profile) + + +def _configure_plugin_builder(payload: dict[str, object]) -> None: + resolved_builder = cast(dict[str, object], payload["builder"]) + builder = cast(dict[str, object], resolved_builder["inline"]) + cast(dict[str, object], builder["data_designer"])["columns"] = [ + {"name": "record_id", "column_type": "sampler", "sampler_type": "uuid", "params": {}}, + {"name": "custom", "column_type": "fake-slurm-column"}, + ] + resolved_builder["content_sha256"] = compute_serialized_json_sha256(builder) + + +def _configure_client_inspection( + client: dict[str, object], + image_distributions: tuple[InstalledDistribution, ...], + installer: Path, +) -> None: + inspection = cast( + dict[str, object], + cast(dict[str, object], cast(dict[str, object], client["image"])["inspection"])["inspection"], + ) + inspection["python_abi"] = f"{interpreter_name()}{interpreter_version()}" + inspection["python_version"] = sys.version.split()[0] + inspection["installer_path"] = installer.as_posix() + inspection["distributions"] = [item.model_dump(mode="json") for item in image_distributions] + dependencies = cast(dict[str, object], cast(dict[str, object], client["authored"])["dependencies"]) + dependencies["requirements"] = ["fake-data-designer-plugin==1.0.0"] + + +def _plugin_dependency_lock( + client: dict[str, object], + image_distributions: tuple[InstalledDistribution, ...], + wheel_path: Path, +) -> ResolvedDependencyLock: + inspection = cast( + dict[str, object], + cast(dict[str, object], cast(dict[str, object], client["image"])["inspection"])["inspection"], + ) + return ResolvedDependencyLock.model_validate( + { + "schema_version": 1, + "resolver_version": "resolver-1", + "python_abi": inspection["python_abi"], + "client_image_sha256": cast(dict[str, object], client["image"])["sha256"], + "authored_requirements": ("fake-data-designer-plugin==1.0.0",), + "authored_source": None, + "source": None, + "image_distributions": image_distributions, + "overlay_packages": ( + { + "name": "fake-data-designer-plugin", + "version": "1.0.0", + "artifact": { + "path": wheel_path.as_posix(), + "sha256": hashlib.sha256(wheel_path.read_bytes()).hexdigest(), + }, + }, + ), + } + ) + + +def _write_dependency_lock(client: dict[str, object], lock: ResolvedDependencyLock) -> None: + lock_path = Path(cast(dict[str, object], client["dependency_lock"])["path"]) + lock_path.parent.mkdir(parents=True, exist_ok=True) + lock_path.write_text(lock.serialize_json()) + cast(dict[str, object], client["dependency_lock"])["sha256"] = lock.compute_sha256() + + +def _build_plugin_wheel(source: Path, destination: Path) -> None: + """Package the installed-layout fixture as the locked wheel under test.""" + destination.parent.mkdir(parents=True, exist_ok=True) + dist_info = "fake_data_designer_plugin-1.0.0.dist-info" + members = { + path.relative_to(source).as_posix(): path.read_bytes() + for path in source.rglob("*") + if path.is_file() and "__pycache__" not in path.parts and path.suffix != ".pyc" + } + members[f"{dist_info}/WHEEL"] = ( + b"Wheel-Version: 1.0\nGenerator: data-designer-tests\nRoot-Is-Purelib: true\nTag: py3-none-any\n" + ) + record_path = f"{dist_info}/RECORD" + output = io.StringIO() + writer = csv.writer(output, lineterminator="\n") + for name in (*sorted(members), record_path): + writer.writerow((name, "", "")) + members[record_path] = output.getvalue().encode() + with zipfile.ZipFile(destination, "w", compression=zipfile.ZIP_DEFLATED) as archive: + for name, content in sorted(members.items()): + archive.writestr(name, content) + + +def _write_test_installer(tmp_path: Path) -> Path: + """Create an offline pip executable without changing the test checkout.""" + environment = tmp_path / "installer-environment" + venv.EnvBuilder(with_pip=True, symlinks=True).create(environment) + return environment / "bin/pip" + + +def _write_inventory_bootstrap( + distributions: tuple[InstalledDistribution, ...], + tmp_path: Path, +) -> Path: + """Present plan-recorded image metadata while imports use the editable checkout.""" + bootstrap_directory = tmp_path / "client-bootstrap" + image_metadata_directory = bootstrap_directory / "image-metadata" + image_metadata_directory.mkdir(parents=True) + for distribution in distributions: + metadata = image_metadata_directory / ( + f"{distribution.name.replace('-', '_')}-{distribution.version}.dist-info/METADATA" + ) + metadata.parent.mkdir() + metadata.write_text(f"Metadata-Version: 2.1\nName: {distribution.name}\nVersion: {distribution.version}\n") + (bootstrap_directory / "sitecustomize.py").write_text( + """from __future__ import annotations + +import importlib.metadata +import os +import sys + +_distributions = importlib.metadata.distributions + + +def _controlled_distributions(**kwargs): + if kwargs.get("path") is not None: + return _distributions(**kwargs) + paths = [os.environ["DATA_DESIGNER_TEST_IMAGE_METADATA"]] + paths.extend(path for path in sys.path if path.endswith("/client-env/site-packages")) + return _distributions(path=paths) + + +importlib.metadata.distributions = _controlled_distributions +""" + ) + return image_metadata_directory + + +def _use_controlled_image_inventory(monkeypatch: pytest.MonkeyPatch, image_metadata_directory: Path) -> None: + """Limit child-process inventory discovery to the plan image and installed overlay.""" + bootstrap_directory = image_metadata_directory.parent + python_path = os.environ.get("PYTHONPATH") + paths = (bootstrap_directory.as_posix(),) if python_path is None else (bootstrap_directory.as_posix(), python_path) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(paths)) + monkeypatch.setenv("DATA_DESIGNER_TEST_IMAGE_METADATA", image_metadata_directory.as_posix()) + + def _phase_arguments( operation: str, runtime_case: RuntimeCase,