Add DOCA_MEMOS preflight for AI Dynamo

conf/cmx-daria/hook/daria_cmx_host_prep.toml

@@ -0,0 +1,6 @@
+name = "daria_cmx_host_prep"
+
+[[Tests]]
+id = "daria-cmx-host-prep"
+test_name = "daria-cmx-host-prep"
+time_limit = "00:10:00"

conf/cmx-daria/hook/test/daria_cmx_host_prep.toml

@@ -0,0 +1,27 @@
+name = "daria-cmx-host-prep"
+description = "Daria CMX host-side pre-test checks for Lustre, RDMA UCM, and cached images"
+test_template_name = "BashCmd"
+
+[cmd_args]
+cmd = """
+set -euo pipefail
+
+echo Daria CMX host prep on $(hostname)
+echo SLURM_JOB_NODELIST=${SLURM_JOB_NODELIST:-}
+
+test -d /mnt/lustre
+test -r /mnt/lustre
+test -d /mnt/lustre/users/spothula/cloudai/.install
+find /mnt/lustre/users/spothula/cloudai/.install -maxdepth 1 -name '*.sqsh' -print -quit | grep -q .
+
+srun --overlap --ntasks=${SLURM_JOB_NUM_NODES} --ntasks-per-node=1 bash -lc '
+  set -euo pipefail
+  echo node=\\$(hostname)
+  sudo -n modprobe rdma_ucm
+  test -e /dev/infiniband/rdma_cm
+  test -d /mnt/lustre
+  test -r /mnt/lustre
+  test -d /dev/hugepages
+  ls -l /dev/ng* 2>/dev/null || true
+'
+"""

conf/cmx-daria/system/daria_b300_scctl_spothula.toml

@@ -0,0 +1,46 @@
+# Daria system configuration for launches through a Daria compute node.
+# Personal-path variant of Juntao Wang's CMS b300_scctl.toml.
+
+name = "daria-b300-scctl-spothula"
+scheduler = "slurm"
+
+install_path = "/mnt/lustre/users/spothula/cloudai/.install"
+output_path = "/mnt/lustre/users/spothula/cloudai/results"
+hf_home_path = "/mnt/lustre/cloudai/huggingface"
+
+default_partition = "daria"
+gpus_per_node = 4
+mpi = "pmix"
+monitor_interval = 10
+extra_srun_args = "--container-remap-root --container-writable"
+extra_sbatch_args = ["--cpus-per-task=256"]
+
+cache_docker_images_locally = true
+
+[[partitions]]
+name = "daria"
+
+  [[partitions.groups]]
+  name = "b300_nodes"
+  nodes = [
+    "daria01",
+    "daria02",
+    "daria03",
+    "daria04",
+    "daria05",
+    "daria06",
+    "daria07",
+    "daria08",
+  ]
+
+[global_env_vars]
+HF_HOME = "/root/.cache/huggingface/"
+HF_HUB_CACHE = "/root/.cache/huggingface/hub"
+
+MELLANOX_VISIBLE_DEVICES = "all"
+
+NCCL_DEBUG = "WARN"
+NCCL_IB_DISABLE = "0"
+NCCL_NET_GDR_LEVEL = "5"
+NCCL_IB_GID_INDEX = "3"
+NCCL_SOCKET_IFNAME = "^lo,docker"

conf/cmx-daria/test/agentic_gpt_oss20_doca_memos_single.toml

@@ -0,0 +1,98 @@
+name = "agentic-gpt-oss20-doca-memos-single"
+description = "GPT-OSS-20B agentic workload: LMCache local CPU plus CMX DOCA_MEMOS, single node"
+test_template_name = "AIDynamo"
+extra_container_mounts = [
+  "/dev:/dev",
+  "/sys/class/nvme:/sys/class/nvme",
+]
+
+[cmd_args]
+docker_image_url = "artifactory.nvidia.com/sw-nbu-swx-storage-devops-docker-local/nixl/dynamo-vllm-nixl-cmx:1.2.0-cuda13-lmcache-389b9cfc-doca-41a1f020-nixl-cd2f6d49-hnogit-vllm0.20.2-build24-x86_64"
+workloads = "aiperf.sh"
+
+  [cmd_args.doca-memos-preflight]
+  enabled = true
+  health-check = true
+  setup-hugepages = true
+  skip-data-path-check = false
+
+  [cmd_args.lmcache]
+  chunk_size = 512
+  local_cpu = true
+  max_local_cpu_size = 25.0
+  local_cpu_use_hugepages = true
+  nixl_buffer_device = "cpu"
+
+    [cmd_args.lmcache.extra_config]
+    enable_nixl_storage = true
+    nixl_presence_cache = true
+    nixl_presence_cache_only = true
+    nixl_async_put = true
+    nixl_sync_mode = "STRICT"
+    nixl_pool_size = 0
+    nixl_backend = "DOCA_MEMOS"
+
+      [cmd_args.lmcache.extra_config.nixl_backend_params]
+      query_mem_mode = "actual"
+
+  [cmd_args.dynamo]
+  backend = "vllm"
+  model = "openai/gpt-oss-20b"
+  ingress-cmd = "python3 -m dynamo.frontend --no-kv-events --router-reset-states --router-mode kv --router-ttl 86400 --kv-overlap-score-weight 1000"
+
+    [cmd_args.dynamo.prefill_worker]
+    num-nodes = 0
+    cmd = "python3 -m dynamo.vllm --is-prefill-worker"
+    worker-initialized-regex = "VllmWorker.*has.been.initialized"
+
+      [cmd_args.dynamo.prefill_worker.args]
+      tensor-parallel-size = 4
+      pipeline-parallel-size = 1
+      gpu-memory-utilization = 0.4
+      kv-transfer-config = '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_both"}'
+
+    [cmd_args.dynamo.decode_worker]
+    num-nodes = 1
+    cmd = "python3 -m dynamo.vllm"
+    worker-initialized-regex = "VllmWorker.*has.been.initialized"
+    extra-args = "--kv-cache-metrics --kv-cache-metrics-sample 0.01"
+
+      [cmd_args.dynamo.decode_worker.args]
+      tensor-parallel-size = 4
+      pipeline-parallel-size = 1
+      gpu-memory-utilization = 0.4
+      kv-transfer-config = '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_both"}'
+
+  [cmd_args.aiperf]
+    [cmd_args.aiperf.args]
+    server-metrics = "auto"
+    endpoint-type = "chat"
+    streaming = true
+    export-level = "records"
+    slice-duration = 30
+    server-metrics-formats = "csv,json,jsonl"
+    extra-inputs = "ignore_eos:true"
+    dataset-sampling-strategy = "sequential"
+    export-http-trace = true
+    user-centric-rate = 0.6666666667
+    num-users = 40
+    session-turns-mean = 15
+    user-context-prompt-length = 64000
+    shared-system-prompt-length = 1000
+    synthetic-input-tokens-mean = 2048
+    synthetic-input-tokens-stddev = 0
+    osl = 64
+    num-dataset-entries = 40
+    conversation-num = 40
+    concurrency = 40
+    max-workers = 100
+    warmup-request-count = 5
+
+[extra_env_vars]
+AIPERF_DATASET_CONFIGURATION_TIMEOUT = "900"
+AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT = "900"
+HF_HUB_OFFLINE = "1"
+PYTHONHASHSEED = "0"
+PROMETHEUS_MULTIPROC_DIR = "/tmp/lmcache_prometheus"
+UCX_TLS = "all"
+UCX_HANDLE_ERRORS = "bt,freeze"

conf/cmx-daria/test_scenario/agentic_gpt_oss20_doca_memos_single.toml

@@ -0,0 +1,12 @@
+name = "agentic-gpt-oss20-doca-memos-single-review"
+
+[[Tests]]
+id = "agentic-gpt-oss20-doca-memos-single-review"
+test_name = "agentic-gpt-oss20-doca-memos-single"
+nodes = ["daria04"]
+time_limit = "24:00:00"
+
+  [Tests.cmd_args]
+
+    [Tests.cmd_args.dynamo.decode_worker]
+    nodes = "daria04"

src/cloudai/workloads/ai_dynamo/ai_dynamo.py

@@ -17,7 +17,7 @@
 import csv
 import logging
 from pathlib import Path
-from typing import Literal, Optional, cast
+from typing import Any, Literal, Optional, cast
 
 from pydantic import (
     AliasChoices,
@@ -162,6 +162,7 @@ class AIDynamoArgs(BaseModel):
 
     model: str = "Qwen/Qwen3-0.6B"
     backend: Literal["vllm", "sglang", "sglang_dsr1"] = "vllm"
+    mode: Literal["disaggregated", "aggregate"] = "disaggregated"
     endpoint: str = Field(default="v1/chat/completions")
     connector: Optional[str | list[str]] = None
 
@@ -381,6 +382,46 @@ class Constraints(BaseModel):
     tp_times_pp_le_gpus_per_node: bool = True
 
 
+class DocaMemosPreflight(BaseModel):
+    """Optional DOCA_MEMOS-specific preflight checks for AI Dynamo Slurm runs."""
+
+    model_config = ConfigDict(extra="forbid", populate_by_name=True)
+
+    enabled: bool = False
+    health_check: bool = Field(
+        default=True,
+        serialization_alias="health-check",
+        validation_alias=AliasChoices("health-check", "health_check"),
+    )
+    setup_hugepages: bool = Field(
+        default=False,
+        serialization_alias="setup-hugepages",
+        validation_alias=AliasChoices("setup-hugepages", "setup_hugepages"),
+    )
+    skip_data_path_check: bool = Field(
+        default=False,
+        serialization_alias="skip-data-path-check",
+        validation_alias=AliasChoices("skip-data-path-check", "skip_data_path_check"),
+    )
+    probe_size_bytes: int = Field(
+        default=6 * 1024 * 1024,
+        gt=0,
+        serialization_alias="probe-size-bytes",
+        validation_alias=AliasChoices("probe-size-bytes", "probe_size_bytes"),
+    )
+    transfer_timeout_sec: float = Field(
+        default=45.0,
+        gt=0,
+        serialization_alias="transfer-timeout-sec",
+        validation_alias=AliasChoices("transfer-timeout-sec", "transfer_timeout_sec"),
+    )
+    srun_timeout_sec: int = Field(
+        default=90,
+        gt=0,
+        serialization_alias="srun-timeout-sec",
+        validation_alias=AliasChoices("srun-timeout-sec", "srun_timeout_sec"),
+    )
+
 class LMCacheController(BaseModel):
     """Optional LMCache controller process to launch on the frontend node."""
 
@@ -402,6 +443,11 @@ class AIDynamoCmdArgs(CmdArgs):
     hicache: dict | None = None
     lmcache: dict | None = None
     lmcache_controller: LMCacheController | None = None
+    doca_memos_preflight: DocaMemosPreflight = Field(
+        default_factory=DocaMemosPreflight,
+        serialization_alias="doca-memos-preflight",
+        validation_alias=AliasChoices("doca-memos-preflight", "doca_memos_preflight"),
+    )
     genai_perf: GenAIPerf = Field(default_factory=GenAIPerf)
     aiperf: AIPerf = Field(default_factory=AIPerf)
     aiperf_phases: list[AIPerfPhase] | None = None
@@ -418,6 +464,46 @@ def validate_workloads(cls, v: str) -> str:
                 raise ValueError(f"Invalid workload: {workload}. Available workloads: {allowed_workloads}")
         return ",".join(values)
 
+    @staticmethod
+    def _is_enabled_value(value: Any) -> bool:
+        if isinstance(value, list):
+            return all(AIDynamoCmdArgs._is_enabled_value(item) for item in value)
+        return value is True
+
+    @staticmethod
+    def _is_doca_memos_backend(value: Any) -> bool:
+        if isinstance(value, list):
+            return all(AIDynamoCmdArgs._is_doca_memos_backend(item) for item in value)
+        return isinstance(value, str) and value.casefold() == "doca_memos"
+
+    @model_validator(mode="after")
+    def validate_doca_memos_preflight(self) -> "AIDynamoCmdArgs":
+        """Validate DOCA_MEMOS preflight is only enabled for DOCA_MEMOS LMCache configurations."""
+        preflight = self.doca_memos_preflight
+        if not preflight.enabled:
+            return self
+        if not preflight.health_check and not preflight.setup_hugepages:
+            raise ValueError("doca_memos_preflight.enabled requires health_check or setup_hugepages")
+        if self.lmcache is None:
+            raise ValueError("doca_memos_preflight requires cmd_args.lmcache")
+
+        extra_config = self.lmcache.get("extra_config")
+        if not isinstance(extra_config, dict):
+            raise ValueError("doca_memos_preflight requires cmd_args.lmcache.extra_config")
+        if not self._is_enabled_value(extra_config.get("enable_nixl_storage")):
+            raise ValueError("doca_memos_preflight requires enable_nixl_storage=true")
+        if not self._is_doca_memos_backend(extra_config.get("nixl_backend")):
+            raise ValueError('doca_memos_preflight requires nixl_backend="DOCA_MEMOS"')
+
+        backend_params = extra_config.get("nixl_backend_params") or {}
+        if not isinstance(backend_params, dict):
+            raise ValueError("doca_memos_preflight requires nixl_backend_params to be a mapping when set")
+        device_name = str(backend_params.get("device_name", "")).strip()
+        if not preflight.health_check and (not device_name or device_name.casefold() == "auto"):
+            raise ValueError("automatic DOCA_MEMOS device discovery requires doca_memos_preflight.health_check=true")
+
+        return self
+
     @property
     def workloads_list(self) -> list[str]:
         return [w.strip() for w in self.workloads.split(",")]

src/cloudai/workloads/ai_dynamo/ai_dynamo.sh

@@ -43,6 +43,7 @@ SHARED_NODE_DISAGG="false"
 
 declare -A dynamo_args
 dynamo_args["backend"]="vllm"
+dynamo_args["mode"]="disaggregated"
 dynamo_args["node-setup-cmd"]=""
 dynamo_args["ingress-cmd"]="python -m dynamo.frontend --router-mode kv"
 dynamo_args["port"]=$((8080 + SLURM_JOBID % 100))
@@ -307,7 +308,11 @@ _apply_connector_settings() {
 
 _patch_dynamo_args() {
   if [[ -z "${dynamo_args["frontend-node"]}" ]]; then
-    dynamo_args["frontend-node"]=$(echo "${decode_config["node-list"]}" | cut -d',' -f1)
+    if [[ -n "${decode_config["node-list"]}" ]]; then
+      dynamo_args["frontend-node"]=$(echo "${decode_config["node-list"]}" | cut -d',' -f1)
+    else
+      dynamo_args["frontend-node"]=$(echo "${prefill_config["node-list"]}" | cut -d',' -f1)
+    fi
   fi
 
   dynamo_args["url"]="http://${dynamo_args["frontend-node"]}:${dynamo_args["port"]}"
@@ -564,6 +569,19 @@ _current_node_name() {
   echo "${SLURMD_NODENAME:-$(hostname)}"
 }
 
+_select_node_lmcache_config() {
+  [[ -n "${SLURM_JOB_ID:-}" ]] || return 0
+
+  local node_name
+  local node_config
+  node_name="$(_current_node_name)"
+  node_config="${RESULTS_DIR}/lmcache-config-${SLURM_JOB_ID}-${node_name}.yaml"
+  if [[ -f "$node_config" ]]; then
+    export LMCACHE_CONFIG_FILE="$node_config"
+    log "Using node-local LMCache config: $LMCACHE_CONFIG_FILE"
+  fi
+}
+
 _is_frontend_node() {
   local name="$(_current_node_name)"
   [[ ",${dynamo_args["frontend-node"]}," == *",$name,"* ]]
@@ -594,6 +612,7 @@ _is_aiperf_accuracy_enabled() {
 _init_runtime_env() {
   if _is_vllm || _is_sglang; then
     export HF_HOME="${HUGGINGFACE_HOME}"
+    export HF_HUB_CACHE="${HUGGINGFACE_HOME%/}/hub"
     hf cache scan || echo "HF cache scan failed"
   fi
   export NATS_SERVER="nats://${dynamo_args["frontend-node"]}:${dynamo_args["nats-port"]}"
@@ -1432,6 +1451,7 @@ function main()
 
   cd "$RESULTS_DIR" || { log "ERROR: Failed to cd to $RESULTS_DIR"; exit 1; }
 
+  _select_node_lmcache_config
   render_lmcache_config
 
   log_gpu_utilization &