diff --git a/docs/training.md b/docs/training.md index 0372fd68c1..0f253ca0d0 100644 --- a/docs/training.md +++ b/docs/training.md @@ -58,7 +58,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli |---|---| | `orchestrator.batch_size` | Tasks per trainer step. | | `orchestrator.group_size` | Rollouts generated per task. | -| `orchestrator.max_off_policy_steps` | How many distinct policies may have contributed to one rollout before it's discarded (default 8). The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. | +| `orchestrator.max_off_policy_steps` | On the vLLM admin backend, how many distinct policies may contribute to one rollout before it is discarded (default 8). Dynamo instead drains all live-policy/eval requests before mutating weights, so this setting does not apply to Dynamo runs. | | `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). | | `[[orchestrator.train.env]]` | Training environments. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Environments](configuration.md#environments-orchestratortrainenv). | | `[[orchestrator.eval.env]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). | diff --git a/k8s/README.md b/k8s/README.md index 01e251a85b..87f9586e58 100644 --- a/k8s/README.md +++ b/k8s/README.md @@ -59,6 +59,34 @@ helm install my-exp ./prime-rl \ --set config.secrets.name=prime-rl-secrets ``` +### Generated DynamoGraph deployments + +`dynamo-dgd` produces internally consistent Helm values for Dynamo disaggregated inference. The source SHAs, image tag, and image digest are caller-supplied evidence: the renderer checks that they agree, but it does not authenticate source or image provenance. Verify those inputs through the release process before rendering and protect the generated values from modification. + +Chart-managed mode requires explicit runnable controller commands; the renderer supplies safe execution defaults of one replica and `autoStart: true` and refuses to emit a sleeper deployment: + +```bash +uv run dynamo-dgd inference.toml \ + --release-name my-exp \ + --namespace my-namespace \ + --image "$IMAGE@$IMAGE_DIGEST" \ + --image-digest "$IMAGE_DIGEST" \ + --prime-sha "$PRIME_SHA" \ + --dynamo-sha "$DYNAMO_SHA" \ + --output-dir ./artifacts \ + --gpu-architecture arm64 \ + --gpu-product NVIDIA-GB200 \ + --gpu-node-pool prime-gpu \ + --orchestrator-command 'uv run orchestrator @ /app/configs/debug/orch.toml --output-dir /data/outputs' \ + --trainer-command 'uv run trainer @ /app/configs/debug/rl/train.toml --output-dir /data/outputs' + +helm install my-exp ./prime-rl \ + --namespace my-namespace \ + -f ./artifacts/dynamo-helm-values.json +``` + +Use `--external-controller` when another system runs orchestration and training; generated values then render no controller StatefulSets. Generated releases are limited to 41 characters so every derived Kubernetes Service name remains valid. The image tag must contain the first 12 characters of both caller-supplied source SHAs, and the digest must match `--image-digest`. + ## Uninstalling ```bash diff --git a/k8s/prime-rl/templates/_helpers.tpl b/k8s/prime-rl/templates/_helpers.tpl index cd6d7183d5..738aed204d 100644 --- a/k8s/prime-rl/templates/_helpers.tpl +++ b/k8s/prime-rl/templates/_helpers.tpl @@ -5,6 +5,42 @@ Expand the name of the chart. {{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }} {{- end }} +{{/* +Resolve the immutable image reference generated for DGD, with the native chart +repository/tag remaining as the fallback for statefulset mode. +*/}} +{{- define "prime-rl.image" -}} +{{- if .Values.image.reference -}} +{{- .Values.image.reference -}} +{{- else -}} +{{- printf "%s:%s" .Values.image.repository .Values.image.tag -}} +{{- end -}} +{{- end }} + +{{/* +Reuse a supplied shared claim or derive the chart-managed claim name. +*/}} +{{- define "prime-rl.storageClaimName" -}} +{{- default (printf "%s-shared-data" .Release.Name) .Values.storage.existingClaim -}} +{{- end }} + +{{- define "prime-rl.inferenceUrls" -}} +{{- if eq .Values.inference.mode "dynamoGraph" -}} +{{- printf "http://%s-frontend.%s.svc.cluster.local:8000/v1" .Release.Name .Values.namespace -}} +{{- else -}} +{{- $releaseName := .Release.Name -}} +{{- $namespace := .Values.namespace -}} +{{- $port := int .Values.inference.service.port -}} +{{- $replicas := int .Values.inference.replicas -}} +{{- $urls := list -}} +{{- range $i := until $replicas -}} +{{- $url := printf "http://%s-inference-%d.%s-inference-headless.%s.svc.cluster.local:%d/v1" $releaseName $i $releaseName $namespace $port -}} +{{- $urls = append $urls $url -}} +{{- end -}} +{{- $urls | join "," -}} +{{- end -}} +{{- end }} + {{/* Create a default fully qualified app name. */}} diff --git a/k8s/prime-rl/templates/deployment.yaml b/k8s/prime-rl/templates/deployment.yaml index 1ed16b9f5e..3f433a41c5 100644 --- a/k8s/prime-rl/templates/deployment.yaml +++ b/k8s/prime-rl/templates/deployment.yaml @@ -1,4 +1,15 @@ +{{- $dgdWorkload := dict }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- $workloadBinding := required "inference.dynamoGraph.workloadBinding is required" .Values.inference.dynamoGraph.workloadBinding }} +{{- $workloadCanonical := required "inference.dynamoGraph.workloadBinding.canonical is required" $workloadBinding.canonical }} +{{- $dgdWorkload = mustFromJson $workloadCanonical }} +{{- end }} {{- if .Values.orchestrator.enabled }} +{{- $orchestratorPlacement := .Values.orchestrator }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- $orchestratorWorkload := required "canonical orchestrator workload is required" (index $dgdWorkload "orchestrator") }} +{{- $orchestratorPlacement = required "canonical orchestrator placement is required" (index $orchestratorWorkload "placement") }} +{{- end }} apiVersion: apps/v1 kind: StatefulSet metadata: @@ -22,18 +33,31 @@ spec: {{- include "prime-rl.componentLabels" . | nindent 8 }} role: orchestrator spec: - {{- with .Values.orchestrator.nodeSelector }} + {{- with .Values.image.pullSecrets }} + imagePullSecrets: + {{- range . }} + - name: {{ . | quote }} + {{- end }} + {{- end }} + {{- with $orchestratorPlacement.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} + {{- if $orchestratorPlacement.runtimeClassName }} + runtimeClassName: {{ $orchestratorPlacement.runtimeClassName }} + {{- end }} + {{- with $orchestratorPlacement.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: prime-rl-orchestrator - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + image: {{ include "prime-rl.image" . | quote }} imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.orchestrator.autoStart }} command: ["/bin/bash", "-c"] args: - - {{ .Values.orchestrator.command }} + - {{ .Values.orchestrator.command | quote }} {{- else }} command: ["sleep", "infinity"] {{- end }} @@ -62,16 +86,17 @@ spec: - name: HEADLESS_SERVICE value: "{{ .Release.Name }}-orchestrator-headless.{{ .Values.namespace }}.svc.cluster.local" - name: INFERENCE_URL - {{- $releaseName := .Release.Name }} - {{- $namespace := .Values.namespace }} - {{- $port := int .Values.inference.service.port }} - {{- $replicas := int .Values.inference.replicas }} - {{- $urls := list }} - {{- range $i := until $replicas }} - {{- $url := printf "http://%s-inference-%d.%s-inference-headless.%s.svc.cluster.local:%d/v1" $releaseName $i $releaseName $namespace $port }} - {{- $urls = append $urls $url }} - {{- end }} - value: {{ $urls | join "," | quote }} + value: {{ include "prime-rl.inferenceUrls" . | quote }} + {{- if .Values.modelCache.enabled }} + - name: HF_HOME + value: {{ .Values.modelCache.mountPath | quote }} + {{- end }} + {{- if eq .Values.inference.mode "dynamoGraph" }} + - name: DYN_RL_DISCOVERY_URL + value: "http://{{ .Release.Name }}-frontend-rl.{{ .Values.namespace }}.svc.cluster.local:8001" + - name: DYN_RL_TOPOLOGY + value: {{ required "inference.dynamoGraph.clientTopology is required" .Values.inference.dynamoGraph.clientTopology | toJson | quote }} + {{- end }} {{- with .Values.orchestrator.env }} {{- toYaml . | nindent 8 }} {{- end }} @@ -82,6 +107,15 @@ spec: name: {{ .Values.config.secrets.name }} key: wandb-api-key optional: true + {{- end }} + {{- if .Values.huggingFace.tokenSecretName }} + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.huggingFace.tokenSecretName }} + key: {{ .Values.huggingFace.tokenSecretKey }} + optional: false + {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: secretKeyRef: @@ -91,20 +125,33 @@ spec: {{- end }} resources: {{- toYaml .Values.orchestrator.resources | nindent 10 }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + mountPath: {{ .Values.modelCache.mountPath }} + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ include "prime-rl.storageClaimName" . }} + {{- end }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + persistentVolumeClaim: + claimName: {{ required "modelCache.existingClaim is required when modelCache.enabled" .Values.modelCache.existingClaim }} + {{- end }} {{- end }} {{- end }} --- -{{- if .Values.inference.enabled }} +{{- if and .Values.inference.enabled (eq .Values.inference.mode "statefulset") }} apiVersion: apps/v1 kind: StatefulSet metadata: @@ -129,12 +176,18 @@ spec: {{- include "prime-rl.componentLabels" . | nindent 8 }} role: inference spec: + {{- with .Values.image.pullSecrets }} + imagePullSecrets: + {{- range . }} + - name: {{ . | quote }} + {{- end }} + {{- end }} {{- if .Values.inference.runtimeClassName }} runtimeClassName: {{ .Values.inference.runtimeClassName }} {{- end }} containers: - name: prime-rl-inference - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + image: {{ include "prime-rl.image" . | quote }} imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.inference.autoStart }} command: ["/bin/bash", "-c"] @@ -165,6 +218,10 @@ spec: value: "{{ .Values.inference.replicas }}" - name: HEADLESS_SERVICE value: "{{ .Release.Name }}-inference-headless.{{ .Values.namespace }}.svc.cluster.local" + {{- if .Values.modelCache.enabled }} + - name: HF_HOME + value: {{ .Values.modelCache.mountPath | quote }} + {{- end }} {{- with .Values.inference.env }} {{- toYaml . | nindent 8 }} {{- end }} @@ -175,6 +232,15 @@ spec: name: {{ .Values.config.secrets.name }} key: wandb-api-key optional: true + {{- end }} + {{- if .Values.huggingFace.tokenSecretName }} + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.huggingFace.tokenSecretName }} + key: {{ .Values.huggingFace.tokenSecretKey }} + optional: false + {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: secretKeyRef: @@ -220,20 +286,38 @@ spec: failureThreshold: {{ .Values.inference.probes.readiness.failureThreshold }} timeoutSeconds: {{ .Values.inference.probes.readiness.timeoutSeconds }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + mountPath: {{ .Values.modelCache.mountPath }} + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ include "prime-rl.storageClaimName" . }} + {{- end }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + persistentVolumeClaim: + claimName: {{ required "modelCache.existingClaim is required when modelCache.enabled" .Values.modelCache.existingClaim }} + {{- end }} {{- end }} {{- end }} --- {{- if .Values.trainer.enabled }} +{{- $trainerPlacement := .Values.trainer }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- $trainerWorkload := required "canonical trainer workload is required" (index $dgdWorkload "trainer") }} +{{- $trainerPlacement = required "canonical trainer placement is required" (index $trainerWorkload "placement") }} +{{- end }} apiVersion: apps/v1 kind: StatefulSet metadata: @@ -258,17 +342,31 @@ spec: {{- include "prime-rl.componentLabels" . | nindent 8 }} role: trainer spec: - {{- if .Values.trainer.runtimeClassName }} - runtimeClassName: {{ .Values.trainer.runtimeClassName }} + {{- with .Values.image.pullSecrets }} + imagePullSecrets: + {{- range . }} + - name: {{ . | quote }} + {{- end }} + {{- end }} + {{- if $trainerPlacement.runtimeClassName }} + runtimeClassName: {{ $trainerPlacement.runtimeClassName }} + {{- end }} + {{- with $trainerPlacement.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with $trainerPlacement.tolerations }} + tolerations: + {{- toYaml . | nindent 8 }} {{- end }} containers: - name: prime-rl-trainer - image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + image: {{ include "prime-rl.image" . | quote }} imagePullPolicy: {{ .Values.image.pullPolicy }} {{- if .Values.trainer.autoStart }} command: ["/bin/bash", "-c"] args: - - {{ .Values.trainer.command }} + - {{ .Values.trainer.command | quote }} {{- else }} command: ["sleep", "infinity"] {{- end }} @@ -301,16 +399,17 @@ spec: value: {{ .Values.trainer.pytorchCudaAllocConf | quote }} {{- end }} - name: INFERENCE_URL - {{- $releaseName := .Release.Name }} - {{- $namespace := .Values.namespace }} - {{- $port := int .Values.inference.service.port }} - {{- $replicas := int .Values.inference.replicas }} - {{- $urls := list }} - {{- range $i := until $replicas }} - {{- $url := printf "http://%s-inference-%d.%s-inference-headless.%s.svc.cluster.local:%d/v1" $releaseName $i $releaseName $namespace $port }} - {{- $urls = append $urls $url }} - {{- end }} - value: {{ $urls | join "," | quote }} + value: {{ include "prime-rl.inferenceUrls" . | quote }} + {{- if .Values.modelCache.enabled }} + - name: HF_HOME + value: {{ .Values.modelCache.mountPath | quote }} + {{- end }} + {{- if eq .Values.inference.mode "dynamoGraph" }} + - name: DYN_RL_DISCOVERY_URL + value: "http://{{ .Release.Name }}-frontend-rl.{{ .Values.namespace }}.svc.cluster.local:8001" + - name: DYN_RL_TOPOLOGY + value: {{ required "inference.dynamoGraph.clientTopology is required" .Values.inference.dynamoGraph.clientTopology | toJson | quote }} + {{- end }} {{- with .Values.trainer.env }} {{- toYaml . | nindent 8 }} {{- end }} @@ -321,6 +420,15 @@ spec: name: {{ .Values.config.secrets.name }} key: wandb-api-key optional: true + {{- end }} + {{- if .Values.huggingFace.tokenSecretName }} + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: {{ .Values.huggingFace.tokenSecretName }} + key: {{ .Values.huggingFace.tokenSecretKey }} + optional: false + {{- else if .Values.config.secrets.enabled }} - name: HF_TOKEN valueFrom: secretKeyRef: @@ -366,15 +474,28 @@ spec: failureThreshold: {{ .Values.trainer.probes.readiness.failureThreshold }} timeoutSeconds: {{ .Values.trainer.probes.readiness.timeoutSeconds }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumeMounts: + {{- if .Values.storage.enabled }} - name: shared-data mountPath: {{ .Values.storage.mountPath }} {{- end }} - {{- if .Values.storage.enabled }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + mountPath: {{ .Values.modelCache.mountPath }} + {{- end }} + {{- end }} + {{- if or .Values.storage.enabled .Values.modelCache.enabled }} volumes: + {{- if .Values.storage.enabled }} - name: shared-data persistentVolumeClaim: - claimName: {{ .Release.Name }}-shared-data + claimName: {{ include "prime-rl.storageClaimName" . }} + {{- end }} + {{- if .Values.modelCache.enabled }} + - name: model-cache + persistentVolumeClaim: + claimName: {{ required "modelCache.existingClaim is required when modelCache.enabled" .Values.modelCache.existingClaim }} + {{- end }} {{- end }} {{- end }} diff --git a/k8s/prime-rl/templates/dynamo-engine-config.yaml b/k8s/prime-rl/templates/dynamo-engine-config.yaml new file mode 100644 index 0000000000..3daf657552 --- /dev/null +++ b/k8s/prime-rl/templates/dynamo-engine-config.yaml @@ -0,0 +1,17 @@ +{{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} +apiVersion: v1 +kind: ConfigMap +metadata: + name: {{ required "inference.dynamoGraph.engineConfig.name is required" .Values.inference.dynamoGraph.engineConfig.name }} + namespace: {{ .Values.namespace }} + labels: + {{- include "prime-rl.labels" . | nindent 4 }} + role: inference + annotations: + {{- toYaml .Values.inference.dynamoGraph.engineConfig.annotations | nindent 4 }} +immutable: true +data: + {{- range $name, $content := .Values.inference.dynamoGraph.engineConfig.data }} + {{ $name }}: {{ $content | toJson }} + {{- end }} +{{- end }} diff --git a/k8s/prime-rl/templates/dynamo-graph-deployment.yaml b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml new file mode 100644 index 0000000000..5e3fc16269 --- /dev/null +++ b/k8s/prime-rl/templates/dynamo-graph-deployment.yaml @@ -0,0 +1,330 @@ +{{- $candidateGraph := default (dict) .Values.inference.dynamoGraph }} +{{- if and (ne .Values.inference.mode "dynamoGraph") (gt (len $candidateGraph) 0) }} +{{- fail "generated DynamoGraph contract requires dynamoGraph mode" }} +{{- end }} +{{- if eq .Values.inference.mode "dynamoGraph" }} +{{- if not (kindIs "bool" .Values.inference.enabled) }} +{{- fail "inference.enabled must be boolean true in dynamoGraph mode" }} +{{- end }} +{{- if not .Values.inference.enabled }} +{{- fail "inference.enabled must be boolean true in dynamoGraph mode" }} +{{- end }} +{{- $graph := required "inference.dynamoGraph is required" .Values.inference.dynamoGraph }} +{{- $resource := required "inference.dynamoGraph.resource is required" $graph.resource }} +{{- $resourceName := required "inference.dynamoGraph.resource.metadata.name is required" $resource.metadata.name }} +{{- $resourceNamespace := required "inference.dynamoGraph.resource.metadata.namespace is required" $resource.metadata.namespace }} +{{- if ne .Release.Name $resourceName }} +{{- fail (printf "Helm Release.Name %q must match embedded DynamoGraphDeployment metadata.name %q" .Release.Name $resourceName) }} +{{- end }} +{{- if ne .Values.namespace $resourceNamespace }} +{{- fail (printf "values namespace %q must match embedded DynamoGraphDeployment metadata.namespace %q" .Values.namespace $resourceNamespace) }} +{{- end }} +{{- if ne .Release.Namespace $resourceNamespace }} +{{- fail (printf "Helm Release.Namespace %q must match values and embedded DynamoGraphDeployment namespace %q" .Release.Namespace $resourceNamespace) }} +{{- end }} +{{- $topology := required "inference.dynamoGraph.clientTopology is required" $graph.clientTopology }} +{{- $baseURLs := required "inference.dynamoGraph.clientTopology.base_url is required" $topology.base_url }} +{{- $rlBaseURLs := required "inference.dynamoGraph.clientTopology.rl_base_url is required" $topology.rl_base_url }} +{{- if or (ne (len $baseURLs) 1) (ne (index $baseURLs 0) (printf "http://%s-frontend.%s.svc.cluster.local:8000/v1" .Release.Name .Values.namespace)) }} +{{- fail "inference.dynamoGraph.clientTopology.base_url must match the generated release frontend URL" }} +{{- end }} +{{- if or (ne (len $rlBaseURLs) 1) (ne (index $rlBaseURLs 0) (printf "http://%s-frontend-rl.%s.svc.cluster.local:8001" .Release.Name .Values.namespace)) }} +{{- fail "inference.dynamoGraph.clientTopology.rl_base_url must match the generated release RL discovery URL" }} +{{- end }} +{{- $engineConfig := required "inference.dynamoGraph.engineConfig is required" $graph.engineConfig }} +{{- $engineCanonical := required "inference.dynamoGraph.engineConfig.canonicalData is required" $engineConfig.canonicalData }} +{{- $engineHash := required "inference.dynamoGraph.engineConfig.sha256 is required" $engineConfig.sha256 }} +{{- $computedEngineHash := sha256sum $engineCanonical }} +{{- if ne $engineHash $computedEngineHash }} +{{- fail "inference.dynamoGraph.engineConfig.sha256 must match its canonical payload" }} +{{- end }} +{{- $canonicalEngineData := mustFromJson $engineCanonical }} +{{- if ne (toJson $canonicalEngineData) (toJson $engineConfig.data) }} +{{- fail "inference.dynamoGraph.engineConfig.data must match its canonical payload" }} +{{- end }} +{{- $expectedConfigName := printf "%s-dynamo-engine-%s" .Release.Name (trunc 12 $computedEngineHash) }} +{{- if ne (required "inference.dynamoGraph.engineConfig.name is required" $engineConfig.name) $expectedConfigName }} +{{- fail "inference.dynamoGraph.engineConfig.name must be content-addressed by engineConfig.sha256" }} +{{- end }} +{{- $resourceAnnotations := required "DynamoGraphDeployment metadata.annotations are required" $resource.metadata.annotations }} +{{- if ne (required "DynamoGraphDeployment config-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/config-sha256")) $engineHash }} +{{- fail "DynamoGraphDeployment config-sha256 must match engineConfig.sha256" }} +{{- end }} +{{- $engineAnnotations := required "inference.dynamoGraph.engineConfig.annotations are required" $engineConfig.annotations }} +{{- if ne (required "engineConfig config-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/config-sha256")) $engineHash }} +{{- fail "engineConfig config-sha256 annotation must match engineConfig.sha256" }} +{{- end }} +{{- $topologyBinding := required "inference.dynamoGraph.topologyBinding is required" $graph.topologyBinding }} +{{- $topologyCanonical := required "inference.dynamoGraph.topologyBinding.canonical is required" $topologyBinding.canonical }} +{{- $topologyHash := required "inference.dynamoGraph.topologyBinding.sha256 is required" $topologyBinding.sha256 }} +{{- if ne (sha256sum $topologyCanonical) $topologyHash }} +{{- fail "topology binding sha256 must match its canonical payload" }} +{{- end }} +{{- if ne (required "DynamoGraphDeployment topology-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/topology-sha256")) $topologyHash }} +{{- fail "DynamoGraphDeployment topology binding annotation must match topologyBinding.sha256" }} +{{- end }} +{{- $canonicalTopology := mustFromJson $topologyCanonical }} +{{- if ne (toJson (index $canonicalTopology "clientTopology")) (toJson $topology) }} +{{- fail "clientTopology must match the renderer topology binding" }} +{{- end }} +{{- $services := required "DynamoGraphDeployment services are required" $resource.spec.services }} +{{- $boundWorkers := required "topology binding workerServices are required" (index $canonicalTopology "workerServices") }} +{{- if or (ne (len $services) 3) (ne (len $boundWorkers) 2) }} +{{- fail "DynamoGraphDeployment services must match the renderer topology binding" }} +{{- end }} +{{- range $serviceName, $expectedWorker := $boundWorkers }} +{{- $service := required (printf "topology binding service %s is required" $serviceName) (index $services $serviceName) }} +{{- if ne (required (printf "%s componentType is required" $serviceName) $service.componentType) "worker" }} +{{- fail (printf "DynamoGraphDeployment service %s must be a worker in the renderer topology binding" $serviceName) }} +{{- end }} +{{- $actualWorker := dict + "role" (required (printf "%s subComponentType is required" $serviceName) $service.subComponentType) + "replicas" (required (printf "%s replicas is required" $serviceName) $service.replicas) + "requestsGpu" (required (printf "%s GPU request is required" $serviceName) $service.resources.requests.gpu) + "limitsGpu" (required (printf "%s GPU limit is required" $serviceName) $service.resources.limits.gpu) }} +{{- if ne (toJson $expectedWorker) (toJson $actualWorker) }} +{{- fail (printf "DynamoGraphDeployment service %s must match the renderer topology binding" $serviceName) }} +{{- end }} +{{- end }} +{{- $prefill := required "DynamoGraphDeployment VllmPrefillWorker is required" (index $services "VllmPrefillWorker") }} +{{- $decode := required "DynamoGraphDeployment VllmDecodeWorker is required" (index $services "VllmDecodeWorker") }} +{{- $expectedRoles := list }} +{{- range $_ := until (int (required "VllmPrefillWorker replicas is required" $prefill.replicas)) }} +{{- $expectedRoles = append $expectedRoles "prefill" }} +{{- end }} +{{- range $_ := until (int (required "VllmDecodeWorker replicas is required" $decode.replicas)) }} +{{- $expectedRoles = append $expectedRoles "decode" }} +{{- end }} +{{- if ne (toJson $expectedRoles) (toJson (required "clientTopology.dynamo_worker_roles is required" $topology.dynamo_worker_roles)) }} +{{- fail "clientTopology.dynamo_worker_roles must match the DGD worker topology binding" }} +{{- end }} +{{- $gpusPerWorker := printf "%v" (required "clientTopology.dynamo_gpus_per_worker is required" $topology.dynamo_gpus_per_worker) }} +{{- range $serviceName, $service := dict "VllmDecodeWorker" $decode "VllmPrefillWorker" $prefill }} +{{- if or (ne (printf "%v" $service.resources.requests.gpu) $gpusPerWorker) (ne (printf "%v" $service.resources.limits.gpu) $gpusPerWorker) }} +{{- fail (printf "DynamoGraphDeployment service %s GPU resources must match the client topology binding" $serviceName) }} +{{- end }} +{{- end }} +{{- $workloadBinding := required "inference.dynamoGraph.workloadBinding is required" $graph.workloadBinding }} +{{- $workloadCanonical := required "inference.dynamoGraph.workloadBinding.canonical is required" $workloadBinding.canonical }} +{{- $workloadHash := required "inference.dynamoGraph.workloadBinding.sha256 is required" $workloadBinding.sha256 }} +{{- if ne (sha256sum $workloadCanonical) $workloadHash }} +{{- fail "workload binding sha256 must match its canonical payload" }} +{{- end }} +{{- if ne (required "DynamoGraphDeployment workload-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/workload-sha256")) $workloadHash }} +{{- fail "DynamoGraphDeployment workload binding annotation must match workloadBinding.sha256" }} +{{- end }} +{{- $workload := mustFromJson $workloadCanonical }} +{{- $workloadKeys := list "config" "controllerMode" "huggingFace" "image" "modelCache" "orchestrator" "storage" "trainer" }} +{{- if ne (len $workload) (len $workloadKeys) }} +{{- fail "workload binding must contain the complete chart runtime contract" }} +{{- end }} +{{- range $key := $workloadKeys }} +{{- if not (hasKey $workload $key) }} +{{- fail (printf "workload binding is missing %s" $key) }} +{{- end }} +{{- end }} +{{- $imageWorkload := index $workload "image" }} +{{- $actualImage := dict "reference" .Values.image.reference "pullPolicy" .Values.image.pullPolicy "pullSecrets" .Values.image.pullSecrets }} +{{- if ne (toJson $imageWorkload) (toJson $actualImage) }} +{{- fail "image configuration must match the workload binding" }} +{{- end }} +{{- $configWorkload := index $workload "config" }} +{{- if ne (toJson $configWorkload) (toJson .Values.config) }} +{{- fail "config configuration must match the workload binding" }} +{{- end }} +{{- $huggingFaceWorkload := index $workload "huggingFace" }} +{{- if ne (toJson $huggingFaceWorkload) (toJson .Values.huggingFace) }} +{{- fail "huggingFace configuration must match the workload binding" }} +{{- end }} +{{- $modelCacheWorkload := index $workload "modelCache" }} +{{- if ne (toJson $modelCacheWorkload) (toJson .Values.modelCache) }} +{{- fail "modelCache configuration must match the workload binding" }} +{{- end }} +{{- $controllerMode := required "canonical controllerMode is required" (index $workload "controllerMode") }} +{{- if and (ne $controllerMode "chartManaged") (ne $controllerMode "external") }} +{{- fail "canonical controllerMode must be chartManaged or external" }} +{{- end }} +{{- if ne (required "inference.dynamoGraph.controllerMode is required" $graph.controllerMode) $controllerMode }} +{{- fail "inference.dynamoGraph.controllerMode must match the workload binding" }} +{{- end }} +{{- $frontend := required "DynamoGraphDeployment Frontend is required" (index $services "Frontend") }} +{{- $frontendPod := required "DynamoGraphDeployment Frontend extraPodSpec is required" $frontend.extraPodSpec }} +{{- $orchestratorWorkload := required "canonical orchestrator workload is required" (index $workload "orchestrator") }} +{{- $trainerWorkload := required "canonical trainer workload is required" (index $workload "trainer") }} +{{- $storageWorkload := required "canonical storage workload is required" (index $workload "storage") }} +{{- $actualOrchestrator := dict + "enabled" .Values.orchestrator.enabled + "replicas" .Values.orchestrator.replicas + "autoStart" .Values.orchestrator.autoStart + "command" .Values.orchestrator.command + "resources" .Values.orchestrator.resources + "service" .Values.orchestrator.service + "env" .Values.orchestrator.env }} +{{- if ne (toJson (omit $orchestratorWorkload "gpu" "placement")) (toJson $actualOrchestrator) }} +{{- fail "orchestrator configuration must match the workload binding" }} +{{- end }} +{{- $actualTrainer := dict + "enabled" .Values.trainer.enabled + "replicas" .Values.trainer.replicas + "autoStart" .Values.trainer.autoStart + "command" .Values.trainer.command + "gpu" .Values.trainer.gpu + "pytorchCudaAllocConf" .Values.trainer.pytorchCudaAllocConf + "resources" .Values.trainer.resources + "service" .Values.trainer.service + "env" .Values.trainer.env + "probes" .Values.trainer.probes }} +{{- if ne (toJson (omit $trainerWorkload "placement")) (toJson $actualTrainer) }} +{{- fail "trainer configuration must match the workload binding" }} +{{- end }} +{{- if ne (toJson $storageWorkload) (toJson .Values.storage) }} +{{- fail "storage configuration must match the workload binding" }} +{{- end }} +{{- if or (not (hasKey $orchestratorWorkload "enabled")) (not (hasKey $trainerWorkload "enabled")) }} +{{- fail "canonical controller workloads must define enabled" }} +{{- end }} +{{- $orchestratorEnabled := index $orchestratorWorkload "enabled" }} +{{- $trainerEnabled := index $trainerWorkload "enabled" }} +{{- if ne $orchestratorEnabled .Values.orchestrator.enabled }} +{{- fail "orchestrator.enabled must match the workload binding" }} +{{- end }} +{{- if ne $trainerEnabled .Values.trainer.enabled }} +{{- fail "trainer.enabled must match the workload binding" }} +{{- end }} +{{- $orchestratorGPU := required "canonical orchestrator gpu contract is required" (index $orchestratorWorkload "gpu") }} +{{- if not (hasKey $orchestratorGPU "enabled") }} +{{- fail "canonical orchestrator gpu contract must define enabled" }} +{{- end }} +{{- if index $orchestratorGPU "enabled" }} +{{- fail "canonical orchestrator GPU capability must be disabled" }} +{{- end }} +{{- range $componentName, $resources := dict "orchestrator" .Values.orchestrator.resources "trainer" .Values.trainer.resources }} +{{- range $resourceScope, $entries := dict "limits" (default (dict) $resources.limits) "requests" (default (dict) $resources.requests) }} +{{- range $resourceName, $_ := $entries }} +{{- if hasPrefix "nvidia.com/" $resourceName }} +{{- fail (printf "%s resources cannot set NVIDIA extended resource %s in %s" $componentName $resourceName $resourceScope) }} +{{- end }} +{{- end }} +{{- end }} +{{- end }} +{{- $trainerGPU := required "canonical trainer gpu contract is required" (index $trainerWorkload "gpu") }} +{{- if or (not (hasKey $trainerGPU "enabled")) (not (hasKey $trainerGPU "count")) }} +{{- fail "canonical trainer gpu contract must define enabled and count" }} +{{- end }} +{{- $trainerGPUEnabled := index $trainerGPU "enabled" }} +{{- $trainerGPUCount := int (index $trainerGPU "count") }} +{{- if eq $controllerMode "chartManaged" }} +{{- if or (not $orchestratorEnabled) (not $trainerEnabled) (not $trainerGPUEnabled) (lt $trainerGPUCount 1) }} +{{- fail "chartManaged mode requires orchestrator, trainer, and trainer GPU capability" }} +{{- end }} +{{- range $componentName, $canonical := dict "orchestrator" $orchestratorWorkload "trainer" $trainerWorkload }} +{{- $command := index $canonical "command" }} +{{- if or (lt (int (index $canonical "replicas")) 1) (not (index $canonical "autoStart")) (not (regexMatch (printf "^uv[[:space:]]+run[[:space:]]+%s([[:space:]]|$)" $componentName) $command)) }} +{{- fail (printf "chartManaged %s execution requires positive replicas, autoStart=true, and an executable uv run %s command" $componentName $componentName) }} +{{- end }} +{{- end }} +{{- else }} +{{- if or $orchestratorEnabled $trainerEnabled $trainerGPUEnabled (ne $trainerGPUCount 0) }} +{{- fail "external mode forbids chart-managed controller workloads and trainer GPU capability" }} +{{- end }} +{{- range $componentName, $canonical := dict "orchestrator" $orchestratorWorkload "trainer" $trainerWorkload }} +{{- if or (ne (int (index $canonical "replicas")) 0) (index $canonical "autoStart") (not (empty (index $canonical "command"))) }} +{{- fail (printf "external mode requires zeroed %s execution" $componentName) }} +{{- end }} +{{- end }} +{{- end }} +{{- if or (not (hasKey $storageWorkload "enabled")) (not (hasKey $storageWorkload "existingClaim")) (not (hasKey $storageWorkload "mountPath")) }} +{{- fail "canonical storage workload must define enabled, existingClaim, and mountPath" }} +{{- end }} +{{- $storageEnabled := index $storageWorkload "enabled" }} +{{- $storageClaim := index $storageWorkload "existingClaim" }} +{{- if and (eq $controllerMode "chartManaged") (not $storageEnabled) }} +{{- fail "chartManaged mode requires chart storage" }} +{{- end }} +{{- if and (eq $controllerMode "external") $storageEnabled (empty $storageClaim) }} +{{- fail "external mode requires an existing claim when chart storage is enabled" }} +{{- end }} +{{- $orchestratorPlacement := required "canonical orchestrator placement is required" (index $orchestratorWorkload "placement") }} +{{- $expectedOrchestratorPlacement := dict + "nodeSelector" (required "Frontend nodeSelector is required" $frontendPod.nodeSelector) + "tolerations" (required "Frontend tolerations are required" $frontendPod.tolerations) }} +{{- if ne (toJson $orchestratorPlacement) (toJson $expectedOrchestratorPlacement) }} +{{- fail "orchestrator placement must match the manifest-bound DGD frontend placement" }} +{{- end }} +{{- $trainerPlacement := required "canonical trainer placement is required" (index $trainerWorkload "placement") }} +{{- range $serviceName, $service := dict "VllmDecodeWorker" $decode "VllmPrefillWorker" $prefill }} +{{- $workerPod := required (printf "%s extraPodSpec is required" $serviceName) $service.extraPodSpec }} +{{- $expectedTrainerPlacement := dict + "nodeSelector" (required (printf "%s nodeSelector is required" $serviceName) $workerPod.nodeSelector) + "tolerations" (required (printf "%s tolerations are required" $serviceName) $workerPod.tolerations) }} +{{- if hasKey $workerPod "runtimeClassName" }} +{{- $_ := set $expectedTrainerPlacement "runtimeClassName" (required (printf "%s runtimeClassName must not be empty" $serviceName) $workerPod.runtimeClassName) }} +{{- end }} +{{- if ne (toJson $trainerPlacement) (toJson $expectedTrainerPlacement) }} +{{- fail "trainer placement must match every manifest-bound DGD worker placement" }} +{{- end }} +{{- end }} +{{- $frontendLabels := required "DynamoGraphDeployment Frontend extraPodMetadata.labels is required" $frontend.extraPodMetadata.labels }} +{{- $expectedFrontendLabels := dict "app.kubernetes.io/name" (include "prime-rl.name" .) "app.kubernetes.io/instance" .Release.Name }} +{{- if ne (toJson $frontendLabels) (toJson $expectedFrontendLabels) }} +{{- fail "DynamoGraphDeployment Frontend labels must match the chart selector labels" }} +{{- end }} +{{- $manifestCanonical := required "inference.dynamoGraph.manifestCanonical is required" $graph.manifestCanonical }} +{{- $manifestHash := required "DynamoGraphDeployment manifest-sha256 annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/manifest-sha256") }} +{{- if ne (sha256sum $manifestCanonical) $manifestHash }} +{{- fail "DynamoGraphDeployment manifest-sha256 must match its canonical payload" }} +{{- end }} +{{- if ne (required "engineConfig manifest-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/manifest-sha256")) $manifestHash }} +{{- fail "engineConfig manifest-sha256 annotation must match the DynamoGraphDeployment manifest" }} +{{- end }} +{{- if ne (required "engineConfig topology-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/topology-sha256")) $topologyHash }} +{{- fail "engineConfig topology-sha256 annotation must match topologyBinding.sha256" }} +{{- end }} +{{- if ne (required "engineConfig workload-sha256 annotation is required" (index $engineAnnotations "prime-rl.nvidia.com/workload-sha256")) $workloadHash }} +{{- fail "engineConfig workload-sha256 annotation must match workloadBinding.sha256" }} +{{- end }} +{{- if ne (toJson $engineAnnotations) (toJson $resourceAnnotations) }} +{{- fail "engineConfig annotations must match DynamoGraphDeployment identity annotations" }} +{{- end }} +{{- $canonicalResource := mustFromJson $manifestCanonical }} +{{- $scopedResource := deepCopy $resource }} +{{- $_ := unset $scopedResource.metadata.annotations "prime-rl.nvidia.com/manifest-sha256" }} +{{- if ne (toJson $canonicalResource) (toJson $scopedResource) }} +{{- fail "DynamoGraphDeployment resource must match its canonical manifest payload" }} +{{- end }} +{{- $protectedEnv := dict "DYN_RL_DISCOVERY_URL" true "DYN_RL_TOPOLOGY" true "INFERENCE_URL" true "HF_HOME" true "HF_TOKEN" true "HUGGING_FACE_HUB_TOKEN" true }} +{{- range $componentName, $component := dict "orchestrator" .Values.orchestrator "trainer" .Values.trainer }} +{{- range $entry := $component.env }} +{{- if hasKey $protectedEnv $entry.name }} +{{- fail (printf "%s.env cannot override generated %s" $componentName $entry.name) }} +{{- end }} +{{- end }} +{{- end }} +{{- $image := required "image.reference is required in dynamoGraph mode" .Values.image.reference }} +{{- if not (regexMatch "^.+@sha256:[0-9a-f]{64}$" $image) }} +{{- fail "image.reference must be pinned to a full sha256 digest in dynamoGraph mode" }} +{{- end }} +{{- $referenceDigest := regexFind "sha256:[0-9a-f]{64}$" $image }} +{{- $annotatedDigest := required "DynamoGraphDeployment image-digest annotation is required" (index $resourceAnnotations "prime-rl.nvidia.com/image-digest") }} +{{- if ne $referenceDigest $annotatedDigest }} +{{- fail "image.reference digest must match the DynamoGraphDeployment image-digest annotation" }} +{{- end }} +{{- $imageWithoutDigest := trimSuffix (printf "@%s" $referenceDigest) $image }} +{{- $imageTag := regexFind "[^/:]+$" $imageWithoutDigest }} +{{- range $label, $annotation := dict "Prime" "prime-rl.nvidia.com/prime-sha" "Dynamo" "prime-rl.nvidia.com/dynamo-sha" }} +{{- $sourceSHA := required (printf "DynamoGraphDeployment %s SHA annotation is required" $label) (index $resourceAnnotations $annotation) }} +{{- if not (regexMatch "^[0-9a-f]{40}$" $sourceSHA) }} +{{- fail (printf "%s SHA annotation must be a full 40-character Git commit SHA" $label) }} +{{- end }} +{{- if not (contains (substr 0 12 $sourceSHA) $imageTag) }} +{{- fail (printf "%s SHA annotation must match a 12-character commit suffix in the image tag" $label) }} +{{- end }} +{{- end }} +{{- range $serviceName, $service := $resource.spec.services }} +{{- $serviceImage := required (printf "DynamoGraphDeployment service %s image is required" $serviceName) $service.extraPodSpec.mainContainer.image }} +{{- if ne $image $serviceImage }} +{{- fail (printf "DynamoGraphDeployment service %s must use the same image.reference as orchestrator and trainer" $serviceName) }} +{{- end }} +{{- end }} +{{- toYaml $resource }} +{{- end }} diff --git a/k8s/prime-rl/templates/dynamo-rl-service.yaml b/k8s/prime-rl/templates/dynamo-rl-service.yaml new file mode 100644 index 0000000000..371f5a3d08 --- /dev/null +++ b/k8s/prime-rl/templates/dynamo-rl-service.yaml @@ -0,0 +1,23 @@ +{{- if and .Values.inference.enabled (eq .Values.inference.mode "dynamoGraph") }} +apiVersion: v1 +kind: Service +metadata: + name: {{ .Release.Name }}-frontend-rl + namespace: {{ .Values.namespace }} + labels: + {{- include "prime-rl.labels" . | nindent 4 }} + role: inference + annotations: + {{- toYaml .Values.inference.dynamoGraph.resource.metadata.annotations | nindent 4 }} +spec: + type: ClusterIP + selector: + nvidia.com/dynamo-graph-deployment-name: {{ .Release.Name }} + nvidia.com/dynamo-component: Frontend + nvidia.com/dynamo-component-type: frontend + ports: + - name: rl + protocol: TCP + port: 8001 + targetPort: rl +{{- end }} diff --git a/k8s/prime-rl/templates/pvc.yaml b/k8s/prime-rl/templates/pvc.yaml index 7afd6ff386..5bd740f7f6 100644 --- a/k8s/prime-rl/templates/pvc.yaml +++ b/k8s/prime-rl/templates/pvc.yaml @@ -1,4 +1,4 @@ -{{- if .Values.storage.enabled }} +{{- if and .Values.storage.enabled (not .Values.storage.existingClaim) }} apiVersion: v1 kind: PersistentVolumeClaim metadata: diff --git a/k8s/prime-rl/templates/release-validation.yaml b/k8s/prime-rl/templates/release-validation.yaml new file mode 100644 index 0000000000..3d6a4ba065 --- /dev/null +++ b/k8s/prime-rl/templates/release-validation.yaml @@ -0,0 +1,3 @@ +{{- if gt (len .Release.Name) 41 }} +{{- fail (printf "Helm release name %q must be at most 41 characters so generated Service names are valid" .Release.Name) }} +{{- end }} diff --git a/k8s/prime-rl/templates/service.yaml b/k8s/prime-rl/templates/service.yaml index b783a7816d..cc620ff187 100644 --- a/k8s/prime-rl/templates/service.yaml +++ b/k8s/prime-rl/templates/service.yaml @@ -12,6 +12,7 @@ metadata: spec: type: {{ .Values.orchestrator.service.type }} selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: orchestrator ports: @@ -37,6 +38,7 @@ metadata: spec: clusterIP: None # Headless service for StatefulSet selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: orchestrator ports: @@ -46,7 +48,7 @@ spec: name: nccl {{- end }} --- -{{- if .Values.inference.enabled }} +{{- if and .Values.inference.enabled (eq .Values.inference.mode "statefulset") }} {{- if .Values.inference.service.enabled }} apiVersion: v1 kind: Service @@ -60,6 +62,7 @@ metadata: spec: type: {{ .Values.inference.service.type }} selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: inference ports: @@ -84,6 +87,7 @@ metadata: spec: type: {{ .Values.trainer.service.type }} selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: trainer ports: @@ -98,7 +102,7 @@ spec: {{- end }} {{- end }} --- -{{- if .Values.inference.enabled }} +{{- if and .Values.inference.enabled (eq .Values.inference.mode "statefulset") }} apiVersion: v1 kind: Service metadata: @@ -111,6 +115,7 @@ metadata: spec: clusterIP: None # Headless service for StatefulSet selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: inference ports: @@ -131,6 +136,7 @@ metadata: spec: clusterIP: None # Headless service for StatefulSet selector: + {{- include "prime-rl.selectorLabels" . | nindent 4 }} {{- include "prime-rl.componentLabels" . | nindent 4 }} role: trainer ports: diff --git a/k8s/prime-rl/values.schema.json b/k8s/prime-rl/values.schema.json new file mode 100644 index 0000000000..e90aec8ede --- /dev/null +++ b/k8s/prime-rl/values.schema.json @@ -0,0 +1,390 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "image": { + "type": "object", + "properties": { + "reference": {"type": "string"}, + "pullPolicy": { + "type": "string", + "enum": ["Always", "IfNotPresent", "Never"] + }, + "pullSecrets": { + "type": "array", + "items": {"type": "string", "minLength": 1}, + "uniqueItems": true + } + } + }, + "inference": { + "type": "object", + "properties": { + "mode": { + "type": "string", + "enum": ["statefulset", "dynamoGraph"] + }, + "dynamoGraph": { + "type": "object", + "properties": { + "clientTopology": { + "type": "object", + "properties": { + "schema_version": {"const": 1}, + "admin_api": {"const": "dynamo"}, + "base_url": { + "type": "array", + "minItems": 1, + "items": {"type": "string", "minLength": 1} + }, + "rl_base_url": { + "type": "array", + "minItems": 1, + "items": {"type": "string", "minLength": 1} + }, + "dynamo_worker_roles": { + "type": "array", + "minItems": 1, + "items": {"type": "string", "enum": ["agg", "prefill", "decode"]} + }, + "dynamo_gpus_per_worker": {"type": "integer", "minimum": 1} + }, + "required": [ + "schema_version", + "admin_api", + "base_url", + "rl_base_url", + "dynamo_worker_roles", + "dynamo_gpus_per_worker" + ], + "additionalProperties": false + }, + "engineConfig": { + "type": "object", + "properties": { + "name": { + "type": "string", + "pattern": "^.+-dynamo-engine-[0-9a-f]{12}$" + }, + "sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "canonicalData": {"type": "string", "minLength": 1}, + "annotations": { + "type": "object", + "properties": { + "prime-rl.nvidia.com/config-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "prime-rl.nvidia.com/manifest-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "prime-rl.nvidia.com/topology-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "prime-rl.nvidia.com/workload-sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + } + }, + "required": [ + "prime-rl.nvidia.com/config-sha256", + "prime-rl.nvidia.com/manifest-sha256", + "prime-rl.nvidia.com/topology-sha256", + "prime-rl.nvidia.com/workload-sha256" + ], + "additionalProperties": {"type": "string"} + }, + "data": { + "type": "object", + "minProperties": 2, + "additionalProperties": {"type": "string"} + } + }, + "anyOf": [ + {"maxProperties": 0}, + {"required": ["name", "sha256", "canonicalData", "annotations", "data"]} + ], + "additionalProperties": false + }, + "topologyBinding": { + "type": "object", + "properties": { + "sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "canonical": {"type": "string", "minLength": 1} + }, + "required": ["sha256", "canonical"], + "additionalProperties": false + }, + "controllerMode": { + "type": "string", + "enum": ["chartManaged", "external"] + }, + "workloadBinding": { + "type": "object", + "properties": { + "sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + "canonical": {"type": "string", "minLength": 1} + }, + "required": ["sha256", "canonical"], + "additionalProperties": false + }, + "manifestCanonical": {"type": "string", "minLength": 1}, + "resource": {"type": "object"} + }, + "additionalProperties": false + } + } + }, + "storage": { + "type": "object", + "properties": { + "existingClaim": {"type": "string"} + } + }, + "modelCache": { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + "existingClaim": {"type": "string"}, + "mountPath": {"type": "string", "minLength": 1} + } + }, + "huggingFace": { + "type": "object", + "properties": { + "tokenSecretName": {"type": "string"}, + "tokenSecretKey": {"type": "string", "minLength": 1} + } + }, + "orchestrator": { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + "replicas": {"type": "integer", "minimum": 0}, + "autoStart": {"type": "boolean"}, + "command": {"type": "string"}, + "runtimeClassName": {"type": "string"}, + "nodeSelector": { + "type": "object", + "additionalProperties": {"type": "string"} + }, + "tolerations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "key": {"type": "string", "minLength": 1}, + "operator": {"type": "string", "enum": ["Exists", "Equal"]}, + "value": {"type": "string"}, + "effect": { + "type": "string", + "enum": ["NoSchedule", "PreferNoSchedule", "NoExecute"] + } + }, + "required": ["key", "operator"], + "additionalProperties": false + } + } + } + }, + "trainer": { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + "replicas": {"type": "integer", "minimum": 0}, + "autoStart": {"type": "boolean"}, + "command": {"type": "string"}, + "gpu": { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + "count": {"type": "integer", "minimum": 0} + }, + "required": ["enabled", "count"] + }, + "runtimeClassName": {"type": "string"}, + "nodeSelector": { + "type": "object", + "additionalProperties": {"type": "string"} + }, + "tolerations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "key": {"type": "string", "minLength": 1}, + "operator": {"type": "string", "enum": ["Exists", "Equal"]}, + "value": {"type": "string"}, + "effect": { + "type": "string", + "enum": ["NoSchedule", "PreferNoSchedule", "NoExecute"] + } + }, + "required": ["key", "operator"], + "additionalProperties": false + } + } + } + } + }, + "allOf": [ + { + "if": { + "properties": { + "inference": { + "properties": {"mode": {"const": "dynamoGraph"}}, + "required": ["mode"] + } + }, + "required": ["inference"] + }, + "then": { + "properties": { + "image": { + "properties": { + "reference": {"pattern": "^.+@sha256:[0-9a-f]{64}$"} + }, + "required": ["reference"] + }, + "inference": { + "properties": { + "enabled": {"const": true}, + "dynamoGraph": { + "properties": { + "engineConfig": {"minProperties": 5}, + "resource": {"minProperties": 1} + }, + "required": [ + "clientTopology", + "engineConfig", + "topologyBinding", + "controllerMode", + "workloadBinding", + "manifestCanonical", + "resource" + ] + } + }, + "required": ["enabled", "dynamoGraph"] + } + }, + "oneOf": [ + { + "properties": { + "inference": { + "properties": { + "dynamoGraph": { + "properties": {"controllerMode": {"const": "chartManaged"}} + } + } + }, + "orchestrator": { + "properties": { + "enabled": {"const": true}, + "replicas": {"minimum": 1}, + "autoStart": {"const": true}, + "command": { + "pattern": "^uv[ \\t]+run[ \\t]+orchestrator([ \\t]|$)" + } + }, + "required": ["enabled", "replicas", "autoStart", "command"] + }, + "trainer": { + "properties": { + "enabled": {"const": true}, + "replicas": {"minimum": 1}, + "autoStart": {"const": true}, + "command": { + "pattern": "^uv[ \\t]+run[ \\t]+trainer([ \\t]|$)" + }, + "gpu": { + "properties": { + "enabled": {"const": true}, + "count": {"minimum": 1} + } + } + }, + "required": ["enabled", "replicas", "autoStart", "command", "gpu"] + } + } + }, + { + "properties": { + "inference": { + "properties": { + "dynamoGraph": { + "properties": {"controllerMode": {"const": "external"}} + } + } + }, + "orchestrator": { + "properties": { + "enabled": {"const": false}, + "replicas": {"const": 0}, + "autoStart": {"const": false}, + "command": {"const": ""} + }, + "required": ["enabled", "replicas", "autoStart", "command"] + }, + "trainer": { + "properties": { + "enabled": {"const": false}, + "replicas": {"const": 0}, + "autoStart": {"const": false}, + "command": {"const": ""}, + "gpu": { + "properties": { + "enabled": {"const": false}, + "count": {"const": 0} + } + } + }, + "required": ["enabled", "replicas", "autoStart", "command", "gpu"] + } + } + } + ], + "required": ["image", "inference", "orchestrator", "trainer"] + }, + "else": { + "properties": { + "inference": { + "properties": { + "dynamoGraph": {"maxProperties": 0} + } + } + } + } + }, + { + "if": { + "properties": { + "modelCache": { + "properties": {"enabled": {"const": true}}, + "required": ["enabled"] + } + }, + "required": ["modelCache"] + }, + "then": { + "properties": { + "modelCache": { + "properties": {"existingClaim": {"minLength": 1}}, + "required": ["existingClaim"] + } + } + } + } + ] +} diff --git a/k8s/prime-rl/values.yaml b/k8s/prime-rl/values.yaml index 7fc3b48989..22f0d0f23e 100644 --- a/k8s/prime-rl/values.yaml +++ b/k8s/prime-rl/values.yaml @@ -10,17 +10,33 @@ image: repository: primeintellect/prime-rl pullPolicy: IfNotPresent tag: "main" + # Generated DGD values set the caller-supplied repository:tag@sha256 reference. + reference: "" + pullSecrets: [] # Shared storage configuration storage: enabled: true - # PVC name will be automatically set to {{ .Release.Name }}-shared-data + # Reuse an existing ReadWriteMany claim instead of creating a release-owned PVC. + existingClaim: "" + # A chart-managed PVC is named {{ .Release.Name }}-shared-data when existingClaim is empty. storageClassName: nfs accessModes: - ReadWriteMany size: 1Ti mountPath: /data +# Existing model cache mounted consistently into generated runtime pods. +modelCache: + enabled: false + existingClaim: "" + mountPath: /model-cache + +# Existing Hugging Face token secret. The secret value is never embedded in values. +huggingFace: + tokenSecretName: "" + tokenSecretKey: HF_TOKEN + # Orchestrator component orchestrator: enabled: true @@ -49,12 +65,18 @@ orchestrator: nodeSelector: {} # nvidia.com/gpu.present: "true" # Orchestrator doesn't need GPUs + runtimeClassName: "" + tolerations: [] # Inference component inference: enabled: true + mode: statefulset replicas: 1 + # Generated by `dynamo-dgd` when mode is dynamoGraph. + dynamoGraph: {} + # Auto-start configuration (set to false to use sleep infinity for debugging) autoStart: false command: "" # e.g., "uv run inference @ /app/examples/reverse_text/rl/infer.toml" @@ -123,6 +145,9 @@ trainer: runtimeClassName: nvidia + nodeSelector: {} + tolerations: [] + # Health probes for trainer (requires metrics_server config) probes: enabled: false diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 20ac441af0..26e12f528f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -24,7 +24,7 @@ class ServerConfig(BaseConfig): class ParallelConfig(BaseConfig): - tp: int = 1 + tp: int = Field(1, ge=1) """Tensor parallel size. Forwarded to vLLM as ``--tensor-parallel-size``.""" dp: int = Field(1, ge=1) @@ -251,7 +251,7 @@ def validate_scorers(self): class BaseInferenceDeploymentConfig(BaseConfig): - gpus_per_node: int = 8 + gpus_per_node: int = Field(8, ge=1) """GPUs per node.""" backend_port: int = 8100 @@ -304,10 +304,10 @@ class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig): """Extra environment variables exported only on decode nodes.""" prefill_vllm_overrides: dict[str, Any] = {} - """Extra vLLM config options merged into --vllm-extra only for prefill ranks (SLURM only).""" + """Extra vLLM config options merged into the resolved config only for prefill workers.""" decode_vllm_overrides: dict[str, Any] = {} - """Extra vLLM config options merged into --vllm-extra only for decode ranks (SLURM only).""" + """Extra vLLM config options merged into the resolved config only for decode workers.""" @property def num_prefill_nodes(self) -> int: @@ -328,7 +328,24 @@ def num_nodes(self) -> int: ] +class VllmInferenceBackendConfig(BaseConfig): + type: Literal["vllm"] = "vllm" + + +class DynamoInferenceBackendConfig(BaseConfig): + type: Literal["dynamo"] = "dynamo" + + +InferenceBackendConfig: TypeAlias = Annotated[ + VllmInferenceBackendConfig | DynamoInferenceBackendConfig, + Field(discriminator="type"), +] + + class InferenceConfig(BaseConfig): + backend: InferenceBackendConfig = VllmInferenceBackendConfig() + """Serving backend. Existing configs default to Prime's native vLLM launcher.""" + server: ServerConfig = ServerConfig() model: ModelConfig = Field(default_factory=ModelConfig) @@ -426,12 +443,81 @@ class InferenceConfig(BaseConfig): dry_run: bool = False """Only validate and dump resolved configs, then exit early.""" + @property + def dynamo_worker_roles(self) -> tuple[Literal["agg", "prefill", "decode"], ...]: + """Canonical admin and launch order for Dynamo worker groups.""" + if self.deployment.type == "disaggregated": + return ("prefill",) * self.deployment.num_prefill_replicas + ( + "decode", + ) * self.deployment.num_decode_replicas + return ("agg",) + + @property + def dynamo_gpus_per_worker(self) -> int: + """GPU allocation owned by one independently administered Dynamo worker.""" + if self.deployment.type == "disaggregated": + return self.deployment.gpus_per_node + return self.parallel.tp * self.parallel.dp + + @property + def dynamo_local_dp(self) -> int: + """vLLM data-parallel ranks inside one Dynamo worker.""" + if self.deployment.type == "disaggregated": + return self.deployment.gpus_per_node // self.parallel.tp + return self.parallel.dp + @model_validator(mode="after") def validate_multi_node_requires_slurm(self): - if self.deployment.type in ("multi_node", "disaggregated") and self.slurm is None: + if self.deployment.type == "multi_node" and self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if self.deployment.type == "disaggregated" and self.slurm is None and self.backend.type != "dynamo": raise ValueError("Must use SLURM for multi-node / disaggregated deployment.") return self + @model_validator(mode="after") + def validate_disaggregated_topology(self): + if self.deployment.type != "disaggregated": + return self + if self.deployment.gpus_per_node % self.parallel.tp != 0: + raise ValueError( + "inference.deployment.gpus_per_node must be divisible by inference.parallel.tp " + "so every worker contains whole tensor-parallel groups." + ) + if self.backend.type == "dynamo": + local_dp = self.deployment.gpus_per_node // self.parallel.tp + if "dp" in self.parallel.model_fields_set and self.parallel.dp != local_dp: + raise ValueError( + "inference.parallel.dp must equal inference.deployment.gpus_per_node / " + "inference.parallel.tp for a Dynamo disaggregated worker." + ) + if self.data_parallel_size_local is not None and self.data_parallel_size_local != local_dp: + raise ValueError( + "inference.data_parallel_size_local must equal inference.deployment.gpus_per_node / " + "inference.parallel.tp for a Dynamo disaggregated worker." + ) + return self + + @model_validator(mode="after") + def validate_dynamo_backend(self): + if self.backend.type != "dynamo": + return self + if self.slurm is not None: + raise ValueError( + "Dynamo is launched locally or through a DynamoGraphDeployment, not Prime's SLURM template." + ) + if self.deployment.type == "multi_node": + raise ValueError("Dynamo multi-node inference must use a DynamoGraphDeployment.") + if self.enable_lora: + raise ValueError("The Dynamo backend does not support LoRA weight updates.") + router = getattr(self.deployment, "router", None) + if router is not None and router.type == "llm-d": + raise ValueError("The Dynamo backend owns request routing and cannot use the llm-d router.") + if self.deployment.type == "disaggregated": + if self.enable_prefix_caching is False: + raise ValueError("Dynamo disaggregated inference requires prefix caching for exact KV-aware routing.") + self.enable_prefix_caching = True + return self + @model_validator(mode="after") def validate_llmd_no_routed_experts(self): """Reject routed-expert return with the llm-d router (breaks P/D, unverified for multi-node).""" @@ -463,9 +549,7 @@ def auto_setup_disaggregated(self): self.enable_expert_parallel = True if "enable_eplb" not in self.model_fields_set: self.enable_eplb = False - gpus_per_node = self.deployment.gpus_per_node - tp = self.parallel.tp - dp_per_node = gpus_per_node // tp + dp_per_node = self.dynamo_local_dp if self.data_parallel_size_local is None: self.data_parallel_size_local = dp_per_node if self.parallel.dp == 1: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 2c11a4ddb6..a708dc6f12 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -529,7 +529,9 @@ class OrchestratorConfig(BaseConfig): """Maximum training steps. If None, runs indefinitely.""" max_off_policy_steps: int = Field(8, ge=0) - """Maximum policies allowed to generate a single rollout. Rollouts generated more than ``max_off_policy_steps`` ahead of training are discarded. Higher values yield better throughput at the cost of off-policy noise.""" + """Maximum policies allowed to generate one rollout on the vLLM admin backend. Dynamo uses a strict + application drain before every weight mutation, so live-policy and eval requests never span versions and + this tolerance does not apply there.""" bench: bool = False """Benchmark mode. Sets ``max_steps`` to 5 and disables W&B.""" @@ -728,7 +730,6 @@ def resolve_env_config(self): if env.algo.sampling.source == "policy": env.sampling.extra_body.setdefault("top_k", -1) env.sampling.extra_body.setdefault("min_p", 0.0) - env.sampling.extra_body.setdefault("return_token_ids", True) if env.is_legacy: # v0 env: cap per-turn response tokens to the training budget (the legacy # bridge applies extra_env_kwargs via env.set_kwargs). diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index c92e53889a..0a50e17e38 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -460,6 +460,12 @@ def auto_setup_lora(self): return self + @model_validator(mode="after") + def validate_auto_setup_does_not_enable_dynamo_lora(self): + if self.inference is not None and self.inference.backend.type == "dynamo" and self.inference.enable_lora: + raise ValueError("The Dynamo backend does not support LoRA weight updates.") + return self + @model_validator(mode="after") def auto_setup_router_replay(self): if self.trainer.enable_router_replay: @@ -534,7 +540,23 @@ def auto_setup_deployment(self): # fill up inference capacity with dp ranks if self.inference is not None: num_infer_gpus = self.deployment.num_infer_gpus - if num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp: + is_dynamo_disaggregated = ( + self.inference.backend.type == "dynamo" and self.inference.deployment.type == "disaggregated" + ) + if is_dynamo_disaggregated: + infer_deploy = self.inference.deployment + expected_infer_gpus = infer_deploy.num_nodes * infer_deploy.gpus_per_node + if num_infer_gpus != expected_infer_gpus: + raise ValueError( + "deployment.num_infer_gpus must equal the Dynamo prefill/decode topology GPU count " + f"({expected_infer_gpus}), got {num_infer_gpus}." + ) + if self.weight_broadcast is not None and self.weight_broadcast.type == "nccl": + assert self.trainer.weight_broadcast.type == "nccl" + self.trainer.weight_broadcast.inference_world_size = expected_infer_gpus + assert self.orchestrator.weight_broadcast.type == "nccl" + self.orchestrator.weight_broadcast.inference_world_size = expected_infer_gpus + elif num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp: assert num_infer_gpus % self.inference.parallel.tp == 0, ( "Number of inference GPUs must be divisible by the tensor parallel size" ) @@ -675,15 +697,38 @@ def auto_setup_inference_client(self): if self.inference is None: return self client = self.orchestrator.model.client + client_updates: dict[str, Any] = {} + if "admin_api" in client.model_fields_set and client.admin_api != self.inference.backend.type: + raise ValueError( + "orchestrator.model.client.admin_api conflicts with inference.backend.type; " + "configure the backend only under inference." + ) + client_updates["admin_api"] = self.inference.backend.type if "dp_rank_count" not in client.model_fields_set: - if self.deployment.type == "multi_node": - client.dp_rank_count = 1 + if self.inference.backend.type == "dynamo" or self.deployment.type == "multi_node": + client_updates["dp_rank_count"] = 1 else: - client.dp_rank_count = self.inference.data_parallel_size_local or self.inference.parallel.dp + client_updates["dp_rank_count"] = self.inference.data_parallel_size_local or self.inference.parallel.dp + if self.inference.backend.type == "dynamo": + expected_topology = { + "dynamo_worker_roles": self.inference.dynamo_worker_roles, + "dynamo_gpus_per_worker": self.inference.dynamo_gpus_per_worker, + } + for field, expected in expected_topology.items(): + if field in client.model_fields_set and getattr(client, field) != expected: + raise ValueError( + f"orchestrator.model.client.{field} conflicts with the inference topology; " + "configure the topology only under inference." + ) + client_updates.update(expected_topology) if not self.orchestrator.any_policy_sourced and "base_url" not in client.model_fields_set: host = self.inference.server.host or "localhost" port = self.inference.server.port - client.base_url = [f"http://{host}:{port}/v1"] + client_updates["base_url"] = [f"http://{host}:{port}/v1"] + + updated_client = client.model_copy(update=client_updates) + updated_model = self.orchestrator.model.model_copy(update={"client": updated_client}) + self.orchestrator = self.orchestrator.model_copy(update={"model": updated_model}) return self @model_validator(mode="after") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index c63bcc9c53..225d544013 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -147,6 +147,18 @@ class ClientConfig(BaseConfig): admin_base_url: list[str] | None = None """Separate base URLs for admin operations (weight updates, health checks). When set, admin clients bypass routers and hit each server directly — used in disaggregated P/D deployments where the router must not handle admin traffic.""" + admin_api: Literal["vllm", "dynamo"] = "vllm" + """Admin protocol used for health and weight updates. Auto-derived from the local inference backend.""" + + rl_base_url: list[str] | None = None + """Dynamo RL worker-discovery URLs. When omitted, they are derived from the frontend URL or ``DYN_RL_DISCOVERY_URL``.""" + + dynamo_worker_roles: tuple[Literal["agg", "prefill", "decode"], ...] | None = None + """Exact Dynamo worker roles expected during readiness. Auto-derived from the local inference topology.""" + + dynamo_gpus_per_worker: int | None = Field(None, ge=1) + """GPUs owned by each discovered Dynamo worker. Auto-derived from the local inference topology.""" + elastic: ElasticConfig | None = None """Elastic inference pool config for DNS-based service discovery. When set, ``base_url`` is ignored and inference servers are discovered dynamically via DNS.""" @@ -255,7 +267,8 @@ class MetricsServerConfig(BaseConfig): class BaseTransportConfig(BaseConfig): - pass + send_timeout_seconds: float = Field(300.0, gt=0, allow_inf_nan=False) + """Maximum time to ship one rollout batch before the orchestrator fails closed.""" class FileSystemTransportConfig(BaseTransportConfig): diff --git a/pyproject.toml b/pyproject.toml index 0daa7259ab..440f0f8808 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ dependencies = [ rl = "prime_rl.entrypoints.rl:main" sft = "prime_rl.entrypoints.sft:main" inference = "prime_rl.entrypoints.inference:main" +dynamo-dgd = "prime_rl.inference.dgd:main" trainer = "prime_rl.entrypoints.trainer:main" orchestrator = "prime_rl.entrypoints.orchestrator:main" env-server = "prime_rl.orchestrator.env_server.env_server:main" diff --git a/src/prime_rl/entrypoints/inference.py b/src/prime_rl/entrypoints/inference.py index efc120f96e..3a6a62d637 100644 --- a/src/prime_rl/entrypoints/inference.py +++ b/src/prime_rl/entrypoints/inference.py @@ -148,6 +148,13 @@ def inference_local(config: InferenceConfig): logger = setup_logger(config.log.level, json_logging=config.log.json_logging) if config.dry_run: + if config.backend.type == "dynamo": + from prime_rl.inference.dynamo import build_dry_run_worker_specs, build_frontend_process + + specs = build_dry_run_worker_specs(config) + logger.info(f"Dynamo frontend: {' '.join(build_frontend_process(config).command())}") + for spec in specs: + logger.info(f"Dynamo {spec.name}: {' '.join(spec.process.command())}") logger.success("Dry run complete. To start inference locally, remove --dry-run from your command.") return @@ -162,6 +169,12 @@ def inference_local(config: InferenceConfig): setup_vllm_env(config) + if config.backend.type == "dynamo": + from prime_rl.inference.dynamo import run_dynamo_local + + run_dynamo_local(config) + return + from prime_rl.inference.vllm.server import server # pyright: ignore server(config, vllm_extra=config.vllm_extra) diff --git a/src/prime_rl/inference/dgd.py b/src/prime_rl/inference/dgd.py new file mode 100644 index 0000000000..e667aca7dc --- /dev/null +++ b/src/prime_rl/inference/dgd.py @@ -0,0 +1,716 @@ +"""Compile a Prime Dynamo inference config into Helm DGD values.""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +from prime_rl.configs.inference import DisaggregatedInferenceDeploymentConfig, InferenceConfig +from prime_rl.inference.dgd_controller_contract import build_chart_runtime_contract +from prime_rl.inference.dynamo import ( + CHAT_TEMPLATE_ASSET, + DynamoProcessSpec, + build_frontend_process, + build_worker_process, + resolve_chat_template_content, + write_role_engine_configs, +) + +ENGINE_MOUNT_PATH = "/etc/prime-rl/dynamo" +ENGINE_CONFIG_HASH_ANNOTATION = "prime-rl.nvidia.com/config-sha256" +MANIFEST_HASH_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256" +MANIFEST_HASH_SCOPE_ANNOTATION = "prime-rl.nvidia.com/manifest-sha256-scope" +TOPOLOGY_HASH_ANNOTATION = "prime-rl.nvidia.com/topology-sha256" +WORKLOAD_HASH_ANNOTATION = "prime-rl.nvidia.com/workload-sha256" +MANIFEST_HASH_SCOPE = ( + "resource; json.dumps(sort_keys=true,indent=2)+newline; " + "exclude=/metadata/annotations/prime-rl.nvidia.com~1manifest-sha256" +) +_DGD_RESERVED_ENV_KEYS = frozenset( + { + "CONTAINER_NAME", + "DYNAMO_PORT", + "DYN_COMPONENT", + "DYN_DISCOVERY_BACKEND", + "DYN_ENABLE_RL", + "DYN_ENDPOINT", + "DYN_ENDPOINT_TYPES", + "DYN_ETCD_ENDPOINTS", + "DYN_EVENT_PLANE", + "DYN_FILE_KV", + "DYN_HEALTH_CHECK_ENABLED", + "DYN_HTTP_PORT", + "DYN_KUBE_DISCOVERY_MODE", + "DYN_NAMESPACE", + "DYN_NAMESPACE_PREFIX", + "DYN_NAMESPACE_WORKER_SUFFIX", + "DYN_PARENT_DGD_K8S_NAME", + "DYN_PARENT_DGD_K8S_NAMESPACE", + "DYN_RL_ENDPOINT", + "DYN_RL_PORT", + "POD_NAME", + "POD_NAMESPACE", + "POD_UID", + "VLLM_NIXL_SIDE_CHANNEL_HOST", + "VLLM_NIXL_SIDE_CHANNEL_PORT", + } +) +_DGD_RESERVED_ENV_PREFIXES = ("DYN_HEALTH_CHECK_", "DYN_SYSTEM_") +_MODEL_CACHE_MOUNT_PATH = "/model-cache" +_CREDENTIAL_ENV_KEY_PATTERNS = ( + re.compile(r"(?:^|_)(?:TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIALS?)(?:_|$)"), + re.compile(r"(?:^|_)(?:API|ACCESS|PRIVATE|SECRET)_KEY(?:_|$)"), +) +MAX_RELEASE_NAME_LENGTH = 41 + + +@dataclass(frozen=True, slots=True) +class KubernetesToleration: + key: str + operator: Literal["Exists", "Equal"] = "Exists" + effect: Literal["NoSchedule", "PreferNoSchedule", "NoExecute"] | None = "NoSchedule" + value: str | None = None + + def __post_init__(self) -> None: + if not self.key: + raise ValueError("Kubernetes toleration key must not be empty") + if self.operator not in ("Exists", "Equal"): + raise ValueError(f"Unsupported Kubernetes toleration operator: {self.operator!r}") + if self.effect not in (None, "NoSchedule", "PreferNoSchedule", "NoExecute"): + raise ValueError(f"Unsupported Kubernetes toleration effect: {self.effect!r}") + if self.operator == "Exists" and self.value is not None: + raise ValueError("An Exists toleration cannot define a value") + if self.operator == "Equal" and self.value is None: + raise ValueError("An Equal toleration requires a value") + + def as_manifest(self) -> dict[str, str]: + return { + "key": self.key, + "operator": self.operator, + **({"value": self.value} if self.value is not None else {}), + **({"effect": self.effect} if self.effect is not None else {}), + } + + +def _parse_kubernetes_toleration(value: str) -> KubernetesToleration: + """Parse one CLI toleration from JSON without evaluating shell-like input.""" + try: + payload = json.loads(value) + except json.JSONDecodeError as error: + raise ValueError(f"Kubernetes toleration must be a JSON object: {error.msg}") from error + if not isinstance(payload, dict): + raise ValueError("Kubernetes toleration must be a JSON object") + supported = {"key", "operator", "effect", "value"} + unknown = sorted(payload.keys() - supported) + if unknown: + raise ValueError(f"Kubernetes toleration has unsupported fields: {unknown}") + if "key" not in payload: + raise ValueError("Kubernetes toleration requires a key") + return KubernetesToleration(**payload) + + +def _unique_tolerations( + tolerations: tuple[KubernetesToleration, ...], +) -> tuple[KubernetesToleration, ...]: + unique: list[KubernetesToleration] = [] + seen: set[tuple[tuple[str, str], ...]] = set() + for toleration in tolerations: + identity = tuple(sorted(toleration.as_manifest().items())) + if identity not in seen: + unique.append(toleration) + seen.add(identity) + return tuple(unique) + + +@dataclass(frozen=True, slots=True) +class GPUSchedulingProfile: + """Image placement plus the stricter placement required by GPU consumers.""" + + runtime_class_name: str | None + architecture: str + product: str + node_pool: str + node_pool_label: str = "cloud.google.com/gke-nodepool" + tolerations: tuple[KubernetesToleration, ...] = (KubernetesToleration(key="nvidia.com/gpu"),) + additional_image_tolerations: tuple[KubernetesToleration, ...] = () + additional_gpu_tolerations: tuple[KubernetesToleration, ...] = () + + def __post_init__(self) -> None: + required = { + "architecture": self.architecture, + "product": self.product, + "node_pool": self.node_pool, + "node_pool_label": self.node_pool_label, + } + empty = [name for name, value in required.items() if not value] + if empty: + raise ValueError(f"GPU scheduling fields must not be empty: {empty}") + if self.runtime_class_name == "": + raise ValueError("runtime_class_name must be non-empty or None") + if not self.tolerations: + raise ValueError("GPU scheduling requires at least one toleration") + required_gpu_toleration = KubernetesToleration(key="nvidia.com/gpu") + if required_gpu_toleration not in self.tolerations: + raise ValueError("GPU scheduling requires nvidia.com/gpu Exists NoSchedule") + + @property + def image_node_selector(self) -> dict[str, str]: + return { + "kubernetes.io/arch": self.architecture, + self.node_pool_label: self.node_pool, + } + + @property + def node_selector(self) -> dict[str, str]: + return { + **self.image_node_selector, + "nvidia.com/gpu.product": self.product, + } + + @property + def image_tolerations(self) -> tuple[KubernetesToleration, ...]: + required = ( + KubernetesToleration( + key="kubernetes.io/arch", + operator="Equal", + value=self.architecture, + ), + KubernetesToleration(key="nvidia.com/gpu"), + KubernetesToleration(key="prime-rl", operator="Equal", value="true"), + ) + return _unique_tolerations((*required, *self.additional_image_tolerations)) + + @property + def image_toleration_manifests(self) -> list[dict[str, str]]: + return [toleration.as_manifest() for toleration in self.image_tolerations] + + @property + def image_placement(self) -> dict[str, Any]: + return { + "nodeSelector": self.image_node_selector, + "tolerations": self.image_toleration_manifests, + } + + @property + def gpu_tolerations(self) -> tuple[KubernetesToleration, ...]: + return _unique_tolerations((*self.image_tolerations, *self.tolerations, *self.additional_gpu_tolerations)) + + @property + def toleration_manifests(self) -> list[dict[str, str]]: + return [toleration.as_manifest() for toleration in self.gpu_tolerations] + + @property + def gpu_placement(self) -> dict[str, Any]: + placement = { + "nodeSelector": self.node_selector, + "tolerations": self.toleration_manifests, + } + if self.runtime_class_name is not None: + placement["runtimeClassName"] = self.runtime_class_name + return placement + + +@dataclass(frozen=True) +class DynamoGraphRenderOptions: + release_name: str + namespace: str + image: str + output_dir: Path + prime_sha: str + dynamo_sha: str + image_digest: str + run_name: str + gpu_scheduling: GPUSchedulingProfile + external_controller: bool = False + trainer_gpu_count: int = 1 + orchestrator_replicas: int = 1 + trainer_replicas: int = 1 + orchestrator_command: str | None = None + trainer_command: str | None = None + model_cache_pvc: str | None = None + shared_pvc: str | None = None + image_pull_secrets: tuple[str, ...] = () + hf_token_secret: str | None = None + + def __post_init__(self) -> None: + if re.fullmatch(r"[a-z0-9](?:[-a-z0-9]*[a-z0-9])?", self.release_name) is None: + raise ValueError("release_name must be a lowercase Kubernetes DNS label") + if len(self.release_name) > MAX_RELEASE_NAME_LENGTH: + raise ValueError( + f"release_name must be at most {MAX_RELEASE_NAME_LENGTH} characters so generated Service names are valid" + ) + for name, value in (("prime_sha", self.prime_sha), ("dynamo_sha", self.dynamo_sha)): + if re.fullmatch(r"[0-9a-f]{40}", value) is None: + raise ValueError(f"{name} must be a full 40-character Git commit SHA") + if re.fullmatch(r"sha256:[0-9a-f]{64}", self.image_digest) is None: + raise ValueError("image_digest must be a full sha256 digest") + if not self.image.endswith(f"@{self.image_digest}"): + raise ValueError("DGD image must be pinned to image_digest") + image_name = self.image.rsplit("@", 1)[0] + image_tag = image_name.rsplit("/", 1)[-1].partition(":")[2] + if self.prime_sha[:12] not in image_tag or self.dynamo_sha[:12] not in image_tag: + raise ValueError("DGD image tag must include the Prime and Dynamo commit suffixes") + if self.trainer_gpu_count < 1: + raise ValueError("trainer_gpu_count must be at least one") + if not self.external_controller: + for component, replicas, command in ( + ("orchestrator", self.orchestrator_replicas, self.orchestrator_command), + ("trainer", self.trainer_replicas, self.trainer_command), + ): + if replicas < 1: + raise ValueError(f"{component}_replicas must be at least one in chart-managed mode") + expected_prefix = f"uv run {component}" + if command is None or re.match(rf"^uv\s+run\s+{component}(?:\s|$)", command) is None: + raise ValueError(f"{component}_command must start with {expected_prefix!r} in chart-managed mode") + elif self.orchestrator_command is not None or self.trainer_command is not None: + raise ValueError("external_controller cannot define chart-managed controller commands") + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _canonical_json(value: Any) -> bytes: + return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode() + + +def _resource_manifest_canonical(resource: dict[str, Any]) -> bytes: + annotations = resource["metadata"]["annotations"] + scoped_annotations = {key: value for key, value in annotations.items() if key != MANIFEST_HASH_ANNOTATION} + scoped_resource = { + **resource, + "metadata": { + **resource["metadata"], + "annotations": scoped_annotations, + }, + } + return _canonical_json(scoped_resource) + + +def _worker_topology_binding(services: dict[str, Any]) -> dict[str, Any]: + return { + service_name: { + "role": service["subComponentType"], + "replicas": service["replicas"], + "requestsGpu": service["resources"]["requests"]["gpu"], + "limitsGpu": service["resources"]["limits"]["gpu"], + } + for service_name, service in sorted(services.items()) + if service["componentType"] == "worker" + } + + +def _release_pod_labels(options: DynamoGraphRenderOptions) -> dict[str, str]: + return { + "app.kubernetes.io/name": "prime-rl", + "app.kubernetes.io/instance": options.release_name, + } + + +def _worker_env(process: DynamoProcessSpec) -> list[dict[str, Any]]: + values: list[dict[str, Any]] = [ + {"name": name, "value": value} for name, value in sorted(process.environment().items()) + ] + values.append( + { + "name": "VLLM_NIXL_SIDE_CHANNEL_HOST", + "valueFrom": {"fieldRef": {"fieldPath": "status.podIP"}}, + } + ) + return values + + +def _apply_pod_credentials( + pod_spec: dict[str, Any], + container: dict[str, Any], + options: DynamoGraphRenderOptions, +) -> None: + if options.image_pull_secrets: + pod_spec["imagePullSecrets"] = [{"name": name} for name in options.image_pull_secrets] + if options.model_cache_pvc and not any(item["name"] == "HF_HOME" for item in container.get("env", [])): + container.setdefault("env", []).append({"name": "HF_HOME", "value": "/model-cache"}) + if options.hf_token_secret and not any(item["name"] == "HF_TOKEN" for item in container.get("env", [])): + container.setdefault("env", []).append( + { + "name": "HF_TOKEN", + "valueFrom": { + "secretKeyRef": { + "name": options.hf_token_secret, + "key": "HF_TOKEN", + "optional": False, + } + }, + } + ) + + +def _worker_service( + config: InferenceConfig, + options: DynamoGraphRenderOptions, + *, + role: str, + replicas: int, + config_map_name: str, + engine_file: str, +) -> dict[str, Any]: + assert config.deployment.type == "disaggregated" + process = build_worker_process( + config, + role, + Path(ENGINE_MOUNT_PATH) / engine_file, + nixl_host=None, + nixl_port=20100, + ) + container = { + "image": options.image, + "imagePullPolicy": "IfNotPresent", + "command": ["python3", "-m", process.module], + "args": list(process.arguments), + "env": _worker_env(process), + "volumeMounts": [ + { + "name": "dynamo-engine-config", + "mountPath": ENGINE_MOUNT_PATH, + "readOnly": True, + } + ], + } + pod_spec = { + **options.gpu_scheduling.gpu_placement, + "volumes": [ + { + "name": "dynamo-engine-config", + "configMap": {"name": config_map_name}, + } + ], + "mainContainer": container, + } + _apply_pod_credentials(pod_spec, container, options) + return { + "componentType": "worker", + "subComponentType": role, + "replicas": replicas, + "extraPodMetadata": {"labels": _release_pod_labels(options)}, + "sharedMemory": {"size": "64Gi"}, + "resources": { + "requests": {"gpu": str(config.deployment.gpus_per_node)}, + "limits": {"gpu": str(config.deployment.gpus_per_node)}, + }, + "extraPodSpec": pod_spec, + } + + +def _add_pvc(resource: dict[str, Any], service: dict[str, Any], name: str | None, mount_point: str) -> None: + if not name: + return + pvcs = resource["spec"].setdefault("pvcs", []) + if not any(pvc["name"] == name for pvc in pvcs): + pvcs.append({"name": name, "create": False}) + # Keep PodSpec projection as the single mount source. Older alpha operators + # do not realize service.volumeMounts, while current alpha-to-beta conversion + # appends those mounts to extraPodSpec and would otherwise create duplicates. + pod_spec = service["extraPodSpec"] + container_mount = {"name": name, "mountPath": mount_point} + container_mounts = pod_spec["mainContainer"].setdefault("volumeMounts", []) + if container_mount not in container_mounts: + if any(mount["mountPath"] == mount_point for mount in container_mounts): + raise ValueError(f"PVC {name!r} conflicts with an existing container mount at {mount_point!r}") + container_mounts.append(container_mount) + + pod_volume = { + "name": name, + "persistentVolumeClaim": {"claimName": name}, + } + pod_volumes = pod_spec.setdefault("volumes", []) + if pod_volume not in pod_volumes: + if any(volume["name"] == name for volume in pod_volumes): + raise ValueError(f"PVC {name!r} conflicts with an existing pod volume") + pod_volumes.append(pod_volume) + + +def _validate_dgd_environment(config: InferenceConfig) -> None: + environment_sources = [("global", config.env_vars)] + if config.deployment.type == "disaggregated": + environment_sources.extend( + [ + ("prefill", config.deployment.prefill_env_vars), + ("decode", config.deployment.decode_env_vars), + ] + ) + for source, environment in environment_sources: + conflicts = sorted( + key + for key in environment + if key in _DGD_RESERVED_ENV_KEYS or any(key.startswith(prefix) for prefix in _DGD_RESERVED_ENV_PREFIXES) + ) + if conflicts: + raise ValueError(f"{source} env_vars contains {conflicts}; these DGD keys are operator-owned") + + +def _validate_typed_credentials( + config: InferenceConfig, + options: DynamoGraphRenderOptions, +) -> None: + environment_sources = [("global", config.env_vars)] + if config.deployment.type == "disaggregated": + environment_sources.extend( + [ + ("prefill", config.deployment.prefill_env_vars), + ("decode", config.deployment.decode_env_vars), + ] + ) + for source, environment in environment_sources: + credential_conflicts = sorted( + key for key in environment if any(pattern.search(key.upper()) for pattern in _CREDENTIAL_ENV_KEY_PATTERNS) + ) + if credential_conflicts: + raise ValueError( + f"{source} env_vars contains raw credentials {credential_conflicts}; " + "use a typed Kubernetes SecretKeyRef instead" + ) + hf_home = environment.get("HF_HOME") + if options.model_cache_pvc and hf_home is not None and hf_home != _MODEL_CACHE_MOUNT_PATH: + raise ValueError( + f"{source} env_vars sets HF_HOME={hf_home!r}, but the typed model cache mount " + f"requires {_MODEL_CACHE_MOUNT_PATH!r}" + ) + + +def build_dgd_values(config: InferenceConfig, options: DynamoGraphRenderOptions) -> dict[str, Any]: + if config.backend.type != "dynamo" or config.deployment.type != "disaggregated": + raise ValueError("DGD rendering requires a Dynamo disaggregated inference config") + deployment: DisaggregatedInferenceDeploymentConfig = config.deployment + if deployment.num_prefill_nodes != deployment.num_prefill_replicas: + raise ValueError("DGD rendering currently requires one pod per prefill replica") + if deployment.num_decode_nodes != deployment.num_decode_replicas: + raise ValueError("DGD rendering currently requires one pod per decode replica") + if config.weight_broadcast.type == "filesystem" and not options.shared_pvc: + raise ValueError("Dynamo filesystem weight broadcast requires a shared existing PVC") + _validate_dgd_environment(config) + _validate_typed_credentials(config, options) + + engine_paths = write_role_engine_configs(config, options.output_dir) + prefill_text = engine_paths["prefill"].read_text() + decode_text = engine_paths["decode"].read_text() + engine_data = { + "prefill-engine.json": prefill_text, + "decode-engine.json": decode_text, + } + chat_template_content = resolve_chat_template_content(config) + if chat_template_content is not None: + engine_data[CHAT_TEMPLATE_ASSET] = chat_template_content + engine_canonical = _canonical_json(engine_data) + engine_hash = _sha256_bytes(engine_canonical) + config_map_name = f"{options.release_name}-dynamo-engine-{engine_hash[:12]}" + runtime_chat_template_path = ( + Path(ENGINE_MOUNT_PATH) / CHAT_TEMPLATE_ASSET if chat_template_content is not None else None + ) + frontend_process = build_frontend_process( + config, + host="0.0.0.0", + port=8000, + runtime_chat_template_path=runtime_chat_template_path, + ) + frontend_container = { + "image": options.image, + "imagePullPolicy": "IfNotPresent", + "command": ["python3", "-m", frontend_process.module], + "args": list(frontend_process.arguments), + "env": [{"name": name, "value": value} for name, value in sorted(frontend_process.environment().items())], + "ports": [ + {"containerPort": 8000, "name": "http"}, + {"containerPort": 8001, "name": "rl"}, + ], + } + frontend_pod_spec = { + **options.gpu_scheduling.image_placement, + "mainContainer": frontend_container, + } + if chat_template_content is not None: + frontend_container["volumeMounts"] = [ + { + "name": "dynamo-chat-template", + "mountPath": ENGINE_MOUNT_PATH, + "readOnly": True, + } + ] + frontend_pod_spec["volumes"] = [ + { + "name": "dynamo-chat-template", + "configMap": { + "name": config_map_name, + "items": [{"key": CHAT_TEMPLATE_ASSET, "path": CHAT_TEMPLATE_ASSET}], + }, + } + ] + _apply_pod_credentials(frontend_pod_spec, frontend_container, options) + frontend = { + "componentType": "frontend", + "replicas": 1, + "extraPodMetadata": {"labels": _release_pod_labels(options)}, + "extraPodSpec": frontend_pod_spec, + } + prefill = _worker_service( + config, + options, + role="prefill", + replicas=deployment.num_prefill_replicas, + config_map_name=config_map_name, + engine_file="prefill-engine.json", + ) + decode = _worker_service( + config, + options, + role="decode", + replicas=deployment.num_decode_replicas, + config_map_name=config_map_name, + engine_file="decode-engine.json", + ) + client_topology = { + "schema_version": 1, + "admin_api": "dynamo", + "base_url": [f"http://{options.release_name}-frontend.{options.namespace}.svc.cluster.local:8000/v1"], + "rl_base_url": [f"http://{options.release_name}-frontend-rl.{options.namespace}.svc.cluster.local:8001"], + "dynamo_worker_roles": list(config.dynamo_worker_roles), + "dynamo_gpus_per_worker": config.dynamo_gpus_per_worker, + } + worker_services = { + "VllmDecodeWorker": decode, + "VllmPrefillWorker": prefill, + } + topology_binding = { + "clientTopology": client_topology, + "workerServices": _worker_topology_binding(worker_services), + } + topology_canonical = _canonical_json(topology_binding) + topology_hash = _sha256_bytes(topology_canonical) + controller_mode = "external" if options.external_controller else "chartManaged" + workload, chart_values = build_chart_runtime_contract( + controller_mode=controller_mode, + image_reference=options.image, + image_pull_secrets=options.image_pull_secrets, + orchestrator_replicas=options.orchestrator_replicas, + trainer_replicas=options.trainer_replicas, + orchestrator_command=options.orchestrator_command, + trainer_command=options.trainer_command, + trainer_gpu_count=options.trainer_gpu_count, + orchestrator_placement=options.gpu_scheduling.image_placement, + trainer_placement=options.gpu_scheduling.gpu_placement, + shared_pvc=options.shared_pvc, + model_cache_pvc=options.model_cache_pvc, + hf_token_secret=options.hf_token_secret, + ) + workload_canonical = _canonical_json(workload) + workload_hash = _sha256_bytes(workload_canonical) + annotations = { + ENGINE_CONFIG_HASH_ANNOTATION: engine_hash, + "prime-rl.nvidia.com/dynamo-sha": options.dynamo_sha, + "prime-rl.nvidia.com/image-digest": options.image_digest, + "prime-rl.nvidia.com/prime-sha": options.prime_sha, + "prime-rl.nvidia.com/run-name": options.run_name, + MANIFEST_HASH_SCOPE_ANNOTATION: MANIFEST_HASH_SCOPE, + TOPOLOGY_HASH_ANNOTATION: topology_hash, + WORKLOAD_HASH_ANNOTATION: workload_hash, + } + resource: dict[str, Any] = { + "apiVersion": "nvidia.com/v1alpha1", + "kind": "DynamoGraphDeployment", + "metadata": { + "name": options.release_name, + "namespace": options.namespace, + "annotations": annotations, + }, + "spec": { + "backendFramework": "vllm", + "services": { + "Frontend": frontend, + "VllmDecodeWorker": decode, + "VllmPrefillWorker": prefill, + }, + }, + } + for service in (frontend, prefill, decode): + _add_pvc(resource, service, options.model_cache_pvc, "/model-cache") + if config.weight_broadcast.type == "filesystem": + for service in (prefill, decode): + _add_pvc(resource, service, options.shared_pvc, "/data") + + manifest_canonical = _resource_manifest_canonical(resource) + manifest_hash = _sha256_bytes(manifest_canonical) + annotations = {**annotations, MANIFEST_HASH_ANNOTATION: manifest_hash} + resource = { + **resource, + "metadata": { + **resource["metadata"], + "annotations": annotations, + }, + } + values: dict[str, Any] = { + "namespace": options.namespace, + **chart_values, + "inference": { + "enabled": True, + "mode": "dynamoGraph", + "dynamoGraph": { + "controllerMode": controller_mode, + "clientTopology": client_topology, + "engineConfig": { + "name": config_map_name, + "sha256": engine_hash, + "canonicalData": engine_canonical.decode(), + "annotations": annotations, + "data": engine_data, + }, + "topologyBinding": { + "sha256": topology_hash, + "canonical": topology_canonical.decode(), + }, + "workloadBinding": { + "sha256": workload_hash, + "canonical": workload_canonical.decode(), + }, + "manifestCanonical": manifest_canonical.decode(), + "resource": resource, + }, + }, + } + return values + + +def write_dgd_artifacts(config: InferenceConfig, options: DynamoGraphRenderOptions) -> dict[str, Path]: + options.output_dir.mkdir(parents=True, exist_ok=True) + values = build_dgd_values(config, options) + resource = values["inference"]["dynamoGraph"]["resource"] + paths = { + "values": options.output_dir / "dynamo-helm-values.json", + "resource": options.output_dir / "dynamo-graph-deployment.json", + } + paths["values"].write_bytes(_canonical_json(values)) + paths["resource"].write_bytes(_canonical_json(resource)) + manifest_entries = [] + for path in sorted(options.output_dir.glob("*.json")): + manifest_entries.append(f"{_sha256_bytes(path.read_bytes())} {path.name}") + manifest = options.output_dir / "artifact-manifest.sha256" + manifest.write_text("\n".join(manifest_entries) + "\n") + paths["manifest"] = manifest + return paths + + +def _parse_args(): + from prime_rl.inference.dgd_cli import parse_args + + return parse_args() + + +def main() -> None: + from prime_rl.inference.dgd_cli import main as cli_main + + cli_main() + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/inference/dgd_cli.py b/src/prime_rl/inference/dgd_cli.py new file mode 100644 index 0000000000..4efbec20c1 --- /dev/null +++ b/src/prime_rl/inference/dgd_cli.py @@ -0,0 +1,113 @@ +"""Command-line interface for rendering Prime Dynamo graph deployments.""" + +import argparse +from pathlib import Path + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.dgd import ( + DynamoGraphRenderOptions, + GPUSchedulingProfile, + _parse_kubernetes_toleration, + write_dgd_artifacts, +) +from prime_rl.utils.config import cli + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("inference_config", type=Path) + parser.add_argument("--release-name", required=True) + parser.add_argument("--namespace", required=True) + parser.add_argument("--image", required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--prime-sha", required=True) + parser.add_argument("--dynamo-sha", required=True) + parser.add_argument("--image-digest", required=True) + parser.add_argument("--run-name") + runtime_class = parser.add_mutually_exclusive_group() + runtime_class.add_argument("--gpu-runtime-class", default="nvidia") + runtime_class.add_argument( + "--no-gpu-runtime-class", + action="store_true", + help="Request nvidia.com/gpu resources without setting a Kubernetes RuntimeClass", + ) + parser.add_argument("--gpu-architecture", required=True) + parser.add_argument("--gpu-product", required=True) + parser.add_argument("--gpu-node-pool", required=True) + parser.add_argument("--gpu-node-pool-label", default="cloud.google.com/gke-nodepool") + parser.add_argument( + "--external-controller", + action="store_true", + help="Render only DGD inference workloads; an external controller owns orchestration and training", + ) + parser.add_argument( + "--trainer-gpus", + type=int, + default=1, + help="Exact GPU request and limit for the chart-managed trainer", + ) + parser.add_argument("--orchestrator-replicas", type=int, default=1) + parser.add_argument("--trainer-replicas", type=int, default=1) + parser.add_argument( + "--orchestrator-command", + help="Chart-managed command; must start with 'uv run orchestrator'", + ) + parser.add_argument( + "--trainer-command", + help="Chart-managed command; must start with 'uv run trainer'", + ) + parser.add_argument( + "--image-toleration", + action="append", + default=[], + type=_parse_kubernetes_toleration, + help='Additional image-pod toleration as JSON, e.g. \'{"key":"dedicated","operator":"Exists"}\'', + ) + parser.add_argument( + "--gpu-toleration", + action="append", + default=[], + type=_parse_kubernetes_toleration, + help='Additional GPU-pod toleration as JSON, e.g. \'{"key":"capacity","operator":"Exists"}\'', + ) + parser.add_argument("--model-cache-pvc") + parser.add_argument("--shared-pvc") + parser.add_argument("--image-pull-secret", action="append", default=[]) + parser.add_argument("--hf-token-secret") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + config = cli(InferenceConfig, args=["@", str(args.inference_config)]) + options = DynamoGraphRenderOptions( + release_name=args.release_name, + namespace=args.namespace, + image=args.image, + output_dir=args.output_dir, + prime_sha=args.prime_sha, + dynamo_sha=args.dynamo_sha, + image_digest=args.image_digest, + run_name=args.run_name or args.release_name, + gpu_scheduling=GPUSchedulingProfile( + runtime_class_name=None if args.no_gpu_runtime_class else args.gpu_runtime_class, + architecture=args.gpu_architecture, + product=args.gpu_product, + node_pool=args.gpu_node_pool, + node_pool_label=args.gpu_node_pool_label, + additional_image_tolerations=tuple(args.image_toleration), + additional_gpu_tolerations=tuple(args.gpu_toleration), + ), + external_controller=args.external_controller, + trainer_gpu_count=args.trainer_gpus, + orchestrator_replicas=args.orchestrator_replicas, + trainer_replicas=args.trainer_replicas, + orchestrator_command=args.orchestrator_command, + trainer_command=args.trainer_command, + model_cache_pvc=args.model_cache_pvc, + shared_pvc=args.shared_pvc, + image_pull_secrets=tuple(args.image_pull_secret), + hf_token_secret=args.hf_token_secret, + ) + for path in write_dgd_artifacts(config, options).values(): + print(path) diff --git a/src/prime_rl/inference/dgd_controller_contract.py b/src/prime_rl/inference/dgd_controller_contract.py new file mode 100644 index 0000000000..d7be20674a --- /dev/null +++ b/src/prime_rl/inference/dgd_controller_contract.py @@ -0,0 +1,148 @@ +"""Canonical Helm controller and PVC values for DynamoGraph deployments.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from typing import Literal + +ControllerMode = Literal["chartManaged", "external"] + + +def _orchestrator_service() -> dict[str, object]: + return { + "enabled": True, + "type": "ClusterIP", + "port": 8000, + "ncclPort": 29501, + } + + +def _trainer_service() -> dict[str, object]: + return { + "enabled": True, + "type": "ClusterIP", + "port": 8000, + "ncclPort": 29501, + } + + +def _trainer_probes() -> dict[str, object]: + return { + "enabled": False, + "startup": { + "periodSeconds": 10, + "failureThreshold": 60, + "timeoutSeconds": 30, + }, + "liveness": { + "periodSeconds": 30, + "failureThreshold": 6, + "timeoutSeconds": 30, + }, + "readiness": { + "periodSeconds": 10, + "failureThreshold": 3, + "timeoutSeconds": 30, + }, + } + + +def _controller_values(component: Mapping[str, object]) -> dict[str, object]: + return {key: deepcopy(value) for key, value in component.items() if key not in {"gpu", "placement"}} + + +def build_chart_runtime_contract( + *, + controller_mode: ControllerMode, + image_reference: str, + image_pull_secrets: Sequence[str], + orchestrator_replicas: int, + trainer_replicas: int, + orchestrator_command: str | None, + trainer_command: str | None, + trainer_gpu_count: int, + orchestrator_placement: Mapping[str, object], + trainer_placement: Mapping[str, object], + shared_pvc: str | None, + model_cache_pvc: str | None, + hf_token_secret: str | None, +) -> tuple[dict[str, object], dict[str, dict[str, object]]]: + """Return one immutable-by-construction contract and its exact chart values.""" + enabled = controller_mode == "chartManaged" + image = { + "reference": image_reference, + "pullPolicy": "IfNotPresent", + "pullSecrets": list(image_pull_secrets), + } + config = { + "example": "reverse-text", + "secrets": { + "enabled": False, + "name": "prime-rl-secrets", + }, + } + hugging_face = { + "tokenSecretName": hf_token_secret or "", + "tokenSecretKey": "HF_TOKEN", + } + model_cache = { + "enabled": model_cache_pvc is not None, + "existingClaim": model_cache_pvc or "", + "mountPath": "/model-cache", + } + storage = { + "enabled": enabled or shared_pvc is not None, + "existingClaim": shared_pvc or "", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", + "mountPath": "/data", + } + orchestrator = { + "enabled": enabled, + "replicas": orchestrator_replicas if enabled else 0, + "autoStart": enabled, + "command": orchestrator_command if enabled else "", + "gpu": {"enabled": False}, + "placement": deepcopy(orchestrator_placement), + "resources": {"requests": {"memory": "2Gi", "cpu": "1"}}, + "service": _orchestrator_service(), + "env": [], + } + trainer = { + "enabled": enabled, + "replicas": trainer_replicas if enabled else 0, + "autoStart": enabled, + "command": trainer_command if enabled else "", + "gpu": { + "enabled": enabled, + "count": trainer_gpu_count if enabled else 0, + }, + "placement": deepcopy(trainer_placement), + "pytorchCudaAllocConf": "expandable_segments:True", + "resources": {"requests": {"memory": "4Gi", "cpu": "1"}}, + "service": _trainer_service(), + "env": [], + "probes": _trainer_probes(), + } + workload = { + "controllerMode": controller_mode, + "image": image, + "config": config, + "huggingFace": hugging_face, + "modelCache": model_cache, + "orchestrator": orchestrator, + "storage": storage, + "trainer": trainer, + } + values = { + "image": deepcopy(image), + "config": deepcopy(config), + "huggingFace": deepcopy(hugging_face), + "modelCache": deepcopy(model_cache), + "orchestrator": _controller_values(orchestrator), + "storage": deepcopy(storage), + "trainer": _controller_values(trainer) | {"gpu": deepcopy(trainer["gpu"])}, + } + return workload, values diff --git a/src/prime_rl/inference/dynamo.py b/src/prime_rl/inference/dynamo.py new file mode 100644 index 0000000000..ca9f9e0533 --- /dev/null +++ b/src/prime_rl/inference/dynamo.py @@ -0,0 +1,529 @@ +"""Translate Prime inference config into Dynamo worker processes.""" + +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any, Literal + +from prime_rl.configs.inference import DisaggregatedInferenceDeploymentConfig, InferenceConfig +from prime_rl.utils.pathing import get_config_dir + +Role = Literal["agg", "prefill", "decode"] + +ENGINE_CONFIG_DIR = "dynamo" +PREFILL_ENGINE_CONFIG = "prefill-engine.json" +DECODE_ENGINE_CONFIG = "decode-engine.json" +AGG_ENGINE_CONFIG = "agg-engine.json" +CHAT_TEMPLATE_ASSET = "chat-template.jinja" + +_ENGINE_CONFIG_EXCLUDED = frozenset( + { + "api_server_count", + "chat_template", + "enable_auto_tool_choice", + "host", + "liveness_timeout_seconds", + "port", + "reasoning_parser", + "tool_call_parser", + } +) +_RESERVED_ENGINE_KEYS = frozenset( + { + "data_parallel_rpc_port", + "data_parallel_size", + "data_parallel_size_local", + "disaggregation_mode", + "enable_prefix_caching", + "enable_rl", + "kv_events_config", + "kv_transfer_config", + "pipeline_parallel_size", + "tensor_parallel_size", + "worker_extension_cls", + } +) +_WORKER_EXTENSION_CLS = { + "nccl": "prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker", + "filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker", +} +_WORKER_COMPONENT = { + "agg": "backend", + "prefill": "prefill", + "decode": "backend", +} + + +@dataclass(frozen=True) +class DynamoProcessSpec: + module: str + arguments: tuple[str, ...] + environment_items: tuple[tuple[str, str], ...] + + def command(self, executable: str = sys.executable) -> list[str]: + return [executable, "-m", self.module, *self.arguments] + + def environment(self, base: dict[str, str] | None = None) -> dict[str, str]: + return (base or {}) | dict(self.environment_items) + + +@dataclass(frozen=True) +class DynamoWorkerSpec: + name: str + role: Role + gpu_ids: tuple[str, ...] + system_port: int + process: DynamoProcessSpec + + +@dataclass(frozen=True) +class DynamoWorkerPorts: + """Host-local ports reserved by one worker process.""" + + system: int + nixl: int + data_parallel_rpc: int + kv_events: int + + +_LOCAL_WORKER_PORT_BASE = 18_000 +_LOCAL_WORKER_PORT_STRIDE = 4 + + +def _allocate_local_worker_ports(worker_index: int) -> DynamoWorkerPorts: + """Allocate one non-overlapping port block for a same-host worker.""" + base = _LOCAL_WORKER_PORT_BASE + worker_index * _LOCAL_WORKER_PORT_STRIDE + ports = DynamoWorkerPorts( + system=base, + nixl=base + 1, + data_parallel_rpc=base + 2, + kv_events=base + 3, + ) + if ports.kv_events > 65_535: + raise ValueError(f"Local Dynamo worker {worker_index} exceeds the available TCP port range") + return ports + + +def _json_default(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, Enum): + return value.value + raise TypeError(f"Cannot serialize {type(value).__name__}") + + +def _role_overrides(config: InferenceConfig, role: Role) -> dict[str, Any]: + if config.deployment.type != "disaggregated": + return {} + if role == "prefill": + return config.deployment.prefill_vllm_overrides + if role == "decode": + return config.deployment.decode_vllm_overrides + return {} + + +def _validate_overrides(source: str, values: dict[str, Any]) -> None: + conflicts = sorted(_RESERVED_ENGINE_KEYS & values.keys()) + if conflicts: + raise ValueError(f"{source} cannot override Dynamo-managed engine keys: {conflicts}") + wrapper_only = sorted(_ENGINE_CONFIG_EXCLUDED & values.keys()) + if wrapper_only: + raise ValueError(f"{source} keys {wrapper_only} are wrapper/server-only and cannot enter a vLLM engine config") + + +def _environment_items(values: dict[str, str]) -> tuple[tuple[str, str], ...]: + return tuple(sorted(values.items())) + + +def _role_environment(config: InferenceConfig, role: Role) -> dict[str, str]: + if config.deployment.type != "disaggregated": + return {} + if role == "prefill": + return config.deployment.prefill_env_vars + if role == "decode": + return config.deployment.decode_env_vars + return {} + + +def resolve_chat_template_content(config: InferenceConfig) -> str | None: + """Resolve a configured inline or file-backed chat template to immutable content.""" + template = config.model.chat_template + if template is None: + return None + template_source = Path(os.path.expanduser(template)) + return template_source.read_text(encoding="utf-8") if template_source.is_file() else template + + +def _materialize_chat_template( + config: InferenceConfig, + template_content: str, + output_dir: Path | None, +) -> Path: + config_dir = output_dir or (get_config_dir(config.output_dir) / ENGINE_CONFIG_DIR) + template_path = config_dir / CHAT_TEMPLATE_ASSET + template_path.parent.mkdir(parents=True, exist_ok=True) + template_path.write_text(template_content, encoding="utf-8") + return template_path + + +def _frontend_model_arguments( + config: InferenceConfig, + output_dir: Path | None, + runtime_chat_template_path: Path | None, +) -> tuple[str, ...]: + template_content = resolve_chat_template_content(config) + if template_content is None: + return () + template_path = runtime_chat_template_path or _materialize_chat_template(config, template_content, output_dir) + tool_arguments = ( + ("--enable-auto-tool-choice", "--tool-call-parser", config.model.tool_call_parser) + if config.model.tool_call_parser is not None + else () + ) + reasoning_arguments = ( + ("--reasoning-parser", config.model.reasoning_parser) if config.model.reasoning_parser is not None else () + ) + return ( + *tool_arguments, + *reasoning_arguments, + "--dyn-chat-processor", + "vllm", + "--chat-template", + str(template_path), + ) + + +def build_frontend_process( + config: InferenceConfig, + *, + host: str | None = None, + port: int | None = None, + output_dir: Path | None = None, + runtime_chat_template_path: Path | None = None, +) -> DynamoProcessSpec: + """Build the canonical Dynamo frontend process contract.""" + environment = { + **config.env_vars, + "DYN_ENABLE_RL": "1", + "DYN_RL_PORT": "8001", + } + arguments = ( + "--http-host", + host or config.server.host or "0.0.0.0", + "--http-port", + str(port or config.server.port), + "--router-mode", + "kv", + "--router-reset-states", + "--enable-engine-apis", + *_frontend_model_arguments(config, output_dir, runtime_chat_template_path), + ) + return DynamoProcessSpec( + module="dynamo.frontend", + arguments=arguments, + environment_items=_environment_items(environment), + ) + + +def _worker_parser_arguments(config: InferenceConfig, role: Role) -> tuple[str, ...]: + if role == "prefill": + return () + tool_arguments = ( + ("--dyn-tool-call-parser", config.model.tool_call_parser) if config.model.tool_call_parser is not None else () + ) + reasoning_arguments = ( + ("--dyn-reasoning-parser", config.model.reasoning_parser) if config.model.reasoning_parser is not None else () + ) + return (*tool_arguments, *reasoning_arguments) + + +def _worker_endpoint_contract(namespace: str | None, component: str) -> tuple[dict[str, str], tuple[str, ...]]: + if namespace is None: + return {}, () + endpoint = f"dyn://{namespace}.{component}.generate" + return ( + {"DYN_NAMESPACE": namespace, "DYN_ENDPOINT": endpoint}, + ("--endpoint", endpoint), + ) + + +def build_worker_process( + config: InferenceConfig, + role: Role, + engine_config: Path, + *, + nixl_host: str | None, + nixl_port: int, + namespace: str | None = None, +) -> DynamoProcessSpec: + """Build the canonical Dynamo vLLM worker process contract.""" + resolved_namespace = namespace or config.env_vars.get("DYN_NAMESPACE") + component = _WORKER_COMPONENT[role] + endpoint_environment, endpoint_arguments = _worker_endpoint_contract(resolved_namespace, component) + environment = { + **config.env_vars, + **_role_environment(config, role), + "DYN_ENABLE_RL": "1", + "DYN_COMPONENT": component, + **endpoint_environment, + "VLLM_NIXL_SIDE_CHANNEL_PORT": str(nixl_port), + **({"VLLM_NIXL_SIDE_CHANNEL_HOST": nixl_host} if nixl_host is not None else {}), + "VLLM_PLUGINS": "prime_rl", + } + arguments = ( + "--engine-config-json", + str(engine_config), + *endpoint_arguments, + "--disaggregation-mode", + role, + "--enable-rl", + *_worker_parser_arguments(config, role), + ) + return DynamoProcessSpec( + module="dynamo.vllm", + arguments=arguments, + environment_items=_environment_items(environment), + ) + + +def build_engine_config( + config: InferenceConfig, + role: Role, + *, + kv_events_port: int | None = None, + data_parallel_rpc_port: int | None = None, +) -> dict[str, Any]: + """Build one deterministic vLLM ``AsyncEngineArgs`` object.""" + _validate_overrides("vllm_extra", config.vllm_extra) + overrides = _role_overrides(config, role) + _validate_overrides(f"{role}_vllm_overrides", overrides) + + values = vars(config.to_vllm()).copy() + for key in _ENGINE_CONFIG_EXCLUDED: + values.pop(key, None) + values.update(config.vllm_extra) + values.update(overrides) + + if config.deployment.type == "disaggregated": + # Each generated worker is an independent vLLM server. Preserve local + # DP within a worker, but never turn the P/D worker count into vLLM DP. + local_dp = config.dynamo_local_dp + values["data_parallel_size"] = local_dp + if local_dp == 1: + values.pop("data_parallel_size_local", None) + values.pop("data_parallel_rpc_port", None) + else: + values["data_parallel_size_local"] = local_dp + if data_parallel_rpc_port is not None: + values["data_parallel_rpc_port"] = data_parallel_rpc_port + + if role in ("prefill", "agg") and kv_events_port is not None: + values["kv_events_config"] = { + "publisher": "zmq", + "topic": "kv-events", + "endpoint": f"tcp://*:{kv_events_port}", + "enable_kv_cache_events": True, + } + else: + values.pop("kv_events_config", None) + + values["worker_extension_cls"] = _WORKER_EXTENSION_CLS[config.weight_broadcast.type] + return {key: value for key, value in values.items() if value is not None} + + +def _write_json(path: Path, value: dict[str, Any]) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, default=_json_default, indent=2, sort_keys=True) + "\n") + return path + + +def write_role_engine_configs(config: InferenceConfig, output_dir: Path | None = None) -> dict[Role, Path]: + """Write canonical role configs used by DGD and dry-run inspection.""" + config_dir = output_dir or (get_config_dir(config.output_dir) / ENGINE_CONFIG_DIR) + if config.deployment.type == "disaggregated": + return { + "prefill": _write_json( + config_dir / PREFILL_ENGINE_CONFIG, + build_engine_config(config, "prefill", kv_events_port=20080), + ), + "decode": _write_json(config_dir / DECODE_ENGINE_CONFIG, build_engine_config(config, "decode")), + } + return { + "agg": _write_json( + config_dir / AGG_ENGINE_CONFIG, + build_engine_config(config, "agg", kv_events_port=20080), + ) + } + + +def _visible_gpu_ids() -> list[str]: + configured = os.environ.get("CUDA_VISIBLE_DEVICES") + if configured: + return [gpu.strip() for gpu in configured.split(",") if gpu.strip()] + try: + output = subprocess.check_output( + ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError) as exc: + raise RuntimeError("Cannot discover GPUs; set CUDA_VISIBLE_DEVICES explicitly") from exc + return [line.strip() for line in output.splitlines() if line.strip()] + + +def build_local_worker_specs( + config: InferenceConfig, + output_dir: Path | None = None, + gpu_ids: list[str] | None = None, + namespace: str | None = None, +) -> list[DynamoWorkerSpec]: + """Allocate local workers and write instance-specific engine configs.""" + config_dir = output_dir or (get_config_dir(config.output_dir) / ENGINE_CONFIG_DIR) + available = gpu_ids if gpu_ids is not None else _visible_gpu_ids() + resolved_namespace = namespace or config.env_vars.get("DYN_NAMESPACE") or "dynamo" + + if config.deployment.type == "disaggregated": + deployment: DisaggregatedInferenceDeploymentConfig = config.deployment + if deployment.num_prefill_nodes != deployment.num_prefill_replicas: + raise ValueError("Local Dynamo requires one prefill node per prefill replica") + if deployment.num_decode_nodes != deployment.num_decode_replicas: + raise ValueError("Local Dynamo requires one decode node per decode replica") + roles = list(config.dynamo_worker_roles) + gpus_per_worker = config.dynamo_gpus_per_worker + else: + roles = list(config.dynamo_worker_roles) + gpus_per_worker = config.dynamo_gpus_per_worker + + required = len(roles) * gpus_per_worker + if len(available) < required: + raise ValueError(f"Dynamo topology requires {required} GPUs, but only {len(available)} are visible") + + specs: list[DynamoWorkerSpec] = [] + role_indexes: dict[Role, int] = {"agg": 0, "prefill": 0, "decode": 0} + for worker_index, role in enumerate(roles): + role_index = role_indexes[role] + role_indexes[role] += 1 + start = worker_index * gpus_per_worker + worker_gpus = tuple(available[start : start + gpus_per_worker]) + ports = _allocate_local_worker_ports(worker_index) + kv_events_port = ports.kv_events if role in ("prefill", "agg") else None + name = f"{role}-{role_index}" + engine_path = _write_json( + config_dir / f"{name}-engine.json", + build_engine_config( + config, + role, + kv_events_port=kv_events_port, + data_parallel_rpc_port=ports.data_parallel_rpc, + ), + ) + specs.append( + DynamoWorkerSpec( + name=name, + role=role, + gpu_ids=worker_gpus, + system_port=ports.system, + process=build_worker_process( + config, + role, + engine_path, + nixl_host="127.0.0.1", + nixl_port=ports.nixl, + namespace=resolved_namespace, + ), + ) + ) + return specs + + +def build_dry_run_worker_specs( + config: InferenceConfig, + output_dir: Path | None = None, +) -> list[DynamoWorkerSpec]: + """Build local specs without consulting host GPU hardware.""" + if config.deployment.type == "disaggregated": + gpu_count = len(config.dynamo_worker_roles) * config.dynamo_gpus_per_worker + else: + gpu_count = config.dynamo_gpus_per_worker + return build_local_worker_specs( + config, + output_dir=output_dir, + gpu_ids=[f"" for index in range(gpu_count)], + namespace=config.env_vars.get("DYN_NAMESPACE") or "dynamo", + ) + + +def build_worker_environment( + spec: DynamoWorkerSpec, + base_environment: dict[str, str], +) -> dict[str, str]: + return spec.process.environment(base_environment) | { + "CUDA_VISIBLE_DEVICES": ",".join(spec.gpu_ids), + "DYN_SYSTEM_PORT": str(spec.system_port), + } + + +def _terminate(process: subprocess.Popen) -> None: + if process.poll() is not None: + return + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + try: + process.wait(timeout=15) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + + +def run_dynamo_local(config: InferenceConfig) -> None: + """Run a Dynamo frontend and all configured workers until one exits.""" + environment = os.environ.copy() + environment.setdefault("DYN_DISCOVERY_BACKEND", "file") + environment.setdefault("DYN_EVENT_PLANE", "zmq") + environment.setdefault("DYN_FILE_KV_TTL_SECS", "1800") + namespace = config.env_vars.get("DYN_NAMESPACE") or environment.get("DYN_NAMESPACE") or f"prime-rl-{os.getpid()}" + environment["DYN_NAMESPACE"] = namespace + environment.setdefault("PYTHONHASHSEED", "0") + specs = build_local_worker_specs(config, namespace=namespace) + + def request_stop(_signum, _frame): + raise KeyboardInterrupt + + signal.signal(signal.SIGTERM, request_stop) + processes: list[subprocess.Popen] = [] + with tempfile.TemporaryDirectory(prefix="prime-dynamo-") as temporary_dir: + environment.setdefault("DYN_FILE_KV", str(Path(temporary_dir) / "discovery")) + frontend = build_frontend_process(config) + frontend_env = frontend.environment(environment) | {"CUDA_VISIBLE_DEVICES": ""} + frontend_env.pop("DYN_SYSTEM_PORT", None) + + try: + processes.append(subprocess.Popen(frontend.command(), env=frontend_env, start_new_session=True)) + for spec in specs: + worker_env = build_worker_environment(spec, environment) + processes.append(subprocess.Popen(spec.process.command(), env=worker_env, start_new_session=True)) + + exited_process = next((process for process in processes if process.poll() is not None), None) + while exited_process is None: + time.sleep(0.2) + exited_process = next((process for process in processes if process.poll() is not None), None) + returncode = exited_process.returncode + if returncode is None: + raise RuntimeError("Dynamo child exit was observed without a return code") + # A clean child exit is still a service failure while its siblings are supervised. + raise SystemExit(returncode if returncode != 0 else 1) + except KeyboardInterrupt: + return + finally: + for process in reversed(processes): + _terminate(process) diff --git a/src/prime_rl/inference/dynamo_admin.py b/src/prime_rl/inference/dynamo_admin.py new file mode 100644 index 0000000000..02fcaf8bec --- /dev/null +++ b/src/prime_rl/inference/dynamo_admin.py @@ -0,0 +1,489 @@ +"""Dynamo worker discovery and engine administration.""" + +from __future__ import annotations + +import asyncio +import os +from collections import Counter +from collections.abc import Awaitable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypeAlias +from urllib.parse import urlsplit, urlunsplit + +import httpx +from httpx import AsyncClient +from tenacity import AsyncRetrying, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential + +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.async_utils import gather_shielded +from prime_rl.utils.logger import get_logger + +NCCL_READY_MARKER = "NCCL_READY" +ADMIN_TIMEOUT_S = 300.0 +UPDATE_WEIGHTS_TIMEOUT_S = 720.0 +DISCOVERY_REQUEST_TIMEOUT_S = 10.0 +DISCOVERY_POLL_INTERVAL_S = 1.0 +_RETRYABLE_DISCOVERY_HTTP_STATUS_CODES = frozenset({408, 409, 429}) +_REQUIRED_ROUTES = frozenset( + { + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + } +) + +WorkerRole: TypeAlias = Literal["agg", "prefill", "decode"] + + +class _DiscoveryConvergenceError(RuntimeError): + """A structurally valid discovery state that may converge during startup.""" + + +@dataclass(frozen=True, slots=True) +class DynamoWorker: + """Restart-safe identity and admin capabilities for one Dynamo worker.""" + + instance_id: int + component: str + role: WorkerRole + system_url: str + model: str + routes: frozenset[str] + + +@dataclass(frozen=True, slots=True) +class DynamoTopology: + """The exact worker shape Prime expects to administer.""" + + roles: tuple[WorkerRole, ...] + gpus_per_worker: int + + def __post_init__(self) -> None: + if not self.roles: + raise ValueError("Dynamo topology must contain at least one worker") + invalid_roles = set(self.roles) - {"agg", "prefill", "decode"} + if invalid_roles: + raise ValueError(f"Dynamo topology contains invalid roles: {sorted(invalid_roles)}") + if isinstance(self.gpus_per_worker, bool) or self.gpus_per_worker < 1: + raise ValueError("Dynamo topology gpus_per_worker must be at least one") + + def validate(self, workers: Sequence[DynamoWorker]) -> None: + expected = Counter(self.roles) + observed = Counter(worker.role for worker in workers) + if observed != expected: + raise ValueError( + "Dynamo worker topology does not match the configured roles: " + f"expected {dict(sorted(expected.items()))}, observed {dict(sorted(observed.items()))}" + ) + + def role_for_component(self, component: str) -> WorkerRole: + normalized = component.casefold() + if "prefill" in normalized: + return "prefill" + if "decode" in normalized: + return "decode" + if normalized in {"backend", "vllmworker", "agg", "aggregate", "aggregated", "vllmaggworker"}: + expected = set(self.roles) + if "decode" in expected and "agg" not in expected: + return "decode" + if "agg" in expected and "decode" not in expected: + return "agg" + raise ValueError( + f"Dynamo worker component {component!r} is ambiguous for configured roles {sorted(expected)}" + ) + raise ValueError(f"Dynamo worker component {component!r} has no recognized inference role") + + +def _root_url(url: str) -> str: + return url.rstrip("/").removesuffix("/v1") + + +def discovery_urls(config: ClientConfig) -> list[str]: + if config.rl_base_url: + return [_root_url(url) for url in config.rl_base_url] + + configured = os.getenv("DYN_RL_DISCOVERY_URL") + if configured: + return [_root_url(url.strip()) for url in configured.split(",") if url.strip()] + + port = int(os.getenv("DYN_RL_PORT", "8001")) + urls: list[str] = [] + for base_url in config.base_url: + parsed = urlsplit(_root_url(base_url)) + host = parsed.hostname or "localhost" + if ":" in host: + host = f"[{host}]" + netloc = f"{host}:{port}" + urls.append(urlunsplit((parsed.scheme or "http", netloc, "", "", ""))) + return urls + + +def _worker_sort_key(worker: DynamoWorker) -> tuple[str, str, int, str]: + return (worker.role, worker.component, worker.instance_id, worker.system_url) + + +def _parse_worker(value: object, model_name: str, topology: DynamoTopology) -> DynamoWorker: + if not isinstance(value, Mapping): + raise ValueError("Dynamo worker discovery returned a non-object worker") + component = value.get("component") + instance_id = value.get("instance_id") + system_url = value.get("system_url") + model = value.get("model") + routes = value.get("routes") + if not isinstance(component, str) or not component: + raise ValueError("Dynamo worker discovery response is missing component") + if not isinstance(instance_id, int) or isinstance(instance_id, bool) or instance_id < 0: + raise ValueError(f"Dynamo worker {component!r} has an invalid instance_id") + error = value.get("error") + if error is not None: + if not isinstance(error, str) or not error: + raise ValueError(f"Dynamo worker {component}[{instance_id}] has an invalid error") + raise _DiscoveryConvergenceError(f"Dynamo worker {component}[{instance_id}] is unhealthy: {error}") + if not isinstance(system_url, str) or not system_url: + raise ValueError(f"Dynamo worker {component}[{instance_id}] is missing system_url") + if model != model_name: + raise ValueError( + f"Dynamo worker {component}[{instance_id}] at {system_url} serves {model!r}, expected {model_name!r}" + ) + if not isinstance(routes, list) or not all(isinstance(route, str) for route in routes): + raise ValueError(f"Dynamo worker {component}[{instance_id}] has invalid routes") + route_set = frozenset(routes) + missing = _REQUIRED_ROUTES - route_set + if missing: + raise ValueError(f"Dynamo worker {system_url} is missing RL routes: {sorted(missing)}") + return DynamoWorker( + instance_id=instance_id, + component=component, + role=topology.role_for_component(component), + system_url=_root_url(system_url), + model=model, + routes=route_set, + ) + + +def _parse_snapshot( + payload: object, + model_name: str, + topology: DynamoTopology, +) -> tuple[str, tuple[DynamoWorker, ...]]: + if not isinstance(payload, Mapping) or not isinstance(payload.get("workers"), list): + raise ValueError("Dynamo worker discovery returned an invalid response") + namespace = payload.get("namespace") + if not isinstance(namespace, str) or not namespace: + raise ValueError("Dynamo worker discovery response is missing namespace") + workers = tuple( + sorted( + (_parse_worker(value, model_name, topology) for value in payload["workers"]), + key=_worker_sort_key, + ) + ) + identities = {(worker.component, worker.instance_id) for worker in workers} + if len(identities) != len(workers): + raise ValueError("Dynamo worker discovery returned duplicate worker identities") + if len({worker.system_url for worker in workers}) != len(workers): + raise ValueError("Dynamo worker discovery returned duplicate system URLs") + return namespace, workers + + +def _retryable_discovery_error(error: Exception) -> bool: + if isinstance(error, httpx.HTTPStatusError): + status_code = error.response.status_code + return status_code in _RETRYABLE_DISCOVERY_HTTP_STATUS_CODES or status_code >= 500 + return isinstance( + error, + ( + _DiscoveryConvergenceError, + httpx.TimeoutException, + httpx.TransportError, + TimeoutError, + ), + ) + + +async def discover_workers( + discovery_clients: list[AsyncClient], + timeout: float, + *, + model_name: str, + topology: DynamoTopology, +) -> tuple[DynamoWorker, ...]: + """Wait for all frontends to report one identical, complete worker set. + + Incomplete membership, inconsistent valid snapshots, and worker probe errors + are startup convergence states. A model mismatch or missing RL route is an + incompatible deployment contract and fails immediately, as do malformed + payloads and permanent HTTP errors. + """ + if not discovery_clients: + raise ValueError("Dynamo worker discovery requires at least one frontend") + if timeout <= 0: + raise TimeoutError("Dynamo worker discovery deadline has already expired") + logger = get_logger() + last_error: Exception | None = None + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + + while (remaining := deadline - loop.time()) > 0: + try: + request_timeout = min(remaining, DISCOVERY_REQUEST_TIMEOUT_S) + async with asyncio.timeout(remaining): + results = await asyncio.gather( + *( + client.get("/v1/rl/workers", timeout=httpx.Timeout(request_timeout)) + for client in discovery_clients + ) + ) + snapshots: list[tuple[str, tuple[DynamoWorker, ...]]] = [] + for response in results: + response.raise_for_status() + snapshots.append(_parse_snapshot(response.json(), model_name, topology)) + first = snapshots[0] + if any(snapshot != first for snapshot in snapshots[1:]): + raise _DiscoveryConvergenceError("Dynamo discovery frontends returned inconsistent worker snapshots") + workers = first[1] + try: + topology.validate(workers) + except ValueError as exc: + raise _DiscoveryConvergenceError(str(exc)) from exc + logger.info(f"Discovered {len(workers)} Dynamo inference worker(s)") + return workers + except Exception as exc: + if not _retryable_discovery_error(exc): + raise + last_error = exc + remaining = deadline - loop.time() + if remaining > 0: + await asyncio.sleep(min(DISCOVERY_POLL_INTERVAL_S, remaining)) + + raise TimeoutError(f"Dynamo workers were not ready after {timeout} seconds: {last_error!r}") + + +def validate_worker_membership( + expected: Sequence[DynamoWorker], + discovered: Sequence[DynamoWorker], +) -> None: + expected_workers = tuple(sorted(expected, key=_worker_sort_key)) + discovered_workers = tuple(sorted(discovered, key=_worker_sort_key)) + if discovered_workers != expected_workers: + raise RuntimeError( + "Dynamo worker membership changed after initialization: " + f"expected {expected_workers!r}, discovered {discovered_workers!r}" + ) + + +class DynamoAdminAPI: + """Typed adapter for Dynamo's per-worker ``/engine`` endpoints.""" + + def __init__(self) -> None: + self._distributed_updates = False + self._distributed_initialization_indeterminate = False + self._weight_update_indeterminate = False + + def _require_unambiguous_admin_state(self) -> None: + if self._distributed_initialization_indeterminate: + raise RuntimeError( + "Dynamo distributed weight-group initialization is indeterminate after a prior failure or " + "cancellation; refusing further admin mutation" + ) + if self._weight_update_indeterminate: + raise RuntimeError( + "Dynamo worker weight state is indeterminate after a prior update or resume failure; refusing " + "further admin mutation" + ) + + @staticmethod + def _retryable(exception: BaseException) -> bool: + if isinstance(exception, httpx.HTTPStatusError): + return exception.response.status_code >= 500 + return isinstance(exception, (httpx.TimeoutException, httpx.TransportError)) + + async def _post( + self, + client: AsyncClient, + method: str, + body: dict | None = None, + *, + timeout_s: float = ADMIN_TIMEOUT_S, + retry_transient: bool = False, + ) -> dict: + async def post_once() -> dict: + response = await client.post( + f"/engine/{method}", + json=body or {}, + timeout=httpx.Timeout(connect=10.0, read=timeout_s, write=60.0, pool=10.0), + ) + response.raise_for_status() + payload = response.json() + if not isinstance(payload, dict): + raise ValueError(f"Dynamo /engine/{method} returned a non-object response") + if payload.get("status") != "ok": + raise RuntimeError(payload.get("message", f"Dynamo /engine/{method} failed")) + return payload + + if not retry_transient: + return await post_once() + + async for attempt in AsyncRetrying( + retry=retry_if_exception(self._retryable), + stop=stop_after_delay(2 * timeout_s) | stop_after_attempt(10), + wait=wait_exponential(multiplier=1, min=1, max=10), + reraise=True, + ): + with attempt: + return await post_once() + raise AssertionError("unreachable") + + @staticmethod + async def _settle_fanout(awaitables: Iterable[Awaitable[dict]], operation: str) -> None: + """Await every sibling even when one fails or the caller is cancelled.""" + tasks = [asyncio.create_task(awaitable) for awaitable in awaitables] + if not tasks: + return + results, cancellation = await gather_shielded(*tasks) + + failures = [result for result in results if isinstance(result, BaseException)] + primary: BaseException | None = cancellation or (failures[0] if failures else None) + if primary is None: + return + siblings = failures if cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Dynamo {operation} sibling also failed: {sibling!r}") + raise primary + + async def initialize_nccl( + self, + clients: list[AsyncClient], + *, + host: str, + port: int, + timeout: int, + inference_world_size: int | None, + gpus_per_worker: int, + quantize_in_weight_transfer: bool, + ) -> None: + self._require_unambiguous_admin_state() + if not clients: + raise ValueError("Cannot initialize NCCL without Dynamo workers") + if isinstance(gpus_per_worker, bool) or gpus_per_worker < 1: + raise ValueError("gpus_per_worker must be at least one") + expected_world_size = len(clients) * gpus_per_worker + world_size = expected_world_size if inference_world_size is None else inference_world_size + if world_size != expected_world_size: + raise ValueError( + f"inference_world_size={world_size} does not match {len(clients)} Dynamo workers " + f"with {gpus_per_worker} GPUs each ({expected_world_size})" + ) + try: + await self._settle_fanout( + ( + self._post( + client, + "init_weights_update_group", + { + "host": host, + "port": port, + "rank_offset": index * gpus_per_worker, + "inference_world_size": world_size, + "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, + "engine_rpc": "init_broadcaster", + }, + ) + for index, client in enumerate(clients) + ), + "init_weights_update_group", + ) + except BaseException: + # A lost response or one failed sibling cannot prove whether every + # rank committed the collective group. This object must not choose + # the filesystem path or retry into that ambiguous state. + self._distributed_initialization_indeterminate = True + raise + self._distributed_updates = True + + async def update_weights(self, clients: list[AsyncClient], weight_dir: Path | None, step: int) -> None: + self._require_unambiguous_admin_state() + try: + await self._settle_fanout( + ( + self._post( + client, + "pause_generation", + {"mode": "wait", "clear_cache": False}, + retry_transient=True, + ) + for client in clients + ), + "pause_generation", + ) + except BaseException as exc: + # Pausing cannot change weights. Settle the fanout, undo any + # successful pauses, and preserve the original pre-mutation error. + await self._resume_after_pre_mutation_failure(clients, exc) + raise + + try: + if weight_dir is not None: + marker = weight_dir / NCCL_READY_MARKER + marker.parent.mkdir(parents=True, exist_ok=True) + marker.touch() + + if self._distributed_updates: + body = { + "weight_version": str(step), + "weight_dir": weight_dir.as_posix() if weight_dir is not None else None, + "engine_rpc": "update_weights_from_path", + } + method = "update_weights_from_distributed" + else: + if weight_dir is None: + raise ValueError("Dynamo filesystem weight updates require weight_dir") + body = { + "model_path": str(weight_dir.resolve()), + "weight_version": str(step), + "engine_rpc": "update_weights_from_path", + } + method = "update_weights_from_disk" + except BaseException as exc: + # Validation/filesystem preparation failed before any update RPC + # was issued, so generation can safely resume on the old weights. + await self._resume_after_pre_mutation_failure(clients, exc) + raise + + try: + await self._settle_fanout( + (self._post(client, method, body, timeout_s=UPDATE_WEIGHTS_TIMEOUT_S) for client in clients), + method, + ) + except BaseException: + # At least one mutation RPC was issued. A failure or cancellation + # cannot prove which workers committed, so keep generation paused + # and permanently poison this admin object. + self._weight_update_indeterminate = True + raise + + try: + await self._resume_generation(clients) + except BaseException: + self._weight_update_indeterminate = True + raise + + async def _resume_generation(self, clients: list[AsyncClient]) -> None: + await self._settle_fanout( + (self._post(client, "resume_generation", retry_transient=True) for client in clients), + "resume_generation", + ) + + async def _resume_after_pre_mutation_failure( + self, + clients: list[AsyncClient], + primary_error: BaseException, + ) -> None: + try: + await self._resume_generation(clients) + except BaseException as resume_error: + self._weight_update_indeterminate = True + primary_error.add_note(f"Dynamo resume_generation cleanup also failed: {resume_error!r}") diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 8d1baa60a3..c79090d410 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/src/prime_rl/orchestrator/algo/__init__.py @@ -40,6 +40,7 @@ if TYPE_CHECKING: from prime_rl.configs.algorithm import AlgoConfig + from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.utils.client import InferencePool # Runtime dispatch is keyed on ``algo.type`` — it names the algorithm, and @@ -54,7 +55,12 @@ } -def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm: +def build_algorithm( + config: AlgoConfig, + policy_pool: InferencePool, + *, + policy_gate: MutablePolicyGate | None = None, +) -> Algorithm: cls = ALGORITHM_CLASSES[config.type] assert cls.action_loss_type == config.action_loss_type # config and runtime declare in two places # The Algorithm is the runtime of the algorithm config's training signal @@ -62,7 +68,9 @@ def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm # handed the live policy pool — opsd self-distills against it, others may # judge against it or ignore it. Other models (a frozen teacher, a hint # renderer) are built from the algorithm's own config in setup(). - return cls(config, policy_pool) + algorithm = cls(config, policy_pool) + algorithm.policy_gate = policy_gate + return algorithm __all__ = [ diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index b12d17df62..db7a1de24f 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -44,11 +44,13 @@ from prime_rl.configs.algorithm import ActionLossType, AlgoConfig, FrozenModelConfig from prime_rl.orchestrator.algo.routing import stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.policy_gate import PolicyRequestRejected from prime_rl.utils.logger import get_logger if TYPE_CHECKING: from renderers import RendererConfig + from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.orchestrator.types import Rollout from prime_rl.utils.client import InferencePool @@ -119,8 +121,15 @@ class Algorithm: action_loss_type: ClassVar[ActionLossType] = "rl" - def __init__(self, config: AlgoConfig, policy_pool: InferencePool): + def __init__( + self, + config: AlgoConfig, + policy_pool: InferencePool, + *, + policy_gate: MutablePolicyGate | None = None, + ): self.policy_pool = policy_pool + self.policy_gate = policy_gate self.connected_pools: list[InferencePool] = [] # frozen pools connected in setup(); closed at shutdown async def setup(self) -> None: @@ -152,7 +161,12 @@ async def finalize_rollout(self, rollout: Rollout) -> None: """Arrival phase (non-virtual): rollout-local scoring as each rollout is tokenized.""" if rollout.samples: - await self.score_rollout(rollout) + try: + await self.score_rollout(rollout) + except PolicyRequestRejected as exc: + # Preserve the sink's one-arrival-per-group accounting while + # dropping scoring that cannot use the generating version. + rollout.capture_error(exc) async def finalize_group(self, rollouts: list[Rollout]) -> None: """Group phase (non-virtual): group-relative scoring, then stamp each diff --git a/src/prime_rl/orchestrator/algo/opd.py b/src/prime_rl/orchestrator/algo/opd.py index 3135f2a9b2..e35444cc49 100644 --- a/src/prime_rl/orchestrator/algo/opd.py +++ b/src/prime_rl/orchestrator/algo/opd.py @@ -5,7 +5,7 @@ from prime_rl.configs.algorithm import OPDAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm -from prime_rl.utils.client import StaticInferencePool +from prime_rl.utils.client import FixedInferencePool if TYPE_CHECKING: from prime_rl.orchestrator.types import Rollout @@ -29,12 +29,12 @@ class OPDAlgorithm(Algorithm): def __init__(self, config: OPDAlgoConfig, policy_pool: InferencePool): super().__init__(config, policy_pool) self.teacher = config.teacher - self.teacher_pool: StaticInferencePool | None = None # static teacher endpoint, connected in setup() + self.teacher_pool: FixedInferencePool | None = None # fixed teacher endpoint, connected in setup() async def setup(self) -> None: pool = await self.connect(self.teacher) - if not isinstance(pool, StaticInferencePool): - raise TypeError("opd teacher must be a static endpoint — prefill scoring needs fixed endpoints") + if not isinstance(pool, FixedInferencePool): + raise TypeError("opd teacher must be a fixed endpoint — prefill scoring needs fixed endpoints") self.teacher_pool = pool async def score_rollout(self, rollout: Rollout) -> None: diff --git a/src/prime_rl/orchestrator/algo/opsd.py b/src/prime_rl/orchestrator/algo/opsd.py index 737666bea9..a8c5791218 100644 --- a/src/prime_rl/orchestrator/algo/opsd.py +++ b/src/prime_rl/orchestrator/algo/opsd.py @@ -1,14 +1,15 @@ from __future__ import annotations -import asyncio from typing import TYPE_CHECKING from prime_rl.configs.algorithm import OPSDAlgoConfig from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.utils.async_utils import gather_shielded if TYPE_CHECKING: from renderers.base import Renderer + from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.orchestrator.types import Rollout from prime_rl.transport import TrainingSample from prime_rl.utils.client import InferencePool @@ -30,8 +31,14 @@ class OPSDAlgorithm(Algorithm): action_loss_type = "ref_kl" - def __init__(self, config: OPSDAlgoConfig, policy_pool: InferencePool): - super().__init__(config, policy_pool) + def __init__( + self, + config: OPSDAlgoConfig, + policy_pool: InferencePool, + *, + policy_gate: MutablePolicyGate | None = None, + ): + super().__init__(config, policy_pool, policy_gate=policy_gate) self.demo_key = config.demo_key self.template = config.template self.renderer_config = config.renderer @@ -74,4 +81,18 @@ async def score_sample(sample: TrainingSample) -> None: # sample.token_ids (demo-conditioned, the trainer's ref_kl target). sample.ref_logprobs = full_logprobs[len(hint_block) :] - await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) + async def settle_scores() -> None: + results, cancellation = await gather_shielded(*(score_sample(sample) for sample in rollout.samples)) + failures = [result for result in results if isinstance(result, BaseException)] + primary: BaseException | None = cancellation or (failures[0] if failures else None) + if primary is not None: + siblings = failures if cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Another OPSD score failed: {sibling!r}") + raise primary + + if self.policy_gate is None: + await settle_scores() + return + async with self.policy_gate.request(expected_version=rollout.policy_version): + await settle_scores() diff --git a/src/prime_rl/orchestrator/component_supervision.py b/src/prime_rl/orchestrator/component_supervision.py new file mode 100644 index 0000000000..5d1cff41ab --- /dev/null +++ b/src/prime_rl/orchestrator/component_supervision.py @@ -0,0 +1,107 @@ +"""Failure propagation for orchestrator background components.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Sequence +from typing import TypeVar + +T = TypeVar("T") + +# Cleanup must not extend a user-facing operation timeout indefinitely. A task +# that suppresses cancellation is retained below and observed when it settles. +SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS = 1.0 +_ORPHANED_OPERATIONS: set[asyncio.Task] = set() + + +def _observe_operation(task: asyncio.Task) -> None: + _ORPHANED_OPERATIONS.discard(task) + try: + task.exception() + except BaseException: + pass + + +async def _cancel_operation_with_grace(task: asyncio.Task) -> asyncio.CancelledError | None: + task.cancel() + loop = asyncio.get_running_loop() + grace_elapsed = loop.create_future() + caller_cancellation: asyncio.CancelledError | None = None + + def finish_grace() -> None: + if not grace_elapsed.done(): + grace_elapsed.set_result(None) + + timer = loop.call_later(SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS, finish_grace) + try: + while not task.done() and not grace_elapsed.done(): + try: + await asyncio.wait((task, grace_elapsed), return_when=asyncio.FIRST_COMPLETED) + except asyncio.CancelledError as error: + # Finish bounded cleanup before propagating caller cancellation. + # Repeated cancellation must not make cleanup unbounded. + if caller_cancellation is None: + caller_cancellation = error + continue + finally: + timer.cancel() + + if task.done(): + _observe_operation(task) + else: + _ORPHANED_OPERATIONS.add(task) + task.add_done_callback(_observe_operation) + return caller_cancellation + + +def raise_if_component_failed(tasks: Sequence[asyncio.Task]) -> None: + """Raise when a background loop exits while the main loop is still live.""" + for task in tasks: + if not task.done(): + continue + name = task.get_name() + if task.cancelled(): + raise RuntimeError(f"Orchestrator component {name!r} was cancelled unexpectedly") + error = task.exception() + if error is not None: + raise error + raise RuntimeError(f"Orchestrator component {name!r} stopped unexpectedly") + + +async def run_with_component_supervision( + operation: Callable[[], Awaitable[T]], + component_tasks: Sequence[asyncio.Task], + *, + timeout: float | None = None, + timeout_description: str | None = None, +) -> T | None: + """Run one operation while racing every supervised component. + + A component failure wins when both sides complete in the same event-loop + turn. ``None`` denotes timeout unless ``timeout_description`` is supplied, + in which case timeout raises after cancelling the operation. + """ + raise_if_component_failed(component_tasks) + task = asyncio.create_task(operation()) + component_failure_selected = False + try: + await asyncio.wait( + [task, *component_tasks], + return_when=asyncio.FIRST_COMPLETED, + timeout=timeout, + ) + try: + raise_if_component_failed(component_tasks) + except BaseException: + component_failure_selected = True + raise + if task.done(): + return await task + if timeout_description is not None: + raise TimeoutError(f"{timeout_description} timed out after {timeout} seconds") + return None + finally: + if not task.done(): + caller_cancellation = await _cancel_operation_with_grace(task) + if caller_cancellation is not None and not component_failure_selected: + raise caller_cancellation diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 11a1ab4d34..c384e86686 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -2,22 +2,14 @@ - Capacity (``max_inflight_rollouts``) is shared across train + eval. A group-scoring task that runs N rollouts in one call reserves N permits. -- Optional rate limiting via ``AsyncLimiter(tasks_per_minute, 60)``. -- Emit-everything invariant: every dispatched rollout eventually reaches - ``out_q`` exactly once as a ``Rollout``. Failures - (env error, empty trajectory, task exception, off-policy cancel) carry - ``trace.error`` set; sinks decide drop / partial-train policy. +- Emit-everything invariant: every dispatched rollout reaches ``out_q`` once. + Failures carry ``trace.error``; sinks decide drop / partial-train policy. - ``DispatcherMode.PREFER_TRAIN`` / ``PREFER_EVAL`` controls which kind to schedule next. Transitions are level-triggered (driven by the eval source's emptiness), so in-flight rollouts of the opposite kind drain naturally on either side of an eval boundary. -- ``on_version_pending`` (called by the watcher before the engines pause for - the weight update) bumps ``off_policy_steps`` on in-flight train rollouts and - drops groups past ``max_off_policy_steps``. - Eval rollouts are measurements for the policy version they started with, - so they are allowed to finish even if training advances. Train rollouts - sampled from a frozen model never age — their sampler doesn't change - with policy updates. +- Dynamo fences and settles eval/live-policy work before worker pause; other + backends retain the off-policy window. Frozen-pool rollouts survive. Cancellations surface as synthetic ``Cancelled`` markers so the sink's count-to-``group_size`` finalization still fires. """ @@ -27,15 +19,22 @@ import asyncio import uuid from collections import Counter, defaultdict -from dataclasses import dataclass, field from enum import Enum, auto -from typing import Literal import verifiers.v1 as vf from aiolimiter import AsyncLimiter +from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics +from prime_rl.orchestrator.dispatcher_transactions import ( + EmissionRecord, + EmissionTracker, + emit_policy_cancellation_markers, + settle_transaction_cleanup, +) from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_source import EvalSource +from prime_rl.orchestrator.policy_gate import MutablePolicyGate, PolicyUpdateToken, SchedulingEpoch +from prime_rl.orchestrator.pool_identity import client_may_alias_pool, pools_may_alias from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( GroupState, @@ -44,7 +43,7 @@ Rollout, RolloutKind, ) -from prime_rl.utils.async_utils import safe_cancel, safe_cancel_all +from prime_rl.utils.async_utils import gather_shielded, safe_cancel, safe_cancel_all from prime_rl.utils.client import InferencePool, client_identity from prime_rl.utils.logger import get_logger @@ -56,69 +55,12 @@ class DispatcherMode(Enum): PREFER_EVAL = auto() -@dataclass -class DispatcherMetrics: - """Per-tick drain counters for the orchestrator's periodic log. - ``drained()`` returns the current values and clears them; point-in-time - gauges live on ``RolloutDispatcher.gauges`` instead.""" - - cancelled_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( - default_factory=lambda: defaultdict(int) - ) - errored_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( - default_factory=lambda: defaultdict(int) - ) - - def record_cancellation(self, *, kind: Literal["train", "eval"], env_name: str, n: int = 1) -> None: - self.cancelled_by_kind_env[(kind, env_name)] += n - - def record_error(self, *, kind: Literal["train", "eval"], env_name: str) -> None: - self.errored_by_kind_env[(kind, env_name)] += 1 - - def drained(self, *, train_envs: set[str], eval_envs: set[str]) -> dict[str, float]: - """Return per-tick counters and clear them. Emits the full pre- - registered key set every tick (zero when no activity) so the wandb - time axis stays dense and ``define_metric`` lines up.""" - out: dict[str, float] = {} - for kind in ("train", "eval"): - envs = train_envs if kind == "train" else eval_envs - cancelled_total = sum(self.cancelled_by_kind_env.get((kind, e), 0) for e in envs) - errored_total = sum(self.errored_by_kind_env.get((kind, e), 0) for e in envs) - out[f"dispatcher/cancelled/{kind}"] = float(cancelled_total) - out[f"dispatcher/errored/{kind}"] = float(errored_total) - for env in train_envs | eval_envs: - out[f"dispatcher/cancelled/{env}"] = float( - self.cancelled_by_kind_env.get(("train", env), 0) + self.cancelled_by_kind_env.get(("eval", env), 0) - ) - out[f"dispatcher/errored/{env}"] = float( - self.errored_by_kind_env.get(("train", env), 0) + self.errored_by_kind_env.get(("eval", env), 0) - ) - self.cancelled_by_kind_env.clear() - self.errored_by_kind_env.clear() - return out - - @staticmethod - def drain_keys(*, train_envs: set[str], eval_envs: set[str]) -> list[str]: - """Full set of keys ``drained`` may emit; used by the periodic - logger for ``wandb.define_metric``.""" - keys = [ - "dispatcher/cancelled/train", - "dispatcher/cancelled/eval", - "dispatcher/errored/train", - "dispatcher/errored/eval", - ] - for env in train_envs | eval_envs: - keys.append(f"dispatcher/cancelled/{env}") - keys.append(f"dispatcher/errored/{env}") - return keys - - class RolloutDispatcher: """``await dispatcher.start()`` runs the dispatch loop until ``stop()``. Pulls examples from ``TrainSource`` / ``EvalSource``, schedules rollouts under shared capacity, and emits ``Rollout``\\ s to - ``out_q``. The watcher drives ``on_version_pending`` for off-policy - cancellation; the orchestrator triggers eval epochs.""" + ``out_q``. The watcher drives ``on_version_pending`` for the policy-update + barrier; the orchestrator triggers eval epochs.""" def __init__( self, @@ -132,6 +74,8 @@ def __init__( max_inflight_rollouts: int, tasks_per_minute: float | None, max_off_policy_steps: int, + enforce_policy_update_barrier: bool, + policy_gate: MutablePolicyGate | None = None, ) -> None: self.policy = policy self.train_envs = train_envs @@ -142,6 +86,9 @@ def __init__( self.train_source = train_source self.eval_source = eval_source self.max_off_policy_steps = max_off_policy_steps + self.enforce_policy_update_barrier = enforce_policy_update_barrier + self.policy_gate = policy_gate or MutablePolicyGate(policy, enabled=enforce_policy_update_barrier) + self._policy_update: tuple[int, PolicyUpdateToken] | None = None self.max_inflight = max_inflight_rollouts self.inflight_permits = 0 @@ -151,19 +98,17 @@ def __init__( self.inflight: dict[asyncio.Task, InflightRollout] = {} self.groups: dict[uuid.UUID, GroupState] = {} + self._emissions = EmissionTracker() # Bounded so the dispatcher backpressures on a slow sink self.out_q: asyncio.Queue[Rollout] = asyncio.Queue(maxsize=max(8, self.max_inflight)) self.mode: DispatcherMode = DispatcherMode.PREFER_TRAIN - # Set by the orchestrator after the final train step; pipeline then - # winds down without scheduling new train rollouts + # Set after the final train step to wind down new scheduling. self.train_scheduling_disabled: bool = False self.metrics = DispatcherMetrics() - # Orchestrator-owned gate. When clear, ``fill_inflight`` returns - # without scheduling new groups. The dispatcher itself doesn't know - # *why* — the orchestrator toggles this based on step / policy lead. + # Orchestrator-owned step/policy-lead gate. self.dispatch_allowed = asyncio.Event() self.dispatch_allowed.set() @@ -267,40 +212,178 @@ async def stop(self) -> None: self.task = None async def on_version_pending(self, step: int) -> None: - """Bump off-policy counters and drop groups past - ``max_off_policy_steps`` (drop_group emits ``Cancelled`` markers so - the sink still finalizes the partial group). Eval rollouts are not - aged because they are tied to their start-time policy version. - - Runs *before* the inference engines are paused for the weight update so - the resulting aborts are processed while the engine is still stepping — - otherwise the orphaned KV transfers crash the decode engine on resume - (see ``WeightWatcher.apply_policy_update``).""" + """Prepare for mutation: Dynamo fences/drains mutable-policy requests + before worker pause; other backends retain their off-policy window. + Pre-pause cancellation lets P/D abort and connector cleanup settle.""" + if not self.enforce_policy_update_barrier: + await self._advance_off_policy_window() + return + + token = await self.policy_gate.begin_update(step=step) + self._policy_update = (step, token) + try: + claimed_groups, claimed_tasks, claimed_emissions = self._claim_mutable_policy_work() + results, cancellation = await gather_shielded( + self._settle_policy_requests(claimed_groups, claimed_tasks, claimed_emissions), + self.policy_gate.wait_idle(), + ) + failures = [result for result in results if isinstance(result, BaseException)] + primary: BaseException | None = cancellation or (failures[0] if failures else None) + if primary is not None: + siblings = failures if cancellation is not None else failures[1:] + for sibling in siblings: + primary.add_note(f"Another policy barrier drain failed: {sibling!r}") + raise primary + except BaseException as primary_error: + await settle_transaction_cleanup( + self._reopen_policy_admission(step, token), + primary_error, + "Policy transition rollback", + ) + raise + + async def on_new_version(self, step: int) -> None: + """Reopen admission after the new policy is live.""" + if self.enforce_policy_update_barrier: + await self._reopen_policy_admission(step) + + async def on_version_update_failed(self, step: int, error: BaseException) -> None: + """Roll back only the transition fence; the policy stays unchanged.""" + if self.enforce_policy_update_barrier: + await self._reopen_policy_admission(step) + + @property + def policy_update_pending(self) -> bool: + return self.policy_gate.pending + + async def _reopen_policy_admission(self, step: int, token: PolicyUpdateToken | None = None) -> None: + owned = self._policy_update + if owned is None or owned[0] != step or (token is not None and owned[1] is not token): + raise RuntimeError(f"Dispatcher does not own the pending policy transition for step {step}") + await self.policy_gate.finish_update(owned[1]) + self._policy_update = None + + async def _advance_off_policy_window(self) -> None: + """Retain the established non-Dynamo in-flight tolerance policy.""" stale_groups: set[uuid.UUID] = set() - cancelled = 0 for meta in self.inflight.values(): if meta.kind != "train": continue - # Frozen-sourced rollouts never go stale — their sampler doesn't - # change with policy updates. - if not self.train_envs.get(meta.env_name).sampler.samples_from_live_policy: + meta.uses_mutable_policy = meta.uses_mutable_policy or self._uses_mutable_policy(meta.kind, meta.env_name) + if not meta.uses_mutable_policy: continue meta.off_policy_steps += 1 if meta.off_policy_steps > self.max_off_policy_steps: stale_groups.add(meta.group_id) - for gid in stale_groups: - removed = await self.drop_group(gid) - cancelled += removed - + cancelled = 0 + for group_id in stale_groups: + cancelled += await self.drop_group(group_id) if cancelled: get_logger().warning( f"Cancelled {cancelled} train rollouts past max_off_policy_steps={self.max_off_policy_steps}. " "Consider increasing it to avoid this." ) - async def on_new_version(self, step: int) -> None: - """No-op: the dispatcher drains in ``on_version_pending`` (pre-pause).""" + def _uses_mutable_policy(self, kind: RolloutKind, env_name: str) -> bool: + if kind == "eval": + return True + pool, _model_name, samples_from_live_policy = self._train_pool_for(env_name) + # A separately-constructed "frozen" pool is safe only when its model + # and request/admin endpoints do not alias the mutable policy service. + return samples_from_live_policy or pools_may_alias(pool, self.policy_pool) + + def _claim_mutable_policy_work( + self, + ) -> tuple[ + dict[uuid.UUID, GroupState], + list[tuple[asyncio.Task, InflightRollout]], + list[EmissionRecord], + ]: + group_ids: set[uuid.UUID] = set() + for group_id, group in self.groups.items(): + group.uses_mutable_policy = group.uses_mutable_policy or self._uses_mutable_policy( + group.kind, group.env_name + ) + if group.uses_mutable_policy: + group_ids.add(group_id) + for meta in self.inflight.values(): + meta.uses_mutable_policy = meta.uses_mutable_policy or self._uses_mutable_policy(meta.kind, meta.env_name) + if meta.uses_mutable_policy: + group_ids.add(meta.group_id) + + claimed_groups = { + group_id: group for group_id in group_ids if (group := self.groups.pop(group_id, None)) is not None + } + claimed_tasks: list[tuple[asyncio.Task, InflightRollout]] = [] + for task, meta in list(self.inflight.items()): + if meta.group_id not in group_ids: + continue + del self.inflight[task] + self.release(meta.rollout_count) + claimed_tasks.append((task, meta)) + claimed_emissions = self._emissions.claim(group_ids) + return claimed_groups, claimed_tasks, claimed_emissions + + async def _settle_policy_requests( + self, + groups: dict[uuid.UUID, GroupState], + claimed: list[tuple[asyncio.Task, InflightRollout]], + emissions: list[EmissionRecord], + ) -> None: + tasks = [task for task, _meta in claimed] + already_settled = {task for task in tasks if task.done()} + for task in tasks: + task.cancel() + + results: list[object] = [] + cancellation: asyncio.CancelledError | None = None + if tasks: + results, cancellation = await gather_shielded(*tasks) + + emission_results: list[object] = [] + emission_cancellation: asyncio.CancelledError | None = None + if emissions: + emission_results, emission_cancellation = await gather_shielded( + *(record.done.wait() for record in emissions) + ) + + metadata_by_group: dict[uuid.UUID, InflightRollout] = {} + for _task, meta in claimed: + metadata_by_group.setdefault(meta.group_id, meta) + + marker_results, marker_cancellation = await gather_shielded( + emit_policy_cancellation_markers( + groups, + metadata_by_group, + out_q=self.out_q, + stopped=self.stopped, + metrics=self.metrics, + ) + ) + + failures = [ + result + for task, result in zip(tasks, results, strict=True) + if task not in already_settled + and isinstance(result, BaseException) + and not isinstance(result, asyncio.CancelledError) + ] + failures.extend(result for result in marker_results if isinstance(result, BaseException)) + failures.extend(record.error for record in emissions if record.error is not None) + failures.extend(result for result in emission_results if isinstance(result, BaseException)) + primary: BaseException | None = ( + cancellation or emission_cancellation or marker_cancellation or (failures[0] if failures else None) + ) + if primary is not None: + siblings = ( + failures + if cancellation is not None or emission_cancellation is not None or marker_cancellation is not None + else failures[1:] + ) + for sibling in siblings: + primary.add_note(f"Another policy barrier cleanup failed: {sibling!r}") + raise primary async def fill_inflight(self) -> None: """Schedule new rollouts up to ``max_inflight``, honoring @@ -309,26 +392,24 @@ async def fill_inflight(self) -> None: respects it. When ``PREFER_EVAL``'s source exhausts we flip back to ``PREFER_TRAIN`` so the eval tail drains alongside fresh train.""" while True: - if self.available_permits <= 0: + epoch = await self.policy_gate.scheduling_epoch() + if epoch is None or self.available_permits <= 0: return if self.mode == DispatcherMode.PREFER_EVAL: - # PREFER_EVAL is only entered when the orchestrator triggers - # eval, which requires ``eval_source`` to be configured + # PREFER_EVAL implies a configured eval source. assert self.eval_source is not None if not self.eval_has_work: - # Eval source + all eval groups fully dispatched. Flip - # to PREFER_TRAIN so any remaining permits go to train - # while the in-flight eval tail completes naturally + # Fill remaining permits with train while eval drains. self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained") continue - scheduled = await self.try_schedule("eval") + scheduled = await self.try_schedule("eval", epoch=epoch) if not scheduled: return else: # PREFER_TRAIN — respects the orchestrator's dispatch gate if not self.dispatch_allowed.is_set(): return - scheduled = await self.try_schedule("train") + scheduled = await self.try_schedule("train", epoch=epoch) if not scheduled: return @@ -339,7 +420,7 @@ def switch_mode(self, new_mode: DispatcherMode, *, reason: str) -> None: get_logger().info(f"Switching dispatcher mode to prefer {prefer} rollouts because {reason}") self.mode = new_mode - async def try_schedule(self, kind: RolloutKind) -> bool: + async def try_schedule(self, kind: RolloutKind, *, epoch: SchedulingEpoch) -> bool: """Schedule one rollout of ``kind``: prefer continuing an existing group (keeps prefix-cache hits); otherwise open a fresh group from the corresponding source. Returns False if nothing could be @@ -356,14 +437,20 @@ async def try_schedule(self, kind: RolloutKind) -> bool: env = envs.get(group.env_name) cost = group.rollouts_to_schedule if env.requires_group_scoring else 1 if cost <= self.available_permits: - return await self.schedule_group_rollout(gid, group) + return await self.schedule_group_rollout(gid, group, epoch=epoch) - fresh = self.next_fresh_group(kind, envs) - if fresh is None: - return False - gid = uuid.uuid4() - self.groups[gid] = fresh - return await self.schedule_group_rollout(gid, fresh) + # Pop and publish a fresh group inside the short scheduling commit. + # An update either sees the group in its claim snapshot or invalidates + # this epoch before the source is consumed. + async with self.policy_gate.scheduling_commit(epoch) as admitted: + if not admitted or self.available_permits <= 0: + return False + fresh = self.next_fresh_group(kind, envs) + if fresh is None: + return False + gid = uuid.uuid4() + self.groups[gid] = fresh + return await self.schedule_group_rollout(gid, fresh, epoch=epoch) def next_fresh_group(self, kind: RolloutKind, envs) -> GroupState | None: """Pop the next example from the corresponding source and wrap it in @@ -390,9 +477,16 @@ def next_fresh_group(self, kind: RolloutKind, envs) -> GroupState | None: target_rollouts=group_size, eval_step=eval_step, policy_version_at_start=self.policy.version, + uses_mutable_policy=self._uses_mutable_policy(kind, env_name), ) - async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) -> bool: + async def schedule_group_rollout( + self, + group_id: uuid.UUID, + group: GroupState, + *, + epoch: SchedulingEpoch, + ) -> bool: """Dispatch one ``run_rollout`` / ``run_group`` task for this group. Returns False only if we couldn't even schedule one rollout (no clients @@ -409,7 +503,8 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - else: pool, model_name, live_sourced = self._train_pool_for(group.env_name) - # Pin a single client per group to keep prefix-cache hits + # Resolve a client and rate-limit outside the gate. Both operations may + # wait indefinitely for elastic discovery or quota replenishment. if group.pinned_client is None: if group.kind == "eval": client = await pool.get_eval_client() @@ -420,7 +515,6 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - client = await pool.select_train_client(load) if group_id not in self.groups: return False - group.pinned_client = client else: client = group.pinned_client @@ -428,58 +522,70 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - if env_collection is None: return False env = env_collection.get(group.env_name) - # Frozen-sourced train rollouts hit a frozen pool; salting per policy - # version would invalidate its prefix cache every weight update for - # no reason. - if live_sourced: - cache_salt = str(group.policy_version_at_start) - else: - cache_salt = None - if env.requires_group_scoring: permits = group.rollouts_to_schedule - group.rollouts_to_schedule = 0 - await self.acquire(permits) - task: asyncio.Task = asyncio.create_task( - env.run_group( - client=client, - task_idx=group.task_idx, - model_name=model_name, - group_size=permits, - cache_salt=cache_salt, - ) - ) else: permits = 1 - group.rollouts_to_schedule -= 1 - await self.acquire(permits) - task = asyncio.create_task( - env.run_rollout( - client=client, - task_idx=group.task_idx, - model_name=model_name, - cache_salt=cache_salt, + # Snapshot selected identity before elastic churn can occur at the next await. + group.uses_mutable_policy = ( + group.uses_mutable_policy + or live_sourced + or pools_may_alias(pool, self.policy_pool) + or client_may_alias_pool(client, self.policy_pool) + ) + await self._wait_for_rate_limit(permits) + + async with self.policy_gate.scheduling_commit(epoch) as admitted: + if ( + not admitted + or self.groups.get(group_id) is not group + or permits > self.available_permits + or group.rollouts_to_schedule < permits + or (group.kind == "train" and (self.train_scheduling_disabled or not self.dispatch_allowed.is_set())) + ): + return False + + group.pinned_client = client + cache_salt = str(group.policy_version_at_start) if group.uses_mutable_policy else None + if env.requires_group_scoring: + group.rollouts_to_schedule = 0 + task: asyncio.Task = asyncio.create_task( + env.run_group( + client=client, + task_idx=group.task_idx, + model_name=model_name, + group_size=permits, + cache_salt=cache_salt, + ) + ) + else: + group.rollouts_to_schedule -= 1 + task = asyncio.create_task( + env.run_rollout( + client=client, + task_idx=group.task_idx, + model_name=model_name, + cache_salt=cache_salt, + ) ) - ) - self.inflight[task] = InflightRollout( - kind=group.kind, - env_name=group.env_name, - group_id=group_id, - policy_version=group.policy_version_at_start, - rollout_count=permits, - client_config=client, - eval_step=group.eval_step, - ) - return True + self.inflight_permits += permits + self.inflight[task] = InflightRollout( + kind=group.kind, + env_name=group.env_name, + group_id=group_id, + policy_version=group.policy_version_at_start, + rollout_count=permits, + client_config=client, + eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, + ) + return True - async def acquire(self, n: int) -> None: - """Reserve ``n`` permits + rate-limit each one. Caller must precheck - ``available_permits >= n``; this is not a blocking acquire.""" + async def _wait_for_rate_limit(self, n: int) -> None: for _ in range(n): if self.rate_limiter is not None: await self.rate_limiter.acquire() - self.inflight_permits += 1 def release(self, n: int) -> None: self.inflight_permits -= n @@ -496,6 +602,16 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: return # already handled by drop_group / cancel_inflight_rollouts self.release(meta.rollout_count) group = self.groups.get(meta.group_id) + async with self._emissions.track(meta.group_id): + await self._emit_completed_task(task, meta, group) + + async def _emit_completed_task( + self, + task: asyncio.Task, + meta: InflightRollout, + group: GroupState | None, + ) -> None: + """Convert one settled task into its complete output transaction.""" is_synth_exception = False try: @@ -536,9 +652,6 @@ async def emit_rollout(self, meta: InflightRollout, group: GroupState | None, ro if group is not None: eval_step = group.eval_step policy_version = group.policy_version_at_start - group.emitted += 1 - if group.emitted >= group.target_rollouts: - self.groups.pop(meta.group_id, None) rollout.kind = meta.kind rollout.env_name = meta.env_name @@ -549,6 +662,10 @@ async def emit_rollout(self, meta: InflightRollout, group: GroupState | None, ro assert eval_step is not None, "eval rollout missing eval_step" rollout.eval_step = eval_step await self.out_q.put(rollout) + if group is not None: + group.emitted += 1 + if group.emitted >= group.target_rollouts: + self.groups.pop(meta.group_id, None) async def drop_group(self, group_id: uuid.UUID) -> int: """Cancel remaining in-flight tasks for this group and emit a @@ -599,6 +716,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: policy_version=group.policy_version_at_start, rollout_count=1, eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, ) unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): @@ -619,6 +737,7 @@ async def drop_group(self, group_id: uuid.UUID) -> int: policy_version=group.policy_version_at_start if group else 0, rollout_count=1, eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, ) if group is not None else None diff --git a/src/prime_rl/orchestrator/dispatcher_metrics.py b/src/prime_rl/orchestrator/dispatcher_metrics.py new file mode 100644 index 0000000000..cf4e5608ea --- /dev/null +++ b/src/prime_rl/orchestrator/dispatcher_metrics.py @@ -0,0 +1,58 @@ +"""Drain counters owned by the rollout dispatcher.""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Literal + + +@dataclass +class DispatcherMetrics: + """Per-tick cancellation and error counters for pipeline logging.""" + + cancelled_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( + default_factory=lambda: defaultdict(int) + ) + errored_by_kind_env: dict[tuple[Literal["train", "eval"], str], int] = field( + default_factory=lambda: defaultdict(int) + ) + + def record_cancellation(self, *, kind: Literal["train", "eval"], env_name: str, n: int = 1) -> None: + self.cancelled_by_kind_env[(kind, env_name)] += n + + def record_error(self, *, kind: Literal["train", "eval"], env_name: str) -> None: + self.errored_by_kind_env[(kind, env_name)] += 1 + + def drained(self, *, train_envs: set[str], eval_envs: set[str]) -> dict[str, float]: + """Return the dense counter set for this tick and clear it.""" + out: dict[str, float] = {} + for kind in ("train", "eval"): + envs = train_envs if kind == "train" else eval_envs + out[f"dispatcher/cancelled/{kind}"] = float( + sum(self.cancelled_by_kind_env.get((kind, env), 0) for env in envs) + ) + out[f"dispatcher/errored/{kind}"] = float(sum(self.errored_by_kind_env.get((kind, env), 0) for env in envs)) + for env in train_envs | eval_envs: + out[f"dispatcher/cancelled/{env}"] = float( + self.cancelled_by_kind_env.get(("train", env), 0) + self.cancelled_by_kind_env.get(("eval", env), 0) + ) + out[f"dispatcher/errored/{env}"] = float( + self.errored_by_kind_env.get(("train", env), 0) + self.errored_by_kind_env.get(("eval", env), 0) + ) + self.cancelled_by_kind_env.clear() + self.errored_by_kind_env.clear() + return out + + @staticmethod + def drain_keys(*, train_envs: set[str], eval_envs: set[str]) -> list[str]: + """Return every key :meth:`drained` may emit.""" + keys = [ + "dispatcher/cancelled/train", + "dispatcher/cancelled/eval", + "dispatcher/errored/train", + "dispatcher/errored/eval", + ] + for env in train_envs | eval_envs: + keys.extend((f"dispatcher/cancelled/{env}", f"dispatcher/errored/{env}")) + return keys diff --git a/src/prime_rl/orchestrator/dispatcher_transactions.py b/src/prime_rl/orchestrator/dispatcher_transactions.py new file mode 100644 index 0000000000..fd92956c6c --- /dev/null +++ b/src/prime_rl/orchestrator/dispatcher_transactions.py @@ -0,0 +1,128 @@ +"""Small transactional helpers for dispatcher output ownership.""" + +from __future__ import annotations + +import asyncio +import uuid +from collections import defaultdict +from collections.abc import AsyncIterator, Awaitable, Iterable +from contextlib import asynccontextmanager +from dataclasses import dataclass, field + +import verifiers.v1 as vf + +from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics +from prime_rl.orchestrator.types import GroupState, InflightRollout, Rollout +from prime_rl.utils.async_utils import gather_shielded, safe_cancel +from prime_rl.utils.logger import get_logger + + +@dataclass +class EmissionRecord: + """A completion handler that still owns group-output accounting.""" + + done: asyncio.Event = field(default_factory=asyncio.Event) + error: BaseException | None = None + + +class EmissionTracker: + """Make completion handlers visible after their inflight task is popped.""" + + def __init__(self) -> None: + self._by_group: dict[uuid.UUID, list[EmissionRecord]] = defaultdict(list) + + @asynccontextmanager + async def track(self, group_id: uuid.UUID) -> AsyncIterator[None]: + record = EmissionRecord() + self._by_group[group_id].append(record) + try: + yield + except BaseException as exc: + record.error = exc + raise + finally: + record.done.set() + records = self._by_group[group_id] + records.remove(record) + if not records: + self._by_group.pop(group_id, None) + + def claim(self, group_ids: Iterable[uuid.UUID]) -> list[EmissionRecord]: + return [record for group_id in group_ids for record in self._by_group.get(group_id, ())] + + +async def settle_transaction_cleanup( + cleanup: Awaitable[object], + primary_error: BaseException, + description: str, +) -> None: + """Settle cleanup despite repeated cancellation, preserving the primary error.""" + results, cancellation = await gather_shielded(cleanup) + failures = [result for result in results if isinstance(result, BaseException)] + for cleanup_error in failures: + primary_error.add_note(f"{description} also failed: {cleanup_error!r}") + if cancellation is not None: + primary_error.add_note(f"{description} was cancelled again but settled before propagation") + + +async def _put_before_stop( + out_q: asyncio.Queue[Rollout], + stopped_event: asyncio.Event, + rollout: Rollout, +) -> None: + put = asyncio.create_task(out_q.put(rollout)) + stopped = asyncio.create_task(stopped_event.wait()) + try: + await asyncio.wait((put, stopped), return_when=asyncio.FIRST_COMPLETED) + if put.done(): + put.result() + else: + raise RuntimeError("Dispatcher stopped while emitting policy barrier cancellation markers") + finally: + if not put.done(): + await safe_cancel(put) + if not stopped.done(): + await safe_cancel(stopped) + + +async def emit_policy_cancellation_markers( + groups: dict[uuid.UUID, GroupState], + metadata_by_group: dict[uuid.UUID, InflightRollout], + *, + out_q: asyncio.Queue[Rollout], + stopped: asyncio.Event, + metrics: DispatcherMetrics, +) -> None: + """Finish every claimed group without hanging after dispatcher stop.""" + cancelled = 0 + for group_id, group in groups.items(): + owed = max(0, group.target_rollouts - group.emitted) + if owed == 0: + continue + meta = metadata_by_group.get(group_id) or InflightRollout( + kind=group.kind, + env_name=group.env_name, + group_id=group_id, + policy_version=group.policy_version_at_start, + rollout_count=1, + eval_step=group.eval_step, + uses_mutable_policy=group.uses_mutable_policy, + ) + for _ in range(owed): + rollout = Rollout( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=group.task_idx, prompt=None)), + errors=[vf.Error(type="Cancelled", message="Policy update barrier")], + stop_condition="error", + kind=meta.kind, + env_name=meta.env_name, + group_id=meta.group_id, + policy_version=group.policy_version_at_start, + off_policy_steps=meta.off_policy_steps, + eval_step=group.eval_step if meta.kind == "eval" else None, + ) + await _put_before_stop(out_q, stopped, rollout) + group.emitted += 1 + metrics.record_cancellation(kind=meta.kind, env_name=meta.env_name, n=owed) + cancelled += owed + if cancelled: + get_logger().debug(f"Policy update barrier cancelled {cancelled} mutable-policy rollout(s)") diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 75e1a55d8b..439c02af87 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -280,14 +280,14 @@ class TrainEnvs(Envs[TrainEnv]): :class:`Sampler` and runtime :class:`Algorithm`, built from the env's resolved algorithm config.""" - def __init__(self, configs: Sequence[TrainEnvConfig], *, policy_pool, renderer_config=None): + def __init__(self, configs: Sequence[TrainEnvConfig], *, policy_pool, renderer_config=None, policy_gate=None): self._envs: dict[str, TrainEnv] = {} for config in configs: assert config.algo is not None, "TrainEnvConfig.algo must be resolved before env construction" env = TrainEnv( config, Sampler(config.algo.sampling, policy_pool, renderer_config), - build_algorithm(config.algo, policy_pool), + build_algorithm(config.algo, policy_pool, policy_gate=policy_gate), ) self._envs[env.name] = env diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index c17ed9e2c7..f883b1bcd7 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -39,7 +39,9 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive imports from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.ckpt import setup_ckpt_manager -from prime_rl.orchestrator.dispatcher import DispatcherMetrics, DispatcherMode, RolloutDispatcher +from prime_rl.orchestrator.component_supervision import raise_if_component_failed, run_with_component_supervision +from prime_rl.orchestrator.dispatcher import DispatcherMode, RolloutDispatcher +from prime_rl.orchestrator.dispatcher_metrics import DispatcherMetrics from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_sink import EvalSink from prime_rl.orchestrator.eval_source import EvalSource @@ -50,6 +52,8 @@ monkey_patch_oai_iterable_types, ) from prime_rl.orchestrator.periodic_logger import PeriodicLogger +from prime_rl.orchestrator.policy_gate import MutablePolicyGate +from prime_rl.orchestrator.train_finalization import finalize_train_batch as finalize_train_batch_step from prime_rl.orchestrator.train_sink import TrainSink from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( @@ -71,7 +75,6 @@ from prime_rl.trainer.model import setup_tokenizer from prime_rl.transport import TrainingBatch, setup_training_batch_sender from prime_rl.utils.async_utils import EventLoopLagMonitor, EventLoopLagStats, safe_cancel -from prime_rl.utils.client import init_nccl_broadcast from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.logger import format_time, get_logger, setup_logger from prime_rl.utils.monitor import setup_monitor @@ -90,11 +93,6 @@ # shutdown wedges (env-server ZMQ recv, vLLM admin aclose, etc) SHUTDOWN_TIMEOUT_S = 300 -# Abort after this many consecutive train batches drop all rollouts to -# post-batch filters — usually a misconfigured filter or homogeneous-reward -# dataset; fail loudly instead of spinning -MAX_CONSECUTIVE_EMPTY_BATCHES = 10 - # Maximum batches the orchestrator may run ahead of the trainer. The # dispatcher is paused via ``update_dispatch_gate`` once this is exceeded; # resumed when the watcher advances ``policy.version``. @@ -123,6 +121,7 @@ class Orchestrator: train_source: TrainSource train_sink: TrainSink dispatcher: RolloutDispatcher + policy_gate: MutablePolicyGate watcher: WeightWatcher lag_monitor: EventLoopLagMonitor periodic_logger: PeriodicLogger @@ -210,6 +209,12 @@ async def setup(self) -> None: self.renderer, self.policy_inference = await setup_policy_inference_pool( config=config, tokenizer=self.tokenizer ) + # The deployment boundary may resolve the client from DYN_RL_TOPOLOGY; + # derive the strict gate from the pool that was actually constructed. + self.policy_gate = MutablePolicyGate( + self.policy, + enabled=self.policy_inference.admin_api == "dynamo", + ) self.mm_token_type_ids_mapping = ( getattr(self.renderer, "mm_token_type_id_map", None) if self.renderer is not None else None ) @@ -242,7 +247,10 @@ async def setup(self) -> None: get_logger().info("Loading training environments") self.train_envs = TrainEnvs( - config.train.env, policy_pool=self.policy_inference, renderer_config=config.renderer + config.train.env, + policy_pool=self.policy_inference, + renderer_config=config.renderer, + policy_gate=self.policy_gate, ) get_logger().debug( f"Loaded {len(self.train_envs)} training environment(s) ({', '.join(self.train_envs.names)})" @@ -293,8 +301,7 @@ async def setup(self) -> None: get_logger().info(f"Initializing weight broadcast ({config.weight_broadcast})") if config.weight_broadcast.type == "nccl": - await init_nccl_broadcast( - self.policy_inference.admin_clients, + await self.policy_inference.init_nccl_broadcast( config.weight_broadcast.host, config.weight_broadcast.port, config.weight_broadcast.timeout, @@ -351,6 +358,8 @@ async def setup(self) -> None: max_inflight_rollouts=config.max_inflight_rollouts, tasks_per_minute=config.tasks_per_minute, max_off_policy_steps=config.max_off_policy_steps, + enforce_policy_update_barrier=self.policy_gate.enabled, + policy_gate=self.policy_gate, ) self.train_sink = TrainSink( config, @@ -455,14 +464,14 @@ async def main_loop(self) -> None: to the train / eval sink. Both sinks return a finalized batch (or ``None``) from ``add()``; we just dispatch on the result.""" while not self.stopped.is_set(): + raise_if_component_failed(self.component_tasks) if self.draining and self.dispatcher.is_idle: get_logger().info("Pipeline drained, exiting main loop") self.stopped.set() break - try: - rollout: Rollout = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) - except asyncio.TimeoutError: + rollout = await self._next_rollout() + if rollout is None: continue # Every completed rollout — errored, filtered, or never batched — lands in the @@ -491,136 +500,25 @@ async def main_loop(self) -> None: await self.finalize_train_batch(train_batch) async def finalize_train_batch(self, batch: TrainBatch) -> None: - """Ship one ``TrainBatch`` out to the trainer and handle the I/O - side-effects (ckpt, save_rollouts, reference scoring, sender.send, - metrics, heartbeat, progress, eval trigger). The sink has already - done all data-transformation work.""" - config = self.config - step = self.progress.step - - # Sink-to-sink cycle time — the actual time between batches, not - # including the orchestrator's ship I/O (overlapped with the - # dispatcher producing the next batch) - now = time.perf_counter() - step_time = (now - self.last_batch_at) if self.last_batch_at is not None else 0.0 - self.last_batch_at = now - - if config.max_steps is not None and step > config.max_steps: - self.draining = True - self.dispatcher.disable_train_scheduling() - n_cancelled = await self.dispatcher.cancel_inflight_train_rollouts() - get_logger().info( - f"Draining pipeline (cancelled {n_cancelled} in-flight train rollout(s); " - f"any in-flight evals will complete)" - ) - return - - if not batch.samples: - self.consecutive_empty_batches += 1 - get_logger().warning( - f"Step {step}: empty train batch (0 of {len(batch.rollouts)} generated rollouts shipped — " - f"all errored or filtered out) " - f"(consecutive empty batches: {self.consecutive_empty_batches}/{MAX_CONSECUTIVE_EMPTY_BATCHES})" - ) - if self.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: - raise RuntimeError( - f"{self.consecutive_empty_batches} consecutive empty train batches — " - "check filter config (pre_batch_filters / post_batch_filters) or task difficulty." - ) - return - self.consecutive_empty_batches = 0 - n_trainable = sum(1 for r in batch.rollouts if r.is_trainable) - if n_trainable / len(batch.rollouts) <= 0.1: - get_logger().warning( - f"Only {n_trainable}/{len(batch.rollouts)} generated rollouts are trainable " - f"({n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config" - ) - - # The effective (clean, trained-on) subset lands in the per-step ``effective`` trace file - # at ship time; the full arrival window already streamed into ``all`` on arrival. - # to_record drops the per-node training tensors — they're for training, not the rollout - # record, and can't round-trip json (raw numpy bytes). - effective = batch.rollouts.effective - records = [r.to_record() for r in effective] - await asyncio.to_thread(save_rollouts, records, get_trace_path(config.output_dir, step, "train", "effective")) - - await self.sender.send(TrainingBatch(examples=batch.samples, step=step)) - self.progress.step += 1 - self.update_dispatch_gate() - # Checkpoint the step we just shipped (resume point: continue at step + 1). - save_ckpt_time = await self.maybe_save_ckpt(step) - trim_process_memory() - - # Rollout metrics over the {agg,} × {all,effective} matrix. ``batch.rollouts`` is the - # full arrival window (errored + filtered included); ``.effective`` is the clean subset. - metrics: dict[str, float] = {} - for subset, pool in (("all", batch.rollouts), ("effective", effective)): - metrics |= pool.metrics.to_wandb(prefix="train/agg", subset=subset) - for env_name, env_pool in pool.by_env().items(): - metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) - - # Progress / timing / env-share / pre-filter accounting (assembled here, not in the metrics - # objects). ``num_tokens`` is over the full arrival window; the input/output breakdown is over - # the effective (shipped) subset, summing the same ``vf.Trace`` token properties the metric - # matrix reports. - num_tokens = sum(r.num_total_tokens for r in batch.rollouts) - num_input = sum(r.num_input_tokens for r in effective) - num_output = sum(r.num_output_tokens for r in effective) - num_rollouts = len(batch.rollouts) - num_unique_examples = len({r.group_id for r in batch.rollouts}) - metrics |= { - "progress/tokens": num_tokens, - "progress/input_tokens": num_input, - "progress/output_tokens": num_output, - "progress/rollouts": num_rollouts, - "progress/tasks": num_unique_examples, - "progress/total_tokens": self.progress.total_tokens, - "progress/total_rollouts": self.progress.total_samples, - "progress/total_tasks": self.progress.total_problems, - "time/step": step_time, - "time/save_ckpt": save_ckpt_time, - "time/wait_for_policy": self.wait_for_policy_time, - "step": step, - } - for env_name, env_pool in batch.rollouts.by_env().items(): - metrics[f"batch/{env_name}"] = len(env_pool) / len(batch.rollouts) - if self.train_sink.pre_filter_seen > 0: - metrics["pre_filters/all/dropped_rate"] = ( - self.train_sink.pre_filter_dropped / self.train_sink.pre_filter_seen - ) - for name, count in self.train_sink.pre_filter_dropped_by_name.items(): - metrics[f"pre_filters/all/{name}/rate"] = count / self.train_sink.pre_filter_seen - self.monitor.log(metrics, step=step) - self.wait_for_policy_time = 0.0 - self.monitor.log_samples(effective.rollouts, step=step) - self.monitor.log_distributions( - distributions={ - "rewards": [r.reward for r in effective], - "advantages": [a for r in effective if (a := r.scalar_advantage()) is not None], - }, - step=step, + """Persist, ship, checkpoint, and report one finalized train batch.""" + await finalize_train_batch_step(self, batch) + + async def _next_rollout(self) -> Rollout | None: + """Wait for output while supervising fatal background components.""" + return await run_with_component_supervision( + self.dispatcher.out_q.get, + self.component_tasks, + timeout=0.5, ) - if self.usage_reporter is not None: - run_id = os.getenv("RUN_ID", "") - if run_id: - self.usage_reporter.report_training_usage( - run_id=run_id, - step=step, - tokens=num_input + num_output, - ) - if self.heart is not None: - self.heart.beat() - - self.progress.total_tokens += num_tokens - self.progress.total_samples += num_rollouts - self.progress.total_problems += num_unique_examples - - self.log_train_batch(batch, step=step, step_time=step_time) - - self.train_sink.reset_pre_filter_stats() - self.maybe_trigger_eval(self.progress.step) - trim_process_memory() + async def _send_to_trainer(self, batch: TrainingBatch) -> None: + """Race shipment against fatal components, preferring any failure.""" + await run_with_component_supervision( + lambda: self.sender.send(batch), + self.component_tasks, + timeout=self.config.rollout_transport.send_timeout_seconds, + timeout_description=f"Training batch {batch.step} send", + ) def maybe_trigger_eval(self, step: int) -> None: """Fire eligible eval epochs and flip to ``PREFER_EVAL`` if anything @@ -705,46 +603,6 @@ def collect_pipeline_view(self) -> tuple[str, dict[str, float]]: payload["event_loop_lag/n"] = float(lag_stats.n) return body, payload - def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> None: - """Per-step ``Step …`` success line. Multi-env runs append an indented ``╰─`` line per env. - ``Error`` is the sink-level rate (errored arrivals / total arrivals, over the full window); - the quality metrics are over the effective (clean, trained-on) subset; ``Trainable`` is - relative to all generated rollouts.""" - rollouts = batch.rollouts - effective = rollouts.effective - eff = effective.metrics - n_generated = len(rollouts) - n_trainable = sum(1 for r in rollouts if r.is_trainable) - trainable_rate = (n_trainable / n_generated) if n_generated else 0.0 - max_off_policy = max((r.off_policy_steps for r in effective), default=0) - - head = ( - f"Step {step} | {format_time(step_time):>7} | Reward {eff.reward.mean():.4f} | " - f"Trainable {n_trainable}/{n_generated} ({trainable_rate:.1%}) | " - f"Turns {eff.num_turns.mean():.1f} | Branches {eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max_off_policy} | " - f"Error {rollouts.metrics.has_error.mean():.1%} | Truncation {eff.is_truncated.mean():.1%}" - ) - if len(self.train_envs) <= 1: - get_logger().success(head) - return - - by_env = rollouts.by_env() - name_width = max((len(n) for n in by_env), default=0) - lines = [head] - for env_name in sorted(by_env): - pool = by_env[env_name] - env_eff_pool = pool.effective - env_eff = env_eff_pool.metrics - ratio = (len(pool) / n_generated) if n_generated else 0.0 - lines.append( - f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | Reward {env_eff.reward.mean():.4f} | " - f"Turns {env_eff.num_turns.mean():.1f} | Branches {env_eff.num_branches.mean():.1f} | " - f"Max Off-Policy {max((r.off_policy_steps for r in env_eff_pool), default=0)} | " - f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_eff.is_truncated.mean():.1%}" - ) - get_logger().success("\n\t\t ".join(lines)) - async def finalize_eval_batch(self, batch: EvalBatch) -> None: """Persist + log one completed eval epoch (save_rollouts, monitor.log_eval_samples, monitor.log).""" @@ -849,6 +707,9 @@ async def on_new_version(self, step: int) -> None: re-evaluate the dispatch gate (may resume if the trainer caught up).""" self.update_dispatch_gate() + async def on_version_update_failed(self, step: int, error: BaseException) -> None: + """No transition state is owned here; the policy version is unchanged.""" + async def stop(self) -> None: """Bounded best-effort teardown of all components. Has a global timeout so a wedged peer can't keep the process alive forever — diff --git a/src/prime_rl/orchestrator/policy_gate.py b/src/prime_rl/orchestrator/policy_gate.py new file mode 100644 index 0000000000..8c10b2de5b --- /dev/null +++ b/src/prime_rl/orchestrator/policy_gate.py @@ -0,0 +1,133 @@ +"""Admission and in-flight accounting for mutable-policy inference calls.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass + +from prime_rl.orchestrator.types import Policy + + +class PolicyRequestRejected(RuntimeError): + """A request cannot safely run against the mutable policy version.""" + + +@dataclass(frozen=True) +class SchedulingEpoch: + """Snapshot revalidated at the dispatcher's non-yielding commit point.""" + + value: int + + +@dataclass(frozen=True) +class PolicyUpdateToken: + """Unique ownership of one closed-gate transition.""" + + step: int + epoch: int + + +class MutablePolicyGate: + """Serialize policy mutation with every request that depends on its weights. + + Dispatcher scheduling prepares outside the lock, then revalidates a + :class:`SchedulingEpoch` at its short commit point. Other policy I/O, such + as OPSD prefill scoring, uses :meth:`request` for its full lifetime. + Closing the gate prevents new work; callers can then cancel + dispatcher-owned tasks and await :meth:`wait_idle` before pausing. + """ + + def __init__(self, policy: Policy, *, enabled: bool) -> None: + self.policy = policy + self.enabled = enabled + self._admission_lock = asyncio.Lock() + self._epoch = 0 + self._pending_token: PolicyUpdateToken | None = None + self._active_requests = 0 + self._idle = asyncio.Event() + self._idle.set() + + @property + def pending(self) -> bool: + return self.enabled and self._pending_token is not None + + async def scheduling_epoch(self) -> SchedulingEpoch | None: + """Take a short-lived admission snapshot before slow preparation.""" + if not self.enabled: + return SchedulingEpoch(self._epoch) + async with self._admission_lock: + if self._pending_token is not None: + return None + return SchedulingEpoch(self._epoch) + + @asynccontextmanager + async def scheduling_commit(self, epoch: SchedulingEpoch) -> AsyncIterator[bool]: + """Revalidate and serialize only the scheduling state commit. + + The caller must not await inside the admitted branch. Client discovery + and rate limiting belong before this context so an update can close the + gate promptly. + """ + if not self.enabled: + yield True + return + async with self._admission_lock: + yield self._pending_token is None and epoch.value == self._epoch + + @asynccontextmanager + async def request(self, *, expected_version: int) -> AsyncIterator[None]: + """Register one non-dispatcher policy call for its complete lifetime.""" + if not self.enabled: + yield + return + + async with self._admission_lock: + if self._pending_token is not None: + raise PolicyRequestRejected( + f"Mutable-policy request rejected because a policy update is pending (expected version " + f"{expected_version})" + ) + if expected_version != self.policy.version: + raise PolicyRequestRejected( + f"Mutable-policy request expected policy version {expected_version}, but current version is " + f"{self.policy.version}" + ) + self._active_requests += 1 + self._idle.clear() + + try: + yield + finally: + # No await here: even repeated cancellation must not strand the + # active count and deadlock the mutation barrier. + self._active_requests -= 1 + if self._active_requests == 0: + self._idle.set() + + async def begin_update(self, *, step: int) -> PolicyUpdateToken: + """Close admission after every in-progress scheduling commit.""" + if not self.enabled: + return PolicyUpdateToken(step=step, epoch=self._epoch) + async with self._admission_lock: + if self._pending_token is not None: + raise RuntimeError(f"A policy update is already pending while preparing step {step}") + self._epoch += 1 + token = PolicyUpdateToken(step=step, epoch=self._epoch) + self._pending_token = token + return token + + async def finish_update(self, token: PolicyUpdateToken) -> None: + """Reopen admission after a successful or proven pre-mutation failure.""" + if not self.enabled: + return + async with self._admission_lock: + if self._pending_token is not token: + raise RuntimeError(f"Policy update token for step {token.step} does not own the pending transition") + self._pending_token = None + + async def wait_idle(self) -> None: + """Wait until every already-admitted non-dispatcher request settles.""" + if self.enabled: + await self._idle.wait() diff --git a/src/prime_rl/orchestrator/pool_identity.py b/src/prime_rl/orchestrator/pool_identity.py new file mode 100644 index 0000000000..6429bb6732 --- /dev/null +++ b/src/prime_rl/orchestrator/pool_identity.py @@ -0,0 +1,54 @@ +"""Serving-resource identity used by the mutable-policy barrier.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import verifiers.v1 as vf + + from prime_rl.utils.client import InferencePool + + +def _normalized_url(value: object) -> str: + return str(value).rstrip("/").removesuffix("/v1") + + +def serving_identity(pool: InferencePool) -> tuple[str, frozenset[str], frozenset[str]]: + """Return model, request endpoints, and admin endpoints for one pool.""" + request_endpoints = frozenset(_normalized_url(client.base_url) for client in pool.train_clients) + admin_endpoints = frozenset(_normalized_url(client.base_url) for client in pool.admin_clients) + return str(pool.model_name), request_endpoints, admin_endpoints + + +def client_may_alias_pool(client: vf.ClientConfig, pool: InferencePool) -> bool: + """Whether one selected request client addresses a pool request endpoint.""" + selected_endpoint = _normalized_url(client.base_url) + return any(selected_endpoint == _normalized_url(candidate.base_url) for candidate in pool.train_clients) + + +def pools_may_alias(left: InferencePool, right: InferencePool) -> bool: + """Conservatively detect two pool objects backed by one mutable model. + + Inline frozen references construct a separate Python pool, so object + identity is only the fast path. Equal model names plus an overlapping + request or admin endpoint identify an alias. Missing endpoint information + is ambiguous and therefore treated as mutable. + """ + if left is right: + return True + left_model, left_requests, left_admin = serving_identity(left) + right_model, right_requests, right_admin = serving_identity(right) + # The worker-admin endpoint identifies the mutable engine itself. It + # aliases even when two logical model names or frontend routes differ. + if left_admin and right_admin and left_admin & right_admin: + return True + if left_model != right_model: + return False + if not left_requests or not right_requests: + return True + if left_requests & right_requests: + return True + if not left_admin or not right_admin: + return True + return False diff --git a/src/prime_rl/orchestrator/train_finalization.py b/src/prime_rl/orchestrator/train_finalization.py new file mode 100644 index 0000000000..00932ab901 --- /dev/null +++ b/src/prime_rl/orchestrator/train_finalization.py @@ -0,0 +1,242 @@ +"""Ordered persistence, shipment, and reporting for one train batch.""" + +from __future__ import annotations + +import asyncio +import os +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +from prime_rl.orchestrator.utils import save_rollouts, trim_process_memory +from prime_rl.transport import TrainingBatch +from prime_rl.utils.logger import format_time, get_logger +from prime_rl.utils.pathing import get_trace_path + +if TYPE_CHECKING: + from prime_rl.configs.orchestrator import OrchestratorConfig + from prime_rl.orchestrator.dispatcher import RolloutDispatcher + from prime_rl.orchestrator.envs import TrainEnvs + from prime_rl.orchestrator.metrics import TrainRollouts + from prime_rl.orchestrator.train_sink import TrainSink + from prime_rl.orchestrator.types import Progress, TrainBatch + from prime_rl.utils.heartbeat import Heartbeat + from prime_rl.utils.monitor.base import Monitor + from prime_rl.utils.usage_reporter import UsageReporter + +MAX_CONSECUTIVE_EMPTY_BATCHES = 10 + + +class TrainFinalizationHost(Protocol): + config: OrchestratorConfig + progress: Progress + dispatcher: RolloutDispatcher + train_envs: TrainEnvs + train_sink: TrainSink + monitor: Monitor + usage_reporter: UsageReporter | None + heart: Heartbeat | None + last_batch_at: float | None + consecutive_empty_batches: int + draining: bool + wait_for_policy_time: float + + async def _send_to_trainer(self, batch: TrainingBatch) -> None: ... + + def update_dispatch_gate(self) -> None: ... + + async def maybe_save_ckpt(self, step: int) -> float: ... + + def maybe_trigger_eval(self, step: int) -> None: ... + + +@dataclass(frozen=True) +class TrainStepReport: + metrics: dict[str, float] + num_tokens: int + num_input: int + num_output: int + num_rollouts: int + num_unique_examples: int + + +async def finalize_train_batch(host: TrainFinalizationHost, batch: TrainBatch) -> None: + """Preserve the step's persist → send → checkpoint → report transaction.""" + step = host.progress.step + step_time = _start_step_clock(host) + if await _skip_unshippable_batch(host, batch, step): + return + effective, save_ckpt_time = await _persist_and_ship(host, batch, step) + report = _build_train_step_report(host, batch, effective, step, step_time, save_ckpt_time) + _publish_train_step_report(host, batch, effective, report, step, step_time) + + +def _start_step_clock(host: TrainFinalizationHost) -> float: + now = time.perf_counter() + step_time = (now - host.last_batch_at) if host.last_batch_at is not None else 0.0 + host.last_batch_at = now + return step_time + + +async def _skip_unshippable_batch(host: TrainFinalizationHost, batch: TrainBatch, step: int) -> bool: + if host.config.max_steps is not None and step > host.config.max_steps: + host.draining = True + host.dispatcher.disable_train_scheduling() + n_cancelled = await host.dispatcher.cancel_inflight_train_rollouts() + get_logger().info( + f"Draining pipeline (cancelled {n_cancelled} in-flight train rollout(s); any in-flight evals will complete)" + ) + return True + if not batch.samples: + host.consecutive_empty_batches += 1 + get_logger().warning( + f"Step {step}: empty train batch (0 of {len(batch.rollouts)} generated rollouts shipped — " + f"all errored or filtered out) (consecutive empty batches: " + f"{host.consecutive_empty_batches}/{MAX_CONSECUTIVE_EMPTY_BATCHES})" + ) + if host.consecutive_empty_batches >= MAX_CONSECUTIVE_EMPTY_BATCHES: + raise RuntimeError( + f"{host.consecutive_empty_batches} consecutive empty train batches — " + "check filter config (pre_batch_filters / post_batch_filters) or task difficulty." + ) + return True + host.consecutive_empty_batches = 0 + n_trainable = sum(1 for rollout in batch.rollouts if rollout.is_trainable) + if n_trainable / len(batch.rollouts) <= 0.1: + get_logger().warning( + f"Only {n_trainable}/{len(batch.rollouts)} generated rollouts are trainable " + f"({n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config" + ) + return False + + +async def _persist_and_ship( + host: TrainFinalizationHost, + batch: TrainBatch, + step: int, +) -> tuple[TrainRollouts, float]: + effective = batch.rollouts.effective + records = [rollout.to_record() for rollout in effective] + await asyncio.to_thread( + save_rollouts, + records, + get_trace_path(host.config.output_dir, step, "train", "effective"), + ) + await host._send_to_trainer(TrainingBatch(examples=batch.samples, step=step)) + host.progress.step += 1 + host.update_dispatch_gate() + save_ckpt_time = await host.maybe_save_ckpt(step) + trim_process_memory() + return effective, save_ckpt_time + + +def _build_train_step_report( + host: TrainFinalizationHost, + batch: TrainBatch, + effective: TrainRollouts, + step: int, + step_time: float, + save_ckpt_time: float, +) -> TrainStepReport: + metrics: dict[str, float] = {} + for subset, pool in (("all", batch.rollouts), ("effective", effective)): + metrics |= pool.metrics.to_wandb(prefix="train/agg", subset=subset) + for env_name, env_pool in pool.by_env().items(): + metrics |= env_pool.metrics.to_wandb(prefix=f"train/{env_name}", subset=subset) + num_tokens = sum(rollout.num_total_tokens for rollout in batch.rollouts) + num_input = sum(rollout.num_input_tokens for rollout in effective) + num_output = sum(rollout.num_output_tokens for rollout in effective) + num_rollouts = len(batch.rollouts) + num_unique_examples = len({rollout.group_id for rollout in batch.rollouts}) + metrics |= { + "progress/tokens": num_tokens, + "progress/input_tokens": num_input, + "progress/output_tokens": num_output, + "progress/rollouts": num_rollouts, + "progress/tasks": num_unique_examples, + "progress/total_tokens": host.progress.total_tokens, + "progress/total_rollouts": host.progress.total_samples, + "progress/total_tasks": host.progress.total_problems, + "time/step": step_time, + "time/save_ckpt": save_ckpt_time, + "time/wait_for_policy": host.wait_for_policy_time, + "step": step, + } + for env_name, env_pool in batch.rollouts.by_env().items(): + metrics[f"batch/{env_name}"] = len(env_pool) / num_rollouts + if host.train_sink.pre_filter_seen > 0: + metrics["pre_filters/all/dropped_rate"] = host.train_sink.pre_filter_dropped / host.train_sink.pre_filter_seen + for name, count in host.train_sink.pre_filter_dropped_by_name.items(): + metrics[f"pre_filters/all/{name}/rate"] = count / host.train_sink.pre_filter_seen + return TrainStepReport(metrics, num_tokens, num_input, num_output, num_rollouts, num_unique_examples) + + +def _publish_train_step_report( + host: TrainFinalizationHost, + batch: TrainBatch, + effective: TrainRollouts, + report: TrainStepReport, + step: int, + step_time: float, +) -> None: + host.monitor.log(report.metrics, step=step) + host.wait_for_policy_time = 0.0 + host.monitor.log_samples(effective.rollouts, step=step) + host.monitor.log_distributions( + distributions={ + "rewards": [rollout.reward for rollout in effective], + "advantages": [value for rollout in effective if (value := rollout.scalar_advantage()) is not None], + }, + step=step, + ) + if host.usage_reporter is not None and (run_id := os.getenv("RUN_ID", "")): + host.usage_reporter.report_training_usage( + run_id=run_id, + step=step, + tokens=report.num_input + report.num_output, + ) + if host.heart is not None: + host.heart.beat() + host.progress.total_tokens += report.num_tokens + host.progress.total_samples += report.num_rollouts + host.progress.total_problems += report.num_unique_examples + _log_train_batch(host, batch, step=step, step_time=step_time) + host.train_sink.reset_pre_filter_stats() + host.maybe_trigger_eval(host.progress.step) + trim_process_memory() + + +def _log_train_batch(host: TrainFinalizationHost, batch: TrainBatch, *, step: int, step_time: float) -> None: + rollouts = batch.rollouts + effective = rollouts.effective + metrics = effective.metrics + n_generated = len(rollouts) + n_trainable = sum(1 for rollout in rollouts if rollout.is_trainable) + trainable_rate = (n_trainable / n_generated) if n_generated else 0.0 + max_off_policy = max((rollout.off_policy_steps for rollout in effective), default=0) + head = ( + f"Step {step} | {format_time(step_time):>7} | Reward {metrics.reward.mean():.4f} | " + f"Trainable {n_trainable}/{n_generated} ({trainable_rate:.1%}) | " + f"Turns {metrics.num_turns.mean():.1f} | Branches {metrics.num_branches.mean():.1f} | " + f"Max Off-Policy {max_off_policy} | Error {rollouts.metrics.has_error.mean():.1%} | " + f"Truncation {metrics.is_truncated.mean():.1%}" + ) + if len(host.train_envs) <= 1: + get_logger().success(head) + return + by_env = rollouts.by_env() + name_width = max((len(name) for name in by_env), default=0) + lines = [head] + for env_name in sorted(by_env): + pool = by_env[env_name] + env_effective = pool.effective + env_metrics = env_effective.metrics + ratio = (len(pool) / n_generated) if n_generated else 0.0 + lines.append( + f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | " + f"Reward {env_metrics.reward.mean():.4f} | Turns {env_metrics.num_turns.mean():.1f} | " + f"Branches {env_metrics.num_branches.mean():.1f} | " + f"Max Off-Policy {max((r.off_policy_steps for r in env_effective), default=0)} | " + f"Error {pool.metrics.has_error.mean():.1%} | Truncation {env_metrics.is_truncated.mean():.1%}" + ) + get_logger().success("\n\t\t ".join(lines)) diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index e3d0e93198..0db6e4219d 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -51,6 +51,7 @@ class InflightRollout: client_config: vf.ClientConfig | None = None off_policy_steps: int = 0 eval_step: int | None = None + uses_mutable_policy: bool = False @dataclass @@ -67,6 +68,7 @@ class GroupState: eval_step: int | None = None pinned_client: vf.ClientConfig | None = None policy_version_at_start: int = 0 + uses_mutable_policy: bool = False class Rollout(vf.Trace[DataT], Generic[DataT]): @@ -172,8 +174,14 @@ class VersionObserver(Protocol): ``on_version_pending`` fires *before* the inference engines are paused for the weight update; ``on_new_version`` fires *after* the new weights are live - and ``Policy`` has been mutated.""" + and ``Policy`` has been mutated. ``on_version_update_failed`` rolls back + transition-only state only when failure is known to precede engine + mutation; indeterminate engine-update failures remain fenced. A pending + hook that raises must roll back state it partially entered before raising; + only hooks that return successfully receive a later failure callback.""" async def on_version_pending(self, step: int) -> None: ... async def on_new_version(self, step: int) -> None: ... + + async def on_version_update_failed(self, step: int, error: BaseException) -> None: ... diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 2b23369045..a844d5fe5e 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -12,6 +12,7 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.utils.client import setup_inference_pool from prime_rl.utils.logger import InterceptHandler, get_logger, setup_logger +from prime_rl.utils.policy_client_config import policy_client_config_from_environment from prime_rl.utils.utils import ( get_broadcast_dir, get_ckpt_dir, @@ -31,7 +32,7 @@ async def setup_policy_inference_pool(*, config: OrchestratorConfig, tokenizer): use plain chat-completions.""" from renderers.base import create_renderer - client_config = config.model.client + client_config = policy_client_config_from_environment(config.model.client) model_name = config.model.name renderer = create_renderer(tokenizer, config.renderer) get_logger().info(f"Initialized {type(renderer).__name__} for {model_name}") diff --git a/src/prime_rl/orchestrator/watcher.py b/src/prime_rl/orchestrator/watcher.py index c01d349f40..966d69dfdb 100644 --- a/src/prime_rl/orchestrator/watcher.py +++ b/src/prime_rl/orchestrator/watcher.py @@ -1,6 +1,6 @@ -"""WeightWatcher: polls the broadcast dir, advances ``Policy``, notifies -observers (dispatcher → off-policy cancel). Standalone async task; the -orchestrator's barrier bounds the in-flight lead.""" +"""WeightWatcher: polls broadcasts and applies policy updates behind the +dispatcher admission/drain barrier. Standalone async task; the orchestrator's +lead gate separately bounds sampling ahead of the trainer.""" from __future__ import annotations @@ -9,7 +9,7 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.types import Policy, VersionObserver -from prime_rl.utils.async_utils import safe_cancel +from prime_rl.utils.async_utils import gather_shielded, safe_cancel from prime_rl.utils.client import InferencePool from prime_rl.utils.logger import format_time, get_logger from prime_rl.utils.pathing import get_broadcast_dir, get_step_path, wait_for_path @@ -92,8 +92,11 @@ async def apply_policy_update(self, next_step: int) -> None: f"Orchestrator resumed: checkpoint {next_step} ready (after {format_time(self.last_wait_for_ckpt_time)})" ) - # Drain off-policy rollouts BEFORE pausing the inference engines. - # Aborting a rollout triggers vLLM's KV-connector cleanup (NIXL's + # Establish the backend-specific application transition BEFORE + # pausing inference engines. Dynamo fences new dispatch and drains + # every mutable-policy request; the vLLM admin path retains its + # configured off-policy cancellation window. Aborting a rollout + # triggers vLLM's KV-connector cleanup (NIXL's # ``_reqs_not_processed``), which is only propagated to the workers # while the engine is stepping. If we drain after resume instead, # the aborts race with the flush of KV transfers that completed @@ -102,17 +105,32 @@ async def apply_policy_update(self, next_step: int) -> None: # the engine and cascading to every DP rank. Draining first lets the # aborts settle under normal stepping. ``on_new_version`` (below) # still runs post-update for observers that need the live version. - for observer in self.observers: - try: + entered_observers: list[VersionObserver] = [] + try: + for observer in self.observers: await observer.on_version_pending(next_step) - except Exception as exc: - get_logger().warning( - f"Observer {type(observer).__name__}.on_version_pending({next_step}) raised: {exc!r}" - ) + # A hook owns its own partial-entry rollback. Only a fully + # entered observer may receive a later transition cleanup. + entered_observers.append(observer) + except BaseException as exc: + await self._notify_update_failed(entered_observers, next_step, exc) + raise get_logger().debug(f"Updating weights to step {next_step}") t1 = time.perf_counter() - await self.inference.update_weights(weights_path, lora_name=self.lora_name, step=next_step) + try: + await self.inference.update_weights(weights_path, lora_name=self.lora_name, step=next_step) + except BaseException as exc: + # Once an admin update starts, an error cannot prove that no + # worker committed new weights (a resume failure is even later). + # Keep every transition fence closed and let the component + # failure terminate the run; reopening would stamp old-version + # requests onto mixed or fully-updated workers. + exc.add_note( + f"Policy update {next_step} may have mutated inference workers; mutable-policy admission " + "remains fail-closed" + ) + raise self.last_update_weights_time = time.perf_counter() - t1 self.update_count += 1 get_logger().debug(f"Updated weights to step {next_step} in {format_time(self.last_update_weights_time)}") @@ -123,13 +141,34 @@ async def apply_policy_update(self, next_step: int) -> None: self.inference.update_model_name(self.lora_name) self.policy.model_name = self.lora_name - for observer in self.observers: - try: - await observer.on_new_version(next_step) - except Exception as exc: - get_logger().warning( - f"Observer {type(observer).__name__}.on_new_version({next_step}) raised: {exc!r}" - ) + await self._notify_update_succeeded(entered_observers, next_step) + + @staticmethod + async def _notify_update_succeeded(observers: list[VersionObserver], step: int) -> None: + """Notify observers in dependency order and fail closed on any error.""" + # Complete observers in reverse order: the orchestrator first + # re-evaluates its long-lived lead gate while the dispatcher's short + # transition fence remains closed, then the dispatcher reopens + # admission on the new version. + for observer in reversed(observers): + await observer.on_new_version(step) + + @staticmethod + async def _notify_update_failed( + observers: list[VersionObserver], + step: int, + primary_error: BaseException, + ) -> None: + for observer in reversed(observers): + results, cancellation = await gather_shielded(observer.on_version_update_failed(step, primary_error)) + failures = [result for result in results if isinstance(result, BaseException)] + if cancellation is not None: + failures.append(cancellation) + for cleanup_error in failures: + primary_error.add_note( + f"Observer {type(observer).__name__}.on_version_update_failed({step}) also failed: " + f"{cleanup_error!r}" + ) def gauges(self) -> dict[str, float]: return { diff --git a/src/prime_rl/utils/async_utils.py b/src/prime_rl/utils/async_utils.py index 49e03500e2..79a80d1bf1 100644 --- a/src/prime_rl/utils/async_utils.py +++ b/src/prime_rl/utils/async_utils.py @@ -2,12 +2,30 @@ import asyncio from collections import deque +from collections.abc import Awaitable from time import perf_counter import numpy as np from pydantic import BaseModel +async def gather_shielded( + *awaitables: Awaitable[object], +) -> tuple[list[object], asyncio.CancelledError | None]: + """Settle all awaitables despite repeated cancellation of the caller.""" + settling = asyncio.gather(*awaitables, return_exceptions=True) + cancellation: asyncio.CancelledError | None = None + while not settling.done(): + try: + await asyncio.shield(settling) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + else: + cancellation.add_note("Caller cancelled again while bounded siblings were settling") + return list(settling.result()), cancellation + + async def safe_cancel(task: asyncio.Task) -> None: """Safely cancels and awaits an asyncio.Task.""" task.cancel() diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index ae8c1cde74..19b0e6443a 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -5,17 +5,25 @@ from collections.abc import Mapping from itertools import cycle from pathlib import Path -from typing import Protocol, runtime_checkable +from typing import Literal, Protocol, runtime_checkable import httpx import verifiers.v1 as vf from httpx import AsyncClient -from openai import AsyncOpenAI, NotFoundError +from openai import AsyncOpenAI from renderers import RendererConfig from tenacity import AsyncRetrying, retry, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential from verifiers.v1.clients.config import EvalClientConfig, TrainClientConfig from prime_rl.configs.shared import ClientConfig +from prime_rl.inference.dynamo_admin import ( + DynamoAdminAPI, + DynamoTopology, + DynamoWorker, + discover_workers, + discovery_urls, + validate_worker_membership, +) from prime_rl.utils.logger import get_logger # Identity tuple used by ``select_train_client`` to key load counts. ``base_url`` @@ -29,6 +37,19 @@ def client_identity(client: vf.ClientConfig) -> ClientIdentity: return (client.base_url, client.headers.get("X-data-parallel-rank")) +def _readiness_deadline(timeout: float) -> float: + if timeout <= 0: + raise TimeoutError("Inference readiness timeout must be greater than zero") + return asyncio.get_running_loop().time() + timeout + + +def _remaining_readiness_timeout(deadline: float, phase: str) -> float: + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + raise TimeoutError(f"Inference readiness deadline expired before {phase}") + return remaining + + @runtime_checkable class InferencePool(Protocol): """Protocol for inference pools (static or elastic).""" @@ -38,6 +59,11 @@ def model_name(self) -> str: """Get current model name for inference requests.""" ... + @property + def admin_api(self) -> Literal["vllm", "dynamo"]: + """Administration protocol used by the resolved pool.""" + ... + @property def train_clients(self) -> list[vf.ClientConfig]: """Get inference clients.""" @@ -73,6 +99,17 @@ async def update_weights(self, weight_dir: Path | None, lora_name: str | None = """Update weights on all inference servers.""" ... + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + """Initialize weight broadcast on all inference servers.""" + ... + async def score(self, token_ids: list[int]) -> list[float]: """Prefill-score ``token_ids`` under the pool's model — one logprob per token.""" ... @@ -115,8 +152,8 @@ async def aclose(self) -> None: await asyncio.gather(*(c.close() for c in self._clients.values())) -class StaticInferencePool: - """Static inference pool with fixed client list.""" +class FixedInferencePool: + """Base capability for pools whose endpoint set is fixed for their lifetime.""" def __init__( self, @@ -136,7 +173,7 @@ def __init__( pool_size=pool_size, ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) - self._admin_clients = setup_admin_clients(client_config) + self._client_config = client_config self._skip_model_check = client_config.skip_model_check self._wait_for_ready_timeout = client_config.wait_for_ready_timeout self._eval_cycle = cycle(self._eval_clients) @@ -147,6 +184,10 @@ def __init__( def train_clients(self) -> list[vf.ClientConfig]: return self._train_clients + @property + def admin_api(self) -> Literal["vllm", "dynamo"]: + return self._client_config.admin_api + @property def admin_clients(self) -> list[AsyncClient]: return self._admin_clients @@ -166,22 +207,183 @@ async def select_train_client(self, load: Mapping[ClientIdentity, int]) -> vf.Cl await asyncio.sleep(0.5) return min(self.train_clients, key=lambda c: load[client_identity(c)]) + async def score(self, token_ids: list[int]) -> list[float]: + """Prefill-score tokens under this pool's model.""" + return await self._scorer.score(self.train_clients, self.model_name, token_ids) + + +class StaticInferencePool(FixedInferencePool): + """Static native-vLLM inference pool with fixed clients.""" + + def __init__( + self, + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions", + eval_client_type: str = "openai_chat_completions", + renderer_config: RendererConfig | None = None, + pool_size: int | None = None, + ): + super().__init__( + client_config, + model_name, + train_client_type, + eval_client_type, + renderer_config, + pool_size, + ) + self._admin_clients = setup_admin_clients(client_config) + self._frontend_admin_clients = self._admin_clients + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + ready_timeout = timeout if timeout is not None else self._wait_for_ready_timeout + deadline = _readiness_deadline(ready_timeout) await check_health( - self._admin_clients, timeout=timeout if timeout is not None else self._wait_for_ready_timeout + self._frontend_admin_clients, + timeout=_remaining_readiness_timeout(deadline, "frontend health"), + ) + await maybe_check_has_model( + self._frontend_admin_clients, + model_name, + skip_model_check=self._skip_model_check, + timeout=_remaining_readiness_timeout(deadline, "model registration"), ) - await maybe_check_has_model(self._admin_clients, model_name, skip_model_check=self._skip_model_check) async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: await update_weights(self._admin_clients, weight_dir, lora_name=lora_name, step=step) - async def score(self, token_ids: list[int]) -> list[float]: - """Prefill-score ``token_ids`` under this pool's model (one logprob per - token, 0.0 for the leading token). Delegates to the shared scorer.""" - return await self._scorer.score(self.train_clients, self.model_name, token_ids) + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + await init_nccl_broadcast( + self._admin_clients, + host, + port, + timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) async def stop(self) -> None: await self._scorer.aclose() + clients = {*self._frontend_admin_clients, *self._admin_clients} + await asyncio.gather(*(client.aclose() for client in clients)) + + +class DynamoInferencePool(FixedInferencePool): + """Static Dynamo pool with worker discovery and Dynamo administration.""" + + def __init__( + self, + client_config: ClientConfig, + model_name: str, + train_client_type: str = "openai_chat_completions", + eval_client_type: str = "openai_chat_completions", + renderer_config: RendererConfig | None = None, + pool_size: int | None = None, + ): + if client_config.dynamo_worker_roles is None or client_config.dynamo_gpus_per_worker is None: + raise ValueError( + "Dynamo clients require dynamo_worker_roles and dynamo_gpus_per_worker; " + "local RL configs derive them from the inference topology" + ) + topology = DynamoTopology( + roles=client_config.dynamo_worker_roles, + gpus_per_worker=client_config.dynamo_gpus_per_worker, + ) + super().__init__( + client_config, + model_name, + train_client_type, + eval_client_type, + renderer_config, + pool_size, + ) + self._frontend_admin_clients = setup_admin_clients(client_config, urls=client_config.base_url) + self._discovery_clients = setup_admin_clients(client_config, urls=discovery_urls(client_config)) + self._admin_clients: list[AsyncClient] = [] + self._topology = topology + self._workers: tuple[DynamoWorker, ...] = () + self._admin = DynamoAdminAPI() + + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + ready_timeout = timeout if timeout is not None else self._wait_for_ready_timeout + deadline = _readiness_deadline(ready_timeout) + await check_health( + self._frontend_admin_clients, + timeout=_remaining_readiness_timeout(deadline, "frontend health"), + strict=True, + ) + await maybe_check_has_model( + self._frontend_admin_clients, + model_name, + skip_model_check=self._skip_model_check, + timeout=_remaining_readiness_timeout(deadline, "model registration"), + ) + workers = await discover_workers( + self._discovery_clients, + _remaining_readiness_timeout(deadline, "worker discovery"), + model_name=model_name, + topology=self._topology, + ) + admin_clients = setup_admin_clients(self._client_config, urls=[worker.system_url for worker in workers]) + try: + await check_health( + admin_clients, + timeout=_remaining_readiness_timeout(deadline, "worker health"), + strict=True, + ) + except BaseException: + await asyncio.gather(*(client.aclose() for client in admin_clients)) + raise + previous_clients = self._admin_clients + self._workers = workers + self._admin_clients = admin_clients + await asyncio.gather(*(client.aclose() for client in previous_clients)) + + async def _validate_membership(self) -> None: + discovered = await discover_workers( + self._discovery_clients, + 30, + model_name=self.model_name, + topology=self._topology, + ) + validate_worker_membership(self._workers, discovered) + + async def update_weights(self, weight_dir: Path | None, lora_name: str | None = None, step: int = 0) -> None: + if lora_name is not None: + raise ValueError("Dynamo backend does not yet support Prime LoRA weight updates") + await self._validate_membership() + await self._admin.update_weights(self._admin_clients, weight_dir, step) + + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + await self._validate_membership() + await self._admin.initialize_nccl( + self._admin_clients, + host=host, + port=port, + timeout=timeout, + inference_world_size=inference_world_size, + gpus_per_worker=self._topology.gpus_per_worker, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) + + async def stop(self) -> None: + await self._scorer.aclose() + clients = {*self._frontend_admin_clients, *self._discovery_clients, *self._admin_clients} + await asyncio.gather(*(client.aclose() for client in clients)) async def setup_inference_pool( @@ -193,6 +395,9 @@ async def setup_inference_pool( pool_size: int | None = None, ) -> InferencePool: """Create an inference pool from config (static or elastic).""" + if client_config.admin_api == "dynamo" and client_config.is_elastic: + raise ValueError("Dynamo admin API does not support elastic inference pools") + if client_config.is_elastic: from prime_rl.utils.elastic import ElasticInferencePool @@ -205,7 +410,8 @@ async def setup_inference_pool( pool_size=pool_size, ) - return StaticInferencePool( + pool_type = DynamoInferencePool if client_config.admin_api == "dynamo" else StaticInferencePool + return pool_type( client_config, model_name=model_name, train_client_type=train_client_type, @@ -250,14 +456,15 @@ def setup_clients( return clients -def setup_admin_clients(client_config: ClientConfig) -> list[AsyncClient]: +def setup_admin_clients(client_config: ClientConfig, urls: list[str] | None = None) -> list[AsyncClient]: """Create dedicated admin clients for weight update operations. Uses a separate connection pool to avoid queueing behind streaming requests. When admin_base_url is set, uses those URLs instead of base_url, allowing weight updates to bypass routers in disaggregated P/D deployments. """ - urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url + if urls is None: + urls = client_config.admin_base_url if client_config.admin_base_url else client_config.base_url def _setup_admin_client(base_url: str) -> httpx.AsyncClient: env_headers = { @@ -281,45 +488,144 @@ def _setup_admin_client(base_url: str) -> httpx.AsyncClient: return [_setup_admin_client(base_url) for base_url in urls] +_RETRYABLE_READINESS_STATUS_CODES = frozenset({408, 409, 429}) + + +def _is_retryable_readiness_error(error: BaseException) -> bool: + """Return whether a readiness request can plausibly succeed unchanged. + + HTTP 408 and 429 are request/proxy backpressure. HTTP 409 is also transient + here because inference servers can report a state conflict while their model + lifecycle is transitioning. Other 4xx responses require a caller or routing + change and must fail immediately; server failures and transports are retried. + """ + if isinstance(error, httpx.HTTPStatusError): + status_code = error.response.status_code + return status_code in _RETRYABLE_READINESS_STATUS_CODES or 500 <= status_code < 600 + return isinstance(error, (httpx.TransportError, TimeoutError)) + + +def _model_ids(response: httpx.Response) -> frozenset[str]: + """Validate an OpenAI models response before interpreting model absence.""" + try: + payload = response.json() + except ValueError as error: + raise ValueError("Invalid /v1/models response: expected valid JSON") from error + if not isinstance(payload, dict): + raise ValueError("Invalid /v1/models response: expected a JSON object") + models = payload.get("data") + if not isinstance(models, list): + raise ValueError("Invalid /v1/models response: 'data' must be a list") + + model_ids: set[str] = set() + for index, model in enumerate(models): + if not isinstance(model, dict) or not isinstance(model.get("id"), str): + raise ValueError(f"Invalid /v1/models response: data[{index}].id must be a string") + model_ids.add(model["id"]) + return frozenset(model_ids) + + async def maybe_check_has_model( - admin_clients: list[AsyncClient], model_name: str, skip_model_check: bool = False + admin_clients: list[AsyncClient], + model_name: str, + skip_model_check: bool = False, + timeout: float = 1800, + interval: float = 1, ) -> None: if skip_model_check: return logger = get_logger() + deadline = _readiness_deadline(timeout) logger.debug(f"Checking if model {model_name} is in the inference pool") - results = await asyncio.gather(*[admin_client.get("/v1/models") for admin_client in admin_clients]) - for admin_client, result in zip(admin_clients, results): - models = result.json()["data"] - if not any(model["id"] == model_name for model in models): - raise ValueError(f"Model {model_name} was not found in the inference pool on {admin_client.base_url}") + + async def _check_has_model(admin_client: AsyncClient) -> None: + last_error: Exception | None = None + loop = asyncio.get_running_loop() + while (remaining := deadline - loop.time()) > 0: + try: + async with asyncio.timeout(remaining): + result = await admin_client.get( + "/v1/models", + timeout=httpx.Timeout(min(remaining, 10.0)), + ) + result.raise_for_status() + except Exception as error: + if not _is_retryable_readiness_error(error): + raise + last_error = error + else: + if model_name in _model_ids(result): + return + # A valid 200 response without the requested model is expected + # while registration is still in progress, so it is retryable. + last_error = RuntimeError(f"Model {model_name} is not registered") + + remaining = deadline - loop.time() + if remaining > 0: + await asyncio.sleep(min(interval, remaining)) + + message = ( + f"Model {model_name} was not registered on {admin_client.base_url} " + f"before the {timeout}-second readiness deadline; last error: {last_error!r}" + ) + raise TimeoutError(message) from last_error + + await asyncio.gather(*(_check_has_model(admin_client) for admin_client in admin_clients)) logger.debug(f"Model {model_name} was found in the inference pool") async def check_health( - admin_clients: list[AsyncClient], interval: int = 1, log_interval: int = 10, timeout: int = 1800 + admin_clients: list[AsyncClient], + interval: float = 1, + log_interval: float = 10, + timeout: float = 1800, + strict: bool = False, ) -> None: + """Wait for healthy endpoints, retrying only transient request failures. + + Native endpoints may omit ``/health``; ``strict=True`` requires the route + and is used for Dynamo frontends and workers. All other non-success statuses + are classified by :func:`_is_retryable_readiness_error`. + """ logger = get_logger() async def _check_health(admin_client: AsyncClient) -> None: - wait_time = 0 + loop = asyncio.get_running_loop() + started = loop.time() + deadline = started + timeout + next_log_at = log_interval + last_error: Exception | None = None logger.debug("Starting pinging /health to check health") - while wait_time < timeout: + while (remaining := deadline - loop.time()) > 0: try: - await admin_client.get("/health") - logger.debug(f"Inference pool is ready after {wait_time} seconds") - return - except NotFoundError: - logger.warning("The route /health does not exist. Skipping health check.") + response = await admin_client.get( + "/health", + timeout=httpx.Timeout(min(remaining, 10.0)), + ) + if not strict and response.status_code == 404: + logger.warning("The route /health does not exist. Skipping health check.") + return + response.raise_for_status() + elapsed = loop.time() - started + logger.debug(f"Inference pool is ready after {elapsed:.1f} seconds") return except Exception as e: - if wait_time % log_interval == 0 and wait_time > 0: + if not _is_retryable_readiness_error(e): + raise + last_error = e + elapsed = loop.time() - started + if elapsed >= next_log_at: logger.warning( - f"Inference server was not reached after {wait_time} seconds (Error: {e}) on {admin_client.base_url}" + f"Inference server was not reached after {elapsed:.1f} seconds " + f"(Error: {e}) on {admin_client.base_url}" ) - await asyncio.sleep(interval) - wait_time += interval - msg = f"Inference server is not ready after {wait_time} (>{timeout}) seconds. Aborting..." + next_log_at += log_interval + remaining = deadline - loop.time() + if remaining > 0: + await asyncio.sleep(min(interval, remaining)) + msg = ( + f"Inference server {admin_client.base_url} is not ready after {timeout} seconds; last error: {last_error!r}" + ) logger.error(msg) raise TimeoutError(msg) diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index ef47dee774..208f9ebcd6 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -26,6 +26,7 @@ ClientIdentity, PrefillScorer, client_identity, + init_nccl_broadcast, load_lora_adapter, setup_admin_clients, setup_clients, @@ -224,6 +225,10 @@ def train_clients(self) -> list[vf.ClientConfig]: self._rebuild_clients() return self._train_clients + @property + def admin_api(self) -> Literal["vllm", "dynamo"]: + return self.client_config.admin_api + @property def eval_clients(self) -> list[vf.ClientConfig]: self._rebuild_clients() @@ -511,3 +516,20 @@ async def update_weights(self, weight_dir: Path | None, lora_name: str | None = if lora_name is None: raise ValueError("Elastic inference pool requires LoRA training (lora_name must be set)") await self.sync_weights(weight_dir, lora_name, step) + + async def init_nccl_broadcast( + self, + host: str, + port: int, + timeout: int, + inference_world_size: int | None = None, + quantize_in_weight_transfer: bool = False, + ) -> None: + await init_nccl_broadcast( + self.admin_clients, + host, + port, + timeout, + inference_world_size=inference_world_size, + quantize_in_weight_transfer=quantize_in_weight_transfer, + ) diff --git a/src/prime_rl/utils/policy_client_config.py b/src/prime_rl/utils/policy_client_config.py new file mode 100644 index 0000000000..8250713776 --- /dev/null +++ b/src/prime_rl/utils/policy_client_config.py @@ -0,0 +1,50 @@ +"""Resolve the generated DGD policy-client deployment boundary.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +from prime_rl.configs.shared import ClientConfig + +DYNAMO_TOPOLOGY_ENV = "DYN_RL_TOPOLOGY" + + +class _DynamoClientTopologyEnvironment(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + schema_version: Literal[1] + admin_api: Literal["dynamo"] + base_url: tuple[str, ...] = Field(min_length=1) + rl_base_url: tuple[str, ...] = Field(min_length=1) + dynamo_worker_roles: tuple[Literal["agg", "prefill", "decode"], ...] = Field(min_length=1) + dynamo_gpus_per_worker: int = Field(ge=1) + + +def policy_client_config_from_environment( + client_config: ClientConfig, + environment: Mapping[str, str] | None = None, +) -> ClientConfig: + """Apply the generated policy topology without mutating user config.""" + source = environment if environment is not None else os.environ + serialized = source.get(DYNAMO_TOPOLOGY_ENV) + if serialized is None: + return client_config + topology = _DynamoClientTopologyEnvironment.model_validate_json(serialized) + updates = { + "admin_api": topology.admin_api, + "base_url": list(topology.base_url), + "rl_base_url": list(topology.rl_base_url), + "dynamo_worker_roles": topology.dynamo_worker_roles, + "dynamo_gpus_per_worker": topology.dynamo_gpus_per_worker, + } + for field, expected in updates.items(): + if field in client_config.model_fields_set and getattr(client_config, field) != expected: + raise ValueError( + f"orchestrator.model.client.{field} conflicts with the generated " + f"{DYNAMO_TOPOLOGY_ENV} deployment boundary" + ) + return client_config.model_copy(update=updates) diff --git a/tests/unit/inference/helm_dgd_test_utils.py b/tests/unit/inference/helm_dgd_test_utils.py new file mode 100644 index 0000000000..2d24e77c43 --- /dev/null +++ b/tests/unit/inference/helm_dgd_test_utils.py @@ -0,0 +1,158 @@ +import hashlib +import json +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.dgd import DynamoGraphRenderOptions, GPUSchedulingProfile + +HELM = shutil.which("helm") +CHART = Path(__file__).parents[3] / "k8s" / "prime-rl" +PRIME_SHA = "1" * 40 +DYNAMO_SHA = "2" * 40 +IMAGE_DIGEST = f"sha256:{'3' * 64}" +ORCHESTRATOR_COMMAND = "uv run orchestrator @ /app/configs/debug/orch.toml --output-dir /data/outputs" +TRAINER_COMMAND = "uv run trainer @ /app/configs/debug/rl/train.toml --output-dir /data/outputs" +GPU_SCHEDULING = GPUSchedulingProfile( + runtime_class_name="nvidia", + architecture="arm64", + product="NVIDIA-GB200", + node_pool="customer-gpu-o7v", +) + + +def inference_config( + weight_broadcast: str = "nccl", + *, + chat_template: str | None = None, +) -> InferenceConfig: + model = {"chat_template": chat_template} if chat_template is not None else {} + return InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "model": model, + "weight_broadcast": {"type": weight_broadcast}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + ) + + +def render_options( + tmp_path: Path, + *, + external_controller: bool = False, + release_name: str = "p4-math", + shared_pvc: str | None = None, + trainer_gpu_count: int = 1, +) -> DynamoGraphRenderOptions: + return DynamoGraphRenderOptions( + release_name=release_name, + namespace="bis-vllm", + image=f"nvcr.io/example/prime:prime-{PRIME_SHA[:12]}-dynamo-{DYNAMO_SHA[:12]}@{IMAGE_DIGEST}", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, + external_controller=external_controller, + trainer_gpu_count=trainer_gpu_count, + orchestrator_command=None if external_controller else ORCHESTRATOR_COMMAND, + trainer_command=None if external_controller else TRAINER_COMMAND, + model_cache_pvc="model-cache", + hf_token_secret="hf-token-secret", + shared_pvc=shared_pvc, + image_pull_secrets=("nvcrimagepullsecret",), + ) + + +def helm_template( + *args: str, + release_name: str = "p4-math", + release_namespace: str = "bis-vllm", +) -> str: + if HELM is None: + pytest.skip("helm is not installed") + return subprocess.run( + [HELM, "template", release_name, str(CHART), "--namespace", release_namespace, *args], + check=True, + capture_output=True, + text=True, + ).stdout + + +def rendered_documents(rendered: str) -> list[dict]: + return [document for document in yaml.safe_load_all(rendered) if document] + + +def rendered_resource(rendered: str, kind: str, name: str) -> dict: + return next( + document + for document in rendered_documents(rendered) + if document.get("kind") == kind and document.get("metadata", {}).get("name") == name + ) + + +def labels_match(selector: dict[str, str], labels: dict[str, str]) -> bool: + return selector.items() <= labels.items() + + +def toleration_identity(toleration: dict[str, str]) -> tuple[str, str, str | None, str | None]: + return ( + toleration["key"], + toleration["operator"], + toleration.get("value"), + toleration.get("effect"), + ) + + +def write_values_mutation( + source: Path, + target: Path, + path: tuple[str, ...], + replacement: object, +) -> None: + values = json.loads(source.read_text()) + parent = values + for key in path[:-1]: + parent = parent[key] + parent[path[-1]] = replacement + target.write_text(json.dumps(values)) + + +def canonical_json(value: object) -> str: + return json.dumps(value, indent=2, sort_keys=True) + "\n" + + +def rewrite_valid_integrity(values: dict, workload: dict) -> None: + graph = values["inference"]["dynamoGraph"] + workload_canonical = canonical_json(workload) + workload_hash = hashlib.sha256(workload_canonical.encode()).hexdigest() + graph["workloadBinding"] = { + "canonical": workload_canonical, + "sha256": workload_hash, + } + + resource = graph["resource"] + resource_annotations = resource["metadata"]["annotations"] + resource_annotations["prime-rl.nvidia.com/workload-sha256"] = workload_hash + graph["engineConfig"]["annotations"]["prime-rl.nvidia.com/workload-sha256"] = workload_hash + + scoped_resource = json.loads(json.dumps(resource)) + scoped_resource["metadata"]["annotations"].pop("prime-rl.nvidia.com/manifest-sha256") + manifest_canonical = canonical_json(scoped_resource) + manifest_hash = hashlib.sha256(manifest_canonical.encode()).hexdigest() + graph["manifestCanonical"] = manifest_canonical + resource_annotations["prime-rl.nvidia.com/manifest-sha256"] = manifest_hash + graph["engineConfig"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"] = manifest_hash diff --git a/tests/unit/inference/test_dgd.py b/tests/unit/inference/test_dgd.py new file mode 100644 index 0000000000..67fbbdcb7b --- /dev/null +++ b/tests/unit/inference/test_dgd.py @@ -0,0 +1,784 @@ +import hashlib +import json +import sys +from copy import deepcopy +from dataclasses import replace +from pathlib import Path + +import pytest + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.dgd import ( + DynamoGraphRenderOptions, + GPUSchedulingProfile, + KubernetesToleration, + _add_pvc, + _parse_args, + _parse_kubernetes_toleration, + build_dgd_values, + write_dgd_artifacts, +) +from prime_rl.inference.dynamo import build_frontend_process, build_worker_process + +PRIME_SHA = "1" * 40 +DYNAMO_SHA = "2" * 40 +IMAGE_DIGEST = f"sha256:{'3' * 64}" +ORCHESTRATOR_COMMAND = "uv run orchestrator @ /app/configs/debug/orch.toml --output-dir /data/outputs" +TRAINER_COMMAND = "uv run trainer @ /app/configs/debug/rl/train.toml --output-dir /data/outputs" +GPU_SCHEDULING = GPUSchedulingProfile( + runtime_class_name="nvidia", + architecture="arm64", + product="NVIDIA-GB200", + node_pool="customer-gpu-o7v", +) + + +def inference_config( + weight_broadcast: str = "nccl", + *, + chat_template: str | None = None, +) -> InferenceConfig: + model = {"name": "Qwen/Qwen3-30B-A3B-Thinking-2507"} + if chat_template is not None: + model["chat_template"] = chat_template + return InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "model": model, + "parallel": {"tp": 1}, + "weight_broadcast": {"type": weight_broadcast}, + "env_vars": {"HF_HOME": "/model-cache", "HF_HUB_OFFLINE": "1"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + ) + + +def render_options( + tmp_path: Path, + *, + external_controller: bool = False, + release_name: str = "p4-math", + shared_pvc: str | None = "p4-shared-data", + trainer_gpu_count: int = 1, +) -> DynamoGraphRenderOptions: + return DynamoGraphRenderOptions( + release_name=release_name, + namespace="bis-vllm", + image=f"nvcr.io/example/prime:prime-{PRIME_SHA[:12]}-dynamo-{DYNAMO_SHA[:12]}@{IMAGE_DIGEST}", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, + external_controller=external_controller, + trainer_gpu_count=trainer_gpu_count, + orchestrator_command=None if external_controller else ORCHESTRATOR_COMMAND, + trainer_command=None if external_controller else TRAINER_COMMAND, + model_cache_pvc="model-cache", + shared_pvc=shared_pvc, + image_pull_secrets=("nvcrimagepullsecret",), + hf_token_secret="hf-token-secret", + ) + + +def test_dgd_values_derive_topology_and_role_configs(tmp_path: Path): + options = render_options(tmp_path) + values = build_dgd_values(inference_config(), options) + graph = values["inference"]["dynamoGraph"] + resource = graph["resource"] + services = resource["spec"]["services"] + + assert values["image"] == { + "reference": options.image, + "pullPolicy": "IfNotPresent", + "pullSecrets": ["nvcrimagepullsecret"], + } + assert resource["apiVersion"] == "nvidia.com/v1alpha1" + assert resource["metadata"]["namespace"] == "bis-vllm" + assert services["VllmPrefillWorker"]["replicas"] == 2 + assert services["VllmDecodeWorker"]["replicas"] == 2 + assert services["VllmPrefillWorker"]["resources"]["limits"]["gpu"] == "1" + assert services["VllmDecodeWorker"]["resources"]["limits"]["gpu"] == "1" + assert resource["spec"]["pvcs"] == [ + {"name": "model-cache", "create": False}, + ] + assert values["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", + "mountPath": "/data", + } + assert values["modelCache"] == { + "enabled": True, + "existingClaim": "model-cache", + "mountPath": "/model-cache", + } + assert values["huggingFace"] == { + "tokenSecretName": "hf-token-secret", + "tokenSecretKey": "HF_TOKEN", + } + assert values["inference"]["dynamoGraph"]["clientTopology"] == { + "schema_version": 1, + "admin_api": "dynamo", + "base_url": ["http://p4-math-frontend.bis-vllm.svc.cluster.local:8000/v1"], + "rl_base_url": ["http://p4-math-frontend-rl.bis-vllm.svc.cluster.local:8001"], + "dynamo_worker_roles": ["prefill", "prefill", "decode", "decode"], + "dynamo_gpus_per_worker": 1, + } + expected_cache_mount = {"name": "model-cache", "mountPath": "/model-cache"} + expected_cache_volume = { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + } + for service in services.values(): + assert "volumeMounts" not in service + pod_spec = service["extraPodSpec"] + assert pod_spec["mainContainer"].get("volumeMounts", []).count(expected_cache_mount) == 1 + assert pod_spec.get("volumes", []).count(expected_cache_volume) == 1 + topology_binding = graph["topologyBinding"] + assert hashlib.sha256(topology_binding["canonical"].encode()).hexdigest() == topology_binding["sha256"] + assert json.loads(topology_binding["canonical"])["clientTopology"] == graph["clientTopology"] + assert json.loads(topology_binding["canonical"])["workerServices"] == { + "VllmDecodeWorker": { + "limitsGpu": "1", + "replicas": 2, + "requestsGpu": "1", + "role": "decode", + }, + "VllmPrefillWorker": { + "limitsGpu": "1", + "replicas": 2, + "requestsGpu": "1", + "role": "prefill", + }, + } + assert resource["metadata"]["annotations"]["prime-rl.nvidia.com/topology-sha256"] == topology_binding["sha256"] + workload_binding = graph["workloadBinding"] + assert hashlib.sha256(workload_binding["canonical"].encode()).hexdigest() == workload_binding["sha256"] + assert resource["metadata"]["annotations"]["prime-rl.nvidia.com/workload-sha256"] == workload_binding["sha256"] + for service in services.values(): + pod_spec = service["extraPodSpec"] + assert pod_spec["mainContainer"]["image"] == options.image + assert pod_spec["imagePullSecrets"] == [{"name": "nvcrimagepullsecret"}] + assert any(item["name"] == "HF_TOKEN" for item in pod_spec["mainContainer"]["env"]) + env = {item["name"]: item.get("value") for item in pod_spec["mainContainer"]["env"]} + assert env["HF_HOME"] == "/model-cache" + assert env["HF_HUB_OFFLINE"] == "1" + + image_tolerations = [ + { + "key": "kubernetes.io/arch", + "operator": "Equal", + "value": "arm64", + "effect": "NoSchedule", + }, + { + "key": "nvidia.com/gpu", + "operator": "Exists", + "effect": "NoSchedule", + }, + { + "key": "prime-rl", + "operator": "Equal", + "value": "true", + "effect": "NoSchedule", + }, + ] + gpu_tolerations = image_tolerations + image_selector = { + "kubernetes.io/arch": "arm64", + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + } + gpu_selector = {**image_selector, "nvidia.com/gpu.product": "NVIDIA-GB200"} + frontend_pod = services["Frontend"]["extraPodSpec"] + assert services["Frontend"]["extraPodMetadata"]["labels"] == { + "app.kubernetes.io/name": "prime-rl", + "app.kubernetes.io/instance": "p4-math", + } + assert frontend_pod["nodeSelector"] == image_selector + assert frontend_pod["tolerations"] == image_tolerations + assert "runtimeClassName" not in frontend_pod + assert frontend_pod["mainContainer"]["ports"] == [ + {"containerPort": 8000, "name": "http"}, + {"containerPort": 8001, "name": "rl"}, + ] + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + worker_pod = services[role]["extraPodSpec"] + assert worker_pod["runtimeClassName"] == "nvidia" + assert worker_pod["nodeSelector"] == gpu_selector + assert worker_pod["tolerations"] == gpu_tolerations + workload = json.loads(workload_binding["canonical"]) + assert workload["controllerMode"] == "chartManaged" + assert workload["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", + "mountPath": "/data", + } + assert workload["orchestrator"]["gpu"] == {"enabled": False} + assert workload["orchestrator"]["placement"] == { + "nodeSelector": image_selector, + "tolerations": image_tolerations, + } + assert workload["trainer"]["placement"] == { + "runtimeClassName": "nvidia", + "nodeSelector": gpu_selector, + "tolerations": gpu_tolerations, + } + assert values["orchestrator"] == { + key: value for key, value in workload["orchestrator"].items() if key not in {"gpu", "placement"} + } + assert values["trainer"] == {key: value for key, value in workload["trainer"].items() if key != "placement"} + + prefill_args = services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["args"] + decode_args = services["VllmDecodeWorker"]["extraPodSpec"]["mainContainer"]["args"] + prefill_process = build_worker_process( + inference_config(), + "prefill", + Path("/etc/prime-rl/dynamo/prefill-engine.json"), + nixl_host=None, + nixl_port=20100, + ) + decode_process = build_worker_process( + inference_config(), + "decode", + Path("/etc/prime-rl/dynamo/decode-engine.json"), + nixl_host=None, + nixl_port=20100, + ) + frontend_process = build_frontend_process(inference_config(), host="0.0.0.0", port=8000) + assert prefill_args == list(prefill_process.arguments) + assert decode_args == list(decode_process.arguments) + assert services["Frontend"]["extraPodSpec"]["mainContainer"]["args"] == list(frontend_process.arguments) + + prefill_env = { + item["name"]: item.get("value") + for item in services["VllmPrefillWorker"]["extraPodSpec"]["mainContainer"]["env"] + } + frontend_env = { + item["name"]: item.get("value") for item in services["Frontend"]["extraPodSpec"]["mainContainer"]["env"] + } + assert {key: prefill_env[key] for key in prefill_process.environment()} == prefill_process.environment() + assert {key: frontend_env[key] for key in frontend_process.environment()} == frontend_process.environment() + + assert all("volumeMounts" not in service for service in services.values()) + + prefill = json.loads(graph["engineConfig"]["data"]["prefill-engine.json"]) + decode = json.loads(graph["engineConfig"]["data"]["decode-engine.json"]) + assert prefill["kv_transfer_config"] == decode["kv_transfer_config"] + assert "kv_events_config" in prefill + assert "kv_events_config" not in decode + assert "disaggregation_mode" not in prefill + assert "enable_rl" not in decode + + +def test_dgd_artifacts_are_deterministic_and_manifest_verifies(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + first_values = paths["values"].read_bytes() + write_dgd_artifacts(inference_config(), render_options(tmp_path)) + assert paths["values"].read_bytes() == first_values + + for line in paths["manifest"].read_text().splitlines(): + expected, name = line.split(" ", 1) + assert hashlib.sha256((tmp_path / name).read_bytes()).hexdigest() == expected + + resource = json.loads(paths["resource"].read_text()) + annotations = resource["metadata"]["annotations"] + assert annotations["prime-rl.nvidia.com/manifest-sha256-scope"] == ( + "resource; json.dumps(sort_keys=true,indent=2)+newline; " + "exclude=/metadata/annotations/prime-rl.nvidia.com~1manifest-sha256" + ) + expected_manifest_hash = annotations["prime-rl.nvidia.com/manifest-sha256"] + unhashed_resource = deepcopy(resource) + del unhashed_resource["metadata"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"] + canonical_resource = (json.dumps(unhashed_resource, indent=2, sort_keys=True) + "\n").encode() + values = json.loads(paths["values"].read_text()) + assert values["inference"]["dynamoGraph"]["manifestCanonical"].encode() == canonical_resource + assert hashlib.sha256(canonical_resource).hexdigest() == expected_manifest_hash + + +def test_dgd_embeds_and_mounts_content_addressed_chat_template(tmp_path: Path): + first = build_dgd_values( + inference_config(chat_template="template-v1: {{ messages }}"), + render_options(tmp_path / "first"), + ) + graph = first["inference"]["dynamoGraph"] + engine_config = graph["engineConfig"] + frontend = graph["resource"]["spec"]["services"]["Frontend"]["extraPodSpec"] + + assert engine_config["data"]["chat-template.jinja"] == "template-v1: {{ messages }}" + expected_hash = hashlib.sha256( + (json.dumps(engine_config["data"], indent=2, sort_keys=True) + "\n").encode() + ).hexdigest() + assert engine_config["sha256"] == expected_hash + assert engine_config["canonicalData"] == json.dumps(engine_config["data"], indent=2, sort_keys=True) + "\n" + assert engine_config["name"].endswith(expected_hash[:12]) + assert frontend["mainContainer"]["args"][-1] == "/etc/prime-rl/dynamo/chat-template.jinja" + assert frontend["mainContainer"]["volumeMounts"] == [ + { + "name": "dynamo-chat-template", + "mountPath": "/etc/prime-rl/dynamo", + "readOnly": True, + }, + {"name": "model-cache", "mountPath": "/model-cache"}, + ] + assert frontend["volumes"] == [ + { + "name": "dynamo-chat-template", + "configMap": { + "name": engine_config["name"], + "items": [{"key": "chat-template.jinja", "path": "chat-template.jinja"}], + }, + }, + { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + }, + ] + assert not (tmp_path / "first" / "chat-template.jinja").exists() + + second = build_dgd_values( + inference_config(chat_template="template-v2: {{ messages }}"), + render_options(tmp_path / "second"), + ) + second_config = second["inference"]["dynamoGraph"]["engineConfig"] + assert second_config["name"] != engine_config["name"] + assert second_config["sha256"] != engine_config["sha256"] + + +def test_filesystem_broadcast_requires_one_shared_existing_claim(tmp_path: Path): + options = render_options(tmp_path) + values = build_dgd_values(inference_config("filesystem"), options) + resource = values["inference"]["dynamoGraph"]["resource"] + services = resource["spec"]["services"] + + assert values["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", + "mountPath": "/data", + } + assert resource["spec"]["pvcs"] == [ + {"name": "model-cache", "create": False}, + {"name": "p4-shared-data", "create": False}, + ] + assert all("volumeMounts" not in service for service in services.values()) + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + pod_spec = services[role]["extraPodSpec"] + assert pod_spec["mainContainer"]["volumeMounts"].count({"name": "p4-shared-data", "mountPath": "/data"}) == 1 + assert ( + pod_spec["volumes"].count( + { + "name": "p4-shared-data", + "persistentVolumeClaim": {"claimName": "p4-shared-data"}, + } + ) + == 1 + ) + + +def test_filesystem_broadcast_rejects_missing_shared_claim(tmp_path: Path): + options = replace(render_options(tmp_path), shared_pvc=None) + + with pytest.raises(ValueError, match="shared existing PVC"): + build_dgd_values(inference_config("filesystem"), options) + + +@pytest.mark.parametrize( + "existing_mount", + [ + ({"name": "other", "mountPath": "/model-cache"}, "existing container mount"), + ], +) +def test_add_pvc_rejects_container_mount_path_collision(existing_mount: tuple[dict[str, str], str]): + mount, message = existing_mount + resource = {"spec": {}} + service = { + "extraPodSpec": { + "mainContainer": {"volumeMounts": [mount]}, + "volumes": [], + } + } + + with pytest.raises(ValueError, match=message): + _add_pvc(resource, service, "model-cache", "/model-cache") + + +def test_filesystem_broadcast_can_reuse_model_cache_claim(tmp_path: Path): + options = replace(render_options(tmp_path), shared_pvc="model-cache") + values = build_dgd_values(inference_config("filesystem"), options) + resource = values["inference"]["dynamoGraph"]["resource"] + services = resource["spec"]["services"] + + assert resource["spec"]["pvcs"] == [{"name": "model-cache", "create": False}] + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + pod_spec = services[role]["extraPodSpec"] + mounts = pod_spec["mainContainer"]["volumeMounts"] + assert mounts.count({"name": "model-cache", "mountPath": "/model-cache"}) == 1 + assert mounts.count({"name": "model-cache", "mountPath": "/data"}) == 1 + assert ( + pod_spec["volumes"].count( + { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + } + ) + == 1 + ) + + +def test_add_pvc_rejects_pod_volume_collision(): + resource = {"spec": {}} + service = { + "extraPodSpec": { + "mainContainer": { + "volumeMounts": [{"name": "model-cache", "mountPath": "/model-cache"}], + }, + "volumes": [{"name": "model-cache", "emptyDir": {}}], + } + } + + with pytest.raises(ValueError, match="existing pod volume"): + _add_pvc(resource, service, "model-cache", "/model-cache") + + +@pytest.mark.parametrize( + ("scope", "key"), + [ + ("global", "DYN_DISCOVERY_BACKEND"), + ("global", "DYN_NAMESPACE"), + ("global", "DYN_NAMESPACE_PREFIX"), + ("global", "DYN_NAMESPACE_WORKER_SUFFIX"), + ("global", "DYN_PARENT_DGD_K8S_NAME"), + ("global", "DYN_PARENT_DGD_K8S_NAMESPACE"), + ("global", "DYN_KUBE_DISCOVERY_MODE"), + ("global", "DYN_ENDPOINT_TYPES"), + ("global", "DYN_SYSTEM_ENABLED"), + ("global", "DYN_SYSTEM_HOST"), + ("global", "DYN_SYSTEM_HEALTH_PATH"), + ("global", "DYN_SYSTEM_LIVE_PATH"), + ("global", "DYN_SYSTEM_STARTING_HEALTH_STATUS"), + ("global", "DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS"), + ("global", "DYN_HEALTH_CHECK_ENABLED"), + ("global", "POD_NAME"), + ("global", "POD_NAMESPACE"), + ("global", "POD_UID"), + ("global", "CONTAINER_NAME"), + ("prefill", "DYN_SYSTEM_PORT"), + ("prefill", "DYN_SYSTEM_PORT1"), + ("decode", "DYN_ENDPOINT"), + ], +) +def test_dgd_rejects_operator_owned_environment(scope: str, key: str, tmp_path: Path): + config_data = inference_config().model_dump(mode="python") + if scope == "global": + config_data["env_vars"] = {key: "override"} + else: + config_data["deployment"][f"{scope}_env_vars"] = {key: "override"} + config = InferenceConfig.model_validate(config_data) + + with pytest.raises(ValueError, match=rf"{scope}.*{key}.*operator-owned"): + build_dgd_values(config, render_options(tmp_path)) + + +@pytest.mark.parametrize("scope", ["global", "prefill", "decode"]) +@pytest.mark.parametrize( + "key", + [ + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "AWS_SECRET_ACCESS_KEY", + "NVIDIA_API_KEY", + "WANDB_API_KEY", + "MODEL_REGISTRY_PASSWORD", + "OIDC_CLIENT_SECRET", + ], +) +def test_dgd_rejects_raw_credentials(scope: str, key: str, tmp_path: Path): + config_data = inference_config().model_dump(mode="python") + if scope == "global": + config_data["env_vars"][key] = "plaintext-secret" + else: + config_data["deployment"][f"{scope}_env_vars"] = {key: "plaintext-secret"} + config = InferenceConfig.model_validate(config_data) + + with pytest.raises(ValueError, match=rf"{scope}.*{key}.*SecretKeyRef"): + build_dgd_values(config, render_options(tmp_path)) + + +@pytest.mark.parametrize("scope", ["global", "prefill", "decode"]) +def test_dgd_rejects_hf_home_that_conflicts_with_typed_model_cache(scope: str, tmp_path: Path): + config_data = inference_config().model_dump(mode="python") + if scope == "global": + config_data["env_vars"]["HF_HOME"] = "/wrong-cache" + else: + config_data["deployment"][f"{scope}_env_vars"] = {"HF_HOME": "/wrong-cache"} + config = InferenceConfig.model_validate(config_data) + + with pytest.raises(ValueError, match=rf"{scope}.*HF_HOME.*/model-cache"): + build_dgd_values(config, render_options(tmp_path)) + + +def test_dgd_allows_hf_home_matching_typed_model_cache(tmp_path: Path): + values = build_dgd_values(inference_config(), render_options(tmp_path)) + services = values["inference"]["dynamoGraph"]["resource"]["spec"]["services"] + + for service in services.values(): + env = service["extraPodSpec"]["mainContainer"]["env"] + assert [item for item in env if item["name"] == "HF_HOME"] == [{"name": "HF_HOME", "value": "/model-cache"}] + + +def test_cli_toleration_parser_is_typed_and_rejects_unknown_fields(): + parsed = _parse_kubernetes_toleration( + '{"key":"dedicated","operator":"Equal","value":"prime","effect":"NoSchedule"}' + ) + assert parsed == KubernetesToleration( + key="dedicated", + operator="Equal", + value="prime", + effect="NoSchedule", + ) + + with pytest.raises(ValueError, match="unsupported fields"): + _parse_kubernetes_toleration('{"key":"dedicated","command":"touch /tmp/pwned"}') + + with pytest.raises(ValueError, match="operator"): + _parse_kubernetes_toleration('{"key":"dedicated","operator":"NotARealOperator"}') + + +@pytest.mark.parametrize( + "toleration", + [ + KubernetesToleration(key="nvidia.com/gpu", effect="NoExecute"), + KubernetesToleration(key="nvidia.com/gpu", operator="Equal", value="true"), + ], +) +def test_gpu_scheduling_requires_exact_nodepool_access_toleration( + toleration: KubernetesToleration, +): + with pytest.raises(ValueError, match="nvidia.com/gpu Exists NoSchedule"): + GPUSchedulingProfile( + runtime_class_name="nvidia", + architecture="arm64", + product="NVIDIA-GB200", + node_pool="customer-gpu-o7v", + tolerations=(toleration,), + ) + + +def test_cli_accepts_typed_additional_image_and_gpu_tolerations(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + sys, + "argv", + [ + "dynamo-dgd", + "inference.toml", + "--release-name", + "p4-math", + "--namespace", + "bis-vllm", + "--image", + "registry/image@sha256:digest", + "--output-dir", + "/tmp/artifacts", + "--prime-sha", + PRIME_SHA, + "--dynamo-sha", + DYNAMO_SHA, + "--image-digest", + IMAGE_DIGEST, + "--gpu-architecture", + "arm64", + "--gpu-product", + "NVIDIA-GB200", + "--gpu-node-pool", + "customer-gpu-o7v", + "--no-gpu-runtime-class", + "--external-controller", + "--trainer-gpus", + "4", + "--image-toleration", + '{"key":"image-extra","operator":"Exists"}', + "--gpu-toleration", + '{"key":"gpu-extra","operator":"Equal","value":"true"}', + ], + ) + + args = _parse_args() + + assert args.image_toleration == [KubernetesToleration(key="image-extra")] + assert args.gpu_toleration == [KubernetesToleration(key="gpu-extra", operator="Equal", value="true")] + assert args.external_controller is True + assert args.no_gpu_runtime_class is True + assert args.trainer_gpus == 4 + + +def test_gpu_scheduling_changes_manifest_identity(tmp_path: Path): + first = build_dgd_values(inference_config(), render_options(tmp_path / "first")) + changed_profile = replace(GPU_SCHEDULING, node_pool="customer-gpu-alternate") + changed_options = replace(render_options(tmp_path / "second"), gpu_scheduling=changed_profile) + second = build_dgd_values(inference_config(), changed_options) + + first_resource = first["inference"]["dynamoGraph"]["resource"] + second_resource = second["inference"]["dynamoGraph"]["resource"] + assert ( + first_resource["metadata"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"] + != (second_resource["metadata"]["annotations"]["prime-rl.nvidia.com/manifest-sha256"]) + ) + + +def test_gpu_runtime_class_can_be_explicitly_omitted(tmp_path: Path): + scheduling = replace(GPU_SCHEDULING, runtime_class_name=None) + options = replace(render_options(tmp_path), gpu_scheduling=scheduling) + + values = build_dgd_values(inference_config(), options) + services = values["inference"]["dynamoGraph"]["resource"]["spec"]["services"] + + assert "runtimeClassName" not in services["Frontend"]["extraPodSpec"] + assert "runtimeClassName" not in services["VllmPrefillWorker"]["extraPodSpec"] + assert "runtimeClassName" not in services["VllmDecodeWorker"]["extraPodSpec"] + + +def test_external_controller_binding_disables_chart_workloads(tmp_path: Path): + values = build_dgd_values( + inference_config(), + render_options(tmp_path, external_controller=True, shared_pvc=None), + ) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + + assert graph["controllerMode"] == "external" + assert workload["controllerMode"] == "external" + assert workload["orchestrator"]["enabled"] is False + assert workload["orchestrator"]["replicas"] == 0 + assert workload["orchestrator"]["autoStart"] is False + assert workload["orchestrator"]["command"] == "" + assert workload["orchestrator"]["gpu"] == {"enabled": False} + assert workload["trainer"]["enabled"] is False + assert workload["trainer"]["replicas"] == 0 + assert workload["trainer"]["autoStart"] is False + assert workload["trainer"]["command"] == "" + assert workload["trainer"]["gpu"] == {"enabled": False, "count": 0} + assert workload["storage"] == { + "enabled": False, + "existingClaim": "", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", + "mountPath": "/data", + } + assert values["orchestrator"] == { + key: value for key, value in workload["orchestrator"].items() if key not in {"gpu", "placement"} + } + assert values["trainer"] == {key: value for key, value in workload["trainer"].items() if key != "placement"} + assert values["storage"] == workload["storage"] + + +def test_trainer_gpu_count_must_be_positive(): + with pytest.raises(ValueError, match="trainer_gpu_count"): + replace(render_options(Path("/tmp/not-written")), trainer_gpu_count=0) + + +@pytest.mark.parametrize( + ("change", "error"), + [ + ({"orchestrator_replicas": 0}, "orchestrator_replicas"), + ({"trainer_replicas": 0}, "trainer_replicas"), + ({"orchestrator_command": None}, "orchestrator_command"), + ({"trainer_command": "sleep infinity"}, "trainer_command"), + ({"trainer_command": "uv run trainer-impersonator"}, "trainer_command"), + ], +) +def test_chart_managed_controller_execution_must_be_runnable(change: dict[str, object], error: str): + with pytest.raises(ValueError, match=error): + replace(render_options(Path("/tmp/not-written")), **change) + + +def test_external_controller_rejects_ignored_chart_commands(): + with pytest.raises(ValueError, match="external_controller.*commands"): + replace(render_options(Path("/tmp/not-written")), external_controller=True) + + +def test_release_name_boundary_preserves_generated_service_names(): + boundary = "a" * 41 + + assert replace(render_options(Path("/tmp/not-written")), release_name=boundary).release_name == boundary + + with pytest.raises(ValueError, match="at most 41 characters"): + replace(render_options(Path("/tmp/not-written")), release_name="a" * 42) + + +def test_image_commit_suffixes_must_be_in_the_image_tag(tmp_path: Path): + with pytest.raises(ValueError, match="commit suffixes"): + replace( + render_options(tmp_path, external_controller=True), + image=(f"nvcr.io/prime-{PRIME_SHA[:12]}/dynamo-{DYNAMO_SHA[:12]}/runtime:reviewed@{IMAGE_DIGEST}"), + ) + + +def test_chart_runtime_binding_covers_every_rendered_controller_input(tmp_path: Path): + values = build_dgd_values(inference_config(), render_options(tmp_path)) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + + assert set(workload) == { + "config", + "controllerMode", + "huggingFace", + "image", + "modelCache", + "orchestrator", + "storage", + "trainer", + } + assert workload["image"] == { + "reference": values["image"]["reference"], + "pullPolicy": values["image"]["pullPolicy"], + "pullSecrets": values["image"]["pullSecrets"], + } + assert workload["storage"] == values["storage"] + assert workload["modelCache"] == values["modelCache"] + assert workload["huggingFace"] == values["huggingFace"] + assert workload["config"] == values["config"] + assert workload["orchestrator"] | {"placement": None, "gpu": None} == ( + values["orchestrator"] | {"placement": None, "gpu": None} + ) + assert workload["trainer"] | {"placement": None} == values["trainer"] | {"placement": None} + + +def test_explicit_hf_secret_references_are_required(tmp_path: Path): + values = build_dgd_values(inference_config(), render_options(tmp_path)) + services = values["inference"]["dynamoGraph"]["resource"]["spec"]["services"] + + for service in services.values(): + hf_token = next(item for item in service["extraPodSpec"]["mainContainer"]["env"] if item["name"] == "HF_TOKEN") + assert hf_token["valueFrom"]["secretKeyRef"]["optional"] is False + + +def test_dgd_readme_uses_runtime_image_config_paths_and_states_trust_boundary(): + repository = Path(__file__).parents[3] + readme = (repository / "k8s" / "README.md").read_text() + + for runtime_path in ("/app/configs/debug/orch.toml", "/app/configs/debug/rl/train.toml"): + assert runtime_path in readme + assert (repository / runtime_path.removeprefix("/app/")).is_file() + assert "does not authenticate source or image provenance" in readme + + +def test_dgd_rejects_native_backend(tmp_path: Path): + config = InferenceConfig.model_validate({}) + with pytest.raises(ValueError, match="Dynamo disaggregated"): + build_dgd_values(config, render_options(tmp_path)) diff --git a/tests/unit/inference/test_dynamo.py b/tests/unit/inference/test_dynamo.py new file mode 100644 index 0000000000..b08ead2a6a --- /dev/null +++ b/tests/unit/inference/test_dynamo.py @@ -0,0 +1,385 @@ +import json +from pathlib import Path + +import pytest + +from prime_rl.configs.inference import InferenceConfig +from prime_rl.entrypoints import inference as inference_entrypoint +from prime_rl.inference import dynamo +from prime_rl.inference.dynamo import ( + build_engine_config, + build_frontend_process, + build_local_worker_specs, + build_worker_environment, + build_worker_process, + write_role_engine_configs, +) + + +def disaggregated_config(**overrides) -> InferenceConfig: + data = { + "backend": {"type": "dynamo"}, + "weight_broadcast": {"type": "nccl"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + data.update(overrides) + return InferenceConfig.model_validate(data) + + +def test_role_engine_configs_share_nixl_and_only_prefill_publishes_events(tmp_path: Path): + paths = write_role_engine_configs(disaggregated_config(), tmp_path) + prefill = json.loads(paths["prefill"].read_text()) + decode = json.loads(paths["decode"].read_text()) + + assert prefill["kv_transfer_config"] == decode["kv_transfer_config"] + assert prefill["kv_transfer_config"]["kv_connector"] == "NixlConnector" + assert prefill["kv_events_config"]["enable_kv_cache_events"] is True + assert "kv_events_config" not in decode + assert prefill["worker_extension_cls"].endswith("NCCLWeightUpdateWorker") + assert decode["worker_extension_cls"] == prefill["worker_extension_cls"] + + +def test_role_overrides_are_isolated(): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + "prefill_vllm_overrides": {"max_num_batched_tokens": 8192}, + "decode_vllm_overrides": {"max_num_seqs": 64}, + } + ) + + prefill = build_engine_config(config, "prefill", kv_events_port=20080) + decode = build_engine_config(config, "decode") + + assert prefill["max_num_batched_tokens"] == 8192 + assert "max_num_batched_tokens" not in decode + assert decode["max_num_seqs"] == 64 + assert "max_num_seqs" not in prefill + + +@pytest.mark.parametrize( + "key", + [ + "data_parallel_rpc_port", + "data_parallel_size", + "data_parallel_size_local", + "disaggregation_mode", + "enable_prefix_caching", + "enable_rl", + "kv_transfer_config", + "kv_events_config", + "pipeline_parallel_size", + "tensor_parallel_size", + "worker_extension_cls", + ], +) +def test_reserved_engine_override_is_rejected(key: str): + config = disaggregated_config(vllm_extra={key: {}}) + with pytest.raises(ValueError, match="Dynamo-managed"): + build_engine_config(config, "prefill", kv_events_port=20080) + + +def test_reserved_role_engine_override_is_rejected(): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "prefill_vllm_overrides": {"tensor_parallel_size": 2}, + } + ) + + with pytest.raises(ValueError, match="prefill_vllm_overrides.*tensor_parallel_size"): + build_engine_config(config, "prefill", kv_events_port=20080) + + +@pytest.mark.parametrize("key", sorted(dynamo._ENGINE_CONFIG_EXCLUDED)) +def test_wrapper_only_global_engine_override_is_rejected(key: str): + config = disaggregated_config(vllm_extra={key: "invalid"}) + + with pytest.raises(ValueError, match=rf"vllm_extra.*{key}.*wrapper/server-only"): + build_engine_config(config, "prefill", kv_events_port=20080) + + +@pytest.mark.parametrize("key", sorted(dynamo._ENGINE_CONFIG_EXCLUDED)) +def test_wrapper_only_role_engine_override_is_rejected(key: str): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "decode_vllm_overrides": {key: "invalid"}, + } + ) + + with pytest.raises(ValueError, match=rf"decode_vllm_overrides.*{key}.*wrapper/server-only"): + build_engine_config(config, "decode") + + +def test_local_specs_allocate_four_workers_and_unique_ports(tmp_path: Path): + specs = build_local_worker_specs(disaggregated_config(), tmp_path, gpu_ids=["4", "5", "6", "7"]) + + assert [spec.role for spec in specs] == list(disaggregated_config().dynamo_worker_roles) + assert [spec.gpu_ids for spec in specs] == [("4",), ("5",), ("6",), ("7",)] + assert len({spec.system_port for spec in specs}) == 4 + assert len({spec.process.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"] for spec in specs}) == 4 + prefill_configs = [ + json.loads(Path(spec.process.arguments[1]).read_text()) for spec in specs if spec.role == "prefill" + ] + assert len({config["kv_events_config"]["endpoint"] for config in prefill_configs}) == 2 + assert all("--enable-rl" in spec.process.command() for spec in specs) + + +def test_local_multi_gpu_workers_allocate_globally_unique_coordinator_ports(tmp_path: Path): + config = disaggregated_config( + parallel={"tp": 1}, + deployment={ + "type": "disaggregated", + "gpus_per_node": 2, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + }, + ) + + specs = build_local_worker_specs(config, tmp_path, gpu_ids=["0", "1", "2", "3"]) + engine_configs = [json.loads(Path(spec.process.arguments[1]).read_text()) for spec in specs] + allocated_ports = [ + *(spec.system_port for spec in specs), + *(int(spec.process.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"]) for spec in specs), + *(engine["data_parallel_rpc_port"] for engine in engine_configs), + *( + int(engine["kv_events_config"]["endpoint"].rsplit(":", 1)[1]) + for engine in engine_configs + if "kv_events_config" in engine + ), + ] + + assert len({engine["data_parallel_rpc_port"] for engine in engine_configs}) == len(specs) + assert len(allocated_ports) == len(set(allocated_ports)) + + +def test_wrapper_options_are_not_written_to_engine_json(): + engine = build_engine_config(disaggregated_config(), "prefill", kv_events_port=20080) + assert "disaggregation_mode" not in engine + assert "enable_rl" not in engine + + +def test_process_specs_own_canonical_commands_and_environment(tmp_path: Path): + config = disaggregated_config( + env_vars={"SHARED": "value"}, + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "prefill_env_vars": {"ROLE": "prefill"}, + }, + ) + + frontend = build_frontend_process(config) + prefill = build_worker_process( + config, + "prefill", + tmp_path / "prefill.json", + nixl_host="127.0.0.1", + nixl_port=20100, + ) + + assert frontend.module == "dynamo.frontend" + assert frontend.arguments[-1] == "--enable-engine-apis" + assert frontend.environment()["DYN_ENABLE_RL"] == "1" + assert prefill.module == "dynamo.vllm" + assert prefill.arguments[-3:] == ("--disaggregation-mode", "prefill", "--enable-rl") + assert prefill.environment()["ROLE"] == "prefill" + assert prefill.environment()["VLLM_PLUGINS"] == "prime_rl" + assert prefill.environment()["VLLM_NIXL_SIDE_CHANNEL_PORT"] == "20100" + + +def test_process_specs_preserve_custom_chat_template_and_parsers(tmp_path: Path): + source = tmp_path / "source-template.jinja" + source.write_text("{{ messages | length }}") + config = disaggregated_config( + model={ + "chat_template": str(source), + "tool_call_parser": "hermes", + "reasoning_parser": "qwen3", + } + ) + + frontend = build_frontend_process(config, output_dir=tmp_path / "generated") + worker = build_worker_process( + config, + "decode", + tmp_path / "decode.json", + nixl_host=None, + nixl_port=20100, + ) + + template_path = Path(frontend.arguments[frontend.arguments.index("--chat-template") + 1]) + assert template_path == tmp_path / "generated" / "chat-template.jinja" + assert template_path.read_text() == source.read_text() + assert frontend.arguments[-4:-2] == ("--dyn-chat-processor", "vllm") + assert worker.arguments[-4:] == ( + "--dyn-tool-call-parser", + "hermes", + "--dyn-reasoning-parser", + "qwen3", + ) + + +@pytest.mark.parametrize( + ("role", "component"), + [("prefill", "prefill"), ("decode", "backend"), ("agg", "backend")], +) +def test_worker_process_uses_deterministic_role_endpoint(tmp_path: Path, role: str, component: str): + process = build_worker_process( + disaggregated_config(env_vars={"DYN_NAMESPACE": "prime-test"}), + role, + tmp_path / f"{role}.json", + nixl_host=None, + nixl_port=20100, + ) + + endpoint = f"dyn://prime-test.{component}.generate" + assert process.environment()["DYN_NAMESPACE"] == "prime-test" + assert process.environment()["DYN_COMPONENT"] == component + assert process.environment()["DYN_ENDPOINT"] == endpoint + assert process.arguments[2:4] == ("--endpoint", endpoint) + + +def test_inline_chat_template_is_materialized_verbatim(tmp_path: Path): + config = disaggregated_config(model={"chat_template": "{{ messages }}"}) + + frontend = build_frontend_process(config, output_dir=tmp_path) + + template_path = Path(frontend.arguments[-1]) + assert template_path.read_text() == "{{ messages }}" + + +def test_frontend_runtime_chat_template_path_does_not_materialize_host_file(tmp_path: Path): + config = disaggregated_config(model={"chat_template": "{{ messages }}"}) + output_dir = tmp_path / "render-host" + runtime_path = Path("/etc/prime-rl/dynamo/chat-template.jinja") + + frontend = build_frontend_process( + config, + output_dir=output_dir, + runtime_chat_template_path=runtime_path, + ) + + assert frontend.arguments[-1] == str(runtime_path) + assert not output_dir.exists() + + +def test_worker_environment_applies_only_matching_role_overrides(tmp_path: Path): + config = disaggregated_config( + deployment={ + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + "prefill_env_vars": {"ROLE_SETTING": "prefill"}, + "decode_env_vars": {"ROLE_SETTING": "decode"}, + } + ) + prefill, decode = build_local_worker_specs(config, tmp_path, gpu_ids=["3", "7"]) + + decode_env = build_worker_environment(decode, {"COMMON": "value"}) + prefill_env = build_worker_environment(prefill, {"COMMON": "value"}) + + assert decode_env["ROLE_SETTING"] == "decode" + assert prefill_env["ROLE_SETTING"] == "prefill" + assert prefill_env["CUDA_VISIBLE_DEVICES"] == "3" + assert decode_env["CUDA_VISIBLE_DEVICES"] == "7" + assert decode_env["VLLM_PLUGINS"] == "prime_rl" + assert decode_env["DYN_COMPONENT"] == "backend" + assert prefill_env["DYN_COMPONENT"] == "prefill" + + +def test_aggregated_worker_uses_canonical_component_name(tmp_path: Path): + config = InferenceConfig.model_validate({"backend": {"type": "dynamo"}}) + spec = build_local_worker_specs(config, tmp_path, gpu_ids=["0"])[0] + + assert build_worker_environment(spec, {})["DYN_COMPONENT"] == "backend" + + +def test_dynamo_dry_run_uses_symbolic_gpu_slots(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + config = disaggregated_config(output_dir=tmp_path, dry_run=True) + captured_gpu_ids = [] + original_build_specs = dynamo.build_local_worker_specs + + class FakeLogger: + def info(self, _message): + pass + + def success(self, _message): + pass + + def build_specs(config, output_dir=None, gpu_ids=None, namespace=None): + captured_gpu_ids.extend(gpu_ids or []) + return original_build_specs(config, output_dir=output_dir, gpu_ids=gpu_ids, namespace=namespace) + + monkeypatch.setattr(dynamo, "_visible_gpu_ids", lambda: pytest.fail("dry-run queried physical GPUs")) + monkeypatch.setattr(dynamo, "build_local_worker_specs", build_specs) + monkeypatch.setattr(inference_entrypoint, "setup_logger", lambda *_args, **_kwargs: FakeLogger()) + + inference_entrypoint.inference_local(config) + + assert captured_gpu_ids == ["", "", "", ""] + + +@pytest.mark.parametrize(("child_code", "supervisor_code"), [(7, 7), (0, 1)]) +def test_child_exit_tears_down_complete_process_group( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + child_code: int, + supervisor_code: int, +): + config = disaggregated_config(output_dir=tmp_path) + processes = [] + terminated = [] + + class FakeProcess: + def __init__(self, returncode): + self.pid = 1000 + len(processes) + self.returncode = returncode + + def poll(self): + return self.returncode + + def popen(*_args, **_kwargs): + process = FakeProcess(child_code if not processes else None) + processes.append(process) + return process + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3") + monkeypatch.setattr(dynamo.subprocess, "Popen", popen) + monkeypatch.setattr(dynamo.signal, "signal", lambda *_args: None) + monkeypatch.setattr(dynamo, "_terminate", terminated.append) + + with pytest.raises(SystemExit) as exc: + dynamo.run_dynamo_local(config) + + assert exc.value.code == supervisor_code + assert len(processes) == 5 + assert terminated == list(reversed(processes)) diff --git a/tests/unit/inference/test_dynamo_admin.py b/tests/unit/inference/test_dynamo_admin.py new file mode 100644 index 0000000000..408e08cc17 --- /dev/null +++ b/tests/unit/inference/test_dynamo_admin.py @@ -0,0 +1,586 @@ +import json +from pathlib import Path + +import httpx +import pytest + +from prime_rl.configs.shared import ClientConfig +from prime_rl.inference.dynamo_admin import ( + DynamoAdminAPI, + DynamoTopology, + DynamoWorker, + discover_workers, + discovery_urls, + validate_worker_membership, +) + +ROUTES = frozenset( + { + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + } +) + + +def worker( + instance_id: int, + *, + component: str = "backend", + system_url: str | None = None, + model: str = "test-model", + routes: frozenset[str] = ROUTES, +) -> dict: + return { + "component": component, + "instance_id": instance_id, + "system_url": system_url or f"http://worker-{instance_id}:8081", + "model": model, + "routes": sorted(routes), + } + + +def async_client(handler, base_url: str = "http://worker:8081") -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url=base_url) + + +def test_discovery_url_defaults_to_rl_listener_port(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DYN_RL_DISCOVERY_URL", raising=False) + monkeypatch.delenv("DYN_RL_PORT", raising=False) + config = ClientConfig(base_url=["http://frontend.example:8000/v1"]) + assert discovery_urls(config) == ["http://frontend.example:8001"] + + +@pytest.mark.asyncio +async def test_worker_discovery_validates_and_sorts_system_urls(): + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "namespace": "test", + "workers": [ + worker(2, system_url="http://worker-b:8081"), + worker(1, system_url="http://worker-a:8081"), + ], + }, + ) + + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg", "agg"), gpus_per_worker=2), + ) + finally: + await client.aclose() + + assert [item.system_url for item in discovered] == ["http://worker-a:8081", "http://worker-b:8081"] + assert discovered[0] == DynamoWorker( + instance_id=1, + component="backend", + role="agg", + system_url="http://worker-a:8081", + model="test-model", + routes=ROUTES, + ) + + +@pytest.mark.asyncio +async def test_worker_discovery_waits_for_exact_staggered_topology(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + workers = [worker(1, component="prefill")] + if calls > 1: + # Dynamo advertises decode as `backend`; the expected topology + # disambiguates it from an aggregated `backend` worker. + workers.append(worker(2, component="backend")) + return httpx.Response(200, json={"namespace": "test", "workers": workers}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("prefill", "decode"), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert {item.role for item in discovered} == {"prefill", "decode"} + + +@pytest.mark.asyncio +async def test_worker_discovery_deduplicates_consistent_frontend_snapshots(): + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"namespace": "test", "workers": [worker(7)]}) + + clients = [async_client(handler, f"http://frontend-{index}:8001") for index in range(2)] + try: + discovered = await discover_workers( + clients, + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + for client in clients: + await client.aclose() + + assert len(discovered) == 1 + assert discovered[0].instance_id == 7 + + +@pytest.mark.asyncio +async def test_worker_discovery_rejects_inconsistent_frontend_snapshots(): + def first(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"namespace": "test", "workers": [worker(7)]}) + + def restarted(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"namespace": "test", "workers": [worker(8)]}) + + clients = [async_client(first, "http://frontend-a:8001"), async_client(restarted, "http://frontend-b:8001")] + try: + with pytest.raises(TimeoutError, match="inconsistent worker snapshots"): + await discover_workers( + clients, + timeout=0.01, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + for client in clients: + await client.aclose() + + +@pytest.mark.asyncio +async def test_worker_discovery_bounds_each_get_by_remaining_deadline(): + request_timeouts: list[dict[str, float]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + request_timeouts.append(request.extensions["timeout"]) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + client = async_client(handler, "http://frontend:8001") + try: + await discover_workers( + [client], + timeout=0.5, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert len(request_timeouts) == 1 + assert 0 < request_timeouts[0]["read"] <= 0.5 + + +@pytest.mark.asyncio +async def test_worker_discovery_rejects_incomplete_admin_surface(): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1, routes=frozenset())]}) + + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(ValueError, match="missing RL routes"): + await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_rejects_model_mismatch_without_retry(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1, model="other-model")]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(ValueError, match="serves 'other-model', expected 'test-model'"): + await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [400, 401, 403, 404]) +async def test_worker_discovery_rejects_permanent_http_errors_without_retry( + monkeypatch: pytest.MonkeyPatch, + status_code: int, +): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response(status_code) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(httpx.HTTPStatusError): + await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [408, 409, 429, 500]) +async def test_worker_discovery_retries_transient_http_errors( + monkeypatch: pytest.MonkeyPatch, + status_code: int, +): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + return httpx.Response(status_code) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_retries_transport_errors(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + raise httpx.ConnectError("frontend is starting", request=request) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_retries_request_timeouts(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + raise httpx.ReadTimeout("frontend response timed out", request=request) + return httpx.Response(200, json={"namespace": "test", "workers": [worker(1)]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +async def test_worker_discovery_retries_transient_worker_probe_errors(monkeypatch: pytest.MonkeyPatch): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + discovered_worker = worker(1) + if calls == 1: + discovered_worker["error"] = "worker endpoint has not converged" + return httpx.Response(200, json={"namespace": "test", "workers": [discovered_worker]}) + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + discovered = await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 2 + assert discovered[0].instance_id == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "error_type", "match"), + [ + (httpx.Response(200, content=b"{"), json.JSONDecodeError, "Expecting property name"), + (httpx.Response(200, json={"namespace": "test"}), ValueError, "invalid response"), + ( + httpx.Response( + 200, + json={"namespace": "test", "workers": [{**worker(1), "error": 123}]}, + ), + ValueError, + "invalid error", + ), + ], +) +async def test_worker_discovery_rejects_invalid_payload_without_retry( + monkeypatch: pytest.MonkeyPatch, + response: httpx.Response, + error_type: type[Exception], + match: str, +): + calls = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return response + + monkeypatch.setattr("prime_rl.inference.dynamo_admin.DISCOVERY_POLL_INTERVAL_S", 0) + client = async_client(handler, "http://frontend:8001") + try: + with pytest.raises(error_type, match=match): + await discover_workers( + [client], + timeout=1, + model_name="test-model", + topology=DynamoTopology(roles=("agg",), gpus_per_worker=1), + ) + finally: + await client.aclose() + + assert calls == 1 + + +def test_worker_membership_change_is_rejected(): + expected = (DynamoWorker(1, "backend", "agg", "http://worker:8081", "test-model", ROUTES),) + restarted = [ + DynamoWorker(2, "backend", "agg", "http://worker:8081", "test-model", ROUTES), + ] + with pytest.raises(RuntimeError, match="membership changed"): + validate_worker_membership(expected, restarted) + + +@pytest.mark.asyncio +async def test_nccl_initialization_and_update_use_engine_routes(tmp_path: Path): + requests: list[tuple[str, dict]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append((request.url.path, json.loads(request.content))) + return httpx.Response(200, json={"status": "ok"}) + + clients = [async_client(handler, f"http://worker-{index}:8081") for index in range(2)] + admin = DynamoAdminAPI() + try: + await admin.initialize_nccl( + clients, + host="localhost", + port=29511, + timeout=12000, + inference_world_size=4, + gpus_per_worker=2, + quantize_in_weight_transfer=False, + ) + await admin.update_weights(clients, tmp_path / "step_1", step=1) + finally: + for client in clients: + await client.aclose() + + init_bodies = [body for path, body in requests if path.endswith("/init_weights_update_group")] + assert [body["rank_offset"] for body in init_bodies] == [0, 2] + assert all(body["inference_world_size"] == 4 for body in init_bodies) + + updates = [body for path, body in requests if path.endswith("/update_weights_from_distributed")] + assert len(updates) == 2 + assert all(body["engine_rpc"] == "update_weights_from_path" for body in updates) + assert all(body["weight_version"] == "1" for body in updates) + assert (tmp_path / "step_1" / "NCCL_READY").exists() + + paths = [path for path, _body in requests] + assert paths.count("/engine/pause_generation") == 2 + assert paths.count("/engine/resume_generation") == 2 + pause_bodies = [body for path, body in requests if path.endswith("/pause_generation")] + assert pause_bodies == [{"mode": "wait", "clear_cache": False}] * 2 + + +@pytest.mark.asyncio +async def test_nccl_initialization_does_not_replay_ambiguous_timeout(): + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + raise httpx.ReadTimeout("response lost after collective may have started", request=request) + + client = async_client(handler) + try: + with pytest.raises(httpx.ReadTimeout): + await DynamoAdminAPI().initialize_nccl( + [client], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=1, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + finally: + await client.aclose() + + assert paths == ["/engine/init_weights_update_group"] + + +@pytest.mark.asyncio +async def test_nccl_initialization_rejects_world_size_that_conflicts_with_topology(): + requests = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal requests + requests += 1 + return httpx.Response(200, json={"status": "ok"}) + + client = async_client(handler) + try: + with pytest.raises(ValueError, match="does not match"): + await DynamoAdminAPI().initialize_nccl( + [client], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=1, + gpus_per_worker=2, + quantize_in_weight_transfer=False, + ) + finally: + await client.aclose() + + assert requests == 0 + + +@pytest.mark.asyncio +async def test_weight_collective_does_not_replay_ambiguous_timeout(tmp_path: Path): + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + if request.url.path.endswith("update_weights_from_disk"): + raise httpx.ReadTimeout("response lost after update may have committed", request=request) + return httpx.Response(200, json={"status": "ok"}) + + client = async_client(handler) + try: + with pytest.raises(httpx.ReadTimeout): + await DynamoAdminAPI().update_weights([client], tmp_path / "weights", step=1) + finally: + await client.aclose() + + assert paths == [ + "/engine/pause_generation", + "/engine/update_weights_from_disk", + ] + + +@pytest.mark.asyncio +async def test_engine_status_error_is_not_accepted(): + paths = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + if request.url.path.endswith("resume_generation"): + return httpx.Response(200, json={"status": "ok"}) + return httpx.Response(200, json={"status": "error", "message": "not paused"}) + + client = async_client(handler) + try: + with pytest.raises(RuntimeError, match="not paused"): + await DynamoAdminAPI().update_weights([client], Path("weights"), step=1) + finally: + await client.aclose() + assert paths == ["/engine/pause_generation", "/engine/resume_generation"] + + +@pytest.mark.asyncio +async def test_resume_error_does_not_hide_primary_update_error(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + + async def post(_client, method, *_args, **_kwargs): + if method == "pause_generation": + raise RuntimeError("pause failed") + if method == "resume_generation": + raise RuntimeError("resume failed") + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + + with pytest.raises(RuntimeError, match="pause failed") as exc: + await admin.update_weights([object()], Path("weights"), step=1) + + assert exc.value.__notes__ == ["Dynamo resume_generation cleanup also failed: RuntimeError('resume failed')"] diff --git a/tests/unit/inference/test_dynamo_admin_barriers.py b/tests/unit/inference/test_dynamo_admin_barriers.py new file mode 100644 index 0000000000..d3486af182 --- /dev/null +++ b/tests/unit/inference/test_dynamo_admin_barriers.py @@ -0,0 +1,269 @@ +"""Cancellation and partial-failure contracts for stateful Dynamo admin fanouts.""" + +import asyncio +from pathlib import Path + +import pytest + +from prime_rl.inference.dynamo_admin import DynamoAdminAPI + + +@pytest.mark.asyncio +async def test_nccl_initialization_settles_siblings_and_fails_closed_after_partial_error( + monkeypatch: pytest.MonkeyPatch, +): + admin = DynamoAdminAPI() + delayed_started = asyncio.Event() + release_delayed = asyncio.Event() + delayed_finished = asyncio.Event() + calls = 0 + + async def post(client, method, *_args, **_kwargs): + nonlocal calls + calls += 1 + assert method == "init_weights_update_group" + if client == "failed": + raise RuntimeError("init failed") + delayed_started.set() + await release_delayed.wait() + delayed_finished.set() + return {} + + monkeypatch.setattr(admin, "_post", post) + initialize = asyncio.create_task( + admin.initialize_nccl( + ["failed", "delayed"], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=2, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + ) + await delayed_started.wait() + await asyncio.sleep(0) + assert not initialize.done() + + release_delayed.set() + with pytest.raises(RuntimeError, match="init failed"): + await initialize + assert delayed_finished.is_set() + + with pytest.raises(RuntimeError, match="indeterminate"): + await admin.initialize_nccl( + ["failed", "delayed"], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=2, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + with pytest.raises(RuntimeError, match="indeterminate"): + await admin.update_weights(["failed", "delayed"], Path("weights"), step=1) + assert calls == 2 + + +@pytest.mark.asyncio +async def test_nccl_initialization_settles_before_propagating_repeated_cancellation( + monkeypatch: pytest.MonkeyPatch, +): + admin = DynamoAdminAPI() + init_started = asyncio.Event() + release_init = asyncio.Event() + init_finished = asyncio.Event() + + async def post(_client, method, *_args, **_kwargs): + assert method == "init_weights_update_group" + init_started.set() + await release_init.wait() + init_finished.set() + return {} + + monkeypatch.setattr(admin, "_post", post) + initialize = asyncio.create_task( + admin.initialize_nccl( + ["worker"], + host="localhost", + port=29511, + timeout=12000, + inference_world_size=1, + gpus_per_worker=1, + quantize_in_weight_transfer=False, + ) + ) + await init_started.wait() + initialize.cancel() + await asyncio.sleep(0) + initialize.cancel() + await asyncio.sleep(0) + assert not initialize.done() + + release_init.set() + with pytest.raises(asyncio.CancelledError): + await initialize + assert init_finished.is_set() + + with pytest.raises(RuntimeError, match="indeterminate"): + await admin.update_weights(["worker"], Path("weights"), step=1) + + +@pytest.mark.asyncio +async def test_pause_fanout_settles_delayed_sibling_before_resume(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + delayed_started = asyncio.Event() + release_delayed = asyncio.Event() + delayed_finished = asyncio.Event() + resume_started = asyncio.Event() + + async def post(client, method, *_args, **_kwargs): + if method == "pause_generation": + if client == "fast-failure": + raise RuntimeError("pause failed") + delayed_started.set() + await release_delayed.wait() + delayed_finished.set() + return {} + if method == "resume_generation": + resume_started.set() + assert delayed_finished.is_set() + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + update = asyncio.create_task(admin.update_weights(["fast-failure", "delayed"], Path("weights"), step=1)) + await delayed_started.wait() + await asyncio.sleep(0) + assert not resume_started.is_set() + + release_delayed.set() + with pytest.raises(RuntimeError, match="pause failed"): + await update + + assert resume_started.is_set() + + +@pytest.mark.asyncio +async def test_collective_update_settles_delayed_sibling_without_resuming_indeterminate_workers( + monkeypatch: pytest.MonkeyPatch, +): + admin = DynamoAdminAPI() + delayed_started = asyncio.Event() + release_delayed = asyncio.Event() + delayed_finished = asyncio.Event() + resume_started = asyncio.Event() + + async def post(client, method, *_args, **_kwargs): + if method == "pause_generation": + return {} + if method == "update_weights_from_disk": + if client == "fast-failure": + raise RuntimeError("collective failed") + delayed_started.set() + await release_delayed.wait() + delayed_finished.set() + return {} + if method == "resume_generation": + resume_started.set() + assert delayed_finished.is_set() + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + update = asyncio.create_task(admin.update_weights(["fast-failure", "delayed"], Path("weights"), step=1)) + await delayed_started.wait() + await asyncio.sleep(0) + assert not resume_started.is_set() + + release_delayed.set() + with pytest.raises(RuntimeError, match="collective failed"): + await update + + assert not resume_started.is_set() + with pytest.raises(RuntimeError, match="weight state is indeterminate"): + await admin.update_weights(["fast-failure", "delayed"], Path("weights"), step=2) + + +@pytest.mark.asyncio +async def test_collective_update_is_settled_before_propagating_cancellation(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + update_started = asyncio.Event() + release_update = asyncio.Event() + update_finished = asyncio.Event() + resume_started = asyncio.Event() + + async def post(_client, method, *_args, **_kwargs): + if method == "pause_generation": + return {} + if method == "update_weights_from_disk": + update_started.set() + await release_update.wait() + update_finished.set() + return {} + if method == "resume_generation": + resume_started.set() + assert update_finished.is_set() + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + update = asyncio.create_task(admin.update_weights(["worker"], Path("weights"), step=1)) + await update_started.wait() + update.cancel() + await asyncio.sleep(0) + update.cancel() + await asyncio.sleep(0) + + assert not update.done() + assert not resume_started.is_set() + + release_update.set() + with pytest.raises(asyncio.CancelledError): + await update + + assert update_finished.is_set() + assert not resume_started.is_set() + with pytest.raises(RuntimeError, match="weight state is indeterminate"): + await admin.update_weights(["worker"], Path("weights"), step=2) + + +@pytest.mark.asyncio +async def test_resume_failure_after_successful_mutation_is_reported(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + operations: list[str] = [] + + async def post(_client, method, *_args, **_kwargs): + operations.append(method) + if method == "resume_generation": + raise RuntimeError("resume failed after mutation") + return {} + + monkeypatch.setattr(admin, "_post", post) + + with pytest.raises(RuntimeError, match="resume failed after mutation"): + await admin.update_weights(["worker"], Path("weights"), step=1) + + assert operations == ["pause_generation", "update_weights_from_disk", "resume_generation"] + with pytest.raises(RuntimeError, match="weight state is indeterminate"): + await admin.update_weights(["worker"], Path("weights"), step=2) + assert operations == ["pause_generation", "update_weights_from_disk", "resume_generation"] + + +@pytest.mark.asyncio +async def test_fanout_preserves_primary_error_and_annotates_siblings(monkeypatch: pytest.MonkeyPatch): + admin = DynamoAdminAPI() + + async def post(client, method, *_args, **_kwargs): + if method == "pause_generation": + raise RuntimeError(f"{client} pause failed") + if method == "resume_generation": + return {} + raise AssertionError(method) + + monkeypatch.setattr(admin, "_post", post) + + with pytest.raises(RuntimeError, match="primary pause failed") as exc: + await admin.update_weights(["primary", "sibling"], Path("weights"), step=1) + + assert exc.value.__notes__ == ["Dynamo pause_generation sibling also failed: RuntimeError('sibling pause failed')"] diff --git a/tests/unit/inference/test_helm_dgd.py b/tests/unit/inference/test_helm_dgd.py new file mode 100644 index 0000000000..ae528366ac --- /dev/null +++ b/tests/unit/inference/test_helm_dgd.py @@ -0,0 +1,551 @@ +import hashlib +import json +import subprocess +from dataclasses import replace +from pathlib import Path + +import pytest + +from prime_rl.inference.dgd import DynamoGraphRenderOptions, write_dgd_artifacts +from tests.unit.inference.helm_dgd_test_utils import ( + DYNAMO_SHA, + GPU_SCHEDULING, + IMAGE_DIGEST, + ORCHESTRATOR_COMMAND, + PRIME_SHA, + TRAINER_COMMAND, + helm_template, + inference_config, + labels_match, + render_options, + rendered_documents, + rendered_resource, + rewrite_valid_integrity, + toleration_identity, +) + + +def test_native_chart_still_renders_inference_statefulset(): + rendered = helm_template() + assert "name: p4-math-inference\n" in rendered + assert "kind: DynamoGraphDeployment" not in rendered + assert "p4-math-inference-0.p4-math-inference-headless" in rendered + + +def test_chart_rejects_unknown_inference_mode(): + with pytest.raises(subprocess.CalledProcessError): + helm_template("--set", "inference.mode=typo") + + +def test_native_chart_rejects_invalid_image_pull_policy(): + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("--set", "image.pullPolicy=Sometimes") + + assert "image/pullPolicy" in error.value.stderr + assert "'Always', 'IfNotPresent', 'Never'" in error.value.stderr + + +def test_chart_release_name_boundary_keeps_every_resource_name_valid(): + release_name = "a" * 41 + rendered = helm_template(release_name=release_name, release_namespace="default") + + assert all(len(document["metadata"]["name"]) <= 63 for document in rendered_documents(rendered)) + + +def test_chart_rejects_release_name_that_would_overflow_service_names(): + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template(release_name="a" * 42, release_namespace="default") + + assert "at most 41 characters" in error.value.stderr + + +def test_dgd_chart_renders_generated_graph_without_inference_statefulset(tmp_path: Path): + options = render_options(tmp_path) + paths = write_dgd_artifacts(inference_config(), options) + rendered = helm_template("-f", str(paths["values"])) + graph = json.loads(paths["resource"].read_text()) + rendered_graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + + assert rendered.count("kind: DynamoGraphDeployment") == 1 + assert rendered.count("kind: ConfigMap") == 1 + assert "name: p4-math-inference\n" not in rendered + assert "name: p4-math-frontend-rl" in rendered + assert "http://p4-math-frontend.bis-vllm.svc.cluster.local:8000/v1" in rendered + assert "http://p4-math-frontend-rl.bis-vllm.svc.cluster.local:8001" in rendered + assert graph["spec"]["services"]["VllmPrefillWorker"]["replicas"] == 2 + assert graph["spec"]["services"]["VllmDecodeWorker"]["replicas"] == 2 + assert rendered_graph["spec"]["services"]["Frontend"]["extraPodSpec"]["mainContainer"]["ports"] == [ + {"containerPort": 8000, "name": "http"}, + {"containerPort": 8001, "name": "rl"}, + ] + assert rendered.count(f'image: "{options.image}"') == 2 + assert rendered.count(f"image: {options.image}") == 3 + assert rendered.count("nvcrimagepullsecret") == 5 + assert rendered.count("name: DYN_RL_TOPOLOGY") == 2 + assert rendered.count("claimName: model-cache") == 5 + assert rendered.count("name: HF_TOKEN") == 5 + assert rendered.count("name: HF_HOME") == 5 + chart_pods = { + component: rendered_resource(rendered, "StatefulSet", f"p4-math-{component}")["spec"]["template"]["spec"] + for component in ("orchestrator", "trainer") + } + dgd_pods = { + component: rendered_graph["spec"]["services"][service]["extraPodSpec"] + for component, service in { + "frontend": "Frontend", + "prefill": "VllmPrefillWorker", + "decode": "VllmDecodeWorker", + }.items() + } + pods = {**chart_pods, **dgd_pods} + image_selector = { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + } + gpu_selector = {**image_selector, "nvidia.com/gpu.product": "NVIDIA-GB200"} + required_tolerations = { + ("kubernetes.io/arch", "Equal", "arm64", "NoSchedule"), + ("nvidia.com/gpu", "Exists", None, "NoSchedule"), + ("prime-rl", "Equal", "true", "NoSchedule"), + } + + assert set(pods) == {"orchestrator", "trainer", "frontend", "prefill", "decode"} + for component, pod in pods.items(): + assert {toleration_identity(item) for item in pod["tolerations"]} == required_tolerations + container = pod["containers"][0] if component in chart_pods else pod["mainContainer"] + hf_token = next(item for item in container["env"] if item["name"] == "HF_TOKEN") + assert hf_token["valueFrom"]["secretKeyRef"]["optional"] is False + if component in {"trainer", "prefill", "decode"}: + assert pod["nodeSelector"] == gpu_selector + assert pod["runtimeClassName"] == "nvidia" + else: + assert pod["nodeSelector"] == image_selector + assert "runtimeClassName" not in pod + + assert "nvidia.com/gpu" not in chart_pods["orchestrator"]["containers"][0]["resources"].get("requests", {}) + assert chart_pods["trainer"]["containers"][0]["resources"]["requests"]["nvidia.com/gpu"] == 1 + assert chart_pods["orchestrator"]["containers"][0]["args"] == [ORCHESTRATOR_COMMAND] + assert chart_pods["trainer"]["containers"][0]["args"] == [TRAINER_COMMAND] + assert rendered_resource(rendered, "StatefulSet", "p4-math-orchestrator")["spec"]["replicas"] == 1 + assert rendered_resource(rendered, "StatefulSet", "p4-math-trainer")["spec"]["replicas"] == 1 + assert not any(kind in rendered for kind in ("kind: ClusterRole", "kind: CustomResourceDefinition")) + + +def test_external_controller_mode_renders_only_five_dgd_inference_pods(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, external_controller=True), + ) + rendered = helm_template("-f", str(paths["values"])) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + + assert "kind: StatefulSet" not in rendered + assert "kind: PersistentVolumeClaim" not in rendered + assert "p4-math-shared-data" not in rendered + assert sum(service["replicas"] for service in graph["spec"]["services"].values()) == 5 + assert rendered.count("kind: DynamoGraphDeployment") == 1 + assert rendered.count("name: p4-math-frontend-rl") == 1 + + +@pytest.mark.parametrize("external_controller", [False, True]) +def test_dgd_chart_renders_without_gpu_runtime_class( + tmp_path: Path, + external_controller: bool, +): + options = render_options(tmp_path, external_controller=external_controller) + options = replace( + options, + gpu_scheduling=replace(options.gpu_scheduling, runtime_class_name=None), + ) + paths = write_dgd_artifacts(inference_config(), options) + + rendered = helm_template("-f", str(paths["values"])) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + workload = json.loads( + json.loads(paths["values"].read_text())["inference"]["dynamoGraph"]["workloadBinding"]["canonical"] + ) + + assert "runtimeClassName" not in workload["trainer"]["placement"] + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + assert "runtimeClassName" not in graph["spec"]["services"][role]["extraPodSpec"] + if external_controller: + assert "kind: StatefulSet" not in rendered + else: + trainer = rendered_resource(rendered, "StatefulSet", "p4-math-trainer") + assert "runtimeClassName" not in trainer["spec"]["template"]["spec"] + + +def test_dgd_chart_projects_model_cache_once_into_operator_pod_specs(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + rendered = helm_template("-f", str(paths["values"])) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + expected_mount = {"name": "model-cache", "mountPath": "/model-cache"} + expected_volume = { + "name": "model-cache", + "persistentVolumeClaim": {"claimName": "model-cache"}, + } + + for service in graph["spec"]["services"].values(): + assert "volumeMounts" not in service + pod_spec = service["extraPodSpec"] + assert pod_spec["mainContainer"].get("volumeMounts", []).count(expected_mount) == 1 + assert pod_spec.get("volumes", []).count(expected_volume) == 1 + + +def test_chart_managed_trainer_uses_exact_bound_gpu_resources(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, trainer_gpu_count=4), + ) + values = json.loads(paths["values"].read_text()) + rendered = helm_template("-f", str(paths["values"])) + trainer = rendered_resource(rendered, "StatefulSet", "p4-math-trainer") + resources = trainer["spec"]["template"]["spec"]["containers"][0]["resources"] + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + + assert workload["trainer"]["gpu"] == {"enabled": True, "count": 4} + assert resources["requests"]["nvidia.com/gpu"] == 4 + assert resources["limits"]["nvidia.com/gpu"] == 4 + + +def test_legacy_statefulset_selectors_remain_upgrade_compatible(): + documents = rendered_documents(helm_template()) + controllers = [document for document in documents if document["kind"] == "StatefulSet"] + + assert len(controllers) == 3 + for controller in controllers: + role = controller["metadata"]["labels"]["role"] + assert controller["spec"]["selector"]["matchLabels"] == { + "app": "prime-rl", + "example": "reverse-text", + "role": role, + } + assert controller["spec"]["template"]["metadata"]["labels"]["app.kubernetes.io/instance"] == "p4-math" + + +def test_chart_service_selectors_are_release_disjoint(): + releases = {release: rendered_documents(helm_template(release_name=release)) for release in ("alpha", "beta")} + pod_labels = { + release: [ + document["spec"]["template"]["metadata"]["labels"] + for document in documents + if document["kind"] == "StatefulSet" + ] + for release, documents in releases.items() + } + + for release, documents in releases.items(): + other_release = "beta" if release == "alpha" else "alpha" + services = [document for document in documents if document["kind"] == "Service"] + assert len(services) == 6 + + for service in services: + selector = service["spec"]["selector"] + assert selector["app.kubernetes.io/instance"] == release + assert any(labels_match(selector, labels) for labels in pod_labels[release]) + assert not any(labels_match(selector, labels) for labels in pod_labels[other_release]) + + +def test_dgd_rl_service_selector_is_release_disjoint(tmp_path: Path): + release_pods: dict[str, dict[str, str]] = {} + release_services: dict[str, dict[str, str]] = {} + for release in ("alpha", "beta"): + output_dir = tmp_path / release + paths = write_dgd_artifacts(inference_config(), render_options(output_dir, release_name=release)) + rendered = helm_template("-f", str(paths["values"]), release_name=release) + release_pods[release] = { + # Grove owns and rewrites the conventional app labels on realized + # pods, but Dynamo's identity labels remain stable. + "app.kubernetes.io/name": f"{release}-0-frontend", + "nvidia.com/dynamo-graph-deployment-name": release, + "nvidia.com/dynamo-component": "Frontend", + "nvidia.com/dynamo-component-type": "frontend", + } + release_services[release] = rendered_resource(rendered, "Service", f"{release}-frontend-rl")["spec"]["selector"] + + for release in ("alpha", "beta"): + other_release = "beta" if release == "alpha" else "alpha" + selector = release_services[release] + assert selector == { + "nvidia.com/dynamo-graph-deployment-name": release, + "nvidia.com/dynamo-component": "Frontend", + "nvidia.com/dynamo-component-type": "frontend", + } + assert labels_match(selector, release_pods[release]) + assert not labels_match(selector, release_pods[other_release]) + + +def test_dgd_chart_renders_chat_template_configmap_and_frontend_mount(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(chat_template="template-marker: {{ messages }}"), + render_options(tmp_path), + ) + rendered = helm_template("-f", str(paths["values"])) + values = json.loads(paths["values"].read_text()) + config_map = rendered_resource( + rendered, + "ConfigMap", + values["inference"]["dynamoGraph"]["engineConfig"]["name"], + ) + + assert config_map["data"]["chat-template.jinja"] == "template-marker: {{ messages }}" + assert "/etc/prime-rl/dynamo/chat-template.jinja" in rendered + assert "name: dynamo-chat-template" in rendered + assert "key: chat-template.jinja" in rendered + + +def test_dgd_chart_preserves_exact_content_addressed_configmap_bytes(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config(chat_template="EXACT: {{ messages }}"), + render_options(tmp_path), + ) + values = json.loads(paths["values"].read_text()) + rendered = helm_template("-f", str(paths["values"])) + config_map_name = values["inference"]["dynamoGraph"]["engineConfig"]["name"] + config_map = rendered_resource(rendered, "ConfigMap", config_map_name) + expected_data = values["inference"]["dynamoGraph"]["engineConfig"]["data"] + + assert config_map["data"] == expected_data + expected_hash = values["inference"]["dynamoGraph"]["engineConfig"]["sha256"] + canonical_data = (json.dumps(config_map["data"], indent=2, sort_keys=True) + "\n").encode() + assert hashlib.sha256(canonical_data).hexdigest() == expected_hash + + +def test_dgd_chart_uses_canonical_workload_contract_as_sole_authority(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + workload_binding = values["inference"]["dynamoGraph"]["workloadBinding"] + workload = json.loads(workload_binding["canonical"]) + + assert hashlib.sha256(workload_binding["canonical"].encode()).hexdigest() == workload_binding["sha256"] + assert workload["controllerMode"] == "chartManaged" + for key in ("config", "huggingFace", "image", "modelCache", "storage"): + assert workload[key] == values[key] + assert values["orchestrator"] == { + key: value for key, value in workload["orchestrator"].items() if key not in {"gpu", "placement"} + } + assert values["trainer"] == {key: value for key, value in workload["trainer"].items() if key != "placement"} + assert workload["orchestrator"]["placement"]["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + } + assert workload["trainer"]["placement"]["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + "nvidia.com/gpu.product": "NVIDIA-GB200", + } + + +@pytest.mark.parametrize( + ("component", "selector_key"), + [ + ("orchestrator", "kubernetes.io/arch"), + ("orchestrator", "cloud.google.com/gke-nodepool"), + ("trainer", "kubernetes.io/arch"), + ("trainer", "cloud.google.com/gke-nodepool"), + ("trainer", "nvidia.com/gpu.product"), + ], +) +def test_dgd_chart_rejects_rehashed_chart_selector_mutations( + component: str, + selector_key: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + selector = workload[component]["placement"]["nodeSelector"] + selector[f"tampered.example/{selector_key.rsplit('/', 1)[-1]}"] = selector.pop(selector_key) + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"{component}-selector.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert f"{component} placement must match" in error.value.stderr + + +def test_dgd_chart_rejects_rehashed_runtime_class_mutation(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + workload["trainer"]["placement"]["runtimeClassName"] = "tampered" + rewrite_valid_integrity(values, workload) + mutation = tmp_path / "runtime-class.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert "trainer placement must match" in error.value.stderr + + +@pytest.mark.parametrize("component", ["orchestrator", "trainer"]) +@pytest.mark.parametrize("toleration_key", ["kubernetes.io/arch", "nvidia.com/gpu", "prime-rl"]) +def test_dgd_chart_rejects_every_rehashed_required_toleration_mutation( + component: str, + toleration_key: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + toleration = next(item for item in workload[component]["placement"]["tolerations"] if item["key"] == toleration_key) + toleration["effect"] = "NoExecute" + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"{component}-{toleration_key.replace('/', '-')}.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert f"{component} placement must match" in error.value.stderr + + +def test_dgd_chart_ignores_legacy_component_placement_overlays(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + overlay = tmp_path / "legacy-placement.json" + overlay.write_text( + json.dumps( + { + "orchestrator": { + "nodeSelector": {"tampered": "true"}, + "runtimeClassName": "tampered", + "tolerations": [{"key": "tampered", "operator": "Exists"}], + }, + "trainer": { + "nodeSelector": {"tampered": "true"}, + "runtimeClassName": "tampered", + "tolerations": [{"key": "tampered", "operator": "Exists"}], + }, + } + ) + ) + rendered = helm_template("-f", str(paths["values"]), "-f", str(overlay)) + + orchestrator = rendered_resource(rendered, "StatefulSet", "p4-math-orchestrator")["spec"]["template"]["spec"] + trainer = rendered_resource(rendered, "StatefulSet", "p4-math-trainer")["spec"]["template"]["spec"] + assert orchestrator["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + } + assert "runtimeClassName" not in orchestrator + assert trainer["nodeSelector"] == { + "cloud.google.com/gke-nodepool": "customer-gpu-o7v", + "kubernetes.io/arch": "arm64", + "nvidia.com/gpu.product": "NVIDIA-GB200", + } + assert trainer["runtimeClassName"] == "nvidia" + + +def test_filesystem_broadcast_reuses_existing_claim_without_rendering_pvc(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config("filesystem"), + render_options(tmp_path, shared_pvc="p4-shared-data"), + ) + rendered = helm_template("-f", str(paths["values"])) + + assert "kind: PersistentVolumeClaim" not in rendered + assert rendered.count("claimName: p4-shared-data") == 4 + + +def test_external_filesystem_broadcast_binds_existing_claim_without_rendering_pvc(tmp_path: Path): + paths = write_dgd_artifacts( + inference_config("filesystem"), + render_options( + tmp_path, + external_controller=True, + shared_pvc="p4-shared-data", + ), + ) + rendered = helm_template("-f", str(paths["values"])) + values = json.loads(paths["values"].read_text()) + graph = rendered_resource(rendered, "DynamoGraphDeployment", "p4-math") + services = graph["spec"]["services"] + workload = json.loads(values["inference"]["dynamoGraph"]["workloadBinding"]["canonical"]) + + assert "kind: PersistentVolumeClaim" not in rendered + assert workload["storage"] == { + "enabled": True, + "existingClaim": "p4-shared-data", + "storageClassName": "nfs", + "accessModes": ["ReadWriteMany"], + "size": "1Ti", + "mountPath": "/data", + } + assert graph["spec"]["pvcs"] == [ + {"create": False, "name": "model-cache"}, + {"create": False, "name": "p4-shared-data"}, + ] + assert all("volumeMounts" not in service for service in services.values()) + for role in ("VllmPrefillWorker", "VllmDecodeWorker"): + pod_spec = services[role]["extraPodSpec"] + assert pod_spec["mainContainer"]["volumeMounts"].count({"name": "p4-shared-data", "mountPath": "/data"}) == 1 + assert ( + pod_spec["volumes"].count( + { + "name": "p4-shared-data", + "persistentVolumeClaim": {"claimName": "p4-shared-data"}, + } + ) + == 1 + ) + + +def test_dgd_chart_rejects_mutable_prime_runtime_image(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError): + helm_template( + "-f", + str(paths["values"]), + "--set", + "image.reference=nvcr.io/example/prime:latest", + ) + + +def test_dgd_chart_rejects_runtime_image_that_differs_from_workers(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + other_digest = f"sha256:{'4' * 64}" + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "-f", + str(paths["values"]), + "--set", + f"image.reference=nvcr.io/example/prime:reviewed@{other_digest}", + ) + + assert "image configuration must match the workload binding" in error.value.stderr + + +def test_dgd_rejects_image_without_matching_digest(tmp_path: Path): + with pytest.raises(ValueError, match="must be pinned"): + DynamoGraphRenderOptions( + release_name="p4-math", + namespace="bis-vllm", + image="nvcr.io/example/prime:p4", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, + ) + + +def test_dgd_rejects_image_without_commit_suffixes(tmp_path: Path): + with pytest.raises(ValueError, match="commit suffixes"): + DynamoGraphRenderOptions( + release_name="p4-math", + namespace="bis-vllm", + image=f"nvcr.io/example/prime:p4@{IMAGE_DIGEST}", + output_dir=tmp_path, + prime_sha=PRIME_SHA, + dynamo_sha=DYNAMO_SHA, + image_digest=IMAGE_DIGEST, + run_name="p4-run", + gpu_scheduling=GPU_SCHEDULING, + ) diff --git a/tests/unit/inference/test_helm_dgd_integrity.py b/tests/unit/inference/test_helm_dgd_integrity.py new file mode 100644 index 0000000000..3f9940ff68 --- /dev/null +++ b/tests/unit/inference/test_helm_dgd_integrity.py @@ -0,0 +1,504 @@ +import json +import subprocess +from pathlib import Path + +import pytest + +from prime_rl.inference.dgd import write_dgd_artifacts +from tests.unit.inference.helm_dgd_test_utils import ( + helm_template, + inference_config, + render_options, + rewrite_valid_integrity, + write_values_mutation, +) + + +@pytest.mark.parametrize( + ("case", "path", "replacement", "error_fragment"), + [ + ( + "engine-data", + ("inference", "dynamoGraph", "engineConfig", "data", "prefill-engine.json"), + "tampered", + "engineConfig.data must match its canonical payload", + ), + ( + "engine-sha", + ("inference", "dynamoGraph", "engineConfig", "sha256"), + "0" * 64, + "engineConfig.sha256 must match its canonical payload", + ), + ( + "engine-name", + ("inference", "dynamoGraph", "engineConfig", "name"), + "p4-math-dynamo-engine-000000000000", + "engineConfig.name must be content-addressed", + ), + ( + "dgd-config-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/config-sha256", + ), + "0" * 64, + "DynamoGraphDeployment config-sha256 must match engineConfig.sha256", + ), + ( + "manifest-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/manifest-sha256", + ), + "0" * 64, + "manifest-sha256 must match its canonical payload", + ), + ( + "workload-sha", + ("inference", "dynamoGraph", "workloadBinding", "sha256"), + "0" * 64, + "workload binding sha256 must match its canonical payload", + ), + ( + "dgd-workload-sha", + ( + "inference", + "dynamoGraph", + "resource", + "metadata", + "annotations", + "prime-rl.nvidia.com/workload-sha256", + ), + "0" * 64, + "workload binding annotation must match workloadBinding.sha256", + ), + ], +) +def test_dgd_chart_rejects_content_identity_mutations( + case: str, + path: tuple[str, ...], + replacement: object, + error_fragment: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + mutation = tmp_path / f"{case}.json" + write_values_mutation(paths["values"], mutation, path, replacement) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert error_fragment in error.value.stderr + + +@pytest.mark.parametrize( + ("case", "path", "replacement"), + [ + ( + "client-roles", + ("inference", "dynamoGraph", "clientTopology", "dynamo_worker_roles"), + ["prefill", "decode", "decode", "decode"], + ), + ( + "client-gpus", + ("inference", "dynamoGraph", "clientTopology", "dynamo_gpus_per_worker"), + 2, + ), + ( + "worker-replicas", + ( + "inference", + "dynamoGraph", + "resource", + "spec", + "services", + "VllmDecodeWorker", + "replicas", + ), + 3, + ), + ( + "worker-gpus", + ( + "inference", + "dynamoGraph", + "resource", + "spec", + "services", + "VllmPrefillWorker", + "resources", + "limits", + "gpu", + ), + "2", + ), + ], +) +def test_dgd_chart_rejects_topology_mutations( + case: str, + path: tuple[str, ...], + replacement: object, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + mutation = tmp_path / f"{case}.json" + write_values_mutation(paths["values"], mutation, path, replacement) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert "topology binding" in error.value.stderr + + +def test_dgd_chart_rejects_release_name_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), release_name="other-release") + + assert "must match embedded DynamoGraphDeployment metadata.name" in error.value.stderr + + +def test_dgd_chart_rejects_namespace_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), "--set", "namespace=other-namespace") + + assert "must match embedded DynamoGraphDeployment metadata.namespace" in error.value.stderr + + +def test_dgd_chart_rejects_consistently_rehashed_image_digest_drift(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + drifted_image = values["image"]["reference"].rsplit("@", 1)[0] + f"@sha256:{'4' * 64}" + values["image"]["reference"] = drifted_image + workload["image"]["reference"] = drifted_image + for service in graph["resource"]["spec"]["services"].values(): + service["extraPodSpec"]["mainContainer"]["image"] = drifted_image + rewrite_valid_integrity(values, workload) + mutation = tmp_path / "rehashed-image-drift.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert "image-digest annotation" in error.value.stderr + + +@pytest.mark.parametrize( + ("annotation", "label"), + [ + ("prime-rl.nvidia.com/prime-sha", "Prime"), + ("prime-rl.nvidia.com/dynamo-sha", "Dynamo"), + ], +) +def test_dgd_chart_rejects_rehashed_source_sha_not_present_in_image_tag( + annotation: str, + label: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + graph["resource"]["metadata"]["annotations"][annotation] = "4" * 40 + graph["engineConfig"]["annotations"][annotation] = "4" * 40 + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"rehashed-{label.lower()}-sha.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(mutation)) + + assert f"{label} SHA annotation" in error.value.stderr + + +def test_dgd_chart_rejects_release_namespace_mismatch(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "-f", + str(paths["values"]), + release_namespace="other-namespace", + ) + + assert "Helm Release.Namespace" in error.value.stderr + + +@pytest.mark.parametrize( + ("override_flag", "enabled_value"), + [ + ("--set", "false"), + ("--set-string", "false"), + ("--set-string", "true"), + ], +) +def test_dgd_chart_requires_boolean_true_inference_with_schema_skipped( + override_flag: str, + enabled_value: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "--skip-schema-validation", + "-f", + str(paths["values"]), + override_flag, + f"inference.enabled={enabled_value}", + ) + + assert "inference.enabled must be boolean true in dynamoGraph mode" in error.value.stderr + + +def test_dgd_chart_cannot_switch_mode_and_skip_integrity_validation(tmp_path: Path): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "--skip-schema-validation", + "-f", + str(paths["values"]), + "--set", + "inference.mode=statefulset", + ) + + assert "generated DynamoGraph contract requires dynamoGraph mode" in error.value.stderr + + +@pytest.mark.parametrize("skip_schema", [False, True]) +def test_dgd_chart_cannot_delete_workload_sentinel_and_switch_mode( + skip_schema: bool, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + values["inference"]["mode"] = "statefulset" + del values["inference"]["dynamoGraph"]["workloadBinding"] + mutation = tmp_path / "deleted-sentinel-mode-switch.json" + mutation.write_text(json.dumps(values)) + args = ["-f", str(mutation)] + if skip_schema: + args.insert(0, "--skip-schema-validation") + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template(*args) + + expected = "DynamoGraph contract requires dynamoGraph mode" if skip_schema else "maxProperties" + assert expected in error.value.stderr + + +@pytest.mark.parametrize( + ("external_controller", "overrides", "error_fragment"), + [ + (False, ("--set", "orchestrator.enabled=false"), "orchestrator.enabled must match"), + (False, ("--set", "trainer.enabled=false"), "trainer.enabled must match"), + (False, ("--set", "trainer.gpu.enabled=false"), "trainer GPU configuration must match"), + (False, ("--set", "trainer.gpu.count=2"), "trainer GPU configuration must match"), + (False, ("--set", "orchestrator.replicas=0"), "orchestrator execution must match"), + (False, ("--set", "trainer.replicas=0"), "trainer execution must match"), + (False, ("--set", "orchestrator.autoStart=false"), "orchestrator execution must match"), + (False, ("--set", "trainer.autoStart=false"), "trainer execution must match"), + (False, ("--set-string", "orchestrator.command=sleep infinity"), "orchestrator execution must match"), + (False, ("--set-string", "trainer.command=sleep infinity"), "trainer execution must match"), + ( + False, + ("--set", r"orchestrator.resources.requests.nvidia\.com/gpu=1"), + "orchestrator resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"orchestrator.resources.limits.nvidia\.com/mig-1g\.23gb=1"), + "orchestrator resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"orchestrator.resources.requests.nvidia\.com/gpu\.shared=1"), + "orchestrator resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"trainer.resources.requests.nvidia\.com/gpu=2"), + "trainer resources cannot set NVIDIA extended resource", + ), + ( + False, + ("--set", r"trainer.resources.limits.nvidia\.com/mig-1g\.23gb=1"), + "trainer resources cannot set NVIDIA extended resource", + ), + (True, ("--set", "orchestrator.enabled=true"), "orchestrator.enabled must match"), + (True, ("--set", "trainer.enabled=true"), "trainer.enabled must match"), + (True, ("--set", "trainer.gpu.enabled=true"), "trainer GPU configuration must match"), + (True, ("--set", "storage.enabled=true"), "storage configuration must match"), + ( + True, + ("--set", "inference.dynamoGraph.controllerMode=chartManaged"), + "controllerMode must match", + ), + ], +) +def test_dgd_chart_rejects_workload_contract_overlays_with_schema_skipped( + external_controller: bool, + overrides: tuple[str, ...], + error_fragment: str, + tmp_path: Path, +): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, external_controller=external_controller), + ) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template( + "--skip-schema-validation", + "-f", + str(paths["values"]), + *overrides, + ) + + root = overrides[1].split(".", 1)[0] + complete_contract_error = ( + f"{root} configuration must match the workload binding" + if root in {"orchestrator", "storage", "trainer"} + else "" + ) + assert error_fragment in error.value.stderr or complete_contract_error in error.value.stderr + + +@pytest.mark.parametrize("external_controller", [False, True]) +def test_dgd_chart_rejects_rehashed_workload_mode_contradictions( + external_controller: bool, + tmp_path: Path, +): + paths = write_dgd_artifacts( + inference_config(), + render_options(tmp_path, external_controller=external_controller), + ) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + + if external_controller: + workload["trainer"]["enabled"] = True + workload["trainer"]["gpu"] = {"enabled": True, "count": 1} + values["trainer"] = { + **values["trainer"], + "enabled": True, + "gpu": {"enabled": True, "count": 1}, + } + error_fragment = "external mode forbids chart-managed controller workloads" + else: + workload["orchestrator"]["enabled"] = False + values["orchestrator"]["enabled"] = False + error_fragment = "chartManaged mode requires orchestrator, trainer" + + rewrite_valid_integrity(values, workload) + mutation = tmp_path / "contradictory-workload.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("--skip-schema-validation", "-f", str(mutation)) + + assert error_fragment in error.value.stderr + + +@pytest.mark.parametrize( + ("component", "field", "replacement"), + [ + ("orchestrator", "replicas", 0), + ("trainer", "replicas", 0), + ("orchestrator", "autoStart", False), + ("trainer", "autoStart", False), + ("orchestrator", "command", "sleep infinity"), + ("trainer", "command", "sleep infinity"), + ], +) +def test_dgd_chart_rejects_rehashed_non_runnable_controller_execution( + component: str, + field: str, + replacement: object, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + values = json.loads(paths["values"].read_text()) + graph = values["inference"]["dynamoGraph"] + workload = json.loads(graph["workloadBinding"]["canonical"]) + workload[component][field] = replacement + values[component][field] = replacement + rewrite_valid_integrity(values, workload) + mutation = tmp_path / f"rehashed-{component}-{field}.json" + mutation.write_text(json.dumps(values)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("--skip-schema-validation", "-f", str(mutation)) + + assert f"chartManaged {component} execution requires" in error.value.stderr + + +@pytest.mark.parametrize( + ("component", "name"), + [ + ("orchestrator", "DYN_RL_TOPOLOGY"), + ("orchestrator", "HF_TOKEN"), + ("trainer", "HF_HOME"), + ], +) +def test_dgd_chart_rejects_raw_env_that_overrides_typed_contract( + component: str, + name: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + overlay = tmp_path / f"{component}-{name}.json" + overlay.write_text(json.dumps({component: {"env": [{"name": name, "value": "override"}]}})) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), "-f", str(overlay)) + + assert f"{component} configuration must match the workload binding" in error.value.stderr + + +@pytest.mark.parametrize( + ("overrides", "contract"), + [ + (("--set", "image.pullPolicy=Never"), "image"), + (("--set-json", "image.pullSecrets=[]"), "image"), + (("--set", "storage.storageClassName=other"), "storage"), + (("--set", "storage.size=1Gi"), "storage"), + (("--set", "modelCache.enabled=false"), "modelCache"), + (("--set", "huggingFace.tokenSecretName=other"), "huggingFace"), + (("--set", "config.example=other"), "config"), + (("--set", "config.secrets.enabled=true"), "config"), + (("--set", "orchestrator.resources.requests.memory=1Mi"), "orchestrator"), + (("--set", "orchestrator.service.port=9000"), "orchestrator"), + (("--set-json", 'orchestrator.env=[{"name":"PYTHONPATH","value":"/data/other"}]'), "orchestrator"), + (("--set", "trainer.resources.requests.memory=1Mi"), "trainer"), + (("--set", "trainer.service.ncclPort=9001"), "trainer"), + (("--set", "trainer.probes.enabled=true"), "trainer"), + (("--set", "trainer.pytorchCudaAllocConf=max_split_size_mb:64"), "trainer"), + ], +) +def test_dgd_chart_rejects_complete_controller_contract_drift( + overrides: tuple[str, str], + contract: str, + tmp_path: Path, +): + paths = write_dgd_artifacts(inference_config(), render_options(tmp_path)) + + with pytest.raises(subprocess.CalledProcessError) as error: + helm_template("-f", str(paths["values"]), *overrides) + + assert f"{contract} configuration must match the workload binding" in error.value.stderr diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index 7e7a427112..62b6de4f87 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -1,5 +1,5 @@ import asyncio -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pydantic import pytest @@ -8,10 +8,13 @@ from verifiers.v1.types import AssistantMessage, ToolMessage, UserMessage from prime_rl.configs.algorithm import AlgoConfig, FrozenModelConfig -from prime_rl.orchestrator.algo import EchoAlgorithm, stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.algo import EchoAlgorithm, OPDAlgorithm, OPSDAlgorithm, stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.policy_gate import MutablePolicyGate from prime_rl.orchestrator.trajectories import trace_to_samples -from prime_rl.orchestrator.types import Rollout +from prime_rl.orchestrator.types import Policy, Rollout from prime_rl.transport.types import TrainingSample +from prime_rl.utils.client import DynamoInferencePool, StaticInferencePool +from prime_rl.utils.elastic import ElasticInferencePool FROZEN = {"name": "org/ref-model", "base_url": ["http://ref:8001/v1"]} @@ -78,6 +81,91 @@ def test_opd_teacher_must_be_a_frozen_endpoint(): _build(type="opd", teacher="policy") +@pytest.mark.asyncio +@pytest.mark.parametrize("pool_type", [StaticInferencePool, DynamoInferencePool]) +async def test_opd_accepts_fixed_teacher_pools(pool_type): + pool = object.__new__(pool_type) + algo = OPDAlgorithm(_build(type="opd", teacher=FROZEN), MagicMock()) + algo.connect = AsyncMock(return_value=pool) + + await algo.setup() + + assert algo.teacher_pool is pool + + +@pytest.mark.asyncio +async def test_opd_rejects_elastic_teacher_pool(): + pool = object.__new__(ElasticInferencePool) + algo = OPDAlgorithm(_build(type="opd", teacher=FROZEN), MagicMock()) + algo.connect = AsyncMock(return_value=pool) + + with pytest.raises(TypeError, match="fixed endpoint"): + await algo.setup() + + +@pytest.mark.asyncio +async def test_opsd_rejects_late_live_policy_score_as_rollout_error(): + policy = Policy(version=0, model_name="policy") + gate = MutablePolicyGate(policy, enabled=True) + pool = MagicMock(model_name="policy") + pool.score = AsyncMock(return_value=[0.0, -0.1]) + algo = OPSDAlgorithm(_build(type="opsd"), pool, policy_gate=gate) + algo.renderer = MagicMock() + algo.renderer.render_ids.return_value = [999] + rollout = _make_rollout([_make_sample()]) + rollout.info["demonstration"] = "expert answer" + rollout.policy_version = 0 + + await gate.begin_update(step=1) + await algo.finalize_rollout(rollout) + + assert rollout.has_error + assert rollout.error.type == "PolicyRequestRejected" + pool.score.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_opsd_settles_every_score_sibling_before_releasing_policy_gate(): + policy = Policy(version=0, model_name="policy") + gate = MutablePolicyGate(policy, enabled=True) + sibling_started = asyncio.Event() + release_sibling = asyncio.Event() + calls = 0 + + async def score(_token_ids: list[int]) -> list[float]: + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("score failed") + sibling_started.set() + await release_sibling.wait() + return [0.0] * 7 + + pool = MagicMock(model_name="policy") + pool.score = score + algo = OPSDAlgorithm(_build(type="opsd"), pool, policy_gate=gate) + algo.renderer = MagicMock() + algo.renderer.render_ids.return_value = [999] + rollout = _make_rollout([_make_sample(), _make_sample()]) + rollout.info["demonstration"] = "expert answer" + rollout.policy_version = 0 + + scoring = asyncio.create_task(algo.score_rollout(rollout)) + await sibling_started.wait() + update_token = await gate.begin_update(step=1) + idle = asyncio.create_task(gate.wait_idle()) + await asyncio.sleep(0) + + assert not scoring.done() + assert not idle.done() + + release_sibling.set() + with pytest.raises(RuntimeError, match="score failed"): + await scoring + await idle + await gate.finish_update(update_token) + + def test_sft_requires_teacher(): with pytest.raises(ValueError, match="needs a teacher to sample rollouts from"): _build(type="sft") diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 73768a7698..811480a293 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -1,12 +1,252 @@ import asyncio +from contextlib import suppress from types import SimpleNamespace from unittest.mock import AsyncMock, patch +import pytest from renderers import Qwen3VLRendererConfig +from prime_rl.configs.shared import ClientConfig +from prime_rl.orchestrator import component_supervision +from prime_rl.orchestrator.component_supervision import raise_if_component_failed, run_with_component_supervision from prime_rl.orchestrator.utils import setup_policy_inference_pool +def test_component_failure_is_raised_into_main_loop(): + async def run() -> None: + async def failed_watcher() -> None: + raise RuntimeError("indeterminate policy update") + + watcher = asyncio.create_task(failed_watcher(), name="watcher") + await asyncio.wait({watcher}) + with pytest.raises(RuntimeError, match="indeterminate policy update"): + raise_if_component_failed([watcher]) + + asyncio.run(run()) + + +def test_component_failure_wins_over_queued_output_and_prevents_trainer_send(): + async def run() -> None: + async def failed_watcher() -> None: + raise RuntimeError("indeterminate policy update") + + watcher = asyncio.create_task(failed_watcher(), name="watcher") + await asyncio.wait({watcher}) + operation = AsyncMock(return_value=object()) + + with pytest.raises(RuntimeError, match="indeterminate policy update"): + await run_with_component_supervision(operation, [watcher]) + + operation.assert_not_called() + + asyncio.run(run()) + + +def test_component_failure_cancels_in_progress_external_operation(): + async def run() -> None: + operation_started = asyncio.Event() + operation_cancelled = asyncio.Event() + fail_component = asyncio.Event() + + async def operation() -> None: + operation_started.set() + try: + await asyncio.Future() + finally: + operation_cancelled.set() + + async def watcher() -> None: + await fail_component.wait() + raise RuntimeError("watcher failed during send") + + watcher_task = asyncio.create_task(watcher(), name="watcher") + supervised = asyncio.create_task( + run_with_component_supervision( + operation, + [watcher_task], + timeout=1.0, + timeout_description="training batch send", + ) + ) + await operation_started.wait() + fail_component.set() + + with pytest.raises(RuntimeError, match="watcher failed during send"): + await supervised + assert operation_cancelled.is_set() + + asyncio.run(run()) + + +def test_supervised_operation_timeout_cancels_stuck_send_and_raises(): + async def run() -> None: + operation_cancelled = asyncio.Event() + + async def stuck_send() -> None: + try: + await asyncio.Future() + finally: + operation_cancelled.set() + + with pytest.raises(TimeoutError, match=r"training batch send timed out after 0\.01 seconds"): + await run_with_component_supervision( + stuck_send, + [], + timeout=0.01, + timeout_description="training batch send", + ) + + assert operation_cancelled.is_set() + + asyncio.run(run()) + + +def test_timeout_cleanup_is_bounded_when_operation_suppresses_cancellation(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + operation_started = asyncio.Event() + cancellation_seen = asyncio.Event() + release_orphan = asyncio.Event() + + async def cancellation_suppressing_send() -> None: + operation_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release_orphan.wait() + raise RuntimeError("late orphan failure") + + monkeypatch.setattr( + component_supervision, + "SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS", + 0.01, + raising=False, + ) + supervised = asyncio.create_task( + run_with_component_supervision( + cancellation_suppressing_send, + [], + timeout=0.01, + timeout_description="training batch send", + ) + ) + try: + with pytest.raises(TimeoutError, match=r"training batch send timed out after 0\.01 seconds"): + await asyncio.wait_for(asyncio.shield(supervised), timeout=0.1) + assert cancellation_seen.is_set() + assert len(component_supervision._ORPHANED_OPERATIONS) == 1 + finally: + release_orphan.set() + if not supervised.done(): + with suppress(BaseException): + await supervised + + for _ in range(10): + if not component_supervision._ORPHANED_OPERATIONS: + break + await asyncio.sleep(0) + assert not component_supervision._ORPHANED_OPERATIONS + + asyncio.run(run()) + + +def test_caller_cancellation_during_timeout_cleanup_is_not_swallowed(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + cancellation_seen = asyncio.Event() + release_orphan = asyncio.Event() + + async def cancellation_suppressing_poll() -> None: + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release_orphan.wait() + + monkeypatch.setattr(component_supervision, "SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS", 0.01) + supervised = asyncio.create_task( + run_with_component_supervision(cancellation_suppressing_poll, [], timeout=0.01) + ) + await cancellation_seen.wait() + supervised.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(supervised), timeout=0.1) + assert len(component_supervision._ORPHANED_OPERATIONS) == 1 + + release_orphan.set() + for _ in range(10): + if not component_supervision._ORPHANED_OPERATIONS: + break + await asyncio.sleep(0) + assert not component_supervision._ORPHANED_OPERATIONS + + asyncio.run(run()) + + +def test_component_failure_precedes_caller_cancellation_during_cleanup(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + operation_started = asyncio.Event() + cancellation_seen = asyncio.Event() + release_orphan = asyncio.Event() + fail_watcher = asyncio.Event() + + async def cancellation_suppressing_send() -> None: + operation_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release_orphan.wait() + + async def failed_watcher() -> None: + await fail_watcher.wait() + raise RuntimeError("component failed before cleanup cancellation") + + monkeypatch.setattr(component_supervision, "SUPERVISED_OPERATION_CANCEL_GRACE_SECONDS", 0.01) + watcher = asyncio.create_task(failed_watcher(), name="watcher") + supervised = asyncio.create_task(run_with_component_supervision(cancellation_suppressing_send, [watcher])) + await operation_started.wait() + fail_watcher.set() + await cancellation_seen.wait() + supervised.cancel() + + with pytest.raises(RuntimeError, match="component failed before cleanup cancellation"): + await asyncio.wait_for(asyncio.shield(supervised), timeout=0.1) + assert len(component_supervision._ORPHANED_OPERATIONS) == 1 + + release_orphan.set() + for _ in range(10): + if not component_supervision._ORPHANED_OPERATIONS: + break + await asyncio.sleep(0) + assert not component_supervision._ORPHANED_OPERATIONS + + asyncio.run(run()) + + +def test_component_failure_wins_when_operation_completes_in_same_loop_turn(): + async def run() -> None: + release = asyncio.Event() + + async def operation() -> str: + await release.wait() + return "sent" + + async def watcher() -> None: + await release.wait() + raise RuntimeError("simultaneous watcher failure") + + watcher_task = asyncio.create_task(watcher(), name="watcher") + supervised = asyncio.create_task(run_with_component_supervision(operation, [watcher_task])) + await asyncio.sleep(0) + release.set() + + with pytest.raises(RuntimeError, match="simultaneous watcher failure"): + await supervised + + asyncio.run(run()) + + def test_setup_policy_inference_pool_uses_renderer_when_enabled(): async def run() -> None: tokenizer = object() @@ -96,3 +336,33 @@ async def run() -> None: ) asyncio.run(run()) + + +def test_setup_policy_inference_pool_uses_environment_resolved_client(monkeypatch: pytest.MonkeyPatch): + async def run() -> None: + monkeypatch.setenv( + "DYN_RL_TOPOLOGY", + """{"schema_version":1,"admin_api":"dynamo","base_url":["http://frontend:8000/v1"],"rl_base_url":["http://frontend-rl:8001"],"dynamo_worker_roles":["agg"],"dynamo_gpus_per_worker":1}""", + ) + config = SimpleNamespace( + model=SimpleNamespace(client=ClientConfig(), name="policy-model"), + renderer=Qwen3VLRendererConfig(), + pool_size=None, + any_policy_sourced=True, + ) + + with ( + patch("renderers.base.create_renderer", return_value=object()), + patch( + "prime_rl.orchestrator.utils.setup_inference_pool", + new=AsyncMock(return_value=object()), + ) as setup_pool, + ): + await setup_policy_inference_pool(config=config, tokenizer=object()) + + resolved = setup_pool.await_args.args[0] + assert resolved.admin_api == "dynamo" + assert resolved.base_url == ["http://frontend:8000/v1"] + assert config.model.client.admin_api == "vllm" + + asyncio.run(run()) diff --git a/tests/unit/orchestrator/test_policy_gate_cancellation.py b/tests/unit/orchestrator/test_policy_gate_cancellation.py new file mode 100644 index 0000000000..0738a806ad --- /dev/null +++ b/tests/unit/orchestrator/test_policy_gate_cancellation.py @@ -0,0 +1,68 @@ +import asyncio + +import pytest + +from prime_rl.orchestrator.dispatcher import RolloutDispatcher +from prime_rl.orchestrator.types import Policy + + +def _dispatcher() -> RolloutDispatcher: + pool = type( + "Pool", + (), + { + "model_name": "policy", + "train_clients": [], + "admin_clients": [], + }, + )() + return RolloutDispatcher( + train_envs=object(), + eval_envs=None, + train_source=object(), + eval_source=None, + policy_pool=pool, + policy=Policy(version=0, model_name="policy"), + max_inflight_rollouts=1, + tasks_per_minute=None, + max_off_policy_steps=0, + enforce_policy_update_barrier=True, + ) + + +@pytest.mark.asyncio +async def test_repeated_cancellation_cannot_interrupt_partial_entry_rollback(): + dispatcher = _dispatcher() + settle_started = asyncio.Event() + release_settle = asyncio.Event() + rollback_started = asyncio.Event() + + async def settle_policy_requests(*_args) -> None: + settle_started.set() + await release_settle.wait() + + dispatcher._settle_policy_requests = settle_policy_requests # type: ignore[method-assign] + finish_update = dispatcher.policy_gate.finish_update + + async def observed_finish_update(token) -> None: + rollback_started.set() + await finish_update(token) + + dispatcher.policy_gate.finish_update = observed_finish_update # type: ignore[method-assign] + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await settle_started.wait() + + await dispatcher.policy_gate._admission_lock.acquire() + barrier.cancel() + release_settle.set() + await rollback_started.wait() + barrier.cancel() + await asyncio.sleep(0) + rollback_was_interrupted = barrier.done() + dispatcher.policy_gate._admission_lock.release() + + assert not rollback_was_interrupted + with pytest.raises(asyncio.CancelledError) as exc: + await barrier + assert not dispatcher.policy_update_pending + assert exc.value.__notes__ == ["Policy transition rollback was cancelled again but settled before propagation"] diff --git a/tests/unit/orchestrator/test_pool_identity.py b/tests/unit/orchestrator/test_pool_identity.py new file mode 100644 index 0000000000..e17c4e02df --- /dev/null +++ b/tests/unit/orchestrator/test_pool_identity.py @@ -0,0 +1,52 @@ +from types import SimpleNamespace + +import pytest + +from prime_rl.orchestrator.pool_identity import pools_may_alias + + +def _pool(model: str, request: str | None, admin: str | None): + return SimpleNamespace( + model_name=model, + train_clients=[] if request is None else [SimpleNamespace(base_url=request)], + admin_clients=[] if admin is None else [SimpleNamespace(base_url=admin)], + ) + + +@pytest.mark.parametrize( + ("left", "right", "aliases"), + [ + ( + _pool("policy", "http://frontend/v1", "http://worker:8081"), + _pool("policy", "http://frontend", "http://worker:8081/"), + True, + ), + ( + _pool("policy", "http://policy/v1", "http://policy-worker:8081"), + _pool("policy", "http://frozen/v1", "http://frozen-worker:8081"), + False, + ), + ( + _pool("policy", "http://router-a/v1", "http://shared-worker:8081"), + _pool("policy", "http://router-b/v1", "http://shared-worker:8081"), + True, + ), + ( + _pool("policy", "http://frontend/v1", "http://worker:8081"), + _pool("other-model", "http://other-frontend/v1", "http://worker:8081"), + True, + ), + ( + _pool("policy", None, None), + _pool("policy", "http://frozen/v1", "http://frozen-worker:8081"), + True, + ), + ], +) +def test_pool_aliasing_uses_model_request_and_admin_identity(left, right, aliases: bool): + assert pools_may_alias(left, right) is aliases + + +def test_pool_aliasing_accepts_object_identity(): + pool = _pool("policy", "http://frontend/v1", "http://worker:8081") + assert pools_may_alias(pool, pool) diff --git a/tests/unit/orchestrator/test_train_finalization.py b/tests/unit/orchestrator/test_train_finalization.py new file mode 100644 index 0000000000..6df8627123 --- /dev/null +++ b/tests/unit/orchestrator/test_train_finalization.py @@ -0,0 +1,150 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from prime_rl.orchestrator.train_finalization import finalize_train_batch +from prime_rl.orchestrator.types import Progress + + +class _Stat: + def mean(self) -> float: + return 1.0 + + +class _Metrics: + reward = _Stat() + num_turns = _Stat() + num_branches = _Stat() + is_truncated = _Stat() + has_error = _Stat() + + def to_wandb(self, *, prefix: str, subset: str) -> dict[str, float]: + return {f"{prefix}/{subset}/metric": 1.0} + + +class _Rollout: + is_trainable = True + off_policy_steps = 0 + num_total_tokens = 3 + num_input_tokens = 1 + num_output_tokens = 2 + group_id = "group" + reward = 1.0 + + def to_record(self) -> dict: + return {"group": self.group_id} + + def scalar_advantage(self) -> float: + return 0.5 + + +class _Rollouts(list): + metrics = _Metrics() + + @property + def effective(self): + return self + + @property + def rollouts(self): + return list(self) + + def by_env(self) -> dict: + return {} + + +class _Monitor: + def __init__(self, events: list[str]) -> None: + self.events = events + self.logged_metrics: dict[str, float] = {} + + def log(self, metrics: dict[str, float], *, step: int) -> None: + self.events.append("monitor:log") + self.logged_metrics = metrics + + def log_samples(self, _rollouts, *, step: int) -> None: + self.events.append("monitor:samples") + + def log_distributions(self, *, distributions, step: int) -> None: + self.events.append("monitor:distributions") + + +class _Host: + def __init__(self, tmp_path: Path, events: list[str]) -> None: + self.config = SimpleNamespace(output_dir=tmp_path, max_steps=None) + self.progress = Progress() + self.last_batch_at = 1.0 + self.draining = False + self.consecutive_empty_batches = 0 + self.wait_for_policy_time = 0.25 + self.monitor = _Monitor(events) + self.usage_reporter = None + self.heart = None + self.train_envs = [] + self.train_sink = SimpleNamespace( + pre_filter_seen=0, + pre_filter_dropped=0, + pre_filter_dropped_by_name={}, + reset_pre_filter_stats=lambda: events.append("reset-filters"), + ) + self.events = events + + async def _send_to_trainer(self, _batch) -> None: + self.events.append("send") + + def update_dispatch_gate(self) -> None: + assert self.progress.step == 2 + self.events.append("dispatch-gate") + + async def maybe_save_ckpt(self, step: int) -> float: + assert step == 1 + self.events.append("checkpoint") + return 0.5 + + def maybe_trigger_eval(self, step: int) -> None: + assert step == 2 + self.events.append("eval") + + +@pytest.mark.asyncio +async def test_finalize_train_batch_preserves_ship_and_reporting_order(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + events: list[str] = [] + host = _Host(tmp_path, events) + batch = SimpleNamespace(samples=[object()], rollouts=_Rollouts([_Rollout()])) + + monkeypatch.setattr( + "prime_rl.orchestrator.train_finalization.save_rollouts", + lambda *_args: events.append("save-rollouts"), + ) + monkeypatch.setattr( + "prime_rl.orchestrator.train_finalization.trim_process_memory", + lambda: events.append("trim"), + ) + monkeypatch.setattr( + "prime_rl.orchestrator.train_finalization.get_logger", + lambda: SimpleNamespace(success=lambda _message: events.append("success")), + ) + + await finalize_train_batch(host, batch) + + assert events == [ + "save-rollouts", + "send", + "dispatch-gate", + "checkpoint", + "trim", + "monitor:log", + "monitor:samples", + "monitor:distributions", + "success", + "reset-filters", + "eval", + "trim", + ] + assert host.progress.step == 2 + assert host.progress.total_tokens == 3 + assert host.progress.total_samples == 1 + assert host.progress.total_problems == 1 + assert host.monitor.logged_metrics["progress/total_tokens"] == 0 + assert host.wait_for_policy_time == 0.0 diff --git a/tests/unit/orchestrator/test_weight_update_barrier.py b/tests/unit/orchestrator/test_weight_update_barrier.py new file mode 100644 index 0000000000..637f3b988a --- /dev/null +++ b/tests/unit/orchestrator/test_weight_update_barrier.py @@ -0,0 +1,750 @@ +import asyncio +import uuid +from pathlib import Path +from types import SimpleNamespace + +import pytest +import verifiers.v1 as vf + +from prime_rl.orchestrator.dispatcher import DispatcherMode, RolloutDispatcher +from prime_rl.orchestrator.policy_gate import MutablePolicyGate, PolicyRequestRejected +from prime_rl.orchestrator.types import GroupState, InflightRollout, Policy, Rollout +from prime_rl.orchestrator.watcher import WeightWatcher +from prime_rl.utils.pathing import get_broadcast_dir, get_step_path + + +class _TrainEnvs: + def __init__(self, *, live: bool, pool: object) -> None: + self._live = live + self._pool = pool + + def get(self, _name: str): + return SimpleNamespace(sampler=SimpleNamespace(samples_from_live_policy=self._live, pool=self._pool)) + + +def _dispatcher( + *, + max_inflight: int = 1, + live_train: bool = True, + frozen_uses_policy_pool: bool = False, + frozen_aliases_policy_pool: bool = False, + enforce_policy_update_barrier: bool = True, + max_off_policy_steps: int = 8, + policy_gate: MutablePolicyGate | None = None, +) -> RolloutDispatcher: + def pool(model_name: str, request_url: str, admin_url: str): + return SimpleNamespace( + model_name=model_name, + train_clients=[SimpleNamespace(base_url=request_url, headers={})], + admin_clients=[SimpleNamespace(base_url=admin_url)], + ) + + policy_pool = pool("policy", "http://policy/v1", "http://policy-worker:8081") + if live_train or frozen_uses_policy_pool: + train_pool = policy_pool + elif frozen_aliases_policy_pool: + # A separately-constructed pool can still address the exact mutable + # serving resource. Object identity is not a topology identity. + train_pool = pool("policy", "http://policy/v1", "http://policy-worker:8081") + else: + train_pool = pool("frozen", "http://frozen/v1", "http://frozen-worker:8081") + return RolloutDispatcher( + train_envs=_TrainEnvs(live=live_train, pool=train_pool), + eval_envs=object(), + train_source=object(), + eval_source=object(), + policy_pool=policy_pool, + policy=Policy(version=0, model_name="policy"), + max_inflight_rollouts=max_inflight, + tasks_per_minute=None, + max_off_policy_steps=max_off_policy_steps, + enforce_policy_update_barrier=enforce_policy_update_barrier, + policy_gate=policy_gate, + ) + + +@pytest.mark.asyncio +async def test_policy_barrier_invalidates_slow_scheduling_before_short_commit(): + dispatcher = _dispatcher() + dispatcher.mode = DispatcherMode.PREFER_EVAL + scheduling_started = asyncio.Event() + allow_schedule_to_commit = asyncio.Event() + request_started = asyncio.Event() + + async def active_request() -> None: + request_started.set() + await asyncio.Future() + + async def schedule_one(_kind: str, *, epoch) -> bool: + scheduling_started.set() + await allow_schedule_to_commit.wait() + async with dispatcher.policy_gate.scheduling_commit(epoch) as admitted: + if not admitted: + return False + group_id = uuid.uuid4() + request = asyncio.create_task(active_request()) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[request] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits += 1 + return True + + dispatcher.try_schedule = schedule_one # type: ignore[method-assign] + + fill = asyncio.create_task(dispatcher.fill_inflight()) + await scheduling_started.wait() + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await asyncio.sleep(0) + + await barrier + assert dispatcher.policy_update_pending + + allow_schedule_to_commit.set() + await fill + assert not request_started.is_set() + assert not dispatcher.inflight + + # The transition fence remains closed until the watcher reports either + # success or failure. + await dispatcher.fill_inflight() + assert not dispatcher.inflight + + await dispatcher.on_new_version(1) + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("frozen_uses_policy_pool", "frozen_aliases_policy_pool", "cancelled_by_barrier"), + [(False, False, False), (True, False, True), (False, True, True)], +) +async def test_policy_barrier_only_preserves_frozen_requests_on_a_distinct_pool( + frozen_uses_policy_pool: bool, + frozen_aliases_policy_pool: bool, + cancelled_by_barrier: bool, +): + dispatcher = _dispatcher( + live_train=False, + frozen_uses_policy_pool=frozen_uses_policy_pool, + frozen_aliases_policy_pool=frozen_aliases_policy_pool, + ) + group_id = uuid.uuid4() + request_started = asyncio.Event() + + async def request() -> None: + request_started.set() + await asyncio.Future() + + task = asyncio.create_task(request()) + await request_started.wait() + dispatcher.groups[group_id] = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + policy_version_at_start=0, + ) + dispatcher.inflight[task] = InflightRollout( + kind="train", + env_name="train", + group_id=group_id, + policy_version=0, + rollout_count=1, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(1) + + assert task.cancelled() is cancelled_by_barrier + assert (task in dispatcher.inflight) is not cancelled_by_barrier + + await dispatcher.on_new_version(1) + await dispatcher.cancel_inflight_rollouts() + + +@pytest.mark.asyncio +async def test_failed_policy_barrier_skips_engine_mutation_and_reopens_admission(tmp_path: Path): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + cleanup_attempted = asyncio.Event() + + async def request_with_failed_cleanup() -> None: + try: + await asyncio.Future() + except asyncio.CancelledError: + cleanup_attempted.set() + raise RuntimeError("connector cleanup failed") + + request = asyncio.create_task(request_with_failed_cleanup()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[request] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _Inference: + def __init__(self) -> None: + self.update_calls = 0 + + async def update_weights(self, *_args, **_kwargs) -> None: + self.update_calls += 1 + + inference = _Inference() + policy = dispatcher.policy + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=policy, + inference=inference, + observers=[dispatcher], + lora_name=None, + ) + + with pytest.raises(RuntimeError, match="connector cleanup failed"): + await watcher.apply_policy_update(1) + + assert cleanup_attempted.is_set() + assert inference.update_calls == 0 + assert watcher.ckpt_step == 0 + assert policy.version == 0 + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_accepts_request_error_that_was_already_settled(): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + + async def failed_request() -> None: + raise RuntimeError("ordinary rollout failure") + + request = asyncio.create_task(failed_request()) + await asyncio.wait({request}) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[request] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(1) + + assert dispatcher.policy_update_pending + assert not dispatcher.inflight + await dispatcher.on_new_version(1) + + +@pytest.mark.asyncio +async def test_policy_barrier_settles_cleanup_before_propagating_cancellation(): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + cleanup_finished = asyncio.Event() + + async def request() -> None: + try: + await asyncio.Future() + finally: + cleanup_started.set() + await release_cleanup.wait() + cleanup_finished.set() + + task = asyncio.create_task(request()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.inflight[task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await cleanup_started.wait() + barrier.cancel() + await asyncio.sleep(0) + barrier.cancel() + await asyncio.sleep(0) + + assert not barrier.done() + assert not cleanup_finished.is_set() + + release_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await barrier + + assert cleanup_finished.is_set() + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_shields_marker_enqueue_and_commits_accounting_after_put(): + dispatcher = _dispatcher() + group_id = uuid.uuid4() + request_started = asyncio.Event() + marker_put_started = asyncio.Event() + + class _ObservedQueue(asyncio.Queue): + async def put(self, item) -> None: + marker_put_started.set() + await super().put(item) + + dispatcher.out_q = _ObservedQueue(maxsize=1) + dispatcher.out_q.put_nowait(object()) + + async def request() -> None: + request_started.set() + await asyncio.Future() + + request_task = asyncio.create_task(request()) + await request_started.wait() + group = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + ) + dispatcher.groups[group_id] = group + dispatcher.inflight[request_task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + ) + dispatcher.inflight_permits = 1 + + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await marker_put_started.wait() + assert group.emitted == 0 + + barrier.cancel() + await asyncio.sleep(0) + barrier.cancel() + await asyncio.sleep(0) + assert not barrier.done() + assert group.emitted == 0 + + dispatcher.out_q.get_nowait() + with pytest.raises(asyncio.CancelledError): + await barrier + + marker = dispatcher.out_q.get_nowait() + assert marker.error.type == "Cancelled" + assert group.emitted == 1 + assert not dispatcher.groups + assert not dispatcher.inflight + assert dispatcher.inflight_permits == 0 + + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_waits_for_completed_handler_before_computing_markers(): + dispatcher = _dispatcher(max_inflight=2) + dispatcher.out_q = asyncio.Queue(maxsize=1) + group_id = uuid.uuid4() + + async def completed_group() -> list[Rollout]: + return [ + Rollout( + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt=None)), + errors=[vf.Error(type="Existing", message="result")], + stop_condition="error", + ) + for _ in range(2) + ] + + task = asyncio.create_task(completed_group()) + await task + group = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=2, + eval_step=1, + policy_version_at_start=0, + uses_mutable_policy=True, + ) + dispatcher.groups[group_id] = group + dispatcher.inflight[task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=2, + eval_step=1, + uses_mutable_policy=True, + ) + dispatcher.inflight_permits = 2 + + handler = asyncio.create_task(dispatcher.handle_completed_rollout(task)) + while dispatcher.out_q.qsize() != 1: + await asyncio.sleep(0) + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await asyncio.sleep(0) + + assert not barrier.done() + + first = dispatcher.out_q.get_nowait() + await handler + await barrier + second = dispatcher.out_q.get_nowait() + + assert [first.error.type, second.error.type] == ["Existing", "Existing"] + assert group.emitted == 2 + await dispatcher.on_new_version(1) + + +@pytest.mark.asyncio +async def test_policy_barrier_marker_enqueue_aborts_promptly_on_dispatcher_stop(): + dispatcher = _dispatcher() + marker_put_started = asyncio.Event() + + class _ObservedQueue(asyncio.Queue): + async def put(self, item) -> None: + marker_put_started.set() + await super().put(item) + + dispatcher.out_q = _ObservedQueue(maxsize=1) + dispatcher.out_q.put_nowait(object()) + group_id = uuid.uuid4() + + async def request() -> None: + await asyncio.Future() + + task = asyncio.create_task(request()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="eval", + env_name="eval", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + eval_step=1, + policy_version_at_start=0, + uses_mutable_policy=True, + ) + dispatcher.inflight[task] = InflightRollout( + kind="eval", + env_name="eval", + group_id=group_id, + policy_version=0, + rollout_count=1, + eval_step=1, + uses_mutable_policy=True, + ) + dispatcher.inflight_permits = 1 + + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await marker_put_started.wait() + await dispatcher.stop() + + with pytest.raises(RuntimeError, match="stopped while emitting policy barrier"): + await asyncio.wait_for(barrier, timeout=0.5) + assert not dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_policy_barrier_drains_active_policy_calls_and_rejects_late_or_stale_calls(): + policy = Policy(version=0, model_name="policy") + policy_gate = MutablePolicyGate(policy, enabled=True) + dispatcher = _dispatcher(policy_gate=policy_gate) + # The helper constructs its own Policy, so make both components share the + # exact version object as production does. + dispatcher.policy = policy + + call_started = asyncio.Event() + release_call = asyncio.Event() + + async def active_policy_call() -> None: + async with policy_gate.request(expected_version=0): + call_started.set() + await release_call.wait() + + score = asyncio.create_task(active_policy_call()) + await call_started.wait() + barrier = asyncio.create_task(dispatcher.on_version_pending(1)) + await asyncio.sleep(0) + assert dispatcher.policy_update_pending + assert not barrier.done() + + with pytest.raises(PolicyRequestRejected, match="update is pending"): + async with policy_gate.request(expected_version=0): + pass + + release_call.set() + await score + await barrier + + policy.version = 1 + await dispatcher.on_new_version(1) + with pytest.raises(PolicyRequestRejected, match="expected policy version 0.*current version is 1"): + async with policy_gate.request(expected_version=0): + pass + + +@pytest.mark.asyncio +async def test_non_dynamo_dispatcher_retains_configured_off_policy_window(): + dispatcher = _dispatcher(enforce_policy_update_barrier=False, max_off_policy_steps=1) + group_id = uuid.uuid4() + request_started = asyncio.Event() + + async def request() -> None: + request_started.set() + await asyncio.Future() + + task = asyncio.create_task(request()) + await request_started.wait() + dispatcher.groups[group_id] = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + policy_version_at_start=0, + ) + dispatcher.inflight[task] = InflightRollout( + kind="train", + env_name="train", + group_id=group_id, + policy_version=0, + rollout_count=1, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(1) + assert dispatcher.inflight[task].off_policy_steps == 1 + assert not dispatcher.policy_update_pending + + await dispatcher.on_version_pending(2) + assert task.cancelled() + assert not dispatcher.inflight + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["collective update failed", "resume_generation failed"]) +async def test_indeterminate_engine_failure_keeps_admission_fail_closed_without_advancing_version( + tmp_path: Path, + failure: str, +): + dispatcher = _dispatcher() + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + raise RuntimeError(failure) + + policy = dispatcher.policy + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=policy, + inference=_Inference(), + observers=[dispatcher], + lora_name=None, + ) + + with pytest.raises(RuntimeError, match=failure) as exc: + await watcher.apply_policy_update(1) + + assert watcher.ckpt_step == 0 + assert policy.version == 0 + assert dispatcher.policy_update_pending + assert exc.value.__notes__ == [ + "Policy update 1 may have mutated inference workers; mutable-policy admission remains fail-closed" + ] + + second_weight_path = get_step_path(get_broadcast_dir(tmp_path), 2) + second_weight_path.mkdir(parents=True) + (second_weight_path / "STABLE").touch() + with pytest.raises(RuntimeError, match="already pending"): + await watcher.apply_policy_update(2) + assert dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_success_updates_lead_gate_before_reopening_transition_fence(tmp_path: Path): + events: list[str] = [] + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _Observer: + def __init__(self, name: str) -> None: + self.name = name + + async def on_version_pending(self, _step: int) -> None: + events.append(f"{self.name}:pending") + + async def on_new_version(self, _step: int) -> None: + events.append(f"{self.name}:new") + + async def on_version_update_failed(self, _step: int, _error: BaseException) -> None: + events.append(f"{self.name}:failed") + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + events.append("engine:update") + + policy = Policy(version=0, model_name="policy") + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=policy, + inference=_Inference(), + observers=[_Observer("fence"), _Observer("lead-gate")], + lora_name=None, + ) + + await watcher.apply_policy_update(1) + + assert events == [ + "fence:pending", + "lead-gate:pending", + "engine:update", + "lead-gate:new", + "fence:new", + ] + assert watcher.ckpt_step == 1 + assert policy.version == 1 + + +@pytest.mark.asyncio +async def test_success_callback_cancellation_propagates_and_keeps_transition_fence_closed(tmp_path: Path): + dispatcher = _dispatcher() + callback_started = asyncio.Event() + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _BlockingLeadGate: + async def on_version_pending(self, _step: int) -> None: + return + + async def on_new_version(self, _step: int) -> None: + callback_started.set() + await asyncio.Future() + + async def on_version_update_failed(self, _step: int, _error: BaseException) -> None: + return + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + return + + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=dispatcher.policy, + inference=_Inference(), + observers=[dispatcher, _BlockingLeadGate()], + lora_name=None, + ) + + update = asyncio.create_task(watcher.apply_policy_update(1)) + await callback_started.wait() + update.cancel() + with pytest.raises(asyncio.CancelledError): + await update + + assert watcher.ckpt_step == 1 + assert dispatcher.policy.version == 1 + assert dispatcher.policy_update_pending + + +@pytest.mark.asyncio +async def test_success_callback_error_propagates_and_keeps_transition_fence_closed(tmp_path: Path): + dispatcher = _dispatcher() + weight_path = get_step_path(get_broadcast_dir(tmp_path), 1) + weight_path.mkdir(parents=True) + (weight_path / "STABLE").touch() + + class _BrokenLeadGate: + async def on_version_pending(self, _step: int) -> None: + return + + async def on_new_version(self, _step: int) -> None: + raise RuntimeError("lead gate callback failed") + + async def on_version_update_failed(self, _step: int, _error: BaseException) -> None: + return + + class _Inference: + async def update_weights(self, *_args, **_kwargs) -> None: + return + + watcher = WeightWatcher( + SimpleNamespace(output_dir=tmp_path), + policy=dispatcher.policy, + inference=_Inference(), + observers=[dispatcher, _BrokenLeadGate()], + lora_name=None, + ) + + with pytest.raises(RuntimeError, match="lead gate callback failed"): + await watcher.apply_policy_update(1) + + assert watcher.ckpt_step == 1 + assert dispatcher.policy.version == 1 + assert dispatcher.policy_update_pending diff --git a/tests/unit/orchestrator/test_weight_update_barrier_mutability.py b/tests/unit/orchestrator/test_weight_update_barrier_mutability.py new file mode 100644 index 0000000000..7f0988a84c --- /dev/null +++ b/tests/unit/orchestrator/test_weight_update_barrier_mutability.py @@ -0,0 +1,211 @@ +import asyncio +import uuid +from types import SimpleNamespace + +import pytest + +from prime_rl.orchestrator.dispatcher import RolloutDispatcher +from prime_rl.orchestrator.types import GroupState, InflightRollout, Policy + + +def _pool(model: str, request: str, admin: str): + return SimpleNamespace( + model_name=model, + train_clients=[SimpleNamespace(base_url=request, headers={})], + admin_clients=[SimpleNamespace(base_url=admin)], + ) + + +class _EnvCollection: + def __init__(self, pool, *, run_rollout=None) -> None: + self.env = SimpleNamespace( + sampler=SimpleNamespace(samples_from_live_policy=False, pool=pool), + requires_group_scoring=False, + config=SimpleNamespace(group_size=1), + run_rollout=run_rollout, + ) + + def get(self, _name: str): + return self.env + + +def _dispatcher(*, enforce_barrier: bool, max_off_policy_steps: int = 0, run_rollout=None): + policy_pool = _pool("policy", "http://policy/v1", "http://policy-worker:8081") + distinct_pool = _pool("frozen", "http://frozen/v1", "http://frozen-worker:8081") + train_envs = _EnvCollection(distinct_pool, run_rollout=run_rollout) + dispatcher = RolloutDispatcher( + train_envs=train_envs, + eval_envs=None, + train_source=object(), + eval_source=None, + policy_pool=policy_pool, + policy=Policy(version=4, model_name="policy"), + max_inflight_rollouts=1, + tasks_per_minute=None, + max_off_policy_steps=max_off_policy_steps, + enforce_policy_update_barrier=enforce_barrier, + ) + return dispatcher + + +@pytest.mark.asyncio +async def test_group_mutability_is_monotonic_for_cache_salt_after_pool_identity_churn(): + called: dict[str, object] = {} + + async def run_rollout(**kwargs): + called.update(kwargs) + await asyncio.Future() + + dispatcher = _dispatcher(enforce_barrier=True, run_rollout=run_rollout) + group_id = uuid.uuid4() + formerly_aliased_client = SimpleNamespace(base_url="http://policy/v1", headers={}) + group = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=1, + target_rollouts=1, + pinned_client=formerly_aliased_client, + policy_version_at_start=4, + uses_mutable_policy=False, + ) + dispatcher.groups[group_id] = group + epoch = await dispatcher.policy_gate.scheduling_epoch() + assert epoch is not None + + assert await dispatcher.schedule_group_rollout(group_id, group, epoch=epoch) + await asyncio.sleep(0) + + assert called["cache_salt"] == "4" + assert next(iter(dispatcher.inflight.values())).uses_mutable_policy + await dispatcher.cancel_inflight_rollouts() + + +@pytest.mark.asyncio +async def test_group_mutability_is_monotonic_for_non_dynamo_off_policy_aging(): + dispatcher = _dispatcher(enforce_barrier=False, max_off_policy_steps=0) + group_id = uuid.uuid4() + + async def request() -> None: + await asyncio.Future() + + task = asyncio.create_task(request()) + await asyncio.sleep(0) + dispatcher.groups[group_id] = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=0, + target_rollouts=1, + policy_version_at_start=4, + uses_mutable_policy=True, + ) + dispatcher.inflight[task] = InflightRollout( + kind="train", + env_name="train", + group_id=group_id, + policy_version=4, + rollout_count=1, + uses_mutable_policy=True, + ) + dispatcher.inflight_permits = 1 + + await dispatcher.on_version_pending(5) + + assert task.cancelled() + assert not dispatcher.inflight + + +@pytest.mark.asyncio +async def test_policy_update_does_not_wait_for_slow_client_selection(): + selection_started = asyncio.Event() + release_selection = asyncio.Event() + + class _BlockingPool: + model_name = "policy" + train_clients = [SimpleNamespace(base_url="http://policy/v1", headers={})] + admin_clients = [SimpleNamespace(base_url="http://policy-worker:8081")] + + async def select_train_client(self, _load): + selection_started.set() + await release_selection.wait() + return self.train_clients[0] + + dispatcher = _dispatcher(enforce_barrier=True, run_rollout=None) + dispatcher.train_envs.env.sampler.pool = _BlockingPool() + group_id = uuid.uuid4() + group = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=1, + target_rollouts=1, + policy_version_at_start=4, + uses_mutable_policy=True, + ) + dispatcher.groups[group_id] = group + epoch = await dispatcher.policy_gate.scheduling_epoch() + assert epoch is not None + + scheduling = asyncio.create_task(dispatcher.schedule_group_rollout(group_id, group, epoch=epoch)) + await selection_started.wait() + await asyncio.wait_for(dispatcher.on_version_pending(5), timeout=0.5) + + assert dispatcher.policy_update_pending + assert group_id not in dispatcher.groups + + release_selection.set() + assert not await scheduling + assert not dispatcher.inflight + await dispatcher.on_new_version(5) + + +@pytest.mark.asyncio +async def test_selected_mutable_endpoint_survives_pool_churn_during_rate_limit_wait(): + rate_limit_started = asyncio.Event() + release_rate_limit = asyncio.Event() + selected_client = SimpleNamespace(base_url="http://policy/v1", headers={}) + + class _ChurningPool: + model_name = "frozen" + train_clients = [selected_client] + admin_clients = [SimpleNamespace(base_url="http://frozen-worker:8081")] + + async def select_train_client(self, _load): + # The elastic snapshot loses the selected endpoint before the + # dispatcher reaches its next slow wait. + self.train_clients = [SimpleNamespace(base_url="http://frozen/v1", headers={})] + return selected_client + + dispatcher = _dispatcher(enforce_barrier=True, run_rollout=None) + dispatcher.train_envs.env.sampler.pool = _ChurningPool() + + async def wait_for_rate_limit(_permits: int) -> None: + rate_limit_started.set() + await release_rate_limit.wait() + + dispatcher._wait_for_rate_limit = wait_for_rate_limit # type: ignore[method-assign] + group_id = uuid.uuid4() + group = GroupState( + kind="train", + env_name="train", + task_idx=0, + rollouts_to_schedule=1, + target_rollouts=1, + policy_version_at_start=4, + uses_mutable_policy=False, + ) + dispatcher.groups[group_id] = group + epoch = await dispatcher.policy_gate.scheduling_epoch() + assert epoch is not None + + scheduling = asyncio.create_task(dispatcher.schedule_group_rollout(group_id, group, epoch=epoch)) + await rate_limit_started.wait() + + assert group.uses_mutable_policy + await asyncio.wait_for(dispatcher.on_version_pending(5), timeout=0.5) + assert group_id not in dispatcher.groups + + release_rate_limit.set() + assert not await scheduling + await dispatcher.on_new_version(5) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 49243a9f98..54c46d2075 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -186,6 +186,12 @@ def test_env_algo_overrides_top_level(): assert reloaded.train.env[0].algo is not None and reloaded.train.env[0].algo.type == "grpo" +def test_policy_sampling_does_not_add_response_format_flags(): + config = OrchestratorConfig.model_validate({"train": {"env": [{"id": "math-env"}]}}) + + assert config.train.env[0].sampling.extra_body == {"top_k": -1, "min_p": 0.0} + + def test_trainer_enable_token_export_cli_flag(): assert not cli(TrainerConfig, args=[]).enable_token_export assert cli(TrainerConfig, args=["--enable-token-export"]).enable_token_export @@ -211,6 +217,218 @@ def test_single_node_auto_inference_client_dp_rank_count_matches_local_dp(): assert config.orchestrator.model.client.dp_rank_count == 2 +def test_inference_backend_defaults_to_vllm(): + assert InferenceConfig().backend.type == "vllm" + + +def test_dynamo_disaggregated_config_is_local_and_enables_prefix_caching(): + config = InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + } + ) + + assert config.enable_prefix_caching is True + assert config.use_pd_kv_transfer is True + + +@pytest.mark.parametrize( + "inference", + [ + {"parallel": {"tp": 0}}, + {"deployment": {"type": "single_node", "gpus_per_node": 0}}, + ], +) +def test_inference_topology_rejects_non_positive_gpu_dimensions(inference: dict): + with pytest.raises(ValidationError, match="greater than or equal to 1"): + InferenceConfig.model_validate(inference) + + +def test_dynamo_disaggregated_topology_requires_whole_tp_groups(): + with pytest.raises(ValidationError, match="gpus_per_node must be divisible"): + InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "parallel": {"tp": 2}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 3, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + }, + } + ) + + +@pytest.mark.parametrize( + "inference_override", + [ + {"parallel": {"tp": 2, "dp": 3}}, + {"parallel": {"tp": 2}, "data_parallel_size_local": 1}, + ], +) +def test_dynamo_disaggregated_topology_rejects_conflicting_dp(inference_override: dict): + inference = { + "backend": {"type": "dynamo"}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 4, + "num_prefill_replicas": 1, + "num_decode_replicas": 1, + }, + **inference_override, + } + + with pytest.raises(ValidationError, match="must equal.*gpus_per_node / inference.parallel.tp"): + InferenceConfig.model_validate(inference) + + +def test_dynamo_topology_is_derived_once_from_inference_config(): + config = InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "parallel": {"tp": 2}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 4, + "num_prefill_replicas": 2, + "num_decode_replicas": 1, + }, + } + ) + + assert config.dynamo_worker_roles == ("prefill", "prefill", "decode") + assert config.dynamo_gpus_per_worker == 4 + assert config.dynamo_local_dp == 2 + + +def test_dynamo_disaggregated_config_rejects_disabled_prefix_caching(): + with pytest.raises(ValidationError, match="requires prefix caching"): + InferenceConfig.model_validate( + { + "backend": {"type": "dynamo"}, + "enable_prefix_caching": False, + "deployment": {"type": "disaggregated"}, + } + ) + + +def test_dynamo_backend_rejects_lora(): + with pytest.raises(ValidationError, match="does not support LoRA"): + InferenceConfig.model_validate({"backend": {"type": "dynamo"}, "enable_lora": True}) + + +def test_rl_config_rejects_lora_auto_setup_for_dynamo_backend(): + with pytest.raises(ValidationError, match="does not support LoRA"): + RLConfig.model_validate( + { + "trainer": {"model": {"lora": {}}}, + "orchestrator": {}, + "inference": {"backend": {"type": "dynamo"}}, + "deployment": { + "type": "single_node", + "gpus_per_node": 2, + "num_train_gpus": 1, + "num_infer_gpus": 1, + }, + } + ) + + +def test_native_disaggregated_config_still_requires_slurm(): + with pytest.raises(ValidationError, match="Must use SLURM"): + InferenceConfig.model_validate({"deployment": {"type": "disaggregated"}}) + + +def test_single_node_dynamo_disaggregated_keeps_per_worker_dp_and_sets_nccl_world_size(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": {}, + "inference": { + "backend": {"type": "dynamo"}, + "parallel": {"tp": 1}, + "deployment": { + "type": "disaggregated", + "gpus_per_node": 1, + "prefill_nodes_per_replica": 1, + "decode_nodes_per_replica": 1, + "num_prefill_replicas": 2, + "num_decode_replicas": 2, + }, + }, + "weight_broadcast": {"type": "nccl"}, + "deployment": { + "type": "single_node", + "gpus_per_node": 8, + "num_train_gpus": 4, + "num_infer_gpus": 4, + }, + } + ) + + assert config.inference is not None + assert config.inference.parallel.dp == 1 + assert config.orchestrator.model.client.admin_api == "dynamo" + assert config.orchestrator.model.client.dp_rank_count == 1 + assert config.orchestrator.model.client.dynamo_worker_roles == ("prefill", "prefill", "decode", "decode") + assert config.orchestrator.model.client.dynamo_gpus_per_worker == 1 + assert config.trainer.weight_broadcast.inference_world_size == 4 + assert config.orchestrator.weight_broadcast.inference_world_size == 4 + + +def test_single_node_dynamo_aggregated_derives_one_multi_gpu_worker(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": {}, + "inference": {"backend": {"type": "dynamo"}, "parallel": {"tp": 1}}, + "deployment": { + "type": "single_node", + "gpus_per_node": 4, + "num_train_gpus": 2, + "num_infer_gpus": 2, + }, + } + ) + + assert config.inference is not None + assert config.orchestrator.model.client.dynamo_worker_roles == ("agg",) + assert config.orchestrator.model.client.dynamo_gpus_per_worker == 2 + + +def test_dynamo_topology_metadata_cannot_conflict_with_inference_config(): + with pytest.raises(ValidationError, match="dynamo_worker_roles conflicts"): + RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": { + "model": { + "client": { + "dynamo_worker_roles": ["decode"], + "dynamo_gpus_per_worker": 1, + } + } + }, + "inference": {"backend": {"type": "dynamo"}, "parallel": {"tp": 1}}, + "deployment": { + "type": "single_node", + "gpus_per_node": 2, + "num_train_gpus": 1, + "num_infer_gpus": 1, + }, + } + ) + + def test_multi_node_auto_inference_client_dp_rank_count_uses_router_url(): config = RLConfig.model_validate( { diff --git a/tests/unit/test_deployment_assets.py b/tests/unit/test_deployment_assets.py new file mode 100644 index 0000000000..5768f5bb67 --- /dev/null +++ b/tests/unit/test_deployment_assets.py @@ -0,0 +1,5 @@ +from pathlib import Path + + +def test_chart_does_not_duplicate_prime_runtime_image(): + assert not (Path(__file__).parents[2] / "Dockerfile.dynamo").exists() diff --git a/tests/unit/test_transport_config.py b/tests/unit/test_transport_config.py new file mode 100644 index 0000000000..9a3f60fbaa --- /dev/null +++ b/tests/unit/test_transport_config.py @@ -0,0 +1,21 @@ +import math + +import pytest +from pydantic import ValidationError + +from prime_rl.configs.shared import FileSystemTransportConfig, ZMQTransportConfig + + +@pytest.mark.parametrize("config_type", [FileSystemTransportConfig, ZMQTransportConfig]) +def test_rollout_send_timeout_has_finite_positive_default(config_type): + timeout = config_type().send_timeout_seconds + + assert timeout == 300.0 + assert math.isfinite(timeout) + + +@pytest.mark.parametrize("config_type", [FileSystemTransportConfig, ZMQTransportConfig]) +@pytest.mark.parametrize("timeout", [0, -1, float("inf"), float("-inf"), float("nan")]) +def test_rollout_send_timeout_rejects_nonpositive_or_nonfinite_values(config_type, timeout: float): + with pytest.raises(ValidationError): + config_type(send_timeout_seconds=timeout) diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 40de4cfee6..7053bf69aa 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -1,12 +1,24 @@ import asyncio from pathlib import Path -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import httpx +import pytest from verifiers.v1.clients.config import EvalClientConfig -from prime_rl.configs.shared import ClientConfig -from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients +from prime_rl.configs.shared import ClientConfig, ElasticConfig +from prime_rl.inference.dynamo_admin import DynamoTopology, DynamoWorker +from prime_rl.utils.client import ( + DynamoInferencePool, + StaticInferencePool, + _is_retryable_lora_error, + check_health, + load_lora_adapter, + maybe_check_has_model, + setup_clients, + setup_inference_pool, +) +from prime_rl.utils.policy_client_config import policy_client_config_from_environment def test_is_retryable_lora_error_returns_true_for_404(): @@ -108,3 +120,258 @@ def test_setup_clients_preserves_chat_client_defaults(): headers={}, ) ] + + +def test_native_pool_preserves_admin_base_url_and_reuses_admin_clients(): + pool = StaticInferencePool( + ClientConfig( + base_url=["http://router:8000/v1"], + admin_base_url=["http://worker:8001/v1"], + ), + model_name="test-model", + ) + + assert pool._frontend_admin_clients is pool._admin_clients + assert [str(client.base_url).rstrip("/") for client in pool.admin_clients] == ["http://worker:8001"] + asyncio.run(pool.stop()) + + +def test_setup_inference_pool_selects_dynamo_pool_once(): + pool = asyncio.run( + setup_inference_pool( + ClientConfig( + base_url=["http://frontend:8000/v1"], + admin_api="dynamo", + dynamo_worker_roles=("agg",), + dynamo_gpus_per_worker=1, + ), + model_name="test-model", + ) + ) + + assert isinstance(pool, DynamoInferencePool) + asyncio.run(pool.stop()) + + +def test_generated_dgd_topology_selects_dynamo_pool_without_mutating_config(monkeypatch: pytest.MonkeyPatch): + client_config = ClientConfig() + original_config = client_config.model_copy(deep=True) + monkeypatch.setenv( + "DYN_RL_TOPOLOGY", + """{"schema_version":1,"admin_api":"dynamo","base_url":["http://frontend:8000/v1"],"rl_base_url":["http://frontend-rl:8001"],"dynamo_worker_roles":["prefill","decode"],"dynamo_gpus_per_worker":1}""", + ) + + resolved = policy_client_config_from_environment(client_config) + pool = asyncio.run(setup_inference_pool(resolved, model_name="test-model")) + + assert isinstance(pool, DynamoInferencePool) + assert pool.admin_api == "dynamo" + assert resolved.base_url == ["http://frontend:8000/v1"] + assert resolved.rl_base_url == ["http://frontend-rl:8001"] + assert client_config == original_config + asyncio.run(pool.stop()) + + +def test_generated_dgd_topology_rejects_explicit_client_conflict(monkeypatch: pytest.MonkeyPatch): + client_config = ClientConfig(admin_api="vllm") + monkeypatch.setenv( + "DYN_RL_TOPOLOGY", + """{"schema_version":1,"admin_api":"dynamo","base_url":["http://frontend:8000/v1"],"rl_base_url":["http://frontend-rl:8001"],"dynamo_worker_roles":["agg"],"dynamo_gpus_per_worker":1}""", + ) + + with pytest.raises(ValueError, match="admin_api.*conflicts"): + policy_client_config_from_environment(client_config) + + +@pytest.mark.asyncio +async def test_setup_inference_pool_rejects_dynamo_elastic_before_pool_selection(): + client_config = ClientConfig( + base_url=["http://frontend:8000/v1"], + admin_api="dynamo", + dynamo_worker_roles=("agg",), + dynamo_gpus_per_worker=1, + elastic=ElasticConfig(hostname="inference.example"), + ) + original_config = client_config.model_copy(deep=True) + + with patch("prime_rl.utils.elastic.ElasticInferencePool.from_config", new=AsyncMock()) as from_config: + with pytest.raises(ValueError, match="Dynamo admin API does not support elastic inference pools"): + await setup_inference_pool(client_config, model_name="test-model") + + from_config.assert_not_awaited() + assert client_config == original_config + + +@pytest.mark.asyncio +async def test_model_registration_retries_transient_status_and_empty_models(): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + request = httpx.Request("GET", "http://frontend:8000/v1/models") + conflict = httpx.Response(409, request=request) + unavailable = httpx.Response( + 503, + request=request, + ) + empty = httpx.Response(200, json={"data": []}, request=request) + ready = httpx.Response(200, json={"data": [{"id": "test-model"}]}, request=request) + client.get.side_effect = [httpx.ConnectError("not listening", request=request), conflict, unavailable, empty, ready] + + await maybe_check_has_model([client], "test-model", timeout=1, interval=0) + + assert client.get.await_count == 5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [401, 403, 404]) +async def test_model_registration_fails_permanent_http_status_immediately(status_code: int): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = httpx.Response( + status_code, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ) + + with pytest.raises(httpx.HTTPStatusError): + await maybe_check_has_model([client], "test-model", timeout=1, interval=0) + + client.get.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + httpx.Response( + 200, + content=b"not-json", + headers={"content-type": "application/json"}, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ), + httpx.Response( + 200, + json={"data": {}}, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ), + httpx.Response( + 200, + json={"data": [{"object": "model"}]}, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ), + ], + ids=["invalid-json", "data-not-list", "model-id-missing"], +) +async def test_model_registration_fails_invalid_response_immediately(response: httpx.Response): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = response + + with pytest.raises(ValueError, match=r"Invalid /v1/models response"): + await maybe_check_has_model([client], "test-model", timeout=1, interval=0) + + client.get.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_strict_health_retries_only_transient_failures(): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + request = httpx.Request("GET", "http://frontend:8000/health") + client.get.side_effect = [ + httpx.ConnectError("not listening", request=request), + httpx.Response(429, request=request), + httpx.Response(503, request=request), + httpx.Response(200, request=request), + ] + + await check_health([client], timeout=1, interval=0, strict=True) + + assert client.get.await_count == 4 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [401, 403, 404]) +async def test_strict_health_fails_permanent_http_status_immediately(status_code: int): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = httpx.Response( + status_code, + request=httpx.Request("GET", "http://frontend:8000/health"), + ) + + with pytest.raises(httpx.HTTPStatusError): + await check_health([client], timeout=1, interval=0, strict=True) + + client.get.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_model_registration_timeout_reports_last_error(): + client = AsyncMock() + client.base_url = httpx.URL("http://frontend:8000") + client.get.return_value = httpx.Response( + 503, + request=httpx.Request("GET", "http://frontend:8000/v1/models"), + ) + + with pytest.raises(TimeoutError, match=r"test-model.*frontend:8000.*503 Service Unavailable"): + await maybe_check_has_model([client], "test-model", timeout=0.02, interval=0.001) + + assert client.get.await_count > 1 + + +@pytest.mark.asyncio +async def test_dynamo_readiness_shares_one_monotonic_deadline(monkeypatch: pytest.MonkeyPatch): + pool = DynamoInferencePool( + ClientConfig( + base_url=["http://frontend:8000/v1"], + admin_api="dynamo", + dynamo_worker_roles=("agg",), + dynamo_gpus_per_worker=2, + ), + model_name="test-model", + ) + observed_timeouts: list[float] = [] + worker = DynamoWorker( + instance_id=11, + component="backend", + role="agg", + system_url="http://worker:8081", + model="test-model", + routes=frozenset( + { + "init_weights_update_group", + "pause_generation", + "resume_generation", + "update_weights_from_disk", + "update_weights_from_distributed", + } + ), + ) + + async def health(_clients, *, timeout, strict=False, **_kwargs): + assert strict is True + observed_timeouts.append(timeout) + await asyncio.sleep(0.01) + + async def models(_clients, _model_name, *, skip_model_check, timeout): + assert skip_model_check is False + observed_timeouts.append(timeout) + await asyncio.sleep(0.01) + + async def discover(_clients, timeout, *, model_name, topology): + assert model_name == "test-model" + assert topology == DynamoTopology(roles=("agg",), gpus_per_worker=2) + observed_timeouts.append(timeout) + await asyncio.sleep(0.01) + return (worker,) + + monkeypatch.setattr("prime_rl.utils.client.check_health", health) + monkeypatch.setattr("prime_rl.utils.client.maybe_check_has_model", models) + monkeypatch.setattr("prime_rl.utils.client.discover_workers", discover) + try: + await pool.wait_for_ready("test-model", timeout=1) + finally: + await pool.stop() + + assert len(observed_timeouts) == 4 + assert all(later < earlier for earlier, later in zip(observed_timeouts, observed_timeouts[1:]))