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. diff --git a/fern/versions/latest/pages/slurm/run-configuration.mdx b/fern/versions/latest/pages/slurm/run-configuration.mdx index 0b8186340..cc0efb85f 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. +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 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,87 @@ 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 + +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: + - 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: + 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 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_concurrency: + generator: 64 + judge: 16 +``` + +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 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 + +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 +schema_version: 1 +name: hosted-judging +builder: + source: hosted-builder.yaml +invocation: + num_records: 1000 + dataset_name: hosted-judging +client: + cpus: 4 +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` 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 @@ -169,7 +249,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..fb7b87f92 --- /dev/null +++ b/packages/data-designer-slurm/examples/hosted-only.yaml @@ -0,0 +1,25 @@ +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_concurrency: + judge: 16 +client: + cpus: 4 +deployments: [] +submission: + partition: cpu + memory: 16G 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..6d2b7a11d --- /dev/null +++ b/packages/data-designer-slurm/examples/mixed-providers.yaml @@ -0,0 +1,36 @@ +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_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/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/_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..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 @@ -404,10 +404,22 @@ 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 = 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.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 = [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] @@ -417,23 +429,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 @@ -547,9 +558,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( @@ -650,6 +674,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/config/builder.py b/packages/data-designer-slurm/src/data_designer/slurm/config/builder.py index 21670bb5e..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 @@ -136,8 +136,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/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 dbc4c0fde..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 @@ -333,6 +333,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 @@ -357,7 +362,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) @@ -386,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: - 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/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/src/data_designer/slurm/launcher/renderer.py b/packages/data-designer-slurm/src/data_designer/slurm/launcher/renderer.py index 8b57f5149..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 @@ -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,13 @@ 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.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)) @@ -255,6 +259,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..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 @@ -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,15 @@ 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], 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 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 f107a9e2e..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, @@ -184,7 +185,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/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 10ee11ee4..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 @@ -47,7 +47,12 @@ 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.planning.model_providers import ( + dump_hosted_model_providers, + load_hosted_model_providers, + resolve_hosted_model_providers, +) from data_designer.slurm.types import NetworkPort, PartitionSelection @@ -185,7 +190,8 @@ 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") def validate_input(self) -> ResolvedBuilderInput: @@ -204,6 +210,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 @@ -379,6 +389,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: @@ -394,16 +405,22 @@ 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 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 @@ -420,7 +437,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 +488,28 @@ 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") + 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("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, 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") - 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 +621,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..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 @@ -8,9 +8,10 @@ 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.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 @@ -26,9 +27,13 @@ ) 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 +from data_designer.slurm.planning.model_providers import ( + dump_hosted_model_providers, +) from data_designer.slurm.planning.models import ( ResolvedBuilderInput, ResolvedClientRuntime, @@ -85,9 +90,10 @@ 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 + hosted_model_providers: tuple[dict[str, JsonValue], ...] = () invocation: ResolvedInvocation client_runtime: ResolvedClientRuntime deployment_images: tuple[ResolvedImage, ...] @@ -109,12 +115,18 @@ 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: 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( @@ -131,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, @@ -159,7 +172,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 +245,7 @@ def _resolve_builder( run_root: str, builder_payload: dict[str, JsonValue] | None, ) -> tuple[ResolvedBuilderInput, dict[str, JsonValue] | None]: + 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") @@ -232,6 +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, local_aliases), ), None, ) @@ -250,6 +273,7 @@ def _resolve_builder( source=source, content_sha256=source.sha256, model_aliases=aliases, + external_model_aliases=get_external_model_aliases(validated_payload, local_aliases), ), validated_payload, ) @@ -277,6 +301,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 8a35033c2..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 @@ -10,8 +10,9 @@ 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.model_providers import load_hosted_model_providers, resolve_hosted_model_providers from data_designer.slurm.planning.models import ( ResolvedDependencyLock, ResolvedSlurmRunPlan, @@ -75,6 +76,23 @@ 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, {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 _require(plan.submission.account == expected_account, "resolved account does not match authored/profile input") @@ -89,6 +107,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/src/data_designer/slurm/runtime/bootstrap.py b/packages/data-designer-slurm/src/data_designer/slurm/runtime/bootstrap.py index 214711b7f..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, @@ -155,11 +156,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, + *(ModelProvider.model_validate(value) for value in plan.hosted_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..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 @@ -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 ' @@ -66,13 +65,23 @@ 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] | 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 | to_entries[] 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/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..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 @@ -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,16 @@ 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() + 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 + ), + ) 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 a29a7674c..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 @@ -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 @@ -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 @@ -57,13 +58,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_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 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 @@ -156,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 @@ -164,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( @@ -205,6 +211,27 @@ 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 + 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, + 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, @@ -238,6 +265,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: @@ -276,6 +304,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 @@ -347,7 +377,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: @@ -491,43 +523,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, @@ -849,6 +844,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), @@ -856,6 +852,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, @@ -863,7 +862,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/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/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/client/test_external_model_execution.py b/packages/data-designer-slurm/tests/client/test_external_model_execution.py new file mode 100644 index 000000000..dbf09e128 --- /dev/null +++ b/packages/data-designer-slurm/tests/client/test_external_model_execution.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +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.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 +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 +from data_designer.slurm.state.artifacts import CandidateArtifactVerifier +from data_designer.slurm.state.outputs import CandidateOutputManifest + + +@pytest.fixture +def external_client_case( + hosted_model_providers: tuple[ModelProvider, ...], + external_run: DataDesignerSlurmConfig, + client_worker_case: ClientWorkerCase, +) -> ClientWorkerCase: + case = client_worker_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), + 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("MODEL_API_KEY", "fake-provider-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 == "MODEL_API_KEY" + assert providers["external"].extra_body == {"routing": {"version": 2}} + assert providers["external"].extra_headers == {"X-Project": "test"} + 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", ("MODEL_API_KEY",)) +def test_worker_rejects_unavailable_hosted_credentials( + external_client_case: ClientWorkerCase, + monkeypatch: pytest.MonkeyPatch, + missing: str, +) -> None: + 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"): + 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("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=()) + result = worker.run( + case.plan_path, + prepared=case.prepared, + endpoints={}, + plugins=(), + ) + 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("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" + 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) + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-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.", + }, + ) + 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()) + 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=()) + # 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 + 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: + 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) + 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") + 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..48e17c362 100644 --- a/packages/data-designer-slurm/tests/conftest.py +++ b/packages/data-designer-slurm/tests/conftest.py @@ -19,6 +19,8 @@ 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 from data_designer.slurm.config import ( @@ -85,6 +87,85 @@ 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 +) -> 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_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 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, +) -> ResolvedSlurmRunPlan: + 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), + 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/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/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..a38fc8a3e --- /dev/null +++ b/packages/data-designer-slurm/tests/planning/test_external_providers.py @@ -0,0 +1,289 @@ +# 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.config.models import ModelProvider +from data_designer.slurm.config import ( + BuilderInput, + DataDesignerSlurmConfig, + DataDesignerSlurmConfigBuilder, + 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 SlurmConfigResolutionError, 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 + + +@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 "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() + ) + + +@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 + 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( + 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), + 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) + 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 + 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", + ) + builder.with_submission(memory="16G") + config = builder.build() + assert config.deployments == [] + assert "model_providers" not in config.invocation.model_dump() + assert config.submission.memory == "16G" + + +@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 + + +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", + ( + "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): + validate_hosted_model_providers((ModelProvider(name="external", endpoint=endpoint),)) + + +@pytest.mark.parametrize( + "values", + ( + {"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): + validate_hosted_model_providers( + (ModelProvider.model_validate({"name": "external", "endpoint": "https://models.example/v1", **values}),) + ) + + +def test_external_provider_retains_extension_parameters() -> None: + config = ModelProvider( + 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 = ModelProvider(name="external", endpoint="https://models.example/v1") + with pytest.raises(ValueError, match="model provider names must be unique"): + validate_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 + 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 + + +@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/runtime/test_bootstrap.py b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py index 5c92d5a38..e67a8bb2b 100644 --- a/packages/data-designer-slurm/tests/runtime/test_bootstrap.py +++ b/packages/data-designer-slurm/tests/runtime/test_bootstrap.py @@ -24,6 +24,52 @@ 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 == ("MODEL_API_KEY",) + 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 + == { + "MODEL_API_KEY": "MODEL_API_KEY", + } + for step 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) + 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..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,12 +13,25 @@ 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("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 +48,18 @@ 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["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)) 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 +95,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 +112,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'}" @@ -103,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" diff --git a/packages/data-designer-slurm/tests/services/test_wiring.py b/packages/data-designer-slurm/tests/services/test_wiring.py index 5da340517..df6f66307 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,20 @@ 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) + monkeypatch.setenv("MODEL_API_KEY", "fake-provider-secret") + yield path + + class _Launcher: def __init__( self, @@ -248,6 +265,29 @@ 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( + hosted_settings: Path, + 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, @@ -1121,10 +1161,87 @@ 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}, ] +@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, @@ -1733,3 +1850,132 @@ 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 == [] + + +@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) + + +@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() 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/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), ) 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"): diff --git a/plans/1079/slurm-external-providers.md b/plans/1079/slurm-external-providers.md new file mode 100644 index 000000000..352072e7e --- /dev/null +++ b/plans/1079/slurm-external-providers.md @@ -0,0 +1,59 @@ +# 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 + +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 + +- 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. +- 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. +- 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. +- Restrict the generated dataset directory before publishing it as a candidate, + and validate native DD outputs with the state's real artifact verifier. + +## 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. + +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.