From 6cef6062b55345d803a0d764ed84b72588063464 Mon Sep 17 00:00:00 2001 From: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:47:11 -0700 Subject: [PATCH 01/12] feat(slurm): support external model providers Allow hosted endpoints alongside managed models and run hosted-only recipes in CPU allocations. Resolve credentials at runtime and preserve native Data Designer model settings and asynchronous execution. Reuse shard, retry, resume, and collection contracts while preserving existing local-only plan serialization. Add configuration examples, documentation, and native client and allocation lifecycle coverage. Closes #1079 Signed-off-by: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> --- .../latest/pages/slurm/run-configuration.mdx | 75 ++++++- .../examples/hosted-only.yaml | 30 +++ .../examples/mixed-providers.yaml | 42 ++++ .../src/data_designer/slurm/_errors.py | 7 + .../data_designer/slurm/client/execution.py | 48 +++-- .../data_designer/slurm/config/__init__.py | 2 + .../src/data_designer/slurm/config/builder.py | 5 +- .../src/data_designer/slurm/config/run.py | 59 +++++- .../data_designer/slurm/launcher/renderer.py | 12 +- .../slurm/planning/builder_identity.py | 18 ++ .../data_designer/slurm/planning/compiler.py | 2 +- .../data_designer/slurm/planning/models.py | 40 +++- .../slurm/planning/resolution.py | 25 ++- .../slurm/planning/validation.py | 10 +- .../data_designer/slurm/runtime/bootstrap.py | 8 +- .../data_designer/slurm/runtime/controller.py | 2 + .../data_designer/slurm/runtime/entrypoint.sh | 34 ++-- .../data_designer/slurm/runtime/manifest.py | 5 +- .../slurm/runtime/plan_reader.sh | 5 +- .../src/data_designer/slurm/runtime/ports.py | 2 +- .../data_designer/slurm/services/wiring.py | 2 + .../data_designer/slurm/state/readiness.py | 10 +- .../client/test_external_model_execution.py | 191 ++++++++++++++++++ .../tests/config/test_loading_builder.py | 2 +- .../data-designer-slurm/tests/conftest.py | 59 ++++++ .../tests/planning/test_external_providers.py | 179 ++++++++++++++++ .../tests/runtime/test_bootstrap.py | 47 +++++ .../tests/runtime/test_controller.py | 37 ++++ .../tests/runtime/test_shell_runtime.py | 48 ++++- .../tests/services/test_wiring.py | 22 ++ .../tests/serving/test_resolver.py | 5 + plans/1079/slurm-external-providers.md | 35 ++++ 32 files changed, 993 insertions(+), 75 deletions(-) create mode 100644 packages/data-designer-slurm/examples/hosted-only.yaml create mode 100644 packages/data-designer-slurm/examples/mixed-providers.yaml create mode 100644 packages/data-designer-slurm/tests/client/test_external_model_execution.py create mode 100644 packages/data-designer-slurm/tests/planning/test_external_providers.py create mode 100644 plans/1079/slurm-external-providers.md diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 0b8186340..e656c5e64 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -36,9 +36,9 @@ submission: `builder` accepts exactly one complete serialized Data Designer config through `source` or `inline`. Relative source paths resolve from the run file directory. `invocation.run_config` accepts public `RunConfig` fields except `display_tui`, which is disabled for Slurm jobs. `input_bindings.seed_path` and `managed_assets_path` accept absolute compute-visible paths. -Every `model_alias` in `model_concurrency` and the builder must have one deployment. Deployment aliases must be unique. `model` selects the Hugging Face model ID or absolute model path loaded by the serving container. `served_model_name` is the name exposed to the client and defaults to `model`; set it explicitly when `model` is a container-visible path but the builder uses a stable model name. Shard and collected output files are Parquet. Choose the destination and partition count with `slurm merge` after the run completes. +Every builder model alias must have either one deployment or a provider declared in `invocation.model_providers`. These assignments must not overlap. `model_concurrency` accepts any declared alias, including external models. Deployment aliases must be unique. `model` selects the Hugging Face model ID or absolute model path loaded by the serving container. `served_model_name` is the name exposed to the client and defaults to `model`; set it explicitly when `model` is a container-visible path but the builder uses a stable model name. Shard and collected output files are Parquet. Choose the destination and partition count with `slurm merge` after the run completes. -The client runs from the same shared Python environment as the CLI on the first allocated GPU node. Its step requests no GPUs; the serving steps use the GPUs and their registered images. The `client` block is optional; `client.cpus` defaults to 16. It sets the client step's CPU request and contributes to the batch allocation's CPU request. +The client runs from the same shared Python environment as the CLI on the first allocated node. Its step requests no GPUs; local serving steps use the GPUs and their registered images. The `client` block is optional; `client.cpus` defaults to 16. It sets the client step's CPU request and contributes to the batch allocation's CPU request. Use `submission.partition` to override the profile's GPU partition selection for one run. It accepts a comma-separated value or a YAML list, such as: @@ -47,7 +47,74 @@ submission: partition: [gpu-a, gpu-b] ``` -Slurm schedules the job in one of the selected partitions. With `gpus_per_node: auto`, every selected partition must report GPU resources with one consistent count. +Slurm schedules the job in one of the selected partitions. For runs with local deployments and `gpus_per_node: auto`, every selected partition must report GPU resources with one consistent count. + +## External model providers + +Use `invocation.model_providers` for endpoints managed outside this allocation. The builder keeps its normal model declarations: + +```yaml +data_designer: + model_configs: + - alias: generator + model: example/generator + provider: openai + - alias: judge + model: example/hosted-judge + provider: hosted + # Keep the recipe's columns and inference parameters here. +``` + +The run file declares the external provider and deploys only `generator`: + +```yaml +invocation: + num_records: 1000 + dataset_name: judged-output + model_providers: + - name: hosted + endpoint: https://models.example/v1 + provider_type: openai + api_key: + type: secret + environment: MODEL_API_KEY + model_concurrency: + generator: 64 + judge: 16 +``` + +The provider name must match the builder model's `provider`. Slurm preserves external model names and inference parameters; an explicit `model_concurrency` value overrides that alias's request concurrency. Provider names must not collide with generated local names such as `slurm-generator`. + +`api_key` is an optional `SecretRef`. Set the referenced variable in the submit environment; Slurm forwards it only to native client steps. Provider `extra_body` accepts non-secret request parameters. `extra_headers` maps header names to `type: literal` / `value` or `type: secret` / `environment` bindings. Authentication headers require secret bindings. Credential values are never written into the run configuration or resolved plan. + +Both local and external requests use DD's existing async engine. Ready requests for different records and models can overlap. Async overlap does not guarantee continuous GPU utilization: slow external calls and the final judging tail can still retain an idle GPU allocation. Measure utilization against the cluster's policy when deciding whether to use separate recipe stages. + +### Hosted-only CPU runs + +For a builder whose models all use declared external providers, omit local deployments or set `deployments: []`. Select the cluster's CPU partition: + +```yaml +schema_version: 1 +name: hosted-judging +builder: + source: hosted-builder.yaml +invocation: + num_records: 1000 + dataset_name: hosted-judging + model_providers: + - name: hosted + endpoint: https://models.example/v1 + api_key: + type: secret + environment: MODEL_API_KEY +client: + cpus: 4 +deployments: [] +submission: + partition: cpu +``` + +Each array task requests one node with `--gres=none` and `--cpus-per-task` set to `client.cpus`. It does not discover GPUs, reserve exclusive nodes, require serving images, or start serving containers and proxies. The profile's GPU facts and GPU memory request are unused for this run. Sharding, status, retry, resume, and collection use the same commands as local-model runs. ## Configure vLLM @@ -169,7 +236,7 @@ Increasing `resources.nodes` at TP=1 creates more independent one-GPU servers; i ### Multiple models and images -Add one deployment per builder model alias. Each deployment may use a different registered serving image: +Add one deployment per locally served builder model alias. Each deployment may use a different registered serving image: ```yaml deployments: diff --git a/packages/data-designer-slurm/examples/hosted-only.yaml b/packages/data-designer-slurm/examples/hosted-only.yaml new file mode 100644 index 000000000..9b6850c83 --- /dev/null +++ b/packages/data-designer-slurm/examples/hosted-only.yaml @@ -0,0 +1,30 @@ +schema_version: 1 +name: hosted-judging +builder: + inline: + data_designer: + model_configs: + - alias: judge + model: example/hosted-judge + provider: hosted + columns: + - name: review + column_type: llm-text + model_alias: judge + prompt: Explain one property of a useful product description. +invocation: + num_records: 100 + dataset_name: hosted-judging + model_providers: + - name: hosted + endpoint: https://models.example/v1 + api_key: + type: secret + environment: MODEL_API_KEY + model_concurrency: + judge: 16 +client: + cpus: 4 +deployments: [] +submission: + partition: cpu diff --git a/packages/data-designer-slurm/examples/mixed-providers.yaml b/packages/data-designer-slurm/examples/mixed-providers.yaml new file mode 100644 index 000000000..7e7e40c91 --- /dev/null +++ b/packages/data-designer-slurm/examples/mixed-providers.yaml @@ -0,0 +1,42 @@ +schema_version: 1 +name: local-generation-hosted-review +builder: + inline: + data_designer: + model_configs: + - alias: generator + model: example/generator + provider: openai + - alias: judge + model: example/hosted-judge + provider: hosted + columns: + - name: description + column_type: llm-text + model_alias: generator + prompt: Write a short description of an everyday household product. + - name: review + column_type: llm-text + model_alias: judge + prompt: "Review this product description for clarity: {{ description }}" +invocation: + num_records: 100 + dataset_name: judged-descriptions + model_providers: + - name: hosted + endpoint: https://models.example/v1 + api_key: + type: secret + environment: MODEL_API_KEY + model_concurrency: + generator: 32 + judge: 16 +deployments: + - model_alias: generator + model: example/generator + server: + type: vllm + image: + name: vllm + topology: + tensor_parallel: 8 diff --git a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py index ef4cdbbea..3c2630a77 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/_errors.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/_errors.py @@ -42,6 +42,13 @@ "shared_node requires GRES GPU request mode", "shared_node requires the same resolved serving image", "deployment placement and GPU ownership must match authored resources", + "model concurrency references undeclared aliases", + "external provider names must not collide with managed providers", + "model provider names must be unique", + "client-only runs must resolve zero GPUs", + "server deployments require a positive GPU count", + "resolved external model aliases do not match the inline builder", + "external model aliases require external providers", ) }, "Value error, builder content digest does not match the resolved input": ( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index ff9b4edf5..5b3a0818e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -404,10 +404,29 @@ def _materialize_model_endpoints( if set(endpoints) != set(aliases): raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "runtime endpoints do not match model aliases") configs = {config.alias: config for config in builder.model_configs} - if set(configs) != set(aliases): + if set(configs) != set(plan.builder.model_aliases): raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "builder model aliases do not match the plan") - updated = [] - providers = [] + external_providers = plan.invocation.authored.model_providers + external_names = {provider.name for provider in external_providers} + external_aliases = tuple(config.alias for config in configs.values() if config.provider in external_names) + if external_aliases != plan.builder.external_model_aliases: + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "external model aliases do not match the plan") + updated = dict(configs) + providers = [ + ModelProvider( + name=provider.name, + endpoint=provider.endpoint, + provider_type=provider.provider_type, + api_key=_resolve_secret(provider.api_key) if provider.api_key is not None else None, + extra_body=provider.extra_body, + extra_headers={ + name: binding.value if isinstance(binding, LiteralEnvironmentBinding) else _resolve_secret(binding) + for name, binding in provider.extra_headers.items() + } + or None, + ) + for provider in external_providers + ] for deployment, port in zip(plan.deployments, plan.client.ports, strict=True): alias = deployment.authored.model_alias endpoint = endpoints[alias] @@ -417,23 +436,22 @@ def _materialize_model_endpoints( raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "runtime model endpoint is invalid") provider_name = f"slurm-{alias}" model_config = configs[alias] - inference = model_config.inference_parameters + updated[alias] = model_config.model_copy( + update={"model": deployment.served_model_name, "provider": provider_name} + ) + providers.append(ModelProvider(name=provider_name, endpoint=endpoint, provider_type="openai")) + for model_config in tuple(builder.model_configs): + builder.delete_model_config(model_config.alias) + for alias, model_config in updated.items(): concurrency = plan.invocation.authored.model_concurrency.get(alias) if concurrency is not None: - inference = inference.model_copy(update={"max_parallel_requests": concurrency}) - updated.append( - model_config.model_copy( + model_config = model_config.model_copy( update={ - "model": deployment.served_model_name, - "provider": provider_name, - "inference_parameters": inference, + "inference_parameters": model_config.inference_parameters.model_copy( + update={"max_parallel_requests": concurrency} + ) } ) - ) - providers.append(ModelProvider(name=provider_name, endpoint=endpoint, provider_type="openai")) - for model_config in tuple(builder.model_configs): - builder.delete_model_config(model_config.alias) - for model_config in updated: builder.add_model_config(model_config) return providers diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index 9842577f0..f87309bb9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -60,6 +60,7 @@ DataDesignerSlurmConfig, DeploymentResources, DeploymentTopology, + ExternalModelProviderConfig, InputBindings, InvocationConfig, LocalStdioMCPProviderConfig, @@ -87,6 +88,7 @@ "DEFAULT_PROFILE_FILE_NAME", "DeploymentResources", "DeploymentTopology", + "ExternalModelProviderConfig", "FixedRecordPolicy", "GpuRequestMode", "ImageBuildProfile", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py index 21670bb5e..123ed53bb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -21,6 +21,7 @@ ClientConfig, ClientDependencies, DataDesignerSlurmConfig, + ExternalModelProviderConfig, InputBindings, InvocationConfig, MCPProviderConfig, @@ -72,6 +73,7 @@ def with_invocation( run_config: Mapping[str, object] | None = None, input_bindings: InputBindings | Mapping[str, object] | None = None, mcp_providers: Sequence[MCPProviderConfig | Mapping[str, object]] = (), + model_providers: Sequence[ExternalModelProviderConfig | Mapping[str, object]] = (), model_concurrency: Mapping[str, int] | None = None, ) -> DataDesignerSlurmConfigBuilder: """Set typed Data Designer invocation intent.""" @@ -84,6 +86,7 @@ def with_invocation( "run_config": {} if run_config is None else dict(run_config), "input_bindings": {} if input_bindings is None else input_bindings, "mcp_providers": list(mcp_providers), + "model_providers": list(model_providers), "model_concurrency": {} if model_concurrency is None else dict(model_concurrency), }, ) @@ -136,8 +139,6 @@ def build(self) -> DataDesignerSlurmConfig: missing = [] if self._invocation is None: missing.append("invocation") - if not self._deployments: - missing.append("deployment") if missing: raise SlurmConfigBuilderError(f"Slurm config builder requires: {', '.join(missing)}") assert self._invocation is not None diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py index dbc4c0fde..d110f98b4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -50,6 +50,7 @@ "DeploymentResources", "DeploymentTopology", "EnvironmentBinding", + "ExternalModelProviderConfig", "InputBindings", "InvocationConfig", "LiteralEnvironmentBinding", @@ -106,6 +107,53 @@ def validate_paths(cls, value: str | None) -> str | None: return None if value is None else validate_absolute_path(value) +class ExternalModelProviderConfig(AuthoredConfig): + """A model endpoint managed outside this Slurm allocation. + + Attributes: + name: Provider name referenced by the builder's model configurations. + endpoint: Provider API URL without embedded credentials. + provider_type: Native Data Designer client type, such as ``openai``. + api_key: Optional environment secret reference for authentication. + extra_body: Additional non-secret request parameters. + extra_headers: Header values expressed as literal or secret bindings. + """ + + name: Identifier + endpoint: str + provider_type: Identifier = "openai" + api_key: SecretRef | None = None + extra_body: dict[str, JsonValue] | None = None + extra_headers: dict[str, EnvironmentBinding] = Field(default_factory=dict) + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str) -> str: + validate_url(value, field_name="model provider endpoint") + parsed = urlsplit(value) + if parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment: + raise ValueError("model provider endpoint must not embed credentials, query parameters, or fragments") + return value + + @field_validator("extra_body") + @classmethod + def validate_extra_body(cls, value: dict[str, JsonValue] | None) -> dict[str, JsonValue] | None: + validate_no_plaintext_secrets(value, field_name="model provider extra_body") + return value + + @field_validator("extra_headers") + @classmethod + def validate_headers(cls, values: dict[str, EnvironmentBinding]) -> dict[str, EnvironmentBinding]: + validate_environment_bindings(values) + for name, binding in values.items(): + if re.fullmatch(r"[!#$%&'*+.^_`|~0-9A-Za-z-]+", name) is None: + raise ValueError("model provider headers require valid HTTP field names") + if isinstance(binding, LiteralEnvironmentBinding): + if name.casefold() in {"authorization", "proxy-authorization", "cookie"}: + raise ValueError("authentication headers require external secret references") + return values + + class RemoteMCPProviderConfig(AuthoredConfig): provider_type: Literal["sse", "streamable_http"] name: Identifier @@ -169,6 +217,7 @@ class InvocationConfig(AuthoredConfig): run_config: dict[str, JsonValue] = Field(default_factory=dict) input_bindings: InputBindings = Field(default_factory=InputBindings) mcp_providers: list[MCPProviderConfig] = Field(default_factory=list) + model_providers: list[ExternalModelProviderConfig] = Field(default_factory=list, exclude_if=lambda value: not value) model_concurrency: dict[str, PositiveInt] = Field(default_factory=dict) @field_validator("run_config") @@ -185,6 +234,9 @@ def validate_mcp_providers(self) -> InvocationConfig: names = [provider.name for provider in self.mcp_providers] if len(names) != len(set(names)): raise ValueError("MCP provider names must be unique") + model_names = [provider.name for provider in self.model_providers] + if len(model_names) != len(set(model_names)): + raise ValueError("model provider names must be unique") return self @@ -357,7 +409,7 @@ class DataDesignerSlurmConfig(AuthoredConfig): builder: BuilderInput invocation: InvocationConfig client: ClientConfig = Field(default_factory=ClientConfig) - deployments: list[ServerDeploymentConfig] = Field(min_length=1) + deployments: list[ServerDeploymentConfig] = Field(default_factory=list) array_tasks: ArrayTasksConfig = Field(default_factory=ArrayTasksConfig) submission: SubmissionConfig = Field(default_factory=SubmissionConfig) @@ -368,6 +420,9 @@ def validate_run(self) -> DataDesignerSlurmConfig: aliases = [deployment.model_alias for deployment in self.deployments] if len(aliases) != len(set(aliases)): raise ValueError("deployment model aliases must be unique") + provider_names = {provider.name for provider in self.invocation.model_providers} + if provider_names.intersection(f"slurm-{alias}" for alias in aliases): + raise ValueError("external provider names must not collide with managed providers") groups: dict[str, list[ServerDeploymentConfig]] = {} for deployment in self.deployments: if deployment.resources.shared_node is not None: @@ -387,7 +442,7 @@ def validate_run(self) -> DataDesignerSlurmConfig: ): raise ValueError("shared_node requires matching server environment and queue backpressure") unknown_concurrency = set(self.invocation.model_concurrency).difference(aliases) - if unknown_concurrency: + if unknown_concurrency and not self.invocation.model_providers: raise ValueError( f"model concurrency references undeclared aliases: {', '.join(sorted(unknown_concurrency))}" ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index 8b57f5149..639ab9ab9 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -221,14 +221,14 @@ def _build_generation_directives( resolved_array = f"{resolved_array}%{plan.array_tasks.max_concurrent}" profile = plan.selected_profile.profile - if profile.gpu_request_mode == "gres": + if plan.deployments and profile.gpu_request_mode == "gres": # Reserve CPUs with every GPU so socket-local CPU binding cannot hide GPUs from a serving step. cpus_per_gpu = max( 1, (plan.client.authored.cpus + plan.resolved_gpus_per_node - 1) // plan.resolved_gpus_per_node ) cpu_directive = ("cpus-per-gpu", str(cpus_per_gpu)) else: - if profile.scheduler.mem_per_gpu is not None: + if plan.deployments and profile.scheduler.mem_per_gpu is not None: raise SlurmBatchRenderError("mem_per_gpu requires GRES GPU request mode") cpu_directive = ("cpus-per-task", str(plan.client.authored.cpus)) @@ -244,9 +244,11 @@ def _build_generation_directives( ("output", output_path), ("error", error_path), ] - if profile.gpu_request_mode == "gres": + if not plan.deployments: + values.append(("gres", "none")) + elif profile.gpu_request_mode == "gres": values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) - if profile.scheduler.mem_per_gpu is not None: + if plan.deployments and profile.scheduler.mem_per_gpu is not None: values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) if plan.submission.comment is not None: values.append(("comment", plan.submission.comment)) @@ -255,6 +257,8 @@ def _build_generation_directives( def _requires_exclusive_nodes(plan: ResolvedSlurmRunPlan) -> bool: """Avoid remote-node port collisions when batch GPU IDs describe node zero only.""" + if not plan.deployments: + return False if plan.selected_profile.profile.gpu_request_mode == "visible": return True return any(index > 0 for deployment in plan.deployments for index in deployment.node_indices) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py index 21f59e381..11e8ef890 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py @@ -5,6 +5,8 @@ from __future__ import annotations +from collections.abc import Collection + from pydantic import JsonValue from data_designer.slurm.contracts import ModelAlias, Sha256Digest, compute_serialized_json_sha256 @@ -32,3 +34,19 @@ def get_persisted_builder_identity( ) -> tuple[tuple[ModelAlias, ...], Sha256Digest]: """Return declared aliases and the digest of persisted builder JSON.""" return get_declared_model_aliases(builder), compute_serialized_json_sha256(builder) + + +def get_external_model_aliases( + builder: dict[str, JsonValue], provider_names: Collection[str] +) -> tuple[ModelAlias, ...]: + """Return explicitly external aliases in their authored builder order.""" + get_declared_model_aliases(builder) + data_designer = builder.get("data_designer", builder) + assert isinstance(data_designer, dict) + configs = data_designer.get("model_configs") or [] + assert isinstance(configs, list) + return tuple( + config["alias"] + for config in configs + if isinstance(config, dict) and isinstance(config.get("provider"), str) and config["provider"] in provider_names + ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index f107a9e2e..b72c04996 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -184,7 +184,7 @@ def _compile_client( effective: EffectiveDataDesignerSlurmConfig, deployments: tuple[ResolvedDeployment, ...], ) -> ResolvedClient: - host_node_index = deployments[0].node_indices[0] + host_node_index = deployments[0].node_indices[0] if deployments else 0 ports = tuple( PortClaim( name=f"{deployment.deployment_id}-logical-endpoint", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py index 10ee11ee4..5db34eaaa 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -47,7 +47,7 @@ validate_local_config_path, validate_plain_text, ) -from data_designer.slurm.planning.builder_identity import get_declared_model_aliases +from data_designer.slurm.planning.builder_identity import get_declared_model_aliases, get_external_model_aliases from data_designer.slurm.types import NetworkPort, PartitionSelection @@ -186,6 +186,7 @@ class ResolvedBuilderInput(ContractValue): # Both forms use deterministic persisted builder JSON bytes. content_sha256: Sha256Digest model_aliases: tuple[ModelAlias, ...] + external_model_aliases: tuple[ModelAlias, ...] = Field(default=(), exclude_if=lambda value: not value) @model_validator(mode="after") def validate_input(self) -> ResolvedBuilderInput: @@ -204,6 +205,10 @@ def validate_input(self) -> ResolvedBuilderInput: expected_digest = self.source.sha256 if len(self.model_aliases) != len(set(self.model_aliases)): raise ValueError("resolved builder model aliases must be unique") + if len(self.external_model_aliases) != len(set(self.external_model_aliases)) or not set( + self.external_model_aliases + ).issubset(self.model_aliases): + raise ValueError("external model aliases must be a unique subset of builder aliases") if self.content_sha256 != expected_digest: raise ValueError("builder content digest does not match the resolved input") return self @@ -394,11 +399,11 @@ class ResolvedSlurmRunPlan(ContractRecord): authored_config: ArtifactReference selected_profile: SelectedSlurmProfile scheduler_bin_path: str | None = Field(default=None, exclude_if=lambda value: value is None) - resolved_gpus_per_node: PositiveInt + resolved_gpus_per_node: NonNegativeInt builder: ResolvedBuilderInput invocation: ResolvedInvocation client: ResolvedClient - deployments: tuple[ResolvedDeployment, ...] = Field(min_length=1) + deployments: tuple[ResolvedDeployment, ...] array_tasks: ArrayTasksConfig shards: tuple[PlannedShard, ...] = Field(min_length=1) submission: ResolvedSubmission @@ -420,7 +425,15 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: profile = self.selected_profile.profile if profile.scheduler.bin_path is not None and self.scheduler_bin_path not in (None, profile.scheduler.bin_path): raise ValueError("resolved scheduler bin path does not match the selected profile override") - if profile.gpus_per_node != "auto" and profile.gpus_per_node != self.resolved_gpus_per_node: + if not self.deployments and self.resolved_gpus_per_node != 0: + raise ValueError("client-only runs must resolve zero GPUs") + if self.deployments and self.resolved_gpus_per_node == 0: + raise ValueError("server deployments require a positive GPU count") + if ( + self.deployments + and profile.gpus_per_node != "auto" + and profile.gpus_per_node != self.resolved_gpus_per_node + ): raise ValueError("resolved GPU count does not match the selected profile") if profile.gpu_request_mode == "visible" and profile.scheduler.mem_per_gpu is not None: raise ValueError("mem_per_gpu requires GRES GPU request mode") @@ -463,10 +476,23 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: raise ValueError("resolved deployment IDs must use complete ordered zero-based identities") if len(aliases) != len(set(aliases)): raise ValueError("resolved deployment aliases must be unique") - if set(aliases) != set(self.builder.model_aliases): + external = set(self.builder.external_model_aliases) + if external.intersection(aliases) or set(aliases) | external != set(self.builder.model_aliases): raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") + provider_names = {provider.name for provider in self.invocation.authored.model_providers} + if external and not provider_names: + raise ValueError("external model aliases require external providers") + if provider_names.intersection(f"slurm-{alias}" for alias in aliases): + raise ValueError("external provider names must not collide with managed providers") + if self.builder.inline is not None and self.builder.external_model_aliases != get_external_model_aliases( + self.builder.inline, provider_names + ): + raise ValueError("resolved external model aliases do not match the inline builder") + if set(self.invocation.authored.model_concurrency).difference(self.builder.model_aliases): + raise ValueError("model concurrency references undeclared aliases") - if self.client.host_node_index != self.deployments[0].node_indices[0]: + expected_client_node = self.deployments[0].node_indices[0] if self.deployments else 0 + if self.client.host_node_index != expected_client_node: raise ValueError("client must be colocated on the first node of the first deployment") authored_bindings = self.invocation.authored.input_bindings expected_bindings = InputBindings( @@ -578,7 +604,7 @@ def derive_deployment_placement( placements: list[tuple[tuple[int, ...], int, int]] = [] next_node = 0 shared_nodes = {deployment.resources.shared_node for deployment in deployments} - if shared_nodes != {None} and (len(shared_nodes) != 1 or None in shared_nodes): + if shared_nodes and shared_nodes != {None} and (len(shared_nodes) != 1 or None in shared_nodes): raise ValueError("shared_node requires one group covering all deployments") for deployment in deployments: resources = deployment.resources diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py index 5f735142d..e2a35dcb0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -8,7 +8,7 @@ import posixpath from typing import Annotated -from pydantic import JsonValue, PositiveInt, StringConstraints, TypeAdapter, ValidationError +from pydantic import JsonValue, NonNegativeInt, StringConstraints, TypeAdapter, ValidationError from data_designer.config import RunConfig from data_designer.slurm._errors import format_validation_error @@ -26,6 +26,7 @@ ) from data_designer.slurm.planning.builder_identity import ( get_declared_model_aliases, + get_external_model_aliases, get_persisted_builder_identity, ) from data_designer.slurm.planning.errors import SlurmConfigResolutionError @@ -85,7 +86,7 @@ class EffectiveDataDesignerSlurmConfig(ContractValue): authored: DataDesignerSlurmConfig selected_profile: SelectedSlurmProfile scheduler_bin_path: str | None = None - resolved_gpus_per_node: PositiveInt + resolved_gpus_per_node: NonNegativeInt builder: ResolvedBuilderInput builder_payload: dict[str, JsonValue] | None = None invocation: ResolvedInvocation @@ -114,7 +115,12 @@ def resolve_slurm_config( try: run_id = _RUN_ID_ADAPTER.validate_python(run_id, strict=True) package_version = _PACKAGE_VERSION_ADAPTER.validate_python(package_version, strict=True) - gpus_per_node = _resolve_gpu_count(selected_profile, resolved_gpus_per_node) + if not authored.deployments: + if resolved_gpus_per_node not in (None, 0): + raise SlurmConfigResolutionError("client-only runs must resolve zero GPUs") + gpus_per_node = 0 + else: + gpus_per_node = _resolve_gpu_count(selected_profile, resolved_gpus_per_node) workspace_root = selected_profile.profile.workspace_root run_root = posixpath.join(workspace_root, "runs", run_id) builder, resolved_builder_payload = _resolve_builder( @@ -159,7 +165,15 @@ def validate_effective_slurm_config( authored = effective.authored profile = effective.selected_profile.profile workspace_root = profile.workspace_root - if profile.gpus_per_node != "auto" and profile.gpus_per_node != effective.resolved_gpus_per_node: + if not authored.deployments and effective.resolved_gpus_per_node != 0: + raise SlurmConfigResolutionError("client-only runs must resolve zero GPUs") + if authored.deployments and effective.resolved_gpus_per_node <= 0: + raise SlurmConfigResolutionError("server deployments require a positive GPU count") + if ( + authored.deployments + and profile.gpus_per_node != "auto" + and profile.gpus_per_node != effective.resolved_gpus_per_node + ): raise SlurmConfigResolutionError("resolved GPU count does not match the selected profile") if profile.gpu_request_mode == "visible" and profile.scheduler.mem_per_gpu is not None: raise SlurmConfigResolutionError("mem_per_gpu requires GRES GPU request mode") @@ -224,6 +238,7 @@ def _resolve_builder( run_root: str, builder_payload: dict[str, JsonValue] | None, ) -> tuple[ResolvedBuilderInput, dict[str, JsonValue] | None]: + provider_names = {provider.name for provider in authored.invocation.model_providers} if authored.builder.inline is not None: if builder_payload is not None: raise SlurmConfigResolutionError("inline builder input must not provide a separate payload") @@ -232,6 +247,7 @@ def _resolve_builder( inline=authored.builder.inline, content_sha256=compute_serialized_json_sha256(authored.builder.inline), model_aliases=get_declared_model_aliases(authored.builder.inline), + external_model_aliases=get_external_model_aliases(authored.builder.inline, provider_names), ), None, ) @@ -250,6 +266,7 @@ def _resolve_builder( source=source, content_sha256=source.sha256, model_aliases=aliases, + external_model_aliases=get_external_model_aliases(validated_payload, provider_names), ), validated_payload, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py index 8a35033c2..96d7b790f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -10,7 +10,7 @@ from data_designer.config import RunConfig from data_designer.slurm.config.run import DataDesignerSlurmConfig from data_designer.slurm.contracts import derive_managed_assets_path -from data_designer.slurm.planning.builder_identity import get_persisted_builder_identity +from data_designer.slurm.planning.builder_identity import get_external_model_aliases, get_persisted_builder_identity from data_designer.slurm.planning.errors import SlurmPlanContractError from data_designer.slurm.planning.models import ( ResolvedDependencyLock, @@ -75,6 +75,14 @@ def validate_resolved_plan( _require(plan.builder.source.sha256 == digest, "resolved builder digest does not match builder source") _require(plan.builder.content_sha256 == digest, "resolved builder digest does not match builder source") + payload = authored.builder.inline if authored.builder.inline is not None else builder_payload + assert payload is not None + _require( + plan.builder.external_model_aliases + == get_external_model_aliases(payload, {provider.name for provider in authored.invocation.model_providers}), + "resolved external aliases do not match the builder", + ) + expected_account = authored.submission.account or plan.selected_profile.profile.scheduler.account expected_partition = authored.submission.partition or plan.selected_profile.profile.scheduler.partition _require(plan.submission.account == expected_account, "resolved account does not match authored/profile input") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 214711b7f..b1f93d7cc 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -155,11 +155,15 @@ def _build_client_step( *endpoint_arguments, ) secret_names = collect_secret_environment_names( - (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) + ( + plan.client.authored.dependencies.index_credentials, + plan.invocation.authored.mcp_providers, + plan.invocation.authored.model_providers, + ) ) allocation_environment = ( {"SLURM_JOB_GPUS": environment["SLURM_JOB_GPUS"]} - if plan.selected_profile.profile.gpu_request_mode == "gres" + if plan.deployments and plan.selected_profile.profile.gpu_request_mode == "gres" else {} ) return _step( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py index 91330e39e..9a8312a3e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/controller.py @@ -512,6 +512,8 @@ def _wait_for_endpoints( endpoints: tuple[RuntimeEndpoint, ...], required: tuple[ManagedStep, ...], ) -> None: + if not endpoints: + return timeout = max(status.deployment.launch_policy.startup_timeout_seconds for status in self._statuses) deadline = self._clock.monotonic() + timeout while True: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh index edac5241e..45e8b1f1c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/entrypoint.sh @@ -77,20 +77,22 @@ dd_slurm_run_allocation() { ((${#DD_STEP_IDS[@]} == 1)) dd_run_step "${DD_RUNTIME_MANIFEST}" "${DD_STEP_IDS[0]}" - DD_RUNTIME_PHASE=server_preflight - dd_runtime_log 'Running serving preflight...' - dd_run_role_steps server_preflight - DD_RUNTIME_PHASE=start_servers - dd_runtime_log 'Starting vLLM servers...' - dd_start_servers - DD_RUNTIME_PHASE=server_readiness - dd_runtime_log 'Waiting for vLLM servers to become ready...' - dd_wait_for_role_readiness server - dd_runtime_log 'vLLM servers are ready.' - DD_RUNTIME_PHASE=start_endpoints - dd_start_endpoints - DD_RUNTIME_PHASE=endpoint_readiness - dd_wait_for_role_readiness endpoint + if ((DD_EXPECTED_GPUS > 0)); then + DD_RUNTIME_PHASE=server_preflight + dd_runtime_log 'Running serving preflight...' + dd_run_role_steps server_preflight + DD_RUNTIME_PHASE=start_servers + dd_runtime_log 'Starting vLLM servers...' + dd_start_servers + DD_RUNTIME_PHASE=server_readiness + dd_runtime_log 'Waiting for vLLM servers to become ready...' + dd_wait_for_role_readiness server + dd_runtime_log 'vLLM servers are ready.' + DD_RUNTIME_PHASE=start_endpoints + dd_start_endpoints + DD_RUNTIME_PHASE=endpoint_readiness + dd_wait_for_role_readiness endpoint + fi dd_require_running DD_RUNTIME_PHASE=ready dd_run_bound_control_phase ready @@ -170,6 +172,10 @@ dd_resolve_allocation_hosts() { } dd_verify_gpu_count() { + [[ ${DD_EXPECTED_GPUS} =~ ^(0|[1-9][0-9]*)$ ]] || return 65 + if [[ ${DD_EXPECTED_GPUS} == 0 ]]; then + return 0 + fi local visible=${CUDA_VISIBLE_DEVICES:-${SLURM_JOB_GPUS:-}} local count=0 item local -A devices=() diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py index 60da12f7d..5b8a2933d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/manifest.py @@ -127,7 +127,7 @@ class RuntimeBootstrapManifest(ContractRecord): attempt_id: Identifier plan_sha256: Sha256Digest all_secret_environment_names: tuple[EnvironmentName, ...] - steps: tuple[RuntimeStepSpec, ...] = Field(min_length=4) + steps: tuple[RuntimeStepSpec, ...] = Field(min_length=2) @model_validator(mode="after") def validate_steps(self) -> RuntimeBootstrapManifest: @@ -137,7 +137,8 @@ def validate_steps(self) -> RuntimeBootstrapManifest: roles = tuple(step.role for step in self.steps) if roles.count(RuntimeStepRole.CLIENT_PREFLIGHT) != 1 or roles.count(RuntimeStepRole.CLIENT) != 1: raise ValueError("runtime manifest requires one preflight and generation step") - if RuntimeStepRole.SERVER not in roles or RuntimeStepRole.ENDPOINT not in roles: + managed_roles = {RuntimeStepRole.SERVER_PREFLIGHT, RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT} + if managed_roles.intersection(roles) and not {RuntimeStepRole.SERVER, RuntimeStepRole.ENDPOINT}.issubset(roles): raise ValueError("runtime manifest requires server and endpoint steps") return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh index f1c78a462..47c642867 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh @@ -18,8 +18,7 @@ dd_read_control_plan() { local plan=$1 jq -e ' .schema_version == 1 - and ([.deployments[].node_indices[]] | length > 0) - and ([.deployments[].node_indices[]] | all(type == "number" and . >= 0)) + and ([.client.host_node_index, .deployments[].node_indices[]] | all(type == "number" and . >= 0)) and ([.container_mounts[] | (.source + .target)] | all(test("[,:]") | not)) and ([.container_mounts[].target] | all( . != "/run/data-designer-slurm" @@ -30,7 +29,7 @@ dd_read_control_plan() { DD_CLIENT_PYTHON=$(jq -er '.client.runtime.python_executable' "${plan}") DD_CLIENT_CPUS=$(jq -er '.client.authored.cpus | tostring' "${plan}") DD_EXPECTED_GPUS=$(jq -er '.resolved_gpus_per_node | tostring' "${plan}") - DD_EXPECTED_NODES=$(jq -er '[.deployments[].node_indices[]] | max + 1 | tostring' "${plan}") + DD_EXPECTED_NODES=$(jq -er '[.client.host_node_index, .deployments[].node_indices[]] | max + 1 | tostring' "${plan}") DD_CLIENT_NODE_INDEX=$(jq -er '.client.host_node_index | tostring' "${plan}") DD_GPU_REQUEST_MODE=$(jq -er '.selected_profile.profile.gpu_request_mode' "${plan}") DD_CONTAINER_MOUNTS=$(jq -jr ' diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py index 53d6f049a..c6c7e4896 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/ports.py @@ -40,7 +40,7 @@ def resolve_allocation_plan( environment: Mapping[str, str], ) -> ResolvedSlurmRunPlan: """Return a plan with ports bound to the allocation's assigned GPUs.""" - if plan.selected_profile.profile.gpu_request_mode == "visible": + if not plan.deployments or plan.selected_profile.profile.gpu_request_mode == "visible": return plan gpu_ids = _parse_gpu_ids(environment.get("SLURM_JOB_GPUS")) block_start = _PORT_RANGE_START + min(gpu_ids) * _PORTS_PER_GPU 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 a29a7674c..988c811d0 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 @@ -276,6 +276,8 @@ def _resolve_gpu_count( authored: DataDesignerSlurmConfig, operation: SlurmServiceOperation, ) -> int | None: + if not authored.deployments: + return 0 if self._profile.profile.gpus_per_node != "auto": return None partition = authored.submission.partition or self._profile.profile.scheduler.partition diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py index 27fc325a2..81746860a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/readiness.py @@ -88,7 +88,7 @@ class AttemptReadiness(StateRecord): updated_at: datetime started_at: datetime | None = Field(default=None, exclude_if=lambda value: value is None) state: ReadinessState - deployments: tuple[DeploymentReadiness, ...] = Field(min_length=1) + deployments: tuple[DeploymentReadiness, ...] _updated_at_is_utc = field_validator("updated_at")(validate_utc_timestamp) @@ -175,11 +175,15 @@ def _validate_attempt_state(state: ReadinessState, deployments: tuple[Deployment elif state is ReadinessState.STARTING: if any(item in {ReadinessState.FAILED, ReadinessState.STOPPED} for item in deployment_states): raise ValueError("starting attempts cannot contain failed or stopped deployments") - if all(item is ReadinessState.READY for item in deployment_states): + if deployment_states and all(item is ReadinessState.READY for item in deployment_states): raise ValueError("attempts with every deployment ready must use the ready state") elif state is ReadinessState.READY and any(item is not ReadinessState.READY for item in deployment_states): raise ValueError("ready attempts require every deployment to be ready") - elif state is ReadinessState.FAILED and not any(item is ReadinessState.FAILED for item in deployment_states): + elif ( + state is ReadinessState.FAILED + and deployment_states + and not any(item is ReadinessState.FAILED for item in deployment_states) + ): raise ValueError("failed attempts require at least one failed deployment") elif state is ReadinessState.STOPPED and any(item is not ReadinessState.STOPPED for item in deployment_states): raise ValueError("stopped attempts require every deployment to be stopped") diff --git a/packages/data-designer-slurm/tests/client/test_external_model_execution.py b/packages/data-designer-slurm/tests/client/test_external_model_execution.py new file mode 100644 index 000000000..701911978 --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_external_model_execution.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +import pyarrow.parquet as pq +import pytest +from conftest import ClientWorkerCase, FakeDataDesigner +from pytest_httpx import HTTPXMock + +from data_designer.slurm.client.errors import ClientWorkerError +from data_designer.slurm.client.execution import ClientWorker +from data_designer.slurm.config import DataDesignerSlurmConfig +from data_designer.slurm.contracts import compute_serialized_json_sha256 +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.resolution import resolve_slurm_config + + +@pytest.fixture +def external_client_case( + external_run: DataDesignerSlurmConfig, + client_worker_case: ClientWorkerCase, +) -> ClientWorkerCase: + case = client_worker_case + plan = SlurmRunCompiler.compile( + resolve_slurm_config( + external_run, + selected_profile=case.plan.selected_profile, + client_runtime=case.plan.client.runtime, + deployment_images=tuple(case.plan.deployments[0].image for _ in external_run.deployments), + dependency_lock=case.lock, + runtime_bundle=case.plan.runtime_bundle, + run_id=case.plan.run_id, + package_version=case.plan.package_version, + ) + ) + case.plan_path.write_text(plan.serialize_json()) + return ClientWorkerCase( + plan, + case.plan_path, + case.lock, + case.attempt_dir, + case.scratch_root, + case.prepared, + case.endpoints if external_run.deployments else {}, + ) + + +def test_worker_preserves_hosted_models_and_scopes_secrets( + external_client_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + case = external_client_case + monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") + monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + designers: list[FakeDataDesigner] = [] + + def factory(**kwargs: object) -> FakeDataDesigner: + designer = FakeDataDesigner(**kwargs) + designers.append(designer) + return designer + + worker = ClientWorker(data_designer_factory=factory) + worker.preflight(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) + model = next(config for config in designers[0].validated_builder.model_configs if config.alias == "judge") + assert model.model == "example/hosted-judge" + assert model.provider == "external" + assert model.inference_parameters.temperature == 0.3 + assert model.inference_parameters.extra_body == {"reasoning_effort": "medium"} + assert model.inference_parameters.max_parallel_requests == 3 + providers = {provider.name: provider for provider in designers[0].initialization["model_providers"]} + assert providers["external"].endpoint == "https://models.example/v1" + assert providers["external"].api_key == "fake-provider-secret" + assert providers["external"].extra_body == {"routing": {"version": 2}} + assert providers["external"].extra_headers == {"X-Project": "test", "X-Token": "fake-header-secret"} + if case.plan.deployments: + assert providers["slurm-generator"].endpoint == case.endpoints["generator"] + else: + assert set(providers) == {"external"} + result = worker.run(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) + assert result.actual_records == case.plan.shards[0].requested_records + for path in case.plan_path.parent.rglob("*.json"): + content = path.read_text() + assert "fake-provider-secret" not in content + assert "fake-header-secret" not in content + + +@pytest.mark.parametrize("missing", ("EXTERNAL_API_KEY", "EXTERNAL_HEADER_TOKEN")) +def test_worker_rejects_unavailable_hosted_credentials( + external_client_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, + missing: str, +) -> None: + monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") + monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + monkeypatch.delenv(missing) + case = external_client_case + with pytest.raises(ClientWorkerError, match="required client secret is unavailable"): + ClientWorker(data_designer_factory=FakeDataDesigner).preflight( + case.plan_path, + prepared=case.prepared, + endpoints=case.endpoints, + plugins=(), + ) + + +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +def test_hosted_only_worker_needs_no_scheduler_gpu_environment( + external_client_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("SLURM_JOB_GPUS", raising=False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") + monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + case = external_client_case + worker = ClientWorker(data_designer_factory=FakeDataDesigner, environment={}) + worker.preflight(case.plan_path, prepared=case.prepared, endpoints={}, plugins=()) + result = worker.run( + case.plan_path, + prepared=case.prepared, + endpoints={}, + plugins=(), + ) + assert result.actual_records == case.plan.shards[0].requested_records + + +@pytest.mark.httpx_mock(can_send_already_matched_responses=True) +def test_native_dd_generates_through_managed_and_external_endpoints( + external_client_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, + httpx_mock: HTTPXMock, +) -> None: + case = external_client_case + monkeypatch.setattr("data_designer.engine.models.telemetry.TELEMETRY_ENABLED", False) + monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") + monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + payload = case.plan.model_dump(mode="python") + payload["builder"]["inline"]["data_designer"]["columns"] = [ + { + "column_type": "llm-text", + "name": "answer", + "model_alias": "judge", + "prompt": "Review {{ draft }}." if case.plan.deployments else "Say hello.", + } + ] + if case.plan.deployments: + payload["builder"]["inline"]["data_designer"]["columns"].insert( + 0, + { + "column_type": "llm-text", + "name": "draft", + "model_alias": "generator", + "prompt": "Say hello.", + }, + ) + payload["builder"]["content_sha256"] = compute_serialized_json_sha256(payload["builder"]["inline"]) + plan = ResolvedSlurmRunPlan.model_validate(payload) + case.plan_path.write_text(plan.serialize_json()) + response = { + "id": "test-completion", + "object": "chat.completion", + "created": 0, + "model": "example/hosted-judge", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hello"}}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + hosted_url = "https://models.example/v1/chat/completions" + httpx_mock.add_response(url=hosted_url, json=response) + for endpoint in case.endpoints.values(): + httpx_mock.add_response(url=endpoint + "/chat/completions", json=response) + worker = ClientWorker(environment={"SLURM_JOB_GPUS": "0"} if case.plan.deployments else {}) + worker.preflight(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) + result = worker.run(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) + assert result.actual_records == plan.shards[0].requested_records + table = pq.read_table(case.attempt_dir / "dataset/parquet-files") + assert table["answer"].to_pylist() == ["hello"] * result.actual_records + requests = httpx_mock.get_requests(url=hosted_url) + assert requests + assert all(request.headers["Authorization"] == "Bearer fake-provider-secret" for request in requests) + assert all(request.headers["X-Token"] == "fake-header-secret" for request in requests) + assert all(json.loads(request.content)["model"] == "example/hosted-judge" for request in requests) + for endpoint in case.endpoints.values(): + local_requests = httpx_mock.get_requests(url=endpoint + "/chat/completions") + assert local_requests + assert all("fake-provider-secret" not in str(request.headers) for request in local_requests) + assert all("fake-header-secret" not in str(request.headers) for request in local_requests) + assert all(json.loads(request.content)["model"] == "example/generator" for request in local_requests) diff --git a/packages/data-designer-slurm/tests/config/test_loading_builder.py b/packages/data-designer-slurm/tests/config/test_loading_builder.py index 618a091b1..ea1f63f75 100644 --- a/packages/data-designer-slurm/tests/config/test_loading_builder.py +++ b/packages/data-designer-slurm/tests/config/test_loading_builder.py @@ -73,7 +73,7 @@ def test_client_cpu_default_is_16() -> None: def test_builder_requires_complete_authored_intent() -> None: builder = DataDesignerSlurmConfigBuilder.from_builder_source("builder.json") - with pytest.raises(SlurmConfigBuilderError, match="invocation, deployment"): + with pytest.raises(SlurmConfigBuilderError, match="invocation"): builder.build() diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index ec7950ebe..9b928c888 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -85,6 +85,65 @@ def single_node_plan() -> ResolvedSlurmRunPlan: return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json").read_text()) +@pytest.fixture(params=("mixed", "hosted-only")) +def external_run( + request: pytest.FixtureRequest, authored_run_single: DataDesignerSlurmConfig +) -> DataDesignerSlurmConfig: + payload = authored_run_single.model_dump(mode="python") + models = payload["builder"]["inline"]["data_designer"]["model_configs"] + models.append( + { + "alias": "judge", + "model": "example/hosted-judge", + "provider": "external", + "inference_parameters": { + "temperature": 0.3, + "max_parallel_requests": 7, + "extra_body": {"reasoning_effort": "medium"}, + }, + } + ) + payload["invocation"]["model_providers"] = [ + { + "name": "external", + "endpoint": "https://models.example/v1", + "provider_type": "openai", + "api_key": {"type": "secret", "environment": "EXTERNAL_API_KEY"}, + "extra_body": {"routing": {"version": 2}}, + "extra_headers": { + "X-Project": {"type": "literal", "value": "test"}, + "X-Token": {"type": "secret", "environment": "EXTERNAL_HEADER_TOKEN"}, + }, + } + ] + payload["invocation"]["model_concurrency"]["judge"] = 3 + if request.param == "hosted-only": + models.pop(0) + payload["deployments"] = [] + del payload["invocation"]["model_concurrency"]["generator"] + return DataDesignerSlurmConfig.model_validate(payload) + + +@pytest.fixture +def external_plan( + external_run: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, +) -> ResolvedSlurmRunPlan: + return SlurmRunCompiler.compile( + resolve_slurm_config( + external_run, + selected_profile=single_node_plan.selected_profile, + client_runtime=single_node_plan.client.runtime, + deployment_images=tuple(single_node_plan.deployments[0].image for _ in external_run.deployments), + dependency_lock=dependency_lock_single, + runtime_bundle=single_node_plan.runtime_bundle, + run_id=single_node_plan.run_id, + package_version=single_node_plan.package_version, + ) + ) + + @pytest.fixture def shared_node_config(authored_run_single: DataDesignerSlurmConfig) -> DataDesignerSlurmConfig: payload = authored_run_single.model_dump(mode="python") diff --git a/packages/data-designer-slurm/tests/planning/test_external_providers.py b/packages/data-designer-slurm/tests/planning/test_external_providers.py new file mode 100644 index 000000000..a0f9d6022 --- /dev/null +++ b/packages/data-designer-slurm/tests/planning/test_external_providers.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from data_designer.slurm.config import ( + BuilderInput, + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + ExternalModelProviderConfig, + InvocationConfig, + injected_profile, + load_run_config, +) +from data_designer.slurm.launcher.renderer import render_generation_attempt_script +from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan +from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.errors import SlurmPlanCompilationError +from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.runtime.ports import resolve_allocation_plan + + +@pytest.mark.parametrize("name", ("hosted-only.yaml", "mixed-providers.yaml")) +def test_external_provider_examples_load(name: str) -> None: + path = Path(__file__).parents[2] / "examples" / name + config = load_run_config(path) + assert config.invocation.model_providers[0].name == "hosted" + assert {deployment.model_alias for deployment in config.deployments} == ( + {"generator"} if name == "mixed-providers.yaml" else set() + ) + + +@pytest.mark.parametrize("sourced", (False, True)) +def test_external_models_compile_and_round_trip( + external_run: DataDesignerSlurmConfig, + external_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, + sourced: bool, +) -> None: + builder_payload = external_run.builder.inline + config = external_run + if sourced: + config = config.model_copy(update={"builder": BuilderInput(source="builder.json")}) + plan = SlurmRunCompiler.compile( + resolve_slurm_config( + config, + selected_profile=external_plan.selected_profile, + client_runtime=external_plan.client.runtime, + deployment_images=tuple(deployment.image for deployment in external_plan.deployments), + dependency_lock=dependency_lock_single, + runtime_bundle=external_plan.runtime_bundle, + run_id=external_plan.run_id, + package_version=external_plan.package_version, + builder_payload=builder_payload if sourced else None, + ) + ) + assert plan.builder.external_model_aliases == ("judge",) + assert ResolvedSlurmRunPlan.model_validate_json(plan.serialize_json()) == plan + assert plan.invocation.authored.model_concurrency["judge"] == 3 + script = render_generation_attempt_script(plan, attempt_ordinal=1) + if config.deployments: + assert len(plan.deployments) == 1 + assert "#SBATCH --gres=gpu:8" in script + else: + assert plan.resolved_gpus_per_node == 0 + assert plan.client.host_node_index == 0 + assert plan.client.ports == () + assert "#SBATCH --gres=none" in script + assert "#SBATCH --nodes=1" in script + assert f"#SBATCH --cpus-per-task={config.client.cpus}" in script + assert "--cpus-per-gpu" not in script + assert "--exclusive" not in script + assert resolve_allocation_plan(plan, {}) == plan + + +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +def test_builder_accepts_external_providers_without_deployments(external_run: DataDesignerSlurmConfig) -> None: + builder = DataDesignerSlurmConfigBuilder(external_run.builder).with_invocation( + num_records=8, + dataset_name="hosted", + model_providers=external_run.invocation.model_providers, + ) + config = builder.build() + assert config.deployments == [] + assert config.invocation.model_providers == external_run.invocation.model_providers + + +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +def test_hosted_only_runs_ignore_profile_gpu_memory(external_plan: ResolvedSlurmRunPlan) -> None: + payload = external_plan.model_dump(mode="python") + profile = external_plan.selected_profile.profile + profile = profile.model_copy(update={"scheduler": profile.scheduler.model_copy(update={"mem_per_gpu": "80G"})}) + payload["selected_profile"] = injected_profile(profile) + plan = ResolvedSlurmRunPlan.model_validate(payload) + script = render_generation_attempt_script(plan, attempt_ordinal=1) + assert "#SBATCH --gres=none" in script + assert "--mem-per-gpu" not in script + + +@pytest.mark.parametrize( + "endpoint", + ( + "https://user:password@models.example/v1", + "https://models.example/v1?token=hidden", + "https://models.example/v1#secret", + "file:///tmp/model", + ), +) +def test_external_provider_rejects_credential_bearing_endpoints(endpoint: str) -> None: + with pytest.raises(ValueError): + ExternalModelProviderConfig(name="external", endpoint=endpoint) + + +@pytest.mark.parametrize( + "values", + ( + {"api_key": "plaintext"}, + {"extra_body": {"api_key": "plaintext"}}, + {"extra_headers": {"Authorization": {"type": "literal", "value": "Bearer plaintext"}}}, + {"extra_headers": {"X-Api-Key": {"type": "literal", "value": "plaintext"}}}, + {"extra_headers": {"X-Bad\r\nHeader": {"type": "literal", "value": "test"}}}, + ), +) +def test_external_provider_rejects_plaintext_credentials_and_invalid_headers(values: dict[str, object]) -> None: + with pytest.raises(ValueError): + ExternalModelProviderConfig.model_validate( + {"name": "external", "endpoint": "https://models.example/v1", **values} + ) + + +def test_external_provider_retains_extension_parameters() -> None: + config = ExternalModelProviderConfig( + name="custom", + endpoint="https://models.example/v1", + provider_type="future-client", + extra_body={"future_extension": {"policy": ["one", "two"]}}, + ) + assert config.provider_type == "future-client" + assert config.extra_body == {"future_extension": {"policy": ["one", "two"]}} + + +def test_external_provider_names_are_unique() -> None: + provider = ExternalModelProviderConfig(name="external", endpoint="https://models.example/v1") + with pytest.raises(ValueError, match="model provider names must be unique"): + InvocationConfig(num_records=8, dataset_name="hosted", model_providers=[provider, provider]) + + +def test_external_model_binding_cannot_be_forged(external_plan: ResolvedSlurmRunPlan) -> None: + payload = external_plan.model_dump(mode="python") + payload["builder"]["external_model_aliases"] = () + with pytest.raises(ValueError, match="exactly cover|external model aliases"): + ResolvedSlurmRunPlan.model_validate(payload) + + +def test_unknown_hosted_concurrency_is_rejected( + external_run: DataDesignerSlurmConfig, + external_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, +) -> None: + payload = external_run.model_dump(mode="python") + payload["invocation"]["model_concurrency"]["typo"] = 4 + authored = DataDesignerSlurmConfig.model_validate(payload) + with pytest.raises(SlurmPlanCompilationError, match="undeclared aliases"): + SlurmRunCompiler.compile( + resolve_slurm_config( + authored, + selected_profile=external_plan.selected_profile, + client_runtime=external_plan.client.runtime, + deployment_images=tuple(deployment.image for deployment in external_plan.deployments), + dependency_lock=dependency_lock_single, + runtime_bundle=external_plan.runtime_bundle, + run_id=external_plan.run_id, + package_version=external_plan.package_version, + ) + ) diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 5c92d5a38..a26342aab 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -24,6 +24,53 @@ from data_designer.slurm.state import RetryPlan, RetryShard +def test_external_providers_only_reach_native_clients( + runtime_case: RuntimeCase, + external_plan: ResolvedSlurmRunPlan, +) -> None: + plan = relocate_plan(external_plan, runtime_case.workspace) + attempt = runtime_case.context.attempt.model_copy( + update={ + "resolved_plan": ArtifactReference( + path=runtime_case.context.attempt.resolved_plan.path, + sha256=plan.compute_sha256(), + ), + } + ) + context = replace(runtime_case.context, plan=plan, shard=plan.shards[0], attempt=attempt) + manifest = build_runtime_manifest( + context, + {"SLURM_JOB_GPUS": "0"} if plan.deployments else {}, + runtime_root=Path("/tmp/data-designer-slurm-4101-0/runtime"), + log_directory=context.attempt_directory / "logs/execution-00000002", + layout=AllocationLayout(("compute-001",)), + ) + assert manifest.all_secret_environment_names == ("EXTERNAL_API_KEY", "EXTERNAL_HEADER_TOKEN") + clients = tuple( + step + for step in manifest.steps + if step.role + in { + RuntimeStepRole.CLIENT_PREFLIGHT, + RuntimeStepRole.CLIENT, + } + ) + assert len(clients) == 2 + assert all( + step.secret_environment + == { + "EXTERNAL_API_KEY": "EXTERNAL_API_KEY", + "EXTERNAL_HEADER_TOKEN": "EXTERNAL_HEADER_TOKEN", + } + for step in clients + ) + assert all("EXTERNAL_API_KEY" not in step.secret_environment for step in manifest.steps if step not in clients) + if not plan.deployments: + assert len(manifest.steps) == 2 + assert all(step.execution == "native" and not step.gpu_indices for step in manifest.steps) + assert all("SLURM_JOB_GPUS" not in step.literal_environment for step in manifest.steps) + + def test_bootstrap_manifest_builds_typed_one_node_steps_without_secret_values(runtime_case: RuntimeCase) -> None: context = runtime_case.context runtime_root = Path("/tmp/data-designer-slurm-4101-0/runtime") diff --git a/packages/data-designer-slurm/tests/runtime/test_controller.py b/packages/data-designer-slurm/tests/runtime/test_controller.py index ceb4a1e73..e4f0ca5da 100644 --- a/packages/data-designer-slurm/tests/runtime/test_controller.py +++ b/packages/data-designer-slurm/tests/runtime/test_controller.py @@ -177,6 +177,43 @@ def _supervisor( ) +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +def test_controller_finishes_hosted_only_run_without_serving_steps( + runtime_case: RuntimeCase, + external_plan: ResolvedSlurmRunPlan, +) -> None: + plan = relocate_plan(external_plan, runtime_case.workspace) + attempt = runtime_case.context.attempt.model_copy( + update={ + "resolved_plan": ArtifactReference( + path=runtime_case.context.attempt.resolved_plan.path, + sha256=plan.compute_sha256(), + ), + } + ) + context = replace(runtime_case.context, plan=plan, shard=plan.shards[0], attempt=attempt) + case = replace(runtime_case, context=context) + clock = FakeClock(case.created_at.replace(second=10), monotonic_time=100) + state = FakeStateStore(context.attempt) + runner = _FakeRunner(generation_hook=lambda: _write_complete_result(case, clock)) + controller = OneNodeAllocationController( + context, + runtime_proxy_path=context.attempt_directory / "runtime/proxy.py", + state=state, + supervisor=_supervisor(runner, clock), + preflight=FakePreflight(), + client_steps=FakeClientStepBuilder(), + prober=_FakeProber(ready=True, clock=clock), + clock=clock, + environment={}, + ) + result = controller.run() + assert result.state is AttemptLifecycleState.SUCCEEDED + assert [step.role for step in runner.steps] == [RuntimeStepRole.CLIENT_PREFLIGHT, RuntimeStepRole.CLIENT] + assert all(not snapshot.deployments for snapshot in state.readiness) + assert state.winners[0].attempt_id == result.attempt_id + + def test_controller_runs_preflight_servers_endpoint_client_and_cleanup(runtime_case: RuntimeCase) -> None: clock = FakeClock(runtime_case.created_at.replace(second=10), monotonic_time=100) state = FakeStateStore(runtime_case.context.attempt) diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index 7fe9e6b28..5dbfde7df 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -15,10 +15,19 @@ import pytest -@pytest.mark.parametrize("gpu_request_mode", ("gres", "visible")) +@pytest.mark.parametrize( + ("gpu_request_mode", "hosted_only"), + ( + ("gres", False), + ("visible", False), + ("gres", True), + ("visible", True), + ), +) def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( tmp_path: Path, gpu_request_mode: str, + hosted_only: bool, ) -> None: runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" scratch_root = tmp_path / "data-designer-slurm-4101-0" @@ -35,9 +44,24 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( plan = _plan( tmp_path, runtime_copy, artifacts, gpu_request_mode, python_executable=(fake_bin / "python3").as_posix() ) + manifest = _manifest(attempt_directory, artifacts, "a" * 64) + if hosted_only: + plan["deployments"] = [] + plan["resolved_gpus_per_node"] = 0 + plan["invocation"] = { + "authored": { + "model_providers": [ + { + "api_key": {"type": "secret", "environment": "SOURCE_TOKEN"}, + } + ] + } + } + manifest["steps"] = [step for step in manifest["steps"] if step["role"] in {"client_preflight", "client"}] + manifest["all_secret_environment_names"] = ["SOURCE_TOKEN"] plan_path.write_text(json.dumps(plan)) manifest_path = tmp_path / "manifest-source.json" - manifest_path.write_text(json.dumps(_manifest(attempt_directory, artifacts, "a" * 64))) + manifest_path.write_text(json.dumps(manifest)) marker_path = tmp_path / "host-python-ran" _write_executable(fake_bin / "python3", f"#!/usr/bin/env bash\nprintf ran > {marker_path}\nexit 99\n") _write_executable(fake_bin / "curl", "#!/usr/bin/env bash\nexit 0\n") @@ -73,6 +97,9 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( "SLURM_JOB_GPUS": "0", "SLURM_NODEID": "0", } + if hosted_only: + environment.pop("SLURM_JOB_GPUS") + environment.pop("CUDA_VISIBLE_DEVICES") completed = subprocess.run( ("bash", "-c", command), @@ -87,13 +114,16 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( assert not scratch_root.exists() assert not marker_path.exists() assert "supersecret" not in completed.stdout + completed.stderr + manifest_path.read_text() - assert "Starting vLLM servers..." in completed.stdout - assert "Waiting for vLLM servers to become ready..." in completed.stdout - assert "vLLM servers are ready." in completed.stdout - assert ( - f"Follow vLLM logs: tail -F {log_directory / 'deployment-00000-replica-00000-rank-00000.out'} " - f"{log_directory / 'deployment-00000-replica-00000-rank-00000.err'}" - ) in completed.stdout + if hosted_only: + assert "vLLM" not in completed.stdout + else: + assert "Starting vLLM servers..." in completed.stdout + assert "Waiting for vLLM servers to become ready..." in completed.stdout + assert "vLLM servers are ready." in completed.stdout + assert ( + f"Follow vLLM logs: tail -F {log_directory / 'deployment-00000-replica-00000-rank-00000.out'} " + f"{log_directory / 'deployment-00000-replica-00000-rank-00000.err'}" + ) in completed.stdout assert ( f"Generation started. Follow client logs: tail -F {log_directory / 'client-generation.out'} " f"{log_directory / 'client-generation.err'}" diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 5da340517..381ebbd73 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -248,6 +248,28 @@ def test_production_wiring_dry_run_resolves_and_renders_without_submission( assert launcher.submissions == [] +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +def test_hosted_only_planning_needs_no_gpu_discovery_or_registered_images( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + external_run: DataDesignerSlurmConfig, +) -> None: + profile = _profile(tmp_path, profile_catalog).model_copy(update={"gpus_per_node": "auto"}) + launcher = _Launcher() + service = create_slurm_run_service( + profile=profile, + launcher=launcher, + run_id_factory=lambda: "run-hosted", + package_version="0.9.2", + ) + plan = service.plan(external_run, source_root=tmp_path) + assert plan.resolved_gpus_per_node == 0 + assert plan.deployments == () + result = service.execute(external_run, source_root=tmp_path, dry_run=True) + assert "#SBATCH --gres=none" in result.batch_script + assert launcher.submissions == [] + + def test_missing_srun_fails_before_submission( tmp_path: Path, profile_catalog: SlurmProfileCatalog, diff --git a/packages/data-designer-slurm/tests/serving/test_resolver.py b/packages/data-designer-slurm/tests/serving/test_resolver.py index 26532ca01..60c4aa7b5 100644 --- a/packages/data-designer-slurm/tests/serving/test_resolver.py +++ b/packages/data-designer-slurm/tests/serving/test_resolver.py @@ -298,6 +298,11 @@ def _multi_group_plan(plan: ResolvedSlurmRunPlan) -> ResolvedSlurmRunPlan: "model_configs" ][:1] payload["builder"]["model_aliases"] = ["generator"] + payload["invocation"]["authored"]["model_concurrency"] = { + alias: value + for alias, value in payload["invocation"]["authored"]["model_concurrency"].items() + if alias == "generator" + } payload["builder"]["content_sha256"] = compute_serialized_json_sha256(payload["builder"]["inline"]) deployment = payload["deployments"][0] deployment["authored"]["resources"]["nodes"] = 4 diff --git a/plans/1079/slurm-external-providers.md b/plans/1079/slurm-external-providers.md new file mode 100644 index 000000000..2e9016ecf --- /dev/null +++ b/plans/1079/slurm-external-providers.md @@ -0,0 +1,35 @@ +# Plan: External model providers in Slurm runs + +Issue: https://github.com/NVIDIA-NeMo/DataDesigner/issues/1079 + +The implementation follows the reviewed plan: retain the existing async engine, +add explicit external providers, and measure allocation utilization during SQL +qualification before considering separate generation and judging allocations. + +## Design + +- Add external provider declarations to invocation configuration. Reuse native DD + provider settings, with Slurm `SecretRef` values at credential boundaries. +- Partition declared model aliases into managed deployments and external models. + Require complete, disjoint coverage and reject provider-name collisions. +- Materialize endpoints only for managed aliases. Preserve external model names, + sampling parameters, and explicit concurrency settings. +- Allow an empty deployment list. Resolve zero GPUs, one native client host, no + serving images, no GPU discovery, no serving containers, and no proxy ports. +- Reuse existing shard, retry, resume, secret transport, and collection contracts. + Preserve serialized local-only plans where defaults are unchanged. + +## Validation + +Exercise authored configs and builders, sourced and inline plan resolution, +native client provider routing, credential isolation, resource directives, +allocation startup and cleanup, retry/resume, and local-only golden fixtures. +Run the Slurm suite and the full package checks before opening the PR. + +SQL qualification uses a small mixed-provider async recipe and a CPU-only hosted +recipe. Record per-model timing and actual GPU utilization, and compare idle +windows with the selected cluster policy. Preserve existing SQL prompts, models, +sampling, judge gates, datasets, checkpoints, and completed study results. + +Automatic DAG partitioning, new async machinery, model substitutions, and the +separate single-row profiling defect are outside this change. From dfcbbf000e2295af9e9fdbaf2216b656dcbdd608 Mon Sep 17 00:00:00 2001 From: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:48:27 -0700 Subject: [PATCH 02/12] docs(slurm): describe hosted-only prerequisites Signed-off-by: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> --- fern/versions/latest/pages/slurm/overview.mdx | 6 +++--- fern/versions/latest/pages/slurm/profiles.mdx | 4 +++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/fern/versions/latest/pages/slurm/overview.mdx b/fern/versions/latest/pages/slurm/overview.mdx index e6f64deac..a3556ee59 100644 --- a/fern/versions/latest/pages/slurm/overview.mdx +++ b/fern/versions/latest/pages/slurm/overview.mdx @@ -1,6 +1,6 @@ # Data Designer on Slurm -Data Designer Slurm runs Data Designer clients and vLLM model servers in batch allocations. The optional extension owns profile selection, image inspection, submission, status, retry, result collection, and benchmark workflows. +Data Designer Slurm runs Data Designer clients and optional vLLM model servers in batch allocations. Recipes can combine local models with external providers, or use CPU allocations for hosted-only generation. The optional extension owns profile selection, image inspection, submission, status, retry, result collection, and benchmark workflows. Start with [Getting Started](/slurm/getting-started). Use [Dependencies and tools](/slurm/dependencies-and-tools) for plugins, MCP providers, and secrets, and keep the [CLI reference](/slurm/cli-reference) nearby for automation. @@ -12,7 +12,7 @@ The submit environment needs: - `sbatch`, `squeue`, `sacct`, `scancel`, `scontrol`, and `sinfo` - Access to a shared workspace that is visible from the submit host and every eligible compute node -Compute nodes need access to the same shared Python environment at the exact path used by the submit command. Enroot 3.5 or later is required for both digest-pinned OCI imports and existing SQSH inspection. Serving images must contain vLLM. +Compute nodes need access to the same shared Python environment at the exact path used by the submit command. Local serving requires Enroot 3.5 or later for both digest-pinned OCI imports and existing SQSH inspection, and serving images must contain vLLM. Hosted-only runs require access to their external endpoints and do not need serving images or Enroot. Your Slurm administrator can confirm the account, GPU partition, CPU image-build partition, GPU request mode, and shared filesystem paths for your cluster. @@ -34,6 +34,6 @@ Runtime bundle extraction, the client dependency overlay, Data Designer home and ## Supported in v1 -The public v1 surface supports named cluster profiles, serving image aliases, one or more vLLM deployments, independent replicas, multi-node replicas, Slurm arrays, bounded array concurrency, durable status, cancellation, sparse retry, winner-driven collection, dependencies, plugins, MCP providers, and benchmarks. +The public v1 surface supports named cluster profiles, serving image aliases, vLLM deployments, external model providers, hosted-only CPU runs, independent replicas, multi-node replicas, Slurm arrays, bounded array concurrency, durable status, cancellation, sparse retry, winner-driven collection, dependencies, plugins, MCP providers, and benchmarks. See [External model providers](/slurm/run-configuration#external-model-providers) for mixed and hosted-only configuration. Interactive allocations, code-sandbox sidecars, NeMo Run, SGLang, Dynamo, Kubernetes, Ray, and platform-specific orchestration are not part of the v1 Slurm extension. diff --git a/fern/versions/latest/pages/slurm/profiles.mdx b/fern/versions/latest/pages/slurm/profiles.mdx index 30ad8d69e..205c8e03f 100644 --- a/fern/versions/latest/pages/slurm/profiles.mdx +++ b/fern/versions/latest/pages/slurm/profiles.mdx @@ -85,7 +85,9 @@ scheduler: The CLI also accepts `--partition gpu-a,gpu-b` during profile initialization. Slurm chooses a partition when scheduling the job; the CPU `image_build.partition` remains a single partition. -`gpus_per_node` is both the GPU count used for deployment topology and, in `gres` mode, the count requested from Slurm on each job node. `auto` queries eligible nodes and uses their reported count, so an eight-GPU partition produces `--gres=gpu:8`. Set an explicit smaller count, such as `1`, for partial-node allocations when the cluster permits them. Use `auto` only when requesting every GPU on each selected node is intended, and use an explicit count when eligible nodes are heterogeneous. +For runs with local deployments, `gpus_per_node` is both the GPU count used for deployment topology and, in `gres` mode, the count requested from Slurm on each job node. `auto` queries eligible nodes and uses their reported count, so an eight-GPU partition produces `--gres=gpu:8`. Set an explicit smaller count, such as `1`, for partial-node allocations when the cluster permits them. Use `auto` only when requesting every GPU on each selected node is intended, and use an explicit count when eligible nodes are heterogeneous. + +[Hosted-only runs](/slurm/run-configuration#hosted-only-cpu-runs) skip GPU discovery and ignore the profile's GPU count and `mem_per_gpu`. Select a CPU partition with `submission.partition` in the run file. Each task requests one node, `--gres=none`, and `client.cpus` CPUs. Multi-node generation jobs request exclusive nodes, including in `gres` mode. This keeps serving ports isolated because Slurm reports batch-job GPU IDs for the first node only. A partial-node `gpus_per_node` setting still limits the GPUs used per node, but other jobs cannot share those nodes during the run. From fd16a7ae14aa3025a3d683baf89a5af5b89e2079 Mon Sep 17 00:00:00 2001 From: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:24:49 -0700 Subject: [PATCH 03/12] fix(slurm): finalize preclaimed datasets and configure node memory Signed-off-by: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> --- .../latest/pages/slurm/run-configuration.mdx | 3 + .../examples/hosted-only.yaml | 1 + .../data_designer/slurm/client/execution.py | 17 +++++- .../src/data_designer/slurm/config/run.py | 5 ++ .../data_designer/slurm/launcher/renderer.py | 4 +- .../data_designer/slurm/planning/models.py | 1 + .../slurm/planning/resolution.py | 1 + .../slurm/planning/validation.py | 1 + .../client/test_external_model_execution.py | 55 +++++++++++++++++++ .../tests/planning/test_external_providers.py | 44 ++++++++++++++- plans/1079/slurm-external-providers.md | 4 ++ 11 files changed, 131 insertions(+), 5 deletions(-) diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index e656c5e64..3cded4cc3 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -112,10 +112,13 @@ client: deployments: [] submission: partition: cpu + memory: 16G ``` Each array task requests one node with `--gres=none` and `--cpus-per-task` set to `client.cpus`. It does not discover GPUs, reserve exclusive nodes, require serving images, or start serving containers and proxies. The profile's GPU facts and GPU memory request are unused for this run. Sharding, status, retry, resume, and collection use the same commands as local-model runs. +Set `submission.memory` to the memory needed per node, such as `16G`. It emits `--mem` and takes precedence over the profile's `mem_per_gpu` request for any run. If omitted, local deployments retain the profile's GPU memory request; otherwise Slurm's memory default applies. Some clusters default to reserving all memory on a node, even for a small CPU-only job. + ## Configure vLLM Each server supports readiness and launch controls: diff --git a/packages/data-designer-slurm/examples/hosted-only.yaml b/packages/data-designer-slurm/examples/hosted-only.yaml index 9b6850c83..5586b8026 100644 --- a/packages/data-designer-slurm/examples/hosted-only.yaml +++ b/packages/data-designer-slurm/examples/hosted-only.yaml @@ -28,3 +28,4 @@ client: deployments: [] submission: partition: cpu + memory: 16G diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index 5b3a0818e..906847b20 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -565,9 +565,22 @@ def _prepare_dataset_workspace( execution_resume = context.requested_resume if context.retry_resume is None else context.retry_resume if execution_resume is not ResumeMode.NEVER: ensure_private_directory(context.dataset_path) - if execution_resume is ResumeMode.NEVER and context.dataset_path.exists(): - if not context.dataset_path.is_dir() or any(context.dataset_path.iterdir()): + if execution_resume is ResumeMode.NEVER and ( + context.dataset_path.exists() or context.dataset_path.is_symlink() + ): + if ( + context.dataset_path.is_symlink() + or not context.dataset_path.is_dir() + or any(context.dataset_path.iterdir()) + ): raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "attempt dataset workspace is not empty") + # The state writer prepares an empty directory; DD must claim it exclusively. + try: + context.dataset_path.rmdir() + except OSError as error: + raise ClientWorkerError( + ClientErrorCode.OUTPUT_INVALID, "attempt dataset workspace is unavailable" + ) from error @staticmethod def _generate_dataset( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py index d110f98b4..f018d29ec 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -385,6 +385,11 @@ class SubmissionConfig(AuthoredConfig): job_name: Identifier = "data-designer" time_limit: Annotated[str, StringConstraints(pattern=r"^(?:[0-9]+-)?[0-9]{2}:[0-9]{2}:[0-9]{2}$")] = "03:55:00" comment: Annotated[str, StringConstraints(max_length=256)] | None = None + memory: Annotated[str, StringConstraints(pattern=r"^[1-9][0-9]*(?:K|M|G|T)$")] | None = Field( + default=None, + exclude_if=lambda value: value is None, + description="Memory requested per allocated node, such as 16G; overrides profile mem_per_gpu.", + ) @field_validator("time_limit") @classmethod diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index 639ab9ab9..bd2fde58e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py @@ -248,7 +248,9 @@ def _build_generation_directives( values.append(("gres", "none")) elif profile.gpu_request_mode == "gres": values.append(("gres", f"gpu:{plan.resolved_gpus_per_node}")) - if plan.deployments and profile.scheduler.mem_per_gpu is not None: + if plan.submission.memory is not None: + values.append(("mem", plan.submission.memory)) + elif plan.deployments and profile.scheduler.mem_per_gpu is not None: values.append(("mem-per-gpu", profile.scheduler.mem_per_gpu)) if plan.submission.comment is not None: values.append(("comment", plan.submission.comment)) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py index 5db34eaaa..f3c3efdd7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -384,6 +384,7 @@ class ResolvedSubmission(ContractValue): job_name: Identifier time_limit: str comment: str | None = None + memory: str | None = Field(default=None, exclude_if=lambda value: value is None) @model_validator(mode="after") def validate_submission(self) -> ResolvedSubmission: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py index e2a35dcb0..6930c5611 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -294,6 +294,7 @@ def _materialize_submission( job_name=authored.submission.job_name, time_limit=authored.submission.time_limit, comment=authored.submission.comment, + memory=authored.submission.memory, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py index 96d7b790f..61c93467a 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -97,6 +97,7 @@ def validate_resolved_plan( "resolved time limit does not match authored input", ) _require(plan.submission.comment == authored.submission.comment, "resolved comment does not match authored input") + _require(plan.submission.memory == authored.submission.memory, "resolved memory does not match authored input") _require( plan.client.dependency_lock.sha256 == dependency_lock.compute_sha256(), diff --git a/packages/data-designer-slurm/tests/client/test_external_model_execution.py b/packages/data-designer-slurm/tests/client/test_external_model_execution.py index 701911978..1aacacc95 100644 --- a/packages/data-designer-slurm/tests/client/test_external_model_execution.py +++ b/packages/data-designer-slurm/tests/client/test_external_model_execution.py @@ -4,12 +4,15 @@ from __future__ import annotations import json +from pathlib import Path import pyarrow.parquet as pq import pytest from conftest import ClientWorkerCase, FakeDataDesigner from pytest_httpx import HTTPXMock +from data_designer.config import DataDesignerConfigBuilder +from data_designer.config.seed_source import LocalFileSeedSource from data_designer.slurm.client.errors import ClientWorkerError from data_designer.slurm.client.execution import ClientWorker from data_designer.slurm.config import DataDesignerSlurmConfig @@ -128,11 +131,53 @@ def test_hosted_only_worker_needs_no_scheduler_gpu_environment( assert result.actual_records == case.plan.shards[0].requested_records +@pytest.mark.parametrize("existing_kind", ("nonempty", "symlink", "dangling-symlink", "raced")) +def test_worker_preserves_existing_dataset_paths( + external_client_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, + existing_kind: str, +) -> None: + case = external_client_case + monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") + monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + worker = ClientWorker(data_designer_factory=FakeDataDesigner) + worker.preflight(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) + dataset = case.attempt_dir / "dataset" + if existing_kind in {"nonempty", "raced"}: + dataset.mkdir() + if existing_kind == "nonempty": + (dataset / "retained.txt").write_text("retain") + else: + rmdir = type(dataset).rmdir + + def racing_rmdir(path: Path) -> None: + if path == dataset: + (dataset / "retained.txt").write_text("retain") + rmdir(path) + + monkeypatch.setattr(type(dataset), "rmdir", racing_rmdir) + else: + target = case.plan_path.parent / "unrelated" + if existing_kind == "symlink": + target.mkdir() + dataset.symlink_to(target, target_is_directory=True) + message = "unavailable" if existing_kind == "raced" else "not empty" + with pytest.raises(ClientWorkerError, match=f"attempt dataset workspace is {message}"): + worker.run(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) + if existing_kind in {"nonempty", "raced"}: + assert (dataset / "retained.txt").read_text() == "retain" + else: + assert dataset.is_symlink() + assert not target.exists() or list(target.iterdir()) == [] + + @pytest.mark.httpx_mock(can_send_already_matched_responses=True) +@pytest.mark.parametrize("seeded", (False, True)) def test_native_dd_generates_through_managed_and_external_endpoints( external_client_case: ClientWorkerCase, monkeypatch: pytest.MonkeyPatch, httpx_mock: HTTPXMock, + seeded: bool, ) -> None: case = external_client_case monkeypatch.setattr("data_designer.engine.models.telemetry.TELEMETRY_ENABLED", False) @@ -157,6 +202,12 @@ def test_native_dd_generates_through_managed_and_external_endpoints( "prompt": "Say hello.", }, ) + if seeded: + seed_path = case.plan_path.parent / "seed.jsonl" + seed_path.write_text("".join(json.dumps({"seed_id": index}) + "\n" for index in range(8))) + builder = DataDesignerConfigBuilder.from_config(payload["builder"]["inline"]) + builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_path)), sampling_strategy="ordered") + payload["builder"]["inline"]["data_designer"]["seed_config"] = builder.get_seed_config().model_dump(mode="json") payload["builder"]["content_sha256"] = compute_serialized_json_sha256(payload["builder"]["inline"]) plan = ResolvedSlurmRunPlan.model_validate(payload) case.plan_path.write_text(plan.serialize_json()) @@ -174,10 +225,14 @@ def test_native_dd_generates_through_managed_and_external_endpoints( httpx_mock.add_response(url=endpoint + "/chat/completions", json=response) worker = ClientWorker(environment={"SLURM_JOB_GPUS": "0"} if case.plan.deployments else {}) worker.preflight(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) + # The allocation controller claims the dataset directory before running DD. + (case.attempt_dir / "dataset").mkdir(exist_ok=True) result = worker.run(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) assert result.actual_records == plan.shards[0].requested_records table = pq.read_table(case.attempt_dir / "dataset/parquet-files") assert table["answer"].to_pylist() == ["hello"] * result.actual_records + if seeded: + assert sorted(table["seed_id"].to_pylist()) == list(range(8)) requests = httpx_mock.get_requests(url=hosted_url) assert requests assert all(request.headers["Authorization"] == "Bearer fake-provider-secret" for request in requests) diff --git a/packages/data-designer-slurm/tests/planning/test_external_providers.py b/packages/data-designer-slurm/tests/planning/test_external_providers.py index a0f9d6022..d6e6b94d8 100644 --- a/packages/data-designer-slurm/tests/planning/test_external_providers.py +++ b/packages/data-designer-slurm/tests/planning/test_external_providers.py @@ -13,14 +13,16 @@ DataDesignerSlurmConfigBuilder, ExternalModelProviderConfig, InvocationConfig, + SubmissionConfig, injected_profile, load_run_config, ) from data_designer.slurm.launcher.renderer import render_generation_attempt_script from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan from data_designer.slurm.planning.compiler import SlurmRunCompiler -from data_designer.slurm.planning.errors import SlurmPlanCompilationError +from data_designer.slurm.planning.errors import SlurmPlanCompilationError, SlurmPlanContractError from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.planning.validation import validate_resolved_plan from data_designer.slurm.runtime.ports import resolve_allocation_plan @@ -35,14 +37,18 @@ def test_external_provider_examples_load(name: str) -> None: @pytest.mark.parametrize("sourced", (False, True)) +@pytest.mark.parametrize("memory", (None, "16G")) def test_external_models_compile_and_round_trip( external_run: DataDesignerSlurmConfig, external_plan: ResolvedSlurmRunPlan, dependency_lock_single: ResolvedDependencyLock, sourced: bool, + memory: str | None, ) -> None: builder_payload = external_run.builder.inline - config = external_run + payload = external_run.model_dump(mode="python") + payload["submission"]["memory"] = memory + config = DataDesignerSlurmConfig.model_validate(payload) if sourced: config = config.model_copy(update={"builder": BuilderInput(source="builder.json")}) plan = SlurmRunCompiler.compile( @@ -62,6 +68,8 @@ def test_external_models_compile_and_round_trip( assert ResolvedSlurmRunPlan.model_validate_json(plan.serialize_json()) == plan assert plan.invocation.authored.model_concurrency["judge"] == 3 script = render_generation_attempt_script(plan, attempt_ordinal=1) + assert plan.submission.memory == memory + assert ("#SBATCH --mem=16G" in script) == (memory is not None) if config.deployments: assert len(plan.deployments) == 1 assert "#SBATCH --gres=gpu:8" in script @@ -84,9 +92,11 @@ def test_builder_accepts_external_providers_without_deployments(external_run: Da dataset_name="hosted", model_providers=external_run.invocation.model_providers, ) + builder.with_submission(memory="16G") config = builder.build() assert config.deployments == [] assert config.invocation.model_providers == external_run.invocation.model_providers + assert config.submission.memory == "16G" @pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) @@ -101,6 +111,36 @@ def test_hosted_only_runs_ignore_profile_gpu_memory(external_plan: ResolvedSlurm assert "--mem-per-gpu" not in script +def test_submission_memory_overrides_profile_gpu_memory(external_plan: ResolvedSlurmRunPlan) -> None: + payload = external_plan.model_dump(mode="python") + profile = external_plan.selected_profile.profile + profile = profile.model_copy(update={"scheduler": profile.scheduler.model_copy(update={"mem_per_gpu": "80G"})}) + payload["selected_profile"] = injected_profile(profile) + payload["submission"]["memory"] = "16G" + plan = ResolvedSlurmRunPlan.model_validate(payload) + script = render_generation_attempt_script(plan, attempt_ordinal=1) + assert "#SBATCH --mem=16G" in script + assert "--mem-per-gpu" not in script + + +@pytest.mark.parametrize("memory", ("0", "16", "-1G", "16G;command")) +def test_submission_rejects_invalid_memory(memory: str) -> None: + with pytest.raises(ValueError): + SubmissionConfig(memory=memory) + + +def test_submission_memory_is_bound_to_authored_config( + external_run: DataDesignerSlurmConfig, + external_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, +) -> None: + payload = external_plan.model_dump(mode="python") + payload["submission"]["memory"] = "16G" + plan = ResolvedSlurmRunPlan.model_validate(payload) + with pytest.raises(SlurmPlanContractError, match="resolved memory does not match authored input"): + validate_resolved_plan(external_run, dependency_lock_single, plan) + + @pytest.mark.parametrize( "endpoint", ( diff --git a/plans/1079/slurm-external-providers.md b/plans/1079/slurm-external-providers.md index 2e9016ecf..bc4a9b9c3 100644 --- a/plans/1079/slurm-external-providers.md +++ b/plans/1079/slurm-external-providers.md @@ -18,6 +18,10 @@ qualification before considering separate generation and judging allocations. serving images, no GPU discovery, no serving containers, and no proxy ports. - Reuse existing shard, retry, resume, secret transport, and collection contracts. Preserve serialized local-only plans where defaults are unchanged. +- Allow an explicit per-node submission memory request. The live CPU probe found + that the cluster otherwise reserves all node memory for a four-CPU job. +- Let DD exclusively create a fresh dataset after the state writer has prepared + its empty directory. Preserve nonempty directories and reject symlink targets. ## Validation From 4be9ff26249f09f751cff22234f9f84e9f4e3841 Mon Sep 17 00:00:00 2001 From: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:37:43 -0700 Subject: [PATCH 04/12] fix(slurm): restrict generated datasets before winner publication Signed-off-by: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> --- fern/versions/latest/pages/slurm/run-configuration.mdx | 2 +- .../src/data_designer/slurm/client/execution.py | 1 + .../src/data_designer/slurm/launcher/collection.py | 3 +++ .../tests/client/test_external_model_execution.py | 7 +++++++ .../data-designer-slurm/tests/launcher/test_collection.py | 4 ++++ plans/1079/slurm-external-providers.md | 2 ++ 6 files changed, 18 insertions(+), 1 deletion(-) diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 3cded4cc3..bba44901d 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -117,7 +117,7 @@ submission: Each array task requests one node with `--gres=none` and `--cpus-per-task` set to `client.cpus`. It does not discover GPUs, reserve exclusive nodes, require serving images, or start serving containers and proxies. The profile's GPU facts and GPU memory request are unused for this run. Sharding, status, retry, resume, and collection use the same commands as local-model runs. -Set `submission.memory` to the memory needed per node, such as `16G`. It emits `--mem` and takes precedence over the profile's `mem_per_gpu` request for any run. If omitted, local deployments retain the profile's GPU memory request; otherwise Slurm's memory default applies. Some clusters default to reserving all memory on a node, even for a small CPU-only job. +Set `submission.memory` to the memory needed per node, such as `16G`. It emits `--mem` for generation and collection jobs and takes precedence over the profile's `mem_per_gpu` request for generation. If omitted, local deployments retain the profile's GPU memory request; otherwise Slurm's memory default applies. Some clusters default to reserving all memory on a node, even for a small CPU-only job. ## Configure vLLM diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index 906847b20..e25abc69b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -681,6 +681,7 @@ def _place_dataset( dataset_path = creation.expected_path if dataset_path != creation.expected_path: raise ClientWorkerError(ClientErrorCode.OUTPUT_INVALID, "Data Designer dataset path is invalid") + ensure_private_directory(dataset_path) return dataset_path, output_paths @staticmethod diff --git a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py index 79cac21d3..e841c7a2f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/launcher/collection.py @@ -32,6 +32,8 @@ def render_collection_script( collection_plan.collection_id, ) collection_plan_path = posixpath.join(collection_root, "plan.json") + memory = resolved_plan.submission.memory + memory_directives = () if memory is None else (("mem", memory),) directives = render_batch_directives( ( ("job-name", collection_plan.submission_job_name), @@ -41,6 +43,7 @@ def render_collection_script( ("ntasks", "1"), ("cpus-per-task", str(resolved_plan.client.authored.cpus)), ("time", resolved_plan.submission.time_limit), + *memory_directives, ("chdir", collection_root), ("output", f"{collection_root}/slurm-%j.out"), ("error", f"{collection_root}/slurm-%j.err"), diff --git a/packages/data-designer-slurm/tests/client/test_external_model_execution.py b/packages/data-designer-slurm/tests/client/test_external_model_execution.py index 1aacacc95..a5b27923c 100644 --- a/packages/data-designer-slurm/tests/client/test_external_model_execution.py +++ b/packages/data-designer-slurm/tests/client/test_external_model_execution.py @@ -20,6 +20,8 @@ from data_designer.slurm.planning import ResolvedSlurmRunPlan from data_designer.slurm.planning.compiler import SlurmRunCompiler from data_designer.slurm.planning.resolution import resolve_slurm_config +from data_designer.slurm.state.artifacts import CandidateArtifactVerifier +from data_designer.slurm.state.outputs import CandidateOutputManifest @pytest.fixture @@ -229,6 +231,11 @@ def test_native_dd_generates_through_managed_and_external_endpoints( (case.attempt_dir / "dataset").mkdir(exist_ok=True) result = worker.run(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) assert result.actual_records == plan.shards[0].requested_records + assert result.candidate_output_manifest is not None + candidate = CandidateOutputManifest.model_validate_json(Path(result.candidate_output_manifest.path).read_text()) + with CandidateArtifactVerifier().verify(candidate) as artifacts: + assert sum(artifacts.record_counts) == result.actual_records + artifacts.rebind() table = pq.read_table(case.attempt_dir / "dataset/parquet-files") assert table["answer"].to_pylist() == ["hello"] * result.actual_records if seeded: diff --git a/packages/data-designer-slurm/tests/launcher/test_collection.py b/packages/data-designer-slurm/tests/launcher/test_collection.py index 861bc1731..9ec0dba0c 100644 --- a/packages/data-designer-slurm/tests/launcher/test_collection.py +++ b/packages/data-designer-slurm/tests/launcher/test_collection.py @@ -22,8 +22,10 @@ from data_designer.slurm.state.errors import StateConflictError +@pytest.mark.parametrize("memory", (None, "16G")) def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives( multi_node_plan: ResolvedSlurmRunPlan, + memory: str | None, ) -> None: workspace_mount = ContainerMount(source="/workspace", target="/workspace", read_only=False) output_mount = ContainerMount( @@ -46,6 +48,7 @@ def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives ) ) ) + plan = plan.model_copy(update={"submission": plan.submission.model_copy(update={"memory": memory})}) destination = CollectionDestinationResolver().resolve(plan) collection = CollectionPlan( schema_version=1, @@ -77,6 +80,7 @@ def test_collection_renderer_uses_plan_bound_native_python_and_no_gpu_directives assert f"#SBATCH --job-name={collection.submission_job_name}" in script assert "#SBATCH --partition=cpu" in script assert "#SBATCH --partition=batch" not in script + assert ("#SBATCH --mem=16G" in script) == (memory is not None) assert f"dd-collect-{plan.run_id}" not in script assert ( 'readonly DD_COLLECTION_PLAN="/workspace/primary/runs/run-001/collections/collection-0001/plan.json"' in script diff --git a/plans/1079/slurm-external-providers.md b/plans/1079/slurm-external-providers.md index bc4a9b9c3..b3d4bfacb 100644 --- a/plans/1079/slurm-external-providers.md +++ b/plans/1079/slurm-external-providers.md @@ -22,6 +22,8 @@ qualification before considering separate generation and judging allocations. that the cluster otherwise reserves all node memory for a four-CPU job. - Let DD exclusively create a fresh dataset after the state writer has prepared its empty directory. Preserve nonempty directories and reject symlink targets. +- Restrict the generated dataset directory before publishing it as a candidate, + and validate native DD outputs with the state's real artifact verifier. ## Validation From d259e90907df58ffef9849d0af13562d8e37116d Mon Sep 17 00:00:00 2001 From: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:08:45 -0700 Subject: [PATCH 05/12] refactor(slurm): reuse native model-provider settings Signed-off-by: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> --- .../latest/pages/slurm/run-configuration.mdx | 35 +++--- .../examples/hosted-only.yaml | 6 - .../examples/mixed-providers.yaml | 6 - .../examples/model_providers.yaml | 7 ++ .../data_designer/slurm/client/execution.py | 29 ++--- .../data_designer/slurm/config/__init__.py | 2 - .../src/data_designer/slurm/config/builder.py | 3 - .../data_designer/slurm/config/environment.py | 11 +- .../src/data_designer/slurm/config/run.py | 68 ++---------- .../slurm/planning/builder_identity.py | 10 +- .../data_designer/slurm/planning/compiler.py | 1 + .../slurm/planning/model_providers.py | 89 +++++++++++++++ .../data_designer/slurm/planning/models.py | 22 +++- .../slurm/planning/resolution.py | 13 ++- .../slurm/planning/validation.py | 12 +- .../data_designer/slurm/runtime/bootstrap.py | 3 +- .../slurm/runtime/plan_reader.sh | 3 +- .../src/data_designer/slurm/runtime/steps.py | 7 +- .../slurm/services/environment.py | 55 ++++++++++ .../slurm/services/retry_collection.py | 9 +- .../data_designer/slurm/services/wiring.py | 55 +++------- .../client/test_external_model_execution.py | 26 ++--- .../data-designer-slurm/tests/conftest.py | 29 ++--- .../tests/planning/test_external_providers.py | 103 +++++++++++++----- .../tests/runtime/test_bootstrap.py | 7 +- .../tests/runtime/test_shell_runtime.py | 12 +- .../tests/services/test_wiring.py | 99 ++++++++++++++++- 27 files changed, 479 insertions(+), 243 deletions(-) create mode 100644 packages/data-designer-slurm/examples/model_providers.yaml create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/planning/model_providers.py create mode 100644 packages/data-designer-slurm/src/data_designer/slurm/services/environment.py diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index bba44901d..e956fd808 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -36,7 +36,7 @@ submission: `builder` accepts exactly one complete serialized Data Designer config through `source` or `inline`. Relative source paths resolve from the run file directory. `invocation.run_config` accepts public `RunConfig` fields except `display_tui`, which is disabled for Slurm jobs. `input_bindings.seed_path` and `managed_assets_path` accept absolute compute-visible paths. -Every builder model alias must have either one deployment or a provider declared in `invocation.model_providers`. These assignments must not overlap. `model_concurrency` accepts any declared alias, including external models. Deployment aliases must be unique. `model` selects the Hugging Face model ID or absolute model path loaded by the serving container. `served_model_name` is the name exposed to the client and defaults to `model`; set it explicitly when `model` is a container-visible path but the builder uses a stable model name. Shard and collected output files are Parquet. Choose the destination and partition count with `slurm merge` after the run completes. +Each builder model alias listed in `deployments` is served locally. Every remaining alias must reference a provider in DD's normal `model_providers.yaml` settings. `model_concurrency` accepts any declared alias, including external models. Deployment aliases must be unique. `model` selects the Hugging Face model ID or absolute model path loaded by the serving container. `served_model_name` is the name exposed to the client and defaults to `model`; set it explicitly when `model` is a container-visible path but the builder uses a stable model name. Shard and collected output files are Parquet. Choose the destination and partition count with `slurm merge` after the run completes. The client runs from the same shared Python environment as the CLI on the first allocated node. Its step requests no GPUs; local serving steps use the GPUs and their registered images. The `client` block is optional; `client.cpus` defaults to 16. It sets the client step's CPU request and contributes to the batch allocation's CPU request. @@ -51,7 +51,17 @@ Slurm schedules the job in one of the selected partitions. For runs with local d ## External model providers -Use `invocation.model_providers` for endpoints managed outside this allocation. The builder keeps its normal model declarations: +Configure hosted endpoints through DD's normal `model_providers.yaml`, outside the Slurm run file: + +```yaml +providers: + - name: hosted + endpoint: https://models.example/v1 + provider_type: openai + api_key: MODEL_API_KEY +``` + +Export `MODEL_API_KEY` in the submit environment. Keep its value out of configuration files. The recipe references the provider through its normal `ModelConfig`: ```yaml data_designer: @@ -65,33 +75,26 @@ data_designer: # Keep the recipe's columns and inference parameters here. ``` -The run file declares the external provider and deploys only `generator`: +The Slurm run deploys only `generator`. There is no separate hosted-provider block in the run file: ```yaml invocation: num_records: 1000 dataset_name: judged-output - model_providers: - - name: hosted - endpoint: https://models.example/v1 - provider_type: openai - api_key: - type: secret - environment: MODEL_API_KEY model_concurrency: generator: 64 judge: 16 ``` -The provider name must match the builder model's `provider`. Slurm preserves external model names and inference parameters; an explicit `model_concurrency` value overrides that alias's request concurrency. Provider names must not collide with generated local names such as `slurm-generator`. +Submission captures only the required hosted providers in the resolved run inputs, using DD's native provider representation. Later edits to submit-side settings do not change that run, and compute nodes do not need access to the submit host's settings directory. Missing or ambiguous providers and collisions with generated names such as `slurm-generator` fail before submission. Only locally deployed aliases receive generated endpoints; hosted model names and inference settings stay intact. An explicit `model_concurrency` value overrides the selected alias's request concurrency. -`api_key` is an optional `SecretRef`. Set the referenced variable in the submit environment; Slurm forwards it only to native client steps. Provider `extra_body` accepts non-secret request parameters. `extra_headers` maps header names to `type: literal` / `value` or `type: secret` / `environment` bindings. Authentication headers require secret bindings. Credential values are never written into the run configuration or resolved plan. +API keys must reference environment variables. Slurm forwards their values only to native client steps through runtime credential transport, including retries. Values stay out of persisted plans and generated scripts. Plaintext API keys, sensitive headers, and custom secret resolvers are not supported for hosted Slurm providers. Non-secret `extra_body` parameters and static headers use the normal DD provider fields. Both local and external requests use DD's existing async engine. Ready requests for different records and models can overlap. Async overlap does not guarantee continuous GPU utilization: slow external calls and the final judging tail can still retain an idle GPU allocation. Measure utilization against the cluster's policy when deciding whether to use separate recipe stages. ### Hosted-only CPU runs -For a builder whose models all use declared external providers, omit local deployments or set `deployments: []`. Select the cluster's CPU partition: +For a builder whose models all reference configured hosted providers, omit local deployments or set `deployments: []`. Select the cluster's CPU partition: ```yaml schema_version: 1 @@ -101,12 +104,6 @@ builder: invocation: num_records: 1000 dataset_name: hosted-judging - model_providers: - - name: hosted - endpoint: https://models.example/v1 - api_key: - type: secret - environment: MODEL_API_KEY client: cpus: 4 deployments: [] diff --git a/packages/data-designer-slurm/examples/hosted-only.yaml b/packages/data-designer-slurm/examples/hosted-only.yaml index 5586b8026..fb7b87f92 100644 --- a/packages/data-designer-slurm/examples/hosted-only.yaml +++ b/packages/data-designer-slurm/examples/hosted-only.yaml @@ -15,12 +15,6 @@ builder: invocation: num_records: 100 dataset_name: hosted-judging - model_providers: - - name: hosted - endpoint: https://models.example/v1 - api_key: - type: secret - environment: MODEL_API_KEY model_concurrency: judge: 16 client: diff --git a/packages/data-designer-slurm/examples/mixed-providers.yaml b/packages/data-designer-slurm/examples/mixed-providers.yaml index 7e7e40c91..6d2b7a11d 100644 --- a/packages/data-designer-slurm/examples/mixed-providers.yaml +++ b/packages/data-designer-slurm/examples/mixed-providers.yaml @@ -22,12 +22,6 @@ builder: invocation: num_records: 100 dataset_name: judged-descriptions - model_providers: - - name: hosted - endpoint: https://models.example/v1 - api_key: - type: secret - environment: MODEL_API_KEY model_concurrency: generator: 32 judge: 16 diff --git a/packages/data-designer-slurm/examples/model_providers.yaml b/packages/data-designer-slurm/examples/model_providers.yaml new file mode 100644 index 000000000..88526bb0d --- /dev/null +++ b/packages/data-designer-slurm/examples/model_providers.yaml @@ -0,0 +1,7 @@ +# Copy this provider into your normal DATA_DESIGNER_HOME/model_providers.yaml. +# Export MODEL_API_KEY in the submit environment; do not store its value here. +providers: + - name: hosted + endpoint: https://models.example/v1 + provider_type: openai + api_key: MODEL_API_KEY diff --git a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py index e25abc69b..1c7f3d13b 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/client/execution.py @@ -406,27 +406,20 @@ def _materialize_model_endpoints( configs = {config.alias: config for config in builder.model_configs} if set(configs) != set(plan.builder.model_aliases): raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "builder model aliases do not match the plan") - external_providers = plan.invocation.authored.model_providers + external_providers = tuple( + ModelProvider.model_validate_json(json.dumps(snapshot)) for snapshot in plan.hosted_model_providers + ) external_names = {provider.name for provider in external_providers} - external_aliases = tuple(config.alias for config in configs.values() if config.provider in external_names) - if external_aliases != plan.builder.external_model_aliases: + external_aliases = tuple(config.alias for config in configs.values() if config.alias not in aliases) + if external_aliases != plan.builder.external_model_aliases or any( + configs[alias].provider not in external_names for alias in external_aliases + ): raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "external model aliases do not match the plan") updated = dict(configs) - providers = [ - ModelProvider( - name=provider.name, - endpoint=provider.endpoint, - provider_type=provider.provider_type, - api_key=_resolve_secret(provider.api_key) if provider.api_key is not None else None, - extra_body=provider.extra_body, - extra_headers={ - name: binding.value if isinstance(binding, LiteralEnvironmentBinding) else _resolve_secret(binding) - for name, binding in provider.extra_headers.items() - } - or None, - ) - for provider in external_providers - ] + providers = [provider.model_copy(deep=True) for provider in external_providers] + for provider in providers: + if provider.api_key is not None and not os.environ.get(provider.api_key): + raise ClientWorkerError(ClientErrorCode.CONFIG_INVALID, "required client secret is unavailable") for deployment, port in zip(plan.deployments, plan.client.ports, strict=True): alias = deployment.authored.model_alias endpoint = endpoints[alias] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py index f87309bb9..9842577f0 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/__init__.py @@ -60,7 +60,6 @@ DataDesignerSlurmConfig, DeploymentResources, DeploymentTopology, - ExternalModelProviderConfig, InputBindings, InvocationConfig, LocalStdioMCPProviderConfig, @@ -88,7 +87,6 @@ "DEFAULT_PROFILE_FILE_NAME", "DeploymentResources", "DeploymentTopology", - "ExternalModelProviderConfig", "FixedRecordPolicy", "GpuRequestMode", "ImageBuildProfile", diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py index 123ed53bb..fccb674b1 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py @@ -21,7 +21,6 @@ ClientConfig, ClientDependencies, DataDesignerSlurmConfig, - ExternalModelProviderConfig, InputBindings, InvocationConfig, MCPProviderConfig, @@ -73,7 +72,6 @@ def with_invocation( run_config: Mapping[str, object] | None = None, input_bindings: InputBindings | Mapping[str, object] | None = None, mcp_providers: Sequence[MCPProviderConfig | Mapping[str, object]] = (), - model_providers: Sequence[ExternalModelProviderConfig | Mapping[str, object]] = (), model_concurrency: Mapping[str, int] | None = None, ) -> DataDesignerSlurmConfigBuilder: """Set typed Data Designer invocation intent.""" @@ -86,7 +84,6 @@ def with_invocation( "run_config": {} if run_config is None else dict(run_config), "input_bindings": {} if input_bindings is None else input_bindings, "mcp_providers": list(mcp_providers), - "model_providers": list(model_providers), "model_concurrency": {} if model_concurrency is None else dict(model_concurrency), }, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py index 006effec8..3e2a9d24f 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/environment.py @@ -11,6 +11,7 @@ from pydantic import BaseModel, Field, JsonValue, StringConstraints, field_validator +from data_designer.config.models import ModelProvider from data_designer.slurm.contracts import AuthoredConfig, validate_plain_text from data_designer.slurm.types import EnvironmentName @@ -125,9 +126,17 @@ def _secret_name_segments(value: str) -> list[str]: def _collect_secret_environment_names(value: object, names: set[str]) -> None: if isinstance(value, SecretRef): names.add(value.environment) + elif isinstance(value, ModelProvider): + if value.api_key is not None: + names.add(value.api_key) elif isinstance(value, BaseModel): for field_name in type(value).model_fields: - _collect_secret_environment_names(getattr(value, field_name), names) + child = getattr(value, field_name) + if field_name == "hosted_model_providers": + for snapshot in child: + _collect_secret_environment_names(ModelProvider.model_validate(snapshot), names) + else: + _collect_secret_environment_names(child, names) elif isinstance(value, Mapping): for child in value.values(): _collect_secret_environment_names(child, names) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py index f018d29ec..05147f76c 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/config/run.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/config/run.py @@ -50,7 +50,6 @@ "DeploymentResources", "DeploymentTopology", "EnvironmentBinding", - "ExternalModelProviderConfig", "InputBindings", "InvocationConfig", "LiteralEnvironmentBinding", @@ -107,53 +106,6 @@ def validate_paths(cls, value: str | None) -> str | None: return None if value is None else validate_absolute_path(value) -class ExternalModelProviderConfig(AuthoredConfig): - """A model endpoint managed outside this Slurm allocation. - - Attributes: - name: Provider name referenced by the builder's model configurations. - endpoint: Provider API URL without embedded credentials. - provider_type: Native Data Designer client type, such as ``openai``. - api_key: Optional environment secret reference for authentication. - extra_body: Additional non-secret request parameters. - extra_headers: Header values expressed as literal or secret bindings. - """ - - name: Identifier - endpoint: str - provider_type: Identifier = "openai" - api_key: SecretRef | None = None - extra_body: dict[str, JsonValue] | None = None - extra_headers: dict[str, EnvironmentBinding] = Field(default_factory=dict) - - @field_validator("endpoint") - @classmethod - def validate_endpoint(cls, value: str) -> str: - validate_url(value, field_name="model provider endpoint") - parsed = urlsplit(value) - if parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment: - raise ValueError("model provider endpoint must not embed credentials, query parameters, or fragments") - return value - - @field_validator("extra_body") - @classmethod - def validate_extra_body(cls, value: dict[str, JsonValue] | None) -> dict[str, JsonValue] | None: - validate_no_plaintext_secrets(value, field_name="model provider extra_body") - return value - - @field_validator("extra_headers") - @classmethod - def validate_headers(cls, values: dict[str, EnvironmentBinding]) -> dict[str, EnvironmentBinding]: - validate_environment_bindings(values) - for name, binding in values.items(): - if re.fullmatch(r"[!#$%&'*+.^_`|~0-9A-Za-z-]+", name) is None: - raise ValueError("model provider headers require valid HTTP field names") - if isinstance(binding, LiteralEnvironmentBinding): - if name.casefold() in {"authorization", "proxy-authorization", "cookie"}: - raise ValueError("authentication headers require external secret references") - return values - - class RemoteMCPProviderConfig(AuthoredConfig): provider_type: Literal["sse", "streamable_http"] name: Identifier @@ -217,7 +169,6 @@ class InvocationConfig(AuthoredConfig): run_config: dict[str, JsonValue] = Field(default_factory=dict) input_bindings: InputBindings = Field(default_factory=InputBindings) mcp_providers: list[MCPProviderConfig] = Field(default_factory=list) - model_providers: list[ExternalModelProviderConfig] = Field(default_factory=list, exclude_if=lambda value: not value) model_concurrency: dict[str, PositiveInt] = Field(default_factory=dict) @field_validator("run_config") @@ -234,9 +185,6 @@ def validate_mcp_providers(self) -> InvocationConfig: names = [provider.name for provider in self.mcp_providers] if len(names) != len(set(names)): raise ValueError("MCP provider names must be unique") - model_names = [provider.name for provider in self.model_providers] - if len(model_names) != len(set(model_names)): - raise ValueError("model provider names must be unique") return self @@ -425,9 +373,6 @@ def validate_run(self) -> DataDesignerSlurmConfig: aliases = [deployment.model_alias for deployment in self.deployments] if len(aliases) != len(set(aliases)): raise ValueError("deployment model aliases must be unique") - provider_names = {provider.name for provider in self.invocation.model_providers} - if provider_names.intersection(f"slurm-{alias}" for alias in aliases): - raise ValueError("external provider names must not collide with managed providers") groups: dict[str, list[ServerDeploymentConfig]] = {} for deployment in self.deployments: if deployment.resources.shared_node is not None: @@ -446,9 +391,12 @@ def validate_run(self) -> DataDesignerSlurmConfig: for member in members[1:] ): raise ValueError("shared_node requires matching server environment and queue backpressure") - unknown_concurrency = set(self.invocation.model_concurrency).difference(aliases) - if unknown_concurrency and not self.invocation.model_providers: - raise ValueError( - f"model concurrency references undeclared aliases: {', '.join(sorted(unknown_concurrency))}" - ) + if self.builder.inline is not None: + payload = self.builder.inline.get("data_designer", self.builder.inline) + assert isinstance(payload, dict) + models = payload.get("model_configs") or [] + if isinstance(models, list): + declared = set(aliases) | {model.get("alias") for model in models if isinstance(model, dict)} + if set(self.invocation.model_concurrency).difference(declared): + raise ValueError("model concurrency references undeclared aliases") return self diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py index 11e8ef890..bab1c6b83 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/builder_identity.py @@ -36,17 +36,13 @@ def get_persisted_builder_identity( return get_declared_model_aliases(builder), compute_serialized_json_sha256(builder) -def get_external_model_aliases( - builder: dict[str, JsonValue], provider_names: Collection[str] -) -> tuple[ModelAlias, ...]: - """Return explicitly external aliases in their authored builder order.""" +def get_external_model_aliases(builder: dict[str, JsonValue], local_aliases: Collection[str]) -> tuple[ModelAlias, ...]: + """Return aliases not deployed locally in their authored builder order.""" get_declared_model_aliases(builder) data_designer = builder.get("data_designer", builder) assert isinstance(data_designer, dict) configs = data_designer.get("model_configs") or [] assert isinstance(configs, list) return tuple( - config["alias"] - for config in configs - if isinstance(config, dict) and isinstance(config.get("provider"), str) and config["provider"] in provider_names + config["alias"] for config in configs if isinstance(config, dict) and config["alias"] not in local_aliases ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py index b72c04996..26785deb7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/compiler.py @@ -77,6 +77,7 @@ def compile(effective: EffectiveDataDesignerSlurmConfig) -> ResolvedSlurmRunPlan resolved_gpus_per_node=effective.resolved_gpus_per_node, builder=effective.builder, invocation=effective.invocation, + hosted_model_providers=effective.hosted_model_providers, client=client, deployments=deployments, array_tasks=effective.authored.array_tasks, diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/model_providers.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/model_providers.py new file mode 100644 index 000000000..83ffa584f --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/model_providers.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resolve native DD provider settings into non-secret, immutable run inputs.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Collection, Mapping, Sequence +from urllib.parse import urlsplit + +from pydantic import JsonValue + +from data_designer.config.models import ModelProvider +from data_designer.slurm.config.environment import is_secret_bearing_name, validate_no_plaintext_secrets +from data_designer.slurm.contracts import validate_plain_text, validate_url +from data_designer.slurm.planning.builder_identity import get_declared_model_aliases + + +def validate_hosted_model_providers(providers: Sequence[ModelProvider]) -> tuple[ModelProvider, ...]: + """Preserve DD settings while admitting only environment-backed credentials.""" + names = [provider.name for provider in providers] + if len(names) != len(set(names)): + raise ValueError("hosted model provider names must be unique") + for provider in providers: + if not provider.name: + raise ValueError("hosted model providers require a name") + validate_url(provider.endpoint, field_name="hosted model provider endpoint") + parsed = urlsplit(provider.endpoint) + if parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment: + raise ValueError("hosted model provider endpoints must not contain credentials, queries, or fragments") + if provider.api_key is not None and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", provider.api_key) is None: + raise ValueError("hosted model provider API keys must reference environment variables") + validate_no_plaintext_secrets(provider.extra_body, field_name="hosted model provider extra_body") + for name, value in (provider.extra_headers or {}).items(): + if re.fullmatch(r"[!#$%&'*+.^_`|~0-9A-Za-z-]+", name) is None: + raise ValueError("hosted model provider headers require valid HTTP field names") + if is_secret_bearing_name(name) or name.casefold() in {"authorization", "proxy-authorization", "cookie"}: + raise ValueError("hosted model providers do not support sensitive headers") + validate_plain_text(value, field_name="hosted model provider header") + return tuple( + ModelProvider.model_validate_json(json.dumps(provider.model_dump(mode="json"), allow_nan=False)) + for provider in providers + ) + + +def resolve_hosted_model_providers( + builder: dict[str, JsonValue], + local_aliases: Collection[str], + providers: Sequence[ModelProvider], +) -> tuple[ModelProvider, ...]: + """Snapshot only the native providers required by aliases not deployed locally.""" + get_declared_model_aliases(builder) + envelope = builder.get("data_designer", builder) + assert isinstance(envelope, dict) + configs = envelope.get("model_configs") or [] + assert isinstance(configs, list) + required: list[str] = [] + for config in configs: + assert isinstance(config, dict) + if config["alias"] in local_aliases: + continue + name = config.get("provider") + if not isinstance(name, str) or not name: + raise ValueError("every hosted model must reference a native Data Designer provider") + if name not in required: + required.append(name) + captured: list[ModelProvider] = [] + for name in required: + matches = [provider for provider in providers if provider.name == name] + if len(matches) != 1: + raise ValueError("each hosted model must resolve to exactly one Data Designer provider") + if name in {f"slurm-{alias}" for alias in local_aliases}: + raise ValueError("hosted provider names must not collide with managed providers") + captured.append(matches[0]) + return validate_hosted_model_providers(captured) + + +def load_hosted_model_providers(snapshots: Sequence[Mapping[str, JsonValue]]) -> tuple[ModelProvider, ...]: + """Restore DD's provider objects from their captured native JSON representation.""" + return validate_hosted_model_providers( + tuple(ModelProvider.model_validate_json(json.dumps(value, allow_nan=False)) for value in snapshots) + ) + + +def dump_hosted_model_providers(providers: Sequence[ModelProvider]) -> tuple[dict[str, JsonValue], ...]: + """Serialize native settings for freezing by the Slurm run contract.""" + return tuple(provider.model_dump(mode="json") for provider in validate_hosted_model_providers(providers)) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py index f3c3efdd7..d0b1ec21d 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -48,6 +48,11 @@ validate_plain_text, ) from data_designer.slurm.planning.builder_identity import get_declared_model_aliases, get_external_model_aliases +from data_designer.slurm.planning.model_providers import ( + dump_hosted_model_providers, + load_hosted_model_providers, + resolve_hosted_model_providers, +) from data_designer.slurm.types import NetworkPort, PartitionSelection @@ -410,6 +415,12 @@ class ResolvedSlurmRunPlan(ContractRecord): submission: ResolvedSubmission container_mounts: tuple[ContainerMount, ...] = () runtime_bundle: ArtifactReference + hosted_model_providers: tuple[dict[str, JsonValue], ...] = Field(default=(), exclude_if=lambda value: not value) + + @field_validator("hosted_model_providers", mode="before") + @classmethod + def validate_hosted_providers(cls, value: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return dump_hosted_model_providers(load_hosted_model_providers(value)) @field_validator("scheduler_bin_path") @classmethod @@ -480,15 +491,20 @@ def validate_plan(self) -> ResolvedSlurmRunPlan: external = set(self.builder.external_model_aliases) if external.intersection(aliases) or set(aliases) | external != set(self.builder.model_aliases): raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") - provider_names = {provider.name for provider in self.invocation.authored.model_providers} + providers = load_hosted_model_providers(self.hosted_model_providers) + provider_names = {provider.name for provider in providers} if external and not provider_names: - raise ValueError("external model aliases require external providers") + raise ValueError("resolved deployment aliases must exactly cover Data Designer model aliases") if provider_names.intersection(f"slurm-{alias}" for alias in aliases): raise ValueError("external provider names must not collide with managed providers") if self.builder.inline is not None and self.builder.external_model_aliases != get_external_model_aliases( - self.builder.inline, provider_names + self.builder.inline, aliases ): raise ValueError("resolved external model aliases do not match the inline builder") + if self.builder.inline is not None and providers != resolve_hosted_model_providers( + self.builder.inline, aliases, providers + ): + raise ValueError("hosted model provider snapshot does not match the builder") if set(self.invocation.authored.model_concurrency).difference(self.builder.model_aliases): raise ValueError("model concurrency references undeclared aliases") diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py index 6930c5611..6031c1050 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py @@ -11,6 +11,7 @@ from pydantic import JsonValue, NonNegativeInt, StringConstraints, TypeAdapter, ValidationError from data_designer.config import RunConfig +from data_designer.config.models import ModelProvider from data_designer.slurm._errors import format_validation_error from data_designer.slurm.config.images import ImageKind from data_designer.slurm.config.profiles import SelectedSlurmProfile @@ -30,6 +31,9 @@ get_persisted_builder_identity, ) from data_designer.slurm.planning.errors import SlurmConfigResolutionError +from data_designer.slurm.planning.model_providers import ( + dump_hosted_model_providers, +) from data_designer.slurm.planning.models import ( ResolvedBuilderInput, ResolvedClientRuntime, @@ -89,6 +93,7 @@ class EffectiveDataDesignerSlurmConfig(ContractValue): resolved_gpus_per_node: NonNegativeInt builder: ResolvedBuilderInput builder_payload: dict[str, JsonValue] | None = None + hosted_model_providers: tuple[dict[str, JsonValue], ...] = () invocation: ResolvedInvocation client_runtime: ResolvedClientRuntime deployment_images: tuple[ResolvedImage, ...] @@ -110,6 +115,7 @@ def resolve_slurm_config( scheduler_bin_path: str | None = None, resolved_gpus_per_node: int | None = None, builder_payload: dict[str, JsonValue] | None = None, + hosted_model_providers: tuple[ModelProvider, ...] = (), ) -> EffectiveDataDesignerSlurmConfig: """Materialize all non-secret defaults from explicitly supplied resolved inputs.""" try: @@ -137,6 +143,7 @@ def resolve_slurm_config( resolved_gpus_per_node=gpus_per_node, builder=builder, builder_payload=resolved_builder_payload, + hosted_model_providers=dump_hosted_model_providers(hosted_model_providers), invocation=_materialize_invocation(authored, workspace_root), client_runtime=client_runtime, deployment_images=deployment_images, @@ -238,7 +245,7 @@ def _resolve_builder( run_root: str, builder_payload: dict[str, JsonValue] | None, ) -> tuple[ResolvedBuilderInput, dict[str, JsonValue] | None]: - provider_names = {provider.name for provider in authored.invocation.model_providers} + local_aliases = {deployment.model_alias for deployment in authored.deployments} if authored.builder.inline is not None: if builder_payload is not None: raise SlurmConfigResolutionError("inline builder input must not provide a separate payload") @@ -247,7 +254,7 @@ def _resolve_builder( inline=authored.builder.inline, content_sha256=compute_serialized_json_sha256(authored.builder.inline), model_aliases=get_declared_model_aliases(authored.builder.inline), - external_model_aliases=get_external_model_aliases(authored.builder.inline, provider_names), + external_model_aliases=get_external_model_aliases(authored.builder.inline, local_aliases), ), None, ) @@ -266,7 +273,7 @@ def _resolve_builder( source=source, content_sha256=source.sha256, model_aliases=aliases, - external_model_aliases=get_external_model_aliases(validated_payload, provider_names), + external_model_aliases=get_external_model_aliases(validated_payload, local_aliases), ), validated_payload, ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py index 61c93467a..84aecad9e 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/validation.py @@ -12,6 +12,7 @@ from data_designer.slurm.contracts import derive_managed_assets_path from data_designer.slurm.planning.builder_identity import get_external_model_aliases, get_persisted_builder_identity from data_designer.slurm.planning.errors import SlurmPlanContractError +from data_designer.slurm.planning.model_providers import load_hosted_model_providers, resolve_hosted_model_providers from data_designer.slurm.planning.models import ( ResolvedDependencyLock, ResolvedSlurmRunPlan, @@ -79,9 +80,18 @@ def validate_resolved_plan( assert payload is not None _require( plan.builder.external_model_aliases - == get_external_model_aliases(payload, {provider.name for provider in authored.invocation.model_providers}), + == get_external_model_aliases(payload, {deployment.model_alias for deployment in authored.deployments}), "resolved external aliases do not match the builder", ) + _require( + load_hosted_model_providers(plan.hosted_model_providers) + == resolve_hosted_model_providers( + payload, + {deployment.model_alias for deployment in authored.deployments}, + load_hosted_model_providers(plan.hosted_model_providers), + ), + "hosted model provider snapshot does not match the builder", + ) expected_account = authored.submission.account or plan.selected_profile.profile.scheduler.account expected_partition = authored.submission.partition or plan.selected_profile.profile.scheduler.partition diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index b1f93d7cc..4f5d91fc4 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py @@ -10,6 +10,7 @@ from pathlib import Path from typing import Literal +from data_designer.config.models import ModelProvider from data_designer.slurm.config.environment import ( SecretRef, collect_secret_environment_names, @@ -158,7 +159,7 @@ def _build_client_step( ( plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers, - plan.invocation.authored.model_providers, + *(ModelProvider.model_validate(value) for value in plan.hosted_model_providers), ) ) allocation_environment = ( diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh index 47c642867..683399856 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh @@ -66,12 +66,13 @@ dd_read_container_path() { dd_read_plan_secret_names() { local plan=$1 dd_read_null_values DD_REQUIRED_SECRET_NAMES < <( - jq -j '[.. | objects | select(.type? == "secret") | .environment] | unique | .[] | ., "\u0000"' "${plan}" + jq -j '[(.. | objects | select(.type? == "secret") | .environment), (.hosted_model_providers[]?.api_key // empty)] | unique | .[] | ., "\u0000"' "${plan}" ) dd_read_null_values DD_ALL_SECRET_NAMES < <( jq -j ' [ (.. | objects | select(.type? == "secret") | .environment), + (.hosted_model_providers[]?.api_key // empty), ( .deployments[].authored.server.environment | to_entries[] diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py index b1b8f8a3c..a95176f96 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/steps.py @@ -10,6 +10,7 @@ from pathlib import Path from typing import Literal, Protocol +from data_designer.config.models import ModelProvider from data_designer.slurm.config.environment import ( LiteralEnvironmentBinding, SecretRef, @@ -561,7 +562,11 @@ def _base_environment(source_environment: Mapping[str, str]) -> dict[str, str]: def _collect_client_secret_environment_names(plan: ResolvedSlurmRunPlan) -> tuple[str, ...]: return collect_secret_environment_names( - (plan.client.authored.dependencies.index_credentials, plan.invocation.authored.mcp_providers) + ( + plan.client.authored.dependencies.index_credentials, + plan.invocation.authored.mcp_providers, + *(ModelProvider.model_validate(value) for value in plan.hosted_model_providers), + ) ) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/environment.py b/packages/data-designer-slurm/src/data_designer/slurm/services/environment.py new file mode 100644 index 000000000..ad6ebf329 --- /dev/null +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/environment.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime-only credential transport shared by initial submissions and retries.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from data_designer.slurm.config import collect_secret_environment_names +from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation + + +def build_run_export_environment( + plan: ResolvedSlurmRunPlan, + source_environment: Mapping[str, str], + operation: SlurmServiceOperation, +) -> dict[str, str]: + """Forward only required credentials without serializing their values.""" + environment = {"SLURM_EXPORT_ENV": "ALL"} + deployment_type = source_environment.get("NEMO_DEPLOYMENT_TYPE", "slurm") + if type(deployment_type) is not str or "\0" in deployment_type: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "NEMO_DEPLOYMENT_TYPE is invalid", + ) + environment["NEMO_DEPLOYMENT_TYPE"] = deployment_type + if "SLURM_CONF" in source_environment: + slurm_conf = source_environment["SLURM_CONF"] + if type(slurm_conf) is not str or "\0" in slurm_conf: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "SLURM_CONF is invalid", + ) + environment["SLURM_CONF"] = slurm_conf + for name in collect_secret_environment_names(plan): + try: + value = source_environment[name] + except KeyError: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + f"required secret environment {name!r} is unavailable", + ) from None + if type(value) is not str or "\0" in value: + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + f"required secret environment {name!r} is invalid", + ) + environment[name] = value + return environment diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py index 51d01b98d..f7ca19aeb 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py @@ -16,6 +16,7 @@ from data_designer.slurm.launcher.client import SlurmCommandClient from data_designer.slurm.launcher.errors import SlurmLauncherError from data_designer.slurm.planning import ResolvedSlurmRunPlan +from data_designer.slurm.services.environment import build_run_export_environment from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation from data_designer.slurm.services.results import SlurmCollectionExecution, SlurmRetryExecution from data_designer.slurm.state import ( @@ -48,6 +49,7 @@ def __init__( self._workspace_root = workspace_root self._launcher = launcher self._clock = clock + self._source_environment = dict(source_environment) self._export_environment = {"NEMO_DEPLOYMENT_TYPE": source_environment.get("NEMO_DEPLOYMENT_TYPE", "slurm")} def retry( @@ -63,8 +65,13 @@ def retry( run_id = self._resolve_run_reference(run_or_job_id, operation) try: observed_at = self._clock() - coordinator = SlurmRetryCoordinator(self._workspace_root, run_id, self._launcher, self._export_environment) resolved_plan = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan() + export_environment = ( + self._export_environment + if dry_run + else build_run_export_environment(resolved_plan, self._source_environment, operation) + ) + coordinator = SlurmRetryCoordinator(self._workspace_root, run_id, self._launcher, export_environment) effective_resume_mode = self._resolve_retry_resume_mode(resolved_plan, resume) preview = coordinator.preview_active(shard_ids=shard_ids, observed_at=observed_at) if preview is not None and effective_resume_mode not in {None, preview[0].effective_resume_mode}: 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 988c811d0..ae064c497 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 @@ -18,6 +18,7 @@ from pydantic import JsonValue, ValidationError +from data_designer.config.default_model_settings import get_default_providers from data_designer.slurm._errors import format_validation_error from data_designer.slurm.benchmark.execution import SystemBenchmarkBackend from data_designer.slurm.benchmark.observer import PersistedBenchmarkRunObserver @@ -38,7 +39,6 @@ SlurmConfigLoadError, SlurmProfile, SlurmProfileCatalog, - collect_secret_environment_names, load_builder_payload, resolve_profile, ) @@ -57,13 +57,16 @@ from data_designer.slurm.launcher.renderer import render_generation_attempt_script from data_designer.slurm.launcher.scheduler_path import SchedulerPathError, resolve_scheduler_bin_path from data_designer.slurm.planning import ResolvedImage, ResolvedSlurmRunPlan +from data_designer.slurm.planning.builder_identity import get_external_model_aliases from data_designer.slurm.planning.compiler import SlurmRunCompiler +from data_designer.slurm.planning.model_providers import resolve_hosted_model_providers from data_designer.slurm.planning.resolution import resolve_slurm_config from data_designer.slurm.runtime.bundle import stage_runtime_bundle from data_designer.slurm.runtime.errors import SlurmRuntimeError from data_designer.slurm.security import redact_sensitive_diagnostic from data_designer.slurm.services.artifacts import StateRunArtifactPublisher from data_designer.slurm.services.benchmark import SlurmBenchmarkService +from data_designer.slurm.services.environment import build_run_export_environment from data_designer.slurm.services.errors import SlurmServiceError, SlurmServiceErrorCode, SlurmServiceOperation from data_designer.slurm.services.image_lifecycle import SlurmImageLifecycleManager from data_designer.slurm.services.images import SlurmImageService @@ -205,6 +208,14 @@ def _prepare( _safe_builder_load_error(error), ) from None _validate_builder_for_submission(authored, builder_payload, operation) + payload = authored.builder.inline if authored.builder.inline is not None else builder_payload + assert payload is not None + local_aliases = {deployment.model_alias for deployment in authored.deployments} + hosted_providers = resolve_hosted_model_providers( + payload, + local_aliases, + get_default_providers() if get_external_model_aliases(payload, local_aliases) else (), + ) dependencies = stack.enter_context( self._dependencies.resolve( authored.client.dependencies, @@ -238,6 +249,7 @@ def _prepare( scheduler_bin_path=scheduler_bin_path, resolved_gpus_per_node=self._resolve_gpu_count(authored, operation), builder_payload=builder_payload, + hosted_model_providers=hosted_providers, ) plan = SlurmRunCompiler.compile(effective) for deployment in plan.deployments: @@ -349,7 +361,9 @@ def execute( batch_script=prepared.batch_script, ) self._materialize_default_managed_assets(config, plan) - export_environment = self._build_export_environment(config) + export_environment = build_run_export_environment( + plan, self._source_environment, SlurmServiceOperation.EXECUTE_RUN + ) self._initialize_run(config, prepared) self._report_progress("Rendered batch script; invoking sbatch --parsable --hold (script via stdin).") try: @@ -493,43 +507,6 @@ def _materialize_default_managed_assets( "managed assets workspace cannot be prepared", ) from None - def _build_export_environment(self, config: DataDesignerSlurmConfig) -> dict[str, str]: - environment = {"SLURM_EXPORT_ENV": "ALL"} - deployment_type = self._source_environment.get("NEMO_DEPLOYMENT_TYPE", "slurm") - if type(deployment_type) is not str or "\0" in deployment_type: - raise SlurmServiceError( - SlurmServiceErrorCode.INVALID_REQUEST, - SlurmServiceOperation.EXECUTE_RUN, - "NEMO_DEPLOYMENT_TYPE is invalid", - ) - environment["NEMO_DEPLOYMENT_TYPE"] = deployment_type - if "SLURM_CONF" in self._source_environment: - slurm_conf = self._source_environment["SLURM_CONF"] - if type(slurm_conf) is not str or "\0" in slurm_conf: - raise SlurmServiceError( - SlurmServiceErrorCode.INVALID_REQUEST, - SlurmServiceOperation.EXECUTE_RUN, - "SLURM_CONF is invalid", - ) - environment["SLURM_CONF"] = slurm_conf - for name in collect_secret_environment_names(config): - try: - value = self._source_environment[name] - except KeyError: - raise SlurmServiceError( - SlurmServiceErrorCode.INVALID_REQUEST, - SlurmServiceOperation.EXECUTE_RUN, - f"required secret environment {name!r} is unavailable", - ) from None - if type(value) is not str or "\0" in value: - raise SlurmServiceError( - SlurmServiceErrorCode.INVALID_REQUEST, - SlurmServiceOperation.EXECUTE_RUN, - f"required secret environment {name!r} is invalid", - ) - environment[name] = value - return environment - def _initialize_run( self, config: DataDesignerSlurmConfig, diff --git a/packages/data-designer-slurm/tests/client/test_external_model_execution.py b/packages/data-designer-slurm/tests/client/test_external_model_execution.py index a5b27923c..dbf09e128 100644 --- a/packages/data-designer-slurm/tests/client/test_external_model_execution.py +++ b/packages/data-designer-slurm/tests/client/test_external_model_execution.py @@ -12,6 +12,7 @@ from pytest_httpx import HTTPXMock from data_designer.config import DataDesignerConfigBuilder +from data_designer.config.models import ModelProvider from data_designer.config.seed_source import LocalFileSeedSource from data_designer.slurm.client.errors import ClientWorkerError from data_designer.slurm.client.execution import ClientWorker @@ -26,6 +27,7 @@ @pytest.fixture def external_client_case( + hosted_model_providers: tuple[ModelProvider, ...], external_run: DataDesignerSlurmConfig, client_worker_case: ClientWorkerCase, ) -> ClientWorkerCase: @@ -33,6 +35,7 @@ def external_client_case( plan = SlurmRunCompiler.compile( resolve_slurm_config( external_run, + hosted_model_providers=hosted_model_providers, selected_profile=case.plan.selected_profile, client_runtime=case.plan.client.runtime, deployment_images=tuple(case.plan.deployments[0].image for _ in external_run.deployments), @@ -59,8 +62,7 @@ def test_worker_preserves_hosted_models_and_scopes_secrets( monkeypatch: pytest.MonkeyPatch, ) -> None: case = external_client_case - monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") - monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-secret") designers: list[FakeDataDesigner] = [] def factory(**kwargs: object) -> FakeDataDesigner: @@ -78,9 +80,9 @@ def factory(**kwargs: object) -> FakeDataDesigner: assert model.inference_parameters.max_parallel_requests == 3 providers = {provider.name: provider for provider in designers[0].initialization["model_providers"]} assert providers["external"].endpoint == "https://models.example/v1" - assert providers["external"].api_key == "fake-provider-secret" + assert providers["external"].api_key == "MODEL_API_KEY" assert providers["external"].extra_body == {"routing": {"version": 2}} - assert providers["external"].extra_headers == {"X-Project": "test", "X-Token": "fake-header-secret"} + assert providers["external"].extra_headers == {"X-Project": "test"} if case.plan.deployments: assert providers["slurm-generator"].endpoint == case.endpoints["generator"] else: @@ -93,14 +95,13 @@ def factory(**kwargs: object) -> FakeDataDesigner: assert "fake-header-secret" not in content -@pytest.mark.parametrize("missing", ("EXTERNAL_API_KEY", "EXTERNAL_HEADER_TOKEN")) +@pytest.mark.parametrize("missing", ("MODEL_API_KEY",)) def test_worker_rejects_unavailable_hosted_credentials( external_client_case: ClientWorkerCase, monkeypatch: pytest.MonkeyPatch, missing: str, ) -> None: - monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") - monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-secret") monkeypatch.delenv(missing) case = external_client_case with pytest.raises(ClientWorkerError, match="required client secret is unavailable"): @@ -119,8 +120,7 @@ def test_hosted_only_worker_needs_no_scheduler_gpu_environment( ) -> None: monkeypatch.delenv("SLURM_JOB_GPUS", raising=False) monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") - monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-secret") case = external_client_case worker = ClientWorker(data_designer_factory=FakeDataDesigner, environment={}) worker.preflight(case.plan_path, prepared=case.prepared, endpoints={}, plugins=()) @@ -140,8 +140,7 @@ def test_worker_preserves_existing_dataset_paths( existing_kind: str, ) -> None: case = external_client_case - monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") - monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-secret") worker = ClientWorker(data_designer_factory=FakeDataDesigner) worker.preflight(case.plan_path, prepared=case.prepared, endpoints=case.endpoints, plugins=()) dataset = case.attempt_dir / "dataset" @@ -183,8 +182,7 @@ def test_native_dd_generates_through_managed_and_external_endpoints( ) -> None: case = external_client_case monkeypatch.setattr("data_designer.engine.models.telemetry.TELEMETRY_ENABLED", False) - monkeypatch.setenv("EXTERNAL_API_KEY", "fake-provider-secret") - monkeypatch.setenv("EXTERNAL_HEADER_TOKEN", "fake-header-secret") + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-secret") payload = case.plan.model_dump(mode="python") payload["builder"]["inline"]["data_designer"]["columns"] = [ { @@ -243,7 +241,7 @@ def test_native_dd_generates_through_managed_and_external_endpoints( requests = httpx_mock.get_requests(url=hosted_url) assert requests assert all(request.headers["Authorization"] == "Bearer fake-provider-secret" for request in requests) - assert all(request.headers["X-Token"] == "fake-header-secret" for request in requests) + assert all(request.headers["X-Project"] == "test" for request in requests) assert all(json.loads(request.content)["model"] == "example/hosted-judge" for request in requests) for endpoint in case.endpoints.values(): local_requests = httpx_mock.get_requests(url=endpoint + "/chat/completions") diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index 9b928c888..8f6ab75f3 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -19,6 +19,7 @@ FakeVllmBackend, ) +from data_designer.config.models import ModelProvider from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport from data_designer.slurm.client import ClientResult from data_designer.slurm.config import ( @@ -103,19 +104,6 @@ def external_run( }, } ) - payload["invocation"]["model_providers"] = [ - { - "name": "external", - "endpoint": "https://models.example/v1", - "provider_type": "openai", - "api_key": {"type": "secret", "environment": "EXTERNAL_API_KEY"}, - "extra_body": {"routing": {"version": 2}}, - "extra_headers": { - "X-Project": {"type": "literal", "value": "test"}, - "X-Token": {"type": "secret", "environment": "EXTERNAL_HEADER_TOKEN"}, - }, - } - ] payload["invocation"]["model_concurrency"]["judge"] = 3 if request.param == "hosted-only": models.pop(0) @@ -124,8 +112,22 @@ def external_run( return DataDesignerSlurmConfig.model_validate(payload) +@pytest.fixture +def hosted_model_providers() -> tuple[ModelProvider, ...]: + return ( + ModelProvider( + name="external", + endpoint="https://models.example/v1", + api_key="MODEL_API_KEY", + extra_body={"routing": {"version": 2}}, + extra_headers={"X-Project": "test"}, + ), + ) + + @pytest.fixture def external_plan( + hosted_model_providers: tuple[ModelProvider, ...], external_run: DataDesignerSlurmConfig, single_node_plan: ResolvedSlurmRunPlan, dependency_lock_single: ResolvedDependencyLock, @@ -133,6 +135,7 @@ def external_plan( return SlurmRunCompiler.compile( resolve_slurm_config( external_run, + hosted_model_providers=hosted_model_providers, selected_profile=single_node_plan.selected_profile, client_runtime=single_node_plan.client.runtime, deployment_images=tuple(single_node_plan.deployments[0].image for _ in external_run.deployments), diff --git a/packages/data-designer-slurm/tests/planning/test_external_providers.py b/packages/data-designer-slurm/tests/planning/test_external_providers.py index d6e6b94d8..0835aff47 100644 --- a/packages/data-designer-slurm/tests/planning/test_external_providers.py +++ b/packages/data-designer-slurm/tests/planning/test_external_providers.py @@ -7,11 +7,11 @@ import pytest +from data_designer.config.models import ModelProvider from data_designer.slurm.config import ( BuilderInput, DataDesignerSlurmConfig, DataDesignerSlurmConfigBuilder, - ExternalModelProviderConfig, InvocationConfig, SubmissionConfig, injected_profile, @@ -20,7 +20,12 @@ from data_designer.slurm.launcher.renderer import render_generation_attempt_script from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan from data_designer.slurm.planning.compiler import SlurmRunCompiler -from data_designer.slurm.planning.errors import SlurmPlanCompilationError, SlurmPlanContractError +from data_designer.slurm.planning.errors import SlurmPlanContractError +from data_designer.slurm.planning.model_providers import ( + load_hosted_model_providers, + resolve_hosted_model_providers, + validate_hosted_model_providers, +) from data_designer.slurm.planning.resolution import resolve_slurm_config from data_designer.slurm.planning.validation import validate_resolved_plan from data_designer.slurm.runtime.ports import resolve_allocation_plan @@ -30,7 +35,7 @@ def test_external_provider_examples_load(name: str) -> None: path = Path(__file__).parents[2] / "examples" / name config = load_run_config(path) - assert config.invocation.model_providers[0].name == "hosted" + assert "model_providers" not in config.invocation.model_dump() assert {deployment.model_alias for deployment in config.deployments} == ( {"generator"} if name == "mixed-providers.yaml" else set() ) @@ -54,6 +59,7 @@ def test_external_models_compile_and_round_trip( plan = SlurmRunCompiler.compile( resolve_slurm_config( config, + hosted_model_providers=load_hosted_model_providers(external_plan.hosted_model_providers), selected_profile=external_plan.selected_profile, client_runtime=external_plan.client.runtime, deployment_images=tuple(deployment.image for deployment in external_plan.deployments), @@ -90,12 +96,11 @@ def test_builder_accepts_external_providers_without_deployments(external_run: Da builder = DataDesignerSlurmConfigBuilder(external_run.builder).with_invocation( num_records=8, dataset_name="hosted", - model_providers=external_run.invocation.model_providers, ) builder.with_submission(memory="16G") config = builder.build() assert config.deployments == [] - assert config.invocation.model_providers == external_run.invocation.model_providers + assert "model_providers" not in config.invocation.model_dump() assert config.submission.memory == "16G" @@ -152,28 +157,28 @@ def test_submission_memory_is_bound_to_authored_config( ) def test_external_provider_rejects_credential_bearing_endpoints(endpoint: str) -> None: with pytest.raises(ValueError): - ExternalModelProviderConfig(name="external", endpoint=endpoint) + validate_hosted_model_providers((ModelProvider(name="external", endpoint=endpoint),)) @pytest.mark.parametrize( "values", ( - {"api_key": "plaintext"}, - {"extra_body": {"api_key": "plaintext"}}, - {"extra_headers": {"Authorization": {"type": "literal", "value": "Bearer plaintext"}}}, - {"extra_headers": {"X-Api-Key": {"type": "literal", "value": "plaintext"}}}, - {"extra_headers": {"X-Bad\r\nHeader": {"type": "literal", "value": "test"}}}, + {"api_key": "sk-plaintext-key"}, + {"extra_body": {"api_key": "sk-plaintext-key"}}, + {"extra_headers": {"Authorization": "Bearer plaintext"}}, + {"extra_headers": {"X-Api-Key": "plaintext"}}, + {"extra_headers": {"X-Bad\r\nHeader": "test"}}, ), ) def test_external_provider_rejects_plaintext_credentials_and_invalid_headers(values: dict[str, object]) -> None: with pytest.raises(ValueError): - ExternalModelProviderConfig.model_validate( - {"name": "external", "endpoint": "https://models.example/v1", **values} + validate_hosted_model_providers( + (ModelProvider.model_validate({"name": "external", "endpoint": "https://models.example/v1", **values}),) ) def test_external_provider_retains_extension_parameters() -> None: - config = ExternalModelProviderConfig( + config = ModelProvider( name="custom", endpoint="https://models.example/v1", provider_type="future-client", @@ -184,9 +189,9 @@ def test_external_provider_retains_extension_parameters() -> None: def test_external_provider_names_are_unique() -> None: - provider = ExternalModelProviderConfig(name="external", endpoint="https://models.example/v1") + provider = ModelProvider(name="external", endpoint="https://models.example/v1") with pytest.raises(ValueError, match="model provider names must be unique"): - InvocationConfig(num_records=8, dataset_name="hosted", model_providers=[provider, provider]) + validate_hosted_model_providers((provider, provider)) def test_external_model_binding_cannot_be_forged(external_plan: ResolvedSlurmRunPlan) -> None: @@ -203,17 +208,57 @@ def test_unknown_hosted_concurrency_is_rejected( ) -> None: payload = external_run.model_dump(mode="python") payload["invocation"]["model_concurrency"]["typo"] = 4 - authored = DataDesignerSlurmConfig.model_validate(payload) - with pytest.raises(SlurmPlanCompilationError, match="undeclared aliases"): - SlurmRunCompiler.compile( - resolve_slurm_config( - authored, - selected_profile=external_plan.selected_profile, - client_runtime=external_plan.client.runtime, - deployment_images=tuple(deployment.image for deployment in external_plan.deployments), - dependency_lock=dependency_lock_single, - runtime_bundle=external_plan.runtime_bundle, - run_id=external_plan.run_id, - package_version=external_plan.package_version, - ) + with pytest.raises(ValueError, match="undeclared aliases"): + DataDesignerSlurmConfig.model_validate(payload) + + +def test_slurm_invocation_rejects_parallel_provider_configuration() -> None: + with pytest.raises(ValueError): + InvocationConfig.model_validate({"num_records": 2, "dataset_name": "native", "model_providers": []}) + + +def test_provider_snapshot_keeps_only_required_settings_and_is_detached(external_run: DataDesignerSlurmConfig) -> None: + assert external_run.builder.inline is not None + provider = ModelProvider( + name="external", endpoint="https://models.example/v1", api_key="API_KEY", extra_body={"policy": "original"} + ) + unused = ModelProvider(name="unused", endpoint="https://unused.example/v1", api_key="UNUSED_KEY") + captured = resolve_hosted_model_providers( + external_run.builder.inline, {d.model_alias for d in external_run.deployments}, (unused, provider) + ) + provider.extra_body["policy"] = "edited" + assert len(captured) == 1 + assert captured[0].name == "external" + assert captured[0].api_key == "API_KEY" + assert captured[0].extra_body == {"policy": "original"} + + +@pytest.mark.parametrize( + "providers", + ( + (), + ( + ModelProvider(name="external", endpoint="https://one.example/v1"), + ModelProvider(name="external", endpoint="https://two.example/v1"), + ), + ), +) +def test_unresolved_or_ambiguous_hosted_provider_is_rejected( + external_run: DataDesignerSlurmConfig, providers: tuple[ModelProvider, ...] +) -> None: + assert external_run.builder.inline is not None + with pytest.raises(ValueError, match="exactly one"): + resolve_hosted_model_providers( + external_run.builder.inline, {d.model_alias for d in external_run.deployments}, providers ) + + +@pytest.mark.parametrize("external_run", ("mixed",), indirect=True) +def test_local_deployment_wins_when_models_share_a_provider( + external_run: DataDesignerSlurmConfig, hosted_model_providers: tuple[ModelProvider, ...] +) -> None: + assert external_run.builder.inline is not None + payload = external_run.builder.model_dump(mode="json")["inline"] + payload["data_designer"]["model_configs"][0]["provider"] = "external" + captured = resolve_hosted_model_providers(payload, {"generator"}, hosted_model_providers) + assert captured == hosted_model_providers diff --git a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index a26342aab..e67a8bb2b 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -45,7 +45,7 @@ def test_external_providers_only_reach_native_clients( log_directory=context.attempt_directory / "logs/execution-00000002", layout=AllocationLayout(("compute-001",)), ) - assert manifest.all_secret_environment_names == ("EXTERNAL_API_KEY", "EXTERNAL_HEADER_TOKEN") + assert manifest.all_secret_environment_names == ("MODEL_API_KEY",) clients = tuple( step for step in manifest.steps @@ -59,12 +59,11 @@ def test_external_providers_only_reach_native_clients( assert all( step.secret_environment == { - "EXTERNAL_API_KEY": "EXTERNAL_API_KEY", - "EXTERNAL_HEADER_TOKEN": "EXTERNAL_HEADER_TOKEN", + "MODEL_API_KEY": "MODEL_API_KEY", } for step in clients ) - assert all("EXTERNAL_API_KEY" not in step.secret_environment for step in manifest.steps if step not in clients) + assert all("MODEL_API_KEY" not in step.secret_environment for step in manifest.steps if step not in clients) if not plan.deployments: assert len(manifest.steps) == 2 assert all(step.execution == "native" and not step.gpu_indices for step in manifest.steps) diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index 5dbfde7df..f4d3882ed 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -48,15 +48,9 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( if hosted_only: plan["deployments"] = [] plan["resolved_gpus_per_node"] = 0 - plan["invocation"] = { - "authored": { - "model_providers": [ - { - "api_key": {"type": "secret", "environment": "SOURCE_TOKEN"}, - } - ] - } - } + plan["hosted_model_providers"] = [ + {"name": "hosted", "endpoint": "https://models.example/v1", "api_key": "SOURCE_TOKEN"} + ] manifest["steps"] = [step for step in manifest["steps"] if step["role"] in {"client_preflight", "client"}] manifest["all_secret_environment_names"] = ["SOURCE_TOKEN"] plan_path.write_text(json.dumps(plan)) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 381ebbd73..bb20c8059 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -4,7 +4,8 @@ from __future__ import annotations import hashlib -from collections.abc import Mapping +import json +from collections.abc import Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone from pathlib import Path @@ -14,6 +15,8 @@ import data_designer.slurm.services.retry_collection as retry_collection_module import data_designer.slurm.services.wiring as wiring_module +from data_designer.config import default_model_settings +from data_designer.config.models import ModelProvider from data_designer.slurm.benchmark.compiler import BenchmarkCompiler from data_designer.slurm.client.dependencies import ResolvedClientDependencies from data_designer.slurm.client.runtime import ClientRuntimeInspectionError, ClientRuntimeInspector @@ -72,6 +75,19 @@ def _use_persisted_client_runtime( monkeypatch.setattr(wiring_module, "resolve_scheduler_bin_path", lambda configured: configured or "/opt/slurm/bin") +@pytest.fixture +def hosted_settings( + tmp_path: Path, hosted_model_providers: tuple[ModelProvider, ...], monkeypatch: pytest.MonkeyPatch +) -> Iterator[Path]: + path = tmp_path / "dd-home" / "model_providers.yaml" + path.parent.mkdir() + path.write_text( + json.dumps({"providers": [provider.model_dump(mode="json") for provider in hosted_model_providers]}) + ) + monkeypatch.setattr(default_model_settings, "MODEL_PROVIDERS_FILE_PATH", path) + yield path + + class _Launcher: def __init__( self, @@ -250,6 +266,7 @@ def test_production_wiring_dry_run_resolves_and_renders_without_submission( @pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) def test_hosted_only_planning_needs_no_gpu_discovery_or_registered_images( + hosted_settings: Path, tmp_path: Path, profile_catalog: SlurmProfileCatalog, external_run: DataDesignerSlurmConfig, @@ -1143,7 +1160,7 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( assert len(launcher.submissions) == 2 assert launcher.exported_environments == [ {"SLURM_EXPORT_ENV": "ALL", "NEMO_DEPLOYMENT_TYPE": expected_deployment_type}, - {"NEMO_DEPLOYMENT_TYPE": expected_deployment_type}, + {"SLURM_EXPORT_ENV": "ALL", "NEMO_DEPLOYMENT_TYPE": expected_deployment_type}, ] @@ -1755,3 +1772,81 @@ def test_production_image_remove_refuses_to_delete_external_sqsh( assert caught.value.code is SlurmServiceErrorCode.CONFLICT assert "package-owned" in str(caught.value) assert service.get(image.name) == image + + +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +def test_native_provider_settings_are_captured_as_immutable_run_inputs( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + external_run: DataDesignerSlurmConfig, + hosted_settings: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), + run_id_factory=lambda: "run-snapshot", + package_version="0.9.2", + ) + plan = service.plan(external_run, source_root=tmp_path) + assert len(plan.hosted_model_providers) == 1 + assert plan.hosted_model_providers[0]["api_key"] == "MODEL_API_KEY" + original = plan.serialize_json() + hosted_settings.write_text('{"providers": []}') + monkeypatch.setattr( + wiring_module, "get_default_providers", lambda: pytest.fail("persisted plan must not reload provider settings") + ) + assert "#SBATCH --gres=none" in service.render_attempt(plan) + assert plan.serialize_json() == original + with pytest.raises(TypeError, match="frozen"): + plan.hosted_model_providers[0]["endpoint"] = "https://changed.example/v1" + assert "model_providers" not in plan.invocation.authored.model_dump() + + +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +@pytest.mark.parametrize("present", (False, True)) +def test_native_provider_credentials_are_validated_before_publication_and_forwarded_only_at_runtime( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + external_run: DataDesignerSlurmConfig, + hosted_settings: Path, + present: bool, +) -> None: + publisher = _Publisher() + launcher = _Launcher() + source = {"MODEL_API_KEY": "fake-provider-secret"} if present else {} + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + artifact_publisher=publisher, + launcher=launcher, + run_id_factory=lambda: "run-native-key", + package_version="0.9.2", + source_environment=source, + ) + if not present: + with pytest.raises(SlurmServiceError, match="required secret environment"): + service.execute(external_run, source_root=tmp_path) + assert publisher.initializations == [] + assert launcher.submissions == [] + return + service.execute(external_run, source_root=tmp_path) + assert launcher.exported_environments[0]["MODEL_API_KEY"] == "fake-provider-secret" + assert "fake-provider-secret" not in publisher.plans[0].serialize_json() + assert all("fake-provider-secret" not in script for script in launcher.submissions) + + +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +def test_missing_native_hosted_provider_is_rejected_before_submission( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + external_run: DataDesignerSlurmConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(wiring_module, "get_default_providers", list) + launcher = _Launcher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), launcher=launcher, package_version="0.9.2" + ) + with pytest.raises(SlurmServiceError): + service.execute(external_run, source_root=tmp_path) + assert launcher.submissions == [] From b6e0f72c54c3dccfc093ccc00b95ebd4435376e5 Mon Sep 17 00:00:00 2001 From: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:11:11 -0700 Subject: [PATCH 06/12] fix(export): normalize batch column order before Parquet schema casts Signed-off-by: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> --- .../src/data_designer/interface/results.py | 4 ++-- .../tests/interface/test_results.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/packages/data-designer/src/data_designer/interface/results.py b/packages/data-designer/src/data_designer/interface/results.py index 1b0f3589e..3734aad98 100644 --- a/packages/data-designer/src/data_designer/interface/results.py +++ b/packages/data-designer/src/data_designer/interface/results.py @@ -345,6 +345,6 @@ def _export_parquet(batch_files: list[Path], output: Path) -> None: except (OSError, lazy.pa.ArrowException) as e: raise ArtifactStorageError(f"Failed to read dataset artifact {batch_file}: {e}") from e try: - writer.write_table(table.cast(unified_schema)) - except (lazy.pa.ArrowInvalid, ValueError) as e: + writer.write_table(table.select(unified_schema.names).cast(unified_schema)) + except (KeyError, lazy.pa.ArrowInvalid, ValueError) as e: raise InvalidFileFormatError(f"Cannot cast batch {batch_file.name} to unified schema: {e}") from e diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py index 597133762..78c907595 100644 --- a/packages/data-designer/tests/interface/test_results.py +++ b/packages/data-designer/tests/interface/test_results.py @@ -377,6 +377,24 @@ def test_export_parquet_schema_unification(stub_dataset_creation_results, tmp_pa assert list(loaded["value"]) == [1.0, 2.0, 3.0, 4.0] +def test_export_parquet_normalizes_column_order_and_nested_empty_lists( + stub_dataset_creation_results: DatasetCreationResults, tmp_path: Path +) -> None: + batch_dir = tmp_path / "parquet-files" + batch_dir.mkdir() + records = [ + {"candidate_id": "first", "solution": {"sql": "SELECT 1", "reasoning": "one"}, "warnings": []}, + {"warnings": ["check"], "solution": {"reasoning": "two", "sql": "SELECT 2"}, "candidate_id": "second"}, + ] + for index, record in enumerate(records): + lazy.pq.write_table(lazy.pa.Table.from_pylist([record]), batch_dir / f"batch_{index:05d}.parquet") + stub_dataset_creation_results.artifact_storage.final_dataset_path = batch_dir + + exported = stub_dataset_creation_results.export(tmp_path / "out.parquet") + + assert lazy.pq.read_table(exported).to_pylist() == records + + def test_export_unknown_extension_raises(stub_dataset_creation_results, tmp_path) -> None: """export() raises InvalidFileFormatError when the extension is not a supported format.""" with pytest.raises(InvalidFileFormatError, match="Unsupported export format"): From 774774487da61921ca3369f051acddd2562318ae Mon Sep 17 00:00:00 2001 From: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:28:24 -0700 Subject: [PATCH 07/12] fix(slurm): validate hosted credentials before capturing provider settings Signed-off-by: Dhruv Nathawani <128275431+dhruvnathawani@users.noreply.github.com> --- .../data_designer/slurm/services/wiring.py | 17 ++++++++++++-- .../tests/services/test_wiring.py | 23 +++++++++++++++++++ 2 files changed, 38 insertions(+), 2 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 ae064c497..94d00370a 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 @@ -8,7 +8,7 @@ import importlib.metadata import os import shlex -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Collection, Iterator, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime, timezone @@ -39,6 +39,7 @@ SlurmConfigLoadError, SlurmProfile, SlurmProfileCatalog, + collect_secret_environment_names, load_builder_payload, resolve_profile, ) @@ -159,6 +160,7 @@ def __init__( run_id_factory: RunIdFactory, package_version: str, runtime_inspector: ClientRuntimeInspector | None = None, + available_environment: Collection[str] = (), ) -> None: self._profile = selected_profile self._images = image_registry @@ -167,6 +169,7 @@ def __init__( self._run_id_factory = run_id_factory self._package_version = package_version self._runtime_inspector = runtime_inspector or ClientRuntimeInspector() + self._available_environment = frozenset(available_environment) @contextmanager def prepare( @@ -216,6 +219,12 @@ def _prepare( local_aliases, get_default_providers() if get_external_model_aliases(payload, local_aliases) else (), ) + if set(collect_secret_environment_names(hosted_providers)).difference(self._available_environment): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "required secret environment for hosted models is unavailable", + ) dependencies = stack.enter_context( self._dependencies.resolve( authored.client.dependencies, @@ -828,6 +837,7 @@ def _create_slurm_run_service( ) -> SlurmRunService: command_client = launcher or SlurmCommandClient() selected_clock = clock or _utc_now + environment = dict(os.environ if source_environment is None else source_environment) preparer = _RunPreparer( selected, VerifiedImageRegistry(selected.profile.workspace_root), @@ -835,6 +845,9 @@ def _create_slurm_run_service( command_client, run_id_factory or _new_run_id, package_version or importlib.metadata.version("data-designer-slurm"), + available_environment={ + name for name, value in environment.items() if isinstance(value, str) and value and "\0" not in value + }, ) backend = _SystemRunBackend( preparer, @@ -842,7 +855,7 @@ def _create_slurm_run_service( command_client, artifact_publisher or StateRunArtifactPublisher(selected.profile.workspace_root, selected_clock), selected_clock, - dict(os.environ if source_environment is None else source_environment), + environment, progress, ) return SlurmRunService(_SystemRunPlanner(preparer), render_generation_attempt_script, backend) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index bb20c8059..e4a161c26 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -85,6 +85,7 @@ def hosted_settings( json.dumps({"providers": [provider.model_dump(mode="json") for provider in hosted_model_providers]}) ) monkeypatch.setattr(default_model_settings, "MODEL_PROVIDERS_FILE_PATH", path) + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-secret") yield path @@ -1850,3 +1851,25 @@ def test_missing_native_hosted_provider_is_rejected_before_submission( with pytest.raises(SlurmServiceError): service.execute(external_run, source_root=tmp_path) assert launcher.submissions == [] + + +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +@pytest.mark.parametrize("key", ("plaintext", "PLAINTEXT_SECRET", "sk-plaintext-key")) +def test_plaintext_native_provider_keys_are_rejected_without_echoing_or_capturing_the_value( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + external_run: DataDesignerSlurmConfig, + hosted_settings: Path, + key: str, +) -> None: + hosted_settings.write_text( + json.dumps({"providers": [{"name": "external", "endpoint": "https://models.example/v1", "api_key": key}]}) + ) + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=_Launcher(), + source_environment={"MODEL_API_KEY": "fake-provider-secret"}, + ) + with pytest.raises(SlurmServiceError) as error: + service.plan(external_run, source_root=tmp_path) + assert key not in str(error.value) From 17dfe0b130a3ea8268dc2274e4d43ac44ebfc35c Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 26 Oct 2026 15:55:28 -0600 Subject: [PATCH 08/12] docs(slurm): align hosted provider plan Signed-off-by: Nabin Mulepati --- plans/1079/slurm-external-providers.md | 38 +++++++++++++++++++------- 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/plans/1079/slurm-external-providers.md b/plans/1079/slurm-external-providers.md index b3d4bfacb..352072e7e 100644 --- a/plans/1079/slurm-external-providers.md +++ b/plans/1079/slurm-external-providers.md @@ -1,18 +1,33 @@ # Plan: External model providers in Slurm runs Issue: https://github.com/NVIDIA-NeMo/DataDesigner/issues/1079 +PR: https://github.com/NVIDIA-NeMo/DataDesigner/pull/1080 -The implementation follows the reviewed plan: retain the existing async engine, -add explicit external providers, and measure allocation utilization during SQL -qualification before considering separate generation and judging allocations. +Provider ownership follows the [issue feedback](https://github.com/NVIDIA-NeMo/DataDesigner/issues/1079#issuecomment-6369036893). + +The implementation retains the existing async engine and resolves hosted models +through native DD provider settings. Slurm configuration describes only models +deployed locally. Mixed runs retain their GPUs during hosted requests; separate +stages remain appropriate when hosted processing is substantial. ## Design -- Add external provider declarations to invocation configuration. Reuse native DD - provider settings, with Slurm `SecretRef` values at credential boundaries. +- Configure hosted providers in DD's existing `model_providers.yaml`; recipe model + configs reference those provider names. Do not introduce `invocation.model_providers` + or a parallel Slurm provider schema. +- Resolve undeployed aliases against native settings before submission. Capture + only required hosted providers as immutable native JSON run inputs so compute + nodes and retries do not depend on the submit host's settings directory or later + settings edits. +- Keep API-key environment-variable references in snapshots. Validate availability + before capture and submission, and reuse runtime credential transport to forward + values only to client steps, including retries. Initially support environment-backed + API keys and non-secret static headers; reject plaintext keys and sensitive + headers. Custom secret resolvers are outside this change. - Partition declared model aliases into managed deployments and external models. Require complete, disjoint coverage and reject provider-name collisions. -- Materialize endpoints only for managed aliases. Preserve external model names, +- Combine captured hosted providers with generated local providers in the client. + Materialize endpoints only for managed aliases. Preserve external model names, sampling parameters, and explicit concurrency settings. - Allow an empty deployment list. Resolve zero GPUs, one native client host, no serving images, no GPU discovery, no serving containers, and no proxy ports. @@ -32,10 +47,13 @@ native client provider routing, credential isolation, resource directives, allocation startup and cleanup, retry/resume, and local-only golden fixtures. Run the Slurm suite and the full package checks before opening the PR. -SQL qualification uses a small mixed-provider async recipe and a CPU-only hosted -recipe. Record per-model timing and actual GPU utilization, and compare idle -windows with the selected cluster policy. Preserve existing SQL prompts, models, -sampling, judge gates, datasets, checkpoints, and completed study results. +The live qualification reported in PR #1080 uses a small mixed-provider async +recipe and a CPU-only hosted recipe, combined with the independent vLLM +version-gate fix in #1086. It validates routing and publication, not sustained GPU +utilization, Gym SQL execution, production SQL quality, or a full dataset/SFT rerun. +Existing SQL prompts, models, sampling, judge gates, datasets, checkpoints, and +completed study results remain preserved. Measure per-model timing and actual GPU +utilization against cluster policy before deciding to split production stages. Automatic DAG partitioning, new async machinery, model substitutions, and the separate single-row profiling defect are outside this change. From 85a3c4f5017866de6e3e53346dfa8336a961471b Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 26 Oct 2026 16:00:57 -0600 Subject: [PATCH 09/12] fix(slurm): keep provider payloads opaque Exclude native hosted-provider payloads from recursive Slurm secret\nreference discovery while retaining explicit API-key references.\n\nCover mixed and hosted-only plans with SecretRef-shaped request data,\nincluding missing environment fields and nested lists. Signed-off-by: Nabin Mulepati --- .../slurm/runtime/plan_reader.sh | 13 ++++- .../tests/runtime/test_shell_runtime.py | 49 +++++++++++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh index 683399856..eb558bfbc 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh +++ b/packages/data-designer-slurm/src/data_designer/slurm/runtime/plan_reader.sh @@ -65,13 +65,22 @@ dd_read_container_path() { dd_read_plan_secret_names() { local plan=$1 + # Native provider payloads are opaque; only their API-key references are Slurm credentials. dd_read_null_values DD_REQUIRED_SECRET_NAMES < <( - jq -j '[(.. | objects | select(.type? == "secret") | .environment), (.hosted_model_providers[]?.api_key // empty)] | unique | .[] | ., "\u0000"' "${plan}" + jq -j ' + [ + (del(.hosted_model_providers) | .. | objects | select(.type? == "secret") | .environment), + (.hosted_model_providers[]?.api_key // empty) + ] + | unique + | .[] + | ., "\u0000" + ' "${plan}" ) dd_read_null_values DD_ALL_SECRET_NAMES < <( jq -j ' [ - (.. | objects | select(.type? == "secret") | .environment), + (del(.hosted_model_providers) | .. | objects | select(.type? == "secret") | .environment), (.hosted_model_providers[]?.api_key // empty), ( .deployments[].authored.server.environment diff --git a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py index f4d3882ed..c968de769 100644 --- a/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py +++ b/packages/data-designer-slurm/tests/runtime/test_shell_runtime.py @@ -13,6 +13,10 @@ from pathlib import Path import pytest +from pydantic import JsonValue + +from data_designer.slurm.config import collect_secret_environment_names +from data_designer.slurm.planning import ResolvedSlurmRunPlan @pytest.mark.parametrize( @@ -127,6 +131,51 @@ def test_bash_controller_scopes_secrets_cleans_steps_and_never_runs_host_python( assert "Run completed." in completed.stdout +@pytest.mark.parametrize( + "provider_payload", + ( + {"type": "secret", "environment": "REMOTE_SERVICE_ENV"}, + {"type": "secret"}, + [{"nested": {"type": "secret", "environment": "REMOTE_SERVICE_ENV"}}], + ), + ids=("named-environment", "no-environment", "nested-list"), +) +def test_bash_secret_discovery_treats_provider_payloads_as_opaque( + tmp_path: Path, + external_plan: ResolvedSlurmRunPlan, + provider_payload: JsonValue, +) -> None: + payload = external_plan.model_dump(mode="json") + payload["hosted_model_providers"][0]["extra_body"] = {"routing": provider_payload} + plan = ResolvedSlurmRunPlan.model_validate_json(json.dumps(payload)) + required_names = collect_secret_environment_names(plan) + assert required_names == ("MODEL_API_KEY",) + plan_path = tmp_path / "resolved-plan.json" + plan_path.write_text(plan.serialize_json()) + runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" + command = f""" +set -Eeuo pipefail +source {shlex.quote((runtime_root / "plan_reader.sh").as_posix())} +source {shlex.quote((runtime_root / "step_runner.sh").as_posix())} +dd_read_plan_secret_names {shlex.quote(plan_path.as_posix())} +printf '%s\\n' "${{DD_REQUIRED_SECRET_NAMES[@]}}" +printf '%s\\n' "${{DD_ALL_SECRET_NAMES[@]}}" +dd_require_plan_secrets +""" + environment = {**os.environ, "MODEL_API_KEY": "test-provider-key"} + environment.pop("REMOTE_SERVICE_ENV", None) + completed = subprocess.run( + ("bash", "-c", command), + capture_output=True, + text=True, + env=environment, + timeout=5, + ) + + assert completed.returncode == 0, completed.stderr + assert completed.stdout.splitlines() == [*required_names, *required_names] + + def test_staged_shell_modules_parse_as_bash(tmp_path: Path) -> None: del tmp_path runtime_root = Path(__file__).parents[2] / "src/data_designer/slurm/runtime" From 9b4ce45d4576c34a43be1adaf349c201d53df9b3 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 27 Oct 2026 09:16:59 -0600 Subject: [PATCH 10/12] fix(slurm): defer retry credential validation Resolve credentials only when a new retry submission is required, under the coordinator locks and before prepared state is published. Preserve credential-free recovery of active and accepted retries, and cover missing keys, dry runs, and restored credential transport. Signed-off-by: Nabin Mulepati --- .../slurm/services/retry_collection.py | 13 ++-- .../src/data_designer/slurm/state/retry.py | 14 +++- .../tests/services/test_wiring.py | 77 +++++++++++++++++++ .../tests/state/test_retry_collection.py | 12 ++- 4 files changed, 107 insertions(+), 9 deletions(-) diff --git a/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py index f7ca19aeb..5515b6cd7 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/services/retry_collection.py @@ -66,12 +66,15 @@ def retry( try: observed_at = self._clock() resolved_plan = SlurmStateWriter(self._workspace_root, run_id).load_resolved_plan() - export_environment = ( - self._export_environment - if dry_run - else build_run_export_environment(resolved_plan, self._source_environment, operation) + coordinator = SlurmRetryCoordinator( + self._workspace_root, + run_id, + self._launcher, + self._export_environment, + export_environment_factory=lambda: build_run_export_environment( + resolved_plan, self._source_environment, operation + ), ) - coordinator = SlurmRetryCoordinator(self._workspace_root, run_id, self._launcher, export_environment) effective_resume_mode = self._resolve_retry_resume_mode(resolved_plan, resume) preview = coordinator.preview_active(shard_ids=shard_ids, observed_at=observed_at) if preview is not None and effective_resume_mode not in {None, preview[0].effective_resume_mode}: diff --git a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py index 8faef34c3..cf0b44096 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/state/retry.py @@ -5,7 +5,7 @@ from __future__ import annotations -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import ExitStack, contextmanager from datetime import datetime, timezone from pathlib import Path @@ -65,7 +65,12 @@ def query_submissions_by_name( class SlurmRetryCoordinator: - """Select retryable shards and durably submit their next attempts.""" + """Select retryable shards and durably submit their next attempts. + + An optional ``export_environment_factory`` overrides ``export_environment`` + only for new submissions. It runs under the selection locks before prepared + state is persisted; previews and accepted-submission recovery do not call it. + """ def __init__( self, @@ -73,10 +78,13 @@ def __init__( run_id: Identifier, scheduler: RetryScheduler | None = None, export_environment: Mapping[str, str] | None = None, + *, + export_environment_factory: Callable[[], Mapping[str, str]] | None = None, ) -> None: root, normalized_run_id = _validate_location(workspace_root, run_id) self._scheduler = scheduler if scheduler is not None else SlurmCommandClient() self._export_environment = None if export_environment is None else dict(export_environment) + self._export_environment_factory = export_environment_factory self._state = StateStorage(root, normalized_run_id) self._reader = StateReader(self._state, normalized_run_id) self._retries = RetryStorage(self._state) @@ -108,6 +116,8 @@ def retry( script = render_generation_retry_script(self._reader.load_resolved_plan(status.run), plan) with self._acquire_selection_locks(selected): self._require_fresh_selection(status, selected) + if self._export_environment_factory is not None: + self._export_environment = dict(self._export_environment_factory()) self._persist_prepared_retry(plan, timestamp) receipt = self._submit(plan, script, timestamp) attempts, superseded = self._publish_attempts(plan, receipt.job_id, timestamp) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index e4a161c26..ddc89a097 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -1165,6 +1165,83 @@ def test_production_retry_dry_run_and_submission_are_sparse_and_idempotent( ] +@pytest.mark.parametrize("external_run", ("hosted-only",), indirect=True) +@pytest.mark.parametrize("active_state", (SchedulerState.PENDING, SchedulerState.RUNNING)) +def test_hosted_retry_requires_credentials_only_for_new_submissions( + tmp_path: Path, + profile_catalog: SlurmProfileCatalog, + external_run: DataDesignerSlurmConfig, + hosted_settings: Path, + active_state: SchedulerState, +) -> None: + launcher = _Launcher(submission_job_ids=(42, 43, 44)) + now = datetime(2026, 9, 8, tzinfo=timezone.utc) + profile = _profile(tmp_path, profile_catalog) + source_environment = {"MODEL_API_KEY": "fake-provider-secret"} + service = create_slurm_run_service( + profile=profile, + launcher=launcher, + run_id_factory=lambda: "run-hosted-retry", + clock=lambda: now, + package_version="0.9.2", + source_environment=source_environment, + ) + service.execute(external_run, source_root=tmp_path) + initial_failure = SlurmAccountingEntry( + job_identity=SchedulerIdentity(array_job_id=42, array_task_id=0), + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ) + launcher.accounting_entries = (initial_failure,) + without_credentials = create_slurm_run_service( + profile=profile, + launcher=launcher, + clock=lambda: now, + package_version="0.9.2", + source_environment={}, + ) + run_root = tmp_path / "runs/run-hosted-retry" + preview = without_credentials.retry("run-hosted-retry", resume="if_possible", dry_run=True) + assert preview.state == "dry_run" + assert not (run_root / "retries").exists() + assert not (run_root / "retry.lock").exists() + with pytest.raises(SlurmServiceError, match="required secret environment"): + without_credentials.retry("run-hosted-retry", resume="if_possible") + assert not (run_root / "retries").exists() + assert len(launcher.submissions) == 1 + + submitted = service.retry("run-hosted-retry", resume="if_possible") + assert submitted.job_id == 43 + assert launcher.exported_environments[-1]["MODEL_API_KEY"] == "fake-provider-secret" + launcher.queue_entries = ( + SlurmQueueEntry(job_identity=SchedulerIdentity(array_job_id=43, array_task_id=0), state=active_state), + ) + repeated = without_credentials.retry("run-hosted-retry", resume="if_possible") + assert repeated == submitted + assert len(launcher.submissions) == 2 + + launcher.queue_entries = () + launcher.accounting_entries = ( + initial_failure, + SlurmAccountingEntry( + job_identity=SchedulerIdentity(array_job_id=43, array_task_id=0), + state=SchedulerState.FAILED, + process_exit_code=SlurmProcessExitCode(exit_status=1, termination_signal=0), + ), + ) + records = {path: path.read_bytes() for path in (run_root / "retries").rglob("*.json")} + with pytest.raises(SlurmServiceError, match="required secret environment"): + without_credentials.retry("run-hosted-retry", resume="if_possible") + assert {path: path.read_bytes() for path in (run_root / "retries").rglob("*.json")} == records + assert len(launcher.submissions) == 2 + + replacement = service.retry("run-hosted-retry", resume="if_possible") + assert replacement.job_id == 44 + assert replacement.attempt_ids != submitted.attempt_ids + assert launcher.exported_environments[-1]["MODEL_API_KEY"] == "fake-provider-secret" + assert all("fake-provider-secret" not in path.read_text() for path in run_root.rglob("*.json")) + + def test_production_retry_preview_recovers_active_pinned_mode( tmp_path: Path, profile_catalog: SlurmProfileCatalog, diff --git a/packages/data-designer-slurm/tests/state/test_retry_collection.py b/packages/data-designer-slurm/tests/state/test_retry_collection.py index 555ed2488..437ff1998 100644 --- a/packages/data-designer-slurm/tests/state/test_retry_collection.py +++ b/packages/data-designer-slurm/tests/state/test_retry_collection.py @@ -9,6 +9,7 @@ import os import shutil import sys +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -361,7 +362,12 @@ def accept_then_lose_receipt(script: str) -> object: ) runner.script_next("sacct", FakeCommandResponse()) - recovered = SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + def unavailable_credentials() -> Mapping[str, str]: + pytest.fail("accepted retry recovery must not resolve credentials") + + recovered = SlurmRetryCoordinator( + case.workspace, case.plan.run_id, scheduler, export_environment_factory=unavailable_credentials + ).retry( effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=6), ) @@ -370,7 +376,9 @@ def accept_then_lose_receipt(script: str) -> object: assert case.writer.load_attempts(case.shards[0].shard_id) == (first_attempt, recovered[0]) assert storage.read_status("retry-0001").state is RetryState.COMPLETED assert ( - SlurmRetryCoordinator(case.workspace, case.plan.run_id, scheduler).retry( + SlurmRetryCoordinator( + case.workspace, case.plan.run_id, scheduler, export_environment_factory=unavailable_credentials + ).retry( effective_resume_mode="never", observed_at=case.created_at + timedelta(minutes=7), ) From 3d3666fb29f742136c4242bff2da764eb94b9e1f Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 27 Oct 2026 10:43:13 -0600 Subject: [PATCH 11/12] fix(slurm): reject model-free runs before submission Require a nonempty model alias set in resolved builder inputs and reject model-free inline or sourced recipes before staging or submission. Document the limitation and cover omitted, null, and empty model configs while retaining hosted-only CPU run coverage. Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/run-configuration.mdx | 2 ++ .../data_designer/slurm/planning/models.py | 2 +- .../data_designer/slurm/services/wiring.py | 9 +++++- .../data-designer-slurm/tests/conftest.py | 19 ++++++++++++ .../tests/planning/test_external_providers.py | 27 ++++++++++++++++- .../tests/services/test_wiring.py | 29 +++++++++++++++++++ 6 files changed, 85 insertions(+), 3 deletions(-) diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index e956fd808..47896436b 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -94,6 +94,8 @@ Both local and external requests use DD's existing async engine. Ready requests ### Hosted-only CPU runs +Slurm runs require at least one explicitly configured model. Sampler-only recipes are not currently supported. + For a builder whose models all reference configured hosted providers, omit local deployments or set `deployments: []`. Select the cluster's CPU partition: ```yaml diff --git a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py index d0b1ec21d..ee67aeb31 100644 --- a/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py +++ b/packages/data-designer-slurm/src/data_designer/slurm/planning/models.py @@ -190,7 +190,7 @@ class ResolvedBuilderInput(ContractValue): inline: dict[str, JsonValue] | None = None # Both forms use deterministic persisted builder JSON bytes. content_sha256: Sha256Digest - model_aliases: tuple[ModelAlias, ...] + model_aliases: tuple[ModelAlias, ...] = Field(min_length=1) external_model_aliases: tuple[ModelAlias, ...] = Field(default=(), exclude_if=lambda value: not value) @model_validator(mode="after") 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 94d00370a..35d0cb9f3 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 @@ -58,7 +58,7 @@ from data_designer.slurm.launcher.renderer import render_generation_attempt_script from data_designer.slurm.launcher.scheduler_path import SchedulerPathError, resolve_scheduler_bin_path from data_designer.slurm.planning import ResolvedImage, ResolvedSlurmRunPlan -from data_designer.slurm.planning.builder_identity import get_external_model_aliases +from data_designer.slurm.planning.builder_identity import get_declared_model_aliases, get_external_model_aliases from data_designer.slurm.planning.compiler import SlurmRunCompiler from data_designer.slurm.planning.model_providers import resolve_hosted_model_providers from data_designer.slurm.planning.resolution import resolve_slurm_config @@ -213,6 +213,13 @@ def _prepare( _validate_builder_for_submission(authored, builder_payload, operation) payload = authored.builder.inline if authored.builder.inline is not None else builder_payload assert payload is not None + if not get_declared_model_aliases(payload): + raise SlurmServiceError( + SlurmServiceErrorCode.INVALID_REQUEST, + operation, + "Slurm runs require at least one configured model. " + "Sampler-only recipes are not currently supported.", + ) local_aliases = {deployment.model_alias for deployment in authored.deployments} hosted_providers = resolve_hosted_model_providers( payload, diff --git a/packages/data-designer-slurm/tests/conftest.py b/packages/data-designer-slurm/tests/conftest.py index 8f6ab75f3..48e17c362 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -19,6 +19,7 @@ FakeVllmBackend, ) +from data_designer.config import SamplerColumnConfig, UUIDSamplerParams from data_designer.config.models import ModelProvider from data_designer.slurm.benchmark import BenchmarkManifest, BenchmarkReport from data_designer.slurm.client import ClientResult @@ -86,6 +87,24 @@ def single_node_plan() -> ResolvedSlurmRunPlan: return ResolvedSlurmRunPlan.model_validate_json((CONTRACT_GOLDEN_DIRECTORY / "single_node_plan.json").read_text()) +@pytest.fixture(params=("empty", "omitted", "none")) +def model_free_run( + request: pytest.FixtureRequest, authored_run_single: DataDesignerSlurmConfig +) -> DataDesignerSlurmConfig: + payload = authored_run_single.model_dump(mode="json") + builder = payload["builder"]["inline"]["data_designer"] + if request.param == "omitted": + builder.pop("model_configs") + else: + builder["model_configs"] = [] if request.param == "empty" else None + builder["columns"] = [ + SamplerColumnConfig(name="id", sampler_type="uuid", params=UUIDSamplerParams()).model_dump(mode="json") + ] + payload["deployments"] = [] + payload["invocation"]["model_concurrency"] = {} + return DataDesignerSlurmConfig.model_validate_json(json.dumps(payload)) + + @pytest.fixture(params=("mixed", "hosted-only")) def external_run( request: pytest.FixtureRequest, authored_run_single: DataDesignerSlurmConfig diff --git a/packages/data-designer-slurm/tests/planning/test_external_providers.py b/packages/data-designer-slurm/tests/planning/test_external_providers.py index 0835aff47..a38fc8a3e 100644 --- a/packages/data-designer-slurm/tests/planning/test_external_providers.py +++ b/packages/data-designer-slurm/tests/planning/test_external_providers.py @@ -20,7 +20,7 @@ from data_designer.slurm.launcher.renderer import render_generation_attempt_script from data_designer.slurm.planning import ResolvedDependencyLock, ResolvedSlurmRunPlan from data_designer.slurm.planning.compiler import SlurmRunCompiler -from data_designer.slurm.planning.errors import SlurmPlanContractError +from data_designer.slurm.planning.errors import SlurmConfigResolutionError, SlurmPlanContractError from data_designer.slurm.planning.model_providers import ( load_hosted_model_providers, resolve_hosted_model_providers, @@ -262,3 +262,28 @@ def test_local_deployment_wins_when_models_share_a_provider( payload["data_designer"]["model_configs"][0]["provider"] = "external" captured = resolve_hosted_model_providers(payload, {"generator"}, hosted_model_providers) assert captured == hosted_model_providers + + +@pytest.mark.parametrize("sourced", (False, True)) +def test_model_free_runs_are_rejected_during_resolution( + model_free_run: DataDesignerSlurmConfig, + single_node_plan: ResolvedSlurmRunPlan, + dependency_lock_single: ResolvedDependencyLock, + sourced: bool, +) -> None: + payload = model_free_run.builder.inline + config = model_free_run + if sourced: + config = config.model_copy(update={"builder": BuilderInput(source="builder.json")}) + with pytest.raises(SlurmConfigResolutionError, match="model_aliases"): + resolve_slurm_config( + config, + selected_profile=single_node_plan.selected_profile, + client_runtime=single_node_plan.client.runtime, + deployment_images=(), + dependency_lock=dependency_lock_single, + runtime_bundle=single_node_plan.runtime_bundle, + run_id=single_node_plan.run_id, + package_version=single_node_plan.package_version, + builder_payload=payload if sourced else None, + ) diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index ddc89a097..df6f66307 100644 --- a/packages/data-designer-slurm/tests/services/test_wiring.py +++ b/packages/data-designer-slurm/tests/services/test_wiring.py @@ -1950,3 +1950,32 @@ def test_plaintext_native_provider_keys_are_rejected_without_echoing_or_capturin with pytest.raises(SlurmServiceError) as error: service.plan(external_run, source_root=tmp_path) assert key not in str(error.value) + + +@pytest.mark.parametrize("sourced", (False, True)) +def test_model_free_runs_are_rejected_before_submission( + tmp_path: Path, + model_free_run: DataDesignerSlurmConfig, + profile_catalog: SlurmProfileCatalog, + sourced: bool, +) -> None: + config = model_free_run + if sourced: + (tmp_path / "builder.json").write_text(json.dumps(config.builder.inline)) + config = config.model_copy(update={"builder": BuilderInput(source="builder.json")}) + launcher = _Launcher() + publisher = _Publisher() + service = create_slurm_run_service( + profile=_profile(tmp_path, profile_catalog), + launcher=launcher, + artifact_publisher=publisher, + run_id_factory=lambda: "run-model-free", + package_version="0.9.2", + ) + with pytest.raises(SlurmServiceError, match="Slurm runs require at least one configured model") as caught: + service.execute(config, source_root=tmp_path) + assert caught.value.code is SlurmServiceErrorCode.INVALID_REQUEST + assert "Sampler-only recipes are not currently supported" in str(caught.value) + assert launcher.submissions == [] + assert publisher.initializations == [] + assert not (tmp_path / "runs").exists() From f06460c4ea995e34a175b5460f8fc6c65d49e67c Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Tue, 27 Oct 2026 11:05:36 -0600 Subject: [PATCH 12/12] docs(slurm): clarify hosted provider settings Signed-off-by: Nabin Mulepati --- .../latest/pages/slurm/run-configuration.mdx | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 47896436b..cc0efb85f 100644 --- a/fern/versions/latest/pages/slurm/run-configuration.mdx +++ b/fern/versions/latest/pages/slurm/run-configuration.mdx @@ -51,7 +51,18 @@ Slurm schedules the job in one of the selected partitions. For runs with local d ## External model providers -Configure hosted endpoints through DD's normal `model_providers.yaml`, outside the Slurm run file: +Configure hosted endpoints in DD's normal provider settings on the host where you run the Slurm submission command. The default file is `~/.data-designer/model_providers.yaml`. + +`DATA_DESIGNER_HOME` is an optional environment variable that changes DD's settings directory. If it is unset, DD uses `~/.data-designer`. To use a different directory, set it before launching the Slurm CLI: + +```bash +export DATA_DESIGNER_HOME="$HOME/dd-settings" +mkdir -p "$DATA_DESIGNER_HOME" +``` + +With this override, the provider file is `$DATA_DESIGNER_HOME/model_providers.yaml`. See [Default Model Settings](/concepts/models/default-model-settings#how-default-model-providers-and-configurations-work) for more about DD's settings files. + +Add the hosted provider to the `providers` list in that file. Provider definitions are separate from the recipe's `model_configs` and the Slurm run file: ```yaml providers: