Merge pull request #1421 from Dstack-TEE/fix/annotation-cli-duplicates

fix(vmm): refuse repeated annotation keys; carry annotations in the UI

Author: kvinwang

dstack/vmm/src/vmm-cli.py

@@ -372,17 +372,19 @@ def encrypt_env(envs, hex_public_key: str) -> str:
     return result.hex()
 
 
-def parse_annotations(items: Optional[List[str]]) -> Dict[str, str]:
-    """Parse repeated KEY=VALUE annotation arguments."""
-    annotations = {}
-    for item in items or []:
-        key, sep, value = item.partition("=")
+class AnnotationAction(argparse.Action):
+    """Collect repeated KEY=VALUE arguments, refusing a key given twice."""
+
+    def __call__(self, parser, namespace, value, option_string=None):
+        """Add one annotation to the namespace."""
+        key, sep, val = value.partition("=")
         if not sep or not key:
-            raise argparse.ArgumentTypeError(
-                f"invalid annotation {item!r}, expected KEY=VALUE"
-            )
-        annotations[key] = value
-    return annotations
+            parser.error(f"{option_string}: expected KEY=VALUE, got {value!r}")
+        annotations = getattr(namespace, self.dest) or {}
+        if key in annotations:
+            parser.error(f"{option_string}: {key!r} given twice")
+        annotations[key] = val
+        setattr(namespace, self.dest, annotations)
 
 
 def parse_port_mapping(port_str: str) -> Dict:
@@ -996,7 +998,7 @@ def create_vm(self, args) -> None:
             "pin_numa": args.pin_numa,
             "stopped": args.stopped,
             "no_tee": args.no_tee,
-            "annotations": parse_annotations(args.annotation),
+            "annotations": args.annotation or {},
         }
         if args.disk_prealloc:
             params["disk_prealloc"] = args.disk_prealloc
@@ -2049,7 +2051,7 @@ def _patched_format_help():
     )
     deploy_parser.add_argument(
         "--annotation",
-        action="append",
+        action=AnnotationAction,
         metavar="KEY=VALUE",
         help="Annotation for host-side tooling (can be repeated)",
     )
@@ -2326,7 +2328,7 @@ def _patched_format_help():
     annotations_group = update_parser.add_mutually_exclusive_group()
     annotations_group.add_argument(
         "--annotation",
-        action="append",
+        action=AnnotationAction,
         metavar="KEY=VALUE",
         help="Replace the annotations with these (can be repeated)",
     )
@@ -2430,13 +2432,7 @@ def _patched_format_help():
             net_vhost=args.net_vhost,
             net_vhost_inherit=getattr(args, "net_vhost_inherit", False),
             net_queues=args.net_queues,
-            annotations=(
-                {}
-                if args.no_annotations
-                else parse_annotations(args.annotation)
-                if args.annotation
-                else None
-            ),
+            annotations={} if args.no_annotations else args.annotation,
         )
     elif args.command == "kms":
         if not args.kms_action:

dstack/vmm/ui/src/components/CreateVmDialog.ts

@@ -197,6 +197,11 @@ const CreateVmDialogComponent = {
               <textarea id="userConfig" v-model="form.user_config" placeholder="Optional user config placed at /dstack/.user-config in the CVM"></textarea>
             </div>
 
+            <div class="form-group full-width">
+              <label for="annotations">Annotations</label>
+              <textarea id="annotations" v-model="form.annotations" placeholder="Optional KEY=VALUE per line, for host-side tooling; never passed to the CVM"></textarea>
+            </div>
+
             <div class="form-group full-width" v-if="availableGpus.length > 0">
               <gpu-config-editor
                 :available-gpus="availableGpus"

dstack/vmm/ui/src/composables/useVmManager.ts

@@ -36,6 +36,29 @@ type AppCompose = {
   event_log_version?: number;
 };
 
+/** Parses one `KEY=VALUE` annotation per line, refusing a key given twice. */
+function parseAnnotations(text: string): Record<string, string> {
+  const annotations = new Map<string, string>();
+  for (const line of text.split('\n').map((l) => l.trim()).filter(Boolean)) {
+    const eq = line.indexOf('=');
+    if (eq <= 0) {
+      throw new Error(`Invalid annotation "${line}", expected KEY=VALUE`);
+    }
+    const key = line.slice(0, eq);
+    if (annotations.has(key)) {
+      throw new Error(`Annotation "${key}" given twice`);
+    }
+    annotations.set(key, line.slice(eq + 1));
+  }
+  return Object.fromEntries(annotations);
+}
+
+function annotationsText(annotations: Record<string, string> | null | undefined): string {
+  return Object.entries(annotations || {})
+    .map(([key, value]) => `${key}=${value}`)
+    .join('\n');
+}
+
 function initScriptList(initScript: string | string[] | undefined): string[] {
   const scripts = Array.isArray(initScript)
     ? [...initScript]
@@ -167,6 +190,7 @@ type VmFormState = {
   disk_prealloc: string;
   networks: NetworkFormEntry[];
   user_config: string;
+  annotations: string;
   kms_urls: string[];
   gateway_urls: string[];
   stopped: boolean;
@@ -259,6 +283,7 @@ function createVmFormState(preLaunchScript: string): VmFormState {
     disk_prealloc: '',
     networks: [],
     user_config: '',
+    annotations: '',
     kms_urls: [],
     gateway_urls: [],
     stopped: false,
@@ -630,6 +655,7 @@ type CreateVmPayloadSource = {
     kms_urls?: string[];
     gateway_urls?: string[];
     stopped?: boolean;
+    annotations?: Record<string, string>;
   };
 
   function buildCreateVmPayload(source: CreateVmPayloadSource): VmmTypes.IVmConfiguration {
@@ -657,6 +683,7 @@ type CreateVmPayloadSource = {
       kms_urls: source.kms_urls?.filter((url) => url && url.trim().length) ?? [],
       gateway_urls: source.gateway_urls?.filter((url) => url && url.trim().length) ?? [],
       stopped: !!source.stopped,
+      annotations: source.annotations ?? {},
     };
   }
 
@@ -1126,6 +1153,8 @@ type CreateVmPayloadSource = {
     // and inheriting "full" unannounced costs the next deploy a full write of
     // its disk before the VM boots.
     vmForm.value.disk_prealloc = '';
+    // And a clone's annotations are parameters for that VM's host tooling.
+    vmForm.value.annotations = '';
     vmForm.value.swapValue = 0;
     vmForm.value.swapUnit = 'GB';
     vmForm.value.swap_size = 0;
@@ -1265,6 +1294,7 @@ type CreateVmPayloadSource = {
         kms_urls: vmForm.value.kms_urls,
         gateway_urls: vmForm.value.gateway_urls,
         stopped: vmForm.value.stopped,
+        annotations: parseAnnotations(vmForm.value.annotations),
       });
 
       await vmmRpc.createVm(payload);
@@ -1420,6 +1450,7 @@ type CreateVmPayloadSource = {
       no_tee: !!config.no_tee,
       simulated_tee: config.simulated_tee || '',
       user_config: config.user_config || '',
+      annotations: annotationsText(config.annotations),
       stopped: !!config.stopped,
       event_log_version: theVm.appCompose?.event_log_version || 1,
     };

dstack/vmm/ui/src/templates/app.html

@@ -483,6 +483,14 @@ <h4>User Config</h4>
           <pre class="user-config-content">{{ vm.configuration.user_config }}</pre>
         </div>
 
+        <div v-if="Object.keys(vm.configuration?.annotations || {}).length" class="user-config-section">
+          <div class="section-header">
+            <h4>Annotations</h4>
+          </div>
+          <pre class="user-config-content"><template v-for="(value, key) in vm.configuration.annotations">{{ key }}={{ value }}
+</template></pre>
+        </div>
+
         <div class="vm-log-tabs">
           <button class="vm-log-button" @click="showLogs(vm.id, 'serial')">Serial Logs</button>
           <button class="vm-log-button" @click="showLogs(vm.id, 'stdout')">Stdout</button>