Skip to content
Merged
12 changes: 10 additions & 2 deletions components/src/dynamo/profiler/rapid.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
from dynamo.profiler.utils.profile_common import (
derive_backend_image,
needs_profile_data,
resolve_model_path,
)

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -104,6 +105,11 @@ def _generate_dgd_from_pick(
finally:
tc.total_gpus = original_total_gpus

service_cfg = cfg.get("ServiceConfig")
if isinstance(service_cfg, dict):
service_cfg["model_path"] = dgdr.model
service_cfg["served_model_path"] = dgdr.model

artifacts = generate_backend_artifacts(
params=cfg,
backend=tc.backend_name,
Expand Down Expand Up @@ -190,9 +196,10 @@ def _run_autoscale_sim(
"Throughput-based scaling enabled — only disagg mode is supported."
)

local_or_hf_model = resolve_model_path(dgdr)
Comment thread
hhzhang16 marked this conversation as resolved.
task = TaskConfig(
serving_mode="disagg",
model_path=model,
model_path=local_or_hf_model,
Comment thread
hhzhang16 marked this conversation as resolved.
system_name=system,
backend_name=backend,
total_gpus=total_gpus,
Expand Down Expand Up @@ -238,8 +245,9 @@ def _run_default_sim(
picking_mode: str,
) -> dict:
"""Build default task_configs, apply load_match kwargs, run simulation, generate DGD."""
local_or_hf_model = resolve_model_path(dgdr)
task_configs = build_default_task_configs(
model_path=model,
model_path=local_or_hf_model,
Comment thread
hhzhang16 marked this conversation as resolved.
total_gpus=total_gpus,
system=system,
backend=backend,
Expand Down
Loading
Loading