diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py index 246701cf..71adca8d 100644 --- a/src/madengine/orchestration/build_orchestrator.py +++ b/src/madengine/orchestration/build_orchestrator.py @@ -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") @@ -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', '')}'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', '')}'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}") @@ -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', '')}'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', '')}'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: diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index 296af1ee..1f6314e1 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -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() @@ -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)) else: self.rich_console.print("[yellow]Warning: Unable to detect host OS[/yellow]")