no reservation, num_agents->num_nodes

NERSC · Feb 10, 2025 · 35a8320 · 35a8320
1 parent 5c9d799
commit 35a8320
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 38 deletions.
diff --git a/backend/micro/launcher/interactem/launcher/launcher.py b/backend/micro/launcher/interactem/launcher/launcher.py
@@ -114,6 +114,10 @@ async def submit(msg: NATSMsg, js: JetStreamContext) -> None:
     create_task_with_ref(task_refs, msg.in_progress())
     try:
         agent_event = AgentCreateEvent.model_validate_json(msg.data)
+        reservation: str | None = None
+
+        if agent_event.extra:
+            reservation = agent_event.extra.get("reservation", None)
 
         # Convert to JobSubmitEvent because we want to limit the frontend
         # to not knowing what a job is...
@@ -123,8 +127,8 @@ async def submit(msg: NATSMsg, js: JetStreamContext) -> None:
             qos=cfg.SFAPI_QOS,
             constraint=agent_event.compute_type.value,
             walltime=agent_event.duration,
-            reservation=agent_event.reservation,
-            num_nodes=agent_event.num_agents,
+            reservation=reservation,
+            num_nodes=agent_event.num_nodes,
         )
     except ValidationError as e:
         logger.error(f"Failed to parse job request: {e}")

diff --git a/backend/sfapi_models/interactem/sfapi_models/__init__.py b/backend/sfapi_models/interactem/sfapi_models/__init__.py
@@ -1,5 +1,6 @@
 import datetime
 from enum import Enum
+from typing import Any
 
 from pydantic import BaseModel, model_validator
 from sfapi_client._models import StatusValue
@@ -23,8 +24,8 @@ class AgentCreateEvent(BaseModel):
     machine: Machine
     duration: datetime.timedelta
     compute_type: ComputeType
-    num_agents: int
-    reservation: str | None = None
+    num_nodes: int
+    extra: dict[str, Any] | None = None
 
 
 class JobSubmitEvent(BaseModel):

diff --git a/frontend/interactEM/openapi.json b/frontend/interactEM/openapi.json
@@ -809,14 +809,14 @@
             "title": "Duration"
           },
           "compute_type": { "$ref": "#/components/schemas/ComputeType" },
-          "num_agents": { "type": "integer", "title": "Num Agents" },
-          "reservation": {
-            "anyOf": [{ "type": "string" }, { "type": "null" }],
-            "title": "Reservation"
+          "num_nodes": { "type": "integer", "title": "Num Nodes" },
+          "extra": {
+            "anyOf": [{ "type": "object" }, { "type": "null" }],
+            "title": "Extra"
           }
         },
         "type": "object",
-        "required": ["machine", "duration", "compute_type", "num_agents"],
+        "required": ["machine", "duration", "compute_type", "num_nodes"],
         "title": "AgentCreateEvent"
       },
       "Body_login-login_access_token": {

diff --git a/frontend/interactEM/src/client/generated/types.gen.ts b/frontend/interactEM/src/client/generated/types.gen.ts
@@ -4,8 +4,10 @@ export type AgentCreateEvent = {
   machine: PublicHost
   duration: string
   compute_type: ComputeType
-  num_agents: number
-  reservation?: string | null
+  num_nodes: number
+  extra?: {
+    [key: string]: unknown
+  } | null
 }
 
 export type BodyLoginLoginAccessToken = {

diff --git a/frontend/interactEM/src/client/generated/zod.gen.ts b/frontend/interactEM/src/client/generated/zod.gen.ts
@@ -6,8 +6,8 @@ export const zAgentCreateEvent = z.object({
   machine: z.enum(["dtn01", "dtns", "perlmutter"]),
   duration: z.string(),
   compute_type: z.enum(["gpu", "cpu"]),
-  num_agents: z.number().int(),
-  reservation: z.union([z.string(), z.null()]).optional(),
+  num_nodes: z.number().int(),
+  extra: z.union([z.object({}), z.null()]).optional(),
 })
 
 export const zBodyLoginLoginAccessToken = z.object({

diff --git a/frontend/interactEM/src/components/launchagentfab.tsx b/frontend/interactEM/src/components/launchagentfab.tsx
@@ -40,8 +40,7 @@ export const LaunchAgentFab = () => {
       machine: "perlmutter",
       compute_type: "cpu",
       duration: "01:00:00",
-      reservation: "",
-      num_agents: 1,
+      num_nodes: 1,
     },
   })
   const onSubmit: SubmitHandler<AgentCreateEvent> = useCallback(
@@ -161,38 +160,19 @@ export const LaunchAgentFab = () => {
 
             <FormControl fullWidth sx={{ marginBottom: "1rem" }}>
               <Controller
-                name="reservation"
+                name="num_nodes"
                 control={control}
                 render={({ field }) => (
                   <TextField
-                    label="Reservation (optional)"
+                    label="Number of Nodes"
                     variant="outlined"
                     onChange={field.onChange}
                     onBlur={field.onBlur}
                     value={field.value}
                     fullWidth
-                    error={!!errors.reservation}
-                    helperText={errors.reservation?.message}
-                  />
-                )}
-              />
-            </FormControl>
-
-            <FormControl fullWidth sx={{ marginBottom: "1rem" }}>
-              <Controller
-                name="num_agents"
-                control={control}
-                render={({ field }) => (
-                  <TextField
-                    label="Number of Agents"
-                    variant="outlined"
-                    onChange={field.onChange}
-                    onBlur={field.onBlur}
-                    value={field.value}
-                    fullWidth
-                    error={!!errors.num_agents}
+                    error={!!errors.num_nodes}
                     helperText={
-                      errors.num_agents ? "This should be an integer." : ""
+                      errors.num_nodes ? "This should be an integer." : ""
                     }
                   />
                 )}