Skip to content

build_orchestrator: merge model-card slurm/distributed config in the main build path #152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
i-kosarev wants to merge 1 commit into
base: fix/slurm-sglang-disagg-co-located-proxy
Choose a base branch
from
+83 −69
Draft

build_orchestrator: merge model-card slurm/distributed config in the main build path #152

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 76 additions & 62 deletions src/madengine/orchestration/build_orchestrator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,10 @@ def execute(
# Step 6: Save deployment_config to manifest
self._save_deployment_config(manifest_output)

# Step 6b: Merge model's distributed and slurm config into deployment_config
# This ensures launcher and slurm settings are in deployment_config even if not in additional-context
self._merge_model_config_into_deployment(manifest_output, models)

self.rich_console.print(f"[green]✓ Build complete: {manifest_output}[/green]")
self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n")

Expand Down Expand Up @@ -582,70 +586,10 @@ def _execute_with_prebuilt_image(

# Save deployment config
self._save_deployment_config(manifest_output)

# Merge model's distributed and slurm config into deployment_config
# This ensures launcher and slurm settings are in deployment_config even if not in additional-context
if models:
with open(manifest_output, "r") as f:
saved_manifest = json.load(f)

if "deployment_config" not in saved_manifest:
saved_manifest["deployment_config"] = {}

# Merge model's distributed config from the first model.
# If multiple models have differing distributed configs, warn — only the first wins here.
# Use json.dumps for the hash key so nested dicts (e.g. sglang_disagg / vllm_disagg)
# don't trigger TypeError: unhashable type: 'dict' from `tuple(sorted(items()))`.
if len(models) > 1:
distinct_distributed = {
json.dumps(m.get("distributed") or {}, sort_keys=True, default=str)
for m in models
}
if len(distinct_distributed) > 1:
self.rich_console.print(
"[yellow]Warning: discovered models have differing distributed configs; "
f"using {models[0].get('name', '<unknown>')}'s config.[/yellow]"
)
model_distributed = models[0].get("distributed", {})
if model_distributed:
if "distributed" not in saved_manifest["deployment_config"]:
saved_manifest["deployment_config"]["distributed"] = {}

# Copy launcher and other critical fields from model config
for key in ["launcher", "nnodes", "nproc_per_node", "backend", "port", "sglang_disagg", "vllm_disagg"]:
if key in model_distributed and key not in saved_manifest["deployment_config"]["distributed"]:
saved_manifest["deployment_config"]["distributed"][key] = model_distributed[key]

# Merge model's slurm config into deployment_config.slurm from the first model.
# This enables run phase to auto-detect SLURM deployment without --additional-context.
# Warn when multiple models have differing slurm configs (only the first wins here).
# json.dumps key for the same unhashable-nested-dict reason as above.
if len(models) > 1:
distinct_slurm = {
json.dumps(m.get("slurm") or {}, sort_keys=True, default=str)
for m in models
}
if len(distinct_slurm) > 1:
self.rich_console.print(
"[yellow]Warning: discovered models have differing slurm configs; "
f"using {models[0].get('name', '<unknown>')}'s config.[/yellow]"
)
model_slurm = models[0].get("slurm", {})
if model_slurm:
if "slurm" not in saved_manifest["deployment_config"]:
saved_manifest["deployment_config"]["slurm"] = {}

# Copy slurm settings from model config (model card fills in
# values not explicitly set by --additional-context).
# Use _original_user_slurm_keys (captured before ConfigLoader
# applies defaults) so model card values override defaults
# but user's explicit CLI values still win.
for key in ["partition", "nodes", "gpus_per_node", "time", "exclusive", "reservation", "output_dir", "nodelist"]:
if key in model_slurm and key not in self._original_user_slurm_keys:
saved_manifest["deployment_config"]["slurm"][key] = model_slurm[key]

with open(manifest_output, "w") as f:
json.dump(saved_manifest, f, indent=2)
self._merge_model_config_into_deployment(manifest_output, models)

self.rich_console.print(f"[green]✓ Generated manifest: {manifest_output}[/green]")
self.rich_console.print(f" Pre-built image: {use_image}")
Expand Down Expand Up @@ -1284,6 +1228,76 @@ def _save_build_summary(self, manifest_file: str, build_summary: Dict):
except Exception as e:
self.rich_console.print(f"[yellow]Warning: Could not save build summary: {e}[/yellow]")

def _merge_model_config_into_deployment(self, manifest_output: str, models: list):
"""Merge the first discovered model's distributed/slurm config into
manifest deployment_config, so the run phase can auto-detect the
deployment target (e.g. slurm) from the model card alone, without
requiring --additional-context to repeat what's already in models.json.
"""
if not models:
return

with open(manifest_output, "r") as f:
saved_manifest = json.load(f)

if "deployment_config" not in saved_manifest:
saved_manifest["deployment_config"] = {}

# Merge model's distributed config from the first model.
# If multiple models have differing distributed configs, warn — only the first wins here.
# Use json.dumps for the hash key so nested dicts (e.g. sglang_disagg / vllm_disagg)
# don't trigger TypeError: unhashable type: 'dict' from `tuple(sorted(items()))`.
if len(models) > 1:
distinct_distributed = {
json.dumps(m.get("distributed") or {}, sort_keys=True, default=str)
for m in models
}
if len(distinct_distributed) > 1:
self.rich_console.print(
"[yellow]Warning: discovered models have differing distributed configs; "
f"using {models[0].get('name', '<unknown>')}'s config.[/yellow]"
)
model_distributed = models[0].get("distributed", {})
if model_distributed:
if "distributed" not in saved_manifest["deployment_config"]:
saved_manifest["deployment_config"]["distributed"] = {}

# Copy launcher and other critical fields from model config
for key in ["launcher", "nnodes", "nproc_per_node", "backend", "port", "sglang_disagg", "vllm_disagg"]:
if key in model_distributed and key not in saved_manifest["deployment_config"]["distributed"]:
saved_manifest["deployment_config"]["distributed"][key] = model_distributed[key]

# Merge model's slurm config into deployment_config.slurm from the first model.
# This enables run phase to auto-detect SLURM deployment without --additional-context.
# Warn when multiple models have differing slurm configs (only the first wins here).
# json.dumps key for the same unhashable-nested-dict reason as above.
if len(models) > 1:
distinct_slurm = {
json.dumps(m.get("slurm") or {}, sort_keys=True, default=str)
for m in models
}
if len(distinct_slurm) > 1:
self.rich_console.print(
"[yellow]Warning: discovered models have differing slurm configs; "
f"using {models[0].get('name', '<unknown>')}'s config.[/yellow]"
)
model_slurm = models[0].get("slurm", {})
if model_slurm:
if "slurm" not in saved_manifest["deployment_config"]:
saved_manifest["deployment_config"]["slurm"] = {}

# Copy slurm settings from model config (model card fills in
# values not explicitly set by --additional-context).
# Use _original_user_slurm_keys (captured before ConfigLoader
# applies defaults) so model card values override defaults
# but user's explicit CLI values still win.
for key in ["partition", "nodes", "gpus_per_node", "time", "exclusive", "reservation", "output_dir", "nodelist"]:
if key in model_slurm and key not in self._original_user_slurm_keys:
saved_manifest["deployment_config"]["slurm"][key] = model_slurm[key]

with open(manifest_output, "w") as f:
json.dump(saved_manifest, f, indent=2)

def _save_deployment_config(self, manifest_file: str):
"""Save deployment_config from --additional-context to manifest."""
if not self.additional_context:
Expand Down
14 changes: 7 additions & 7 deletions src/madengine/orchestration/run_orchestrator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -517,13 +517,13 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict:
# Load manifest first to check if we have Docker images
with open(manifest_file, "r") as f:
manifest = json.load(f)

has_docker_images = bool(manifest.get("built_images", {}))

if has_docker_images:
# Using Docker containers - containers have GPU support built-in
self.rich_console.print("[dim cyan]Using Docker containers with built-in GPU support[/dim cyan]\n")

# Initialize runtime context (runs full GPU detection on compute nodes)
self._init_runtime_context()

Expand Down Expand Up @@ -765,13 +765,13 @@ def _show_node_info(self):

host_os = self.context.ctx.get("host_os", "")
if "HOST_UBUNTU" in host_os:
print(self.console.sh("apt show rocm-libs -a", canFail=True))
print(self.console.sh("timeout 10 apt show rocm-libs -a", canFail=True))
elif "HOST_CENTOS" in host_os:
print(self.console.sh("yum info rocm-libs", canFail=True))
print(self.console.sh("timeout 10 yum info rocm-libs", canFail=True))
elif "HOST_SLES" in host_os:
print(self.console.sh("zypper info rocm-libs", canFail=True))
print(self.console.sh("timeout 10 zypper info rocm-libs", canFail=True))
elif "HOST_AZURE" in host_os:
print(self.console.sh("tdnf info rocm-libs", canFail=True))
print(self.console.sh("timeout 10 tdnf info rocm-libs", canFail=True))
Comment on lines 767 to +774
else:
self.rich_console.print("[yellow]Warning: Unable to detect host OS[/yellow]")

Expand Down