Creating jobspec via the Python API for cores and GPUs that share a socket

[SteVwonder]

The following script (specifically the insert_socket function) will insert a socket resource just above the slot resource. If you pass non-None values for gpus_per_task and cores_per_task to the JobspecV1.from_command constructor, you will get a jobspec that requests cores and GPUs that share a socket.

#!/usr/bin/flux python

import flux
from flux.job import JobspecV1

def insert_socket(resource_list):
    for idx, resource in enumerate(resource_list):
        try:
            if resource["type"] == "slot":
                resource_list[idx] = {
                    "type": "socket",
                    "count": 1,
                    "with": [resource],
                }
                return
            else:
                insert_socket(resource["with"])
        except (IndexError, KeyError):
            return

jobspec = JobspecV1.from_command(["hostname"],
                                 num_tasks=1,
                                 cores_per_task=2,
                                 gpus_per_task=2,
                                 num_nodes=1)
insert_socket(jobspec.jobspec["resources"])
print(jobspec.dumps())

ƒ(s=1,d=0) fluxuser@23763266b859:/src$ ./core_gpu_samesocket.py | jq .resources
[
  {
    "type": "node",
    "count": 1,
    "with": [
      {
        "type": "socket",
        "count": 1,
        "with": [
          {
            "type": "slot",
            "count": 1,
            "with": [
              {
                "type": "core",
                "count": 2
              },
              {
                "type": "gpu",
                "count": 2
              }
            ],
            "label": "task"
          }
        ]
      }
    ]
  }
]

#3151 prevents this jobspec from being scheduled and executed currently. Once that is closed, this should work for your MuMMi use case @FrankD412.

[SteVwonder]

I confirmed that with the merging of #3151, this is now possible in v0.19.0 with some small modifications to the original example. It turns out that hwloc v1 doesn't put the GPUs and Cores under the same socket but instead under the same Numanode, so after swapping socket out for numanode in the generated jobspec, this appears to work as intended on Lassen.

The magic run line is (note the FLUXION env var at the start which explicitly tells it to allow numanodes in the resource graph):

FLUXION_RESOURCE_OPTIONS="load-allowlist=node,core,gpu,numanode" PMIX_MCA_gds="^ds12,ds21" jsrun -a 1 -c ALL_CPUS -g ALL_GPUS -n 1 --bind=none --smpiargs="-disable_gpu_hooks" flux start ./submit-get-R.sh

Where submit-get-R.sh is just a simple driver of the example demo:

#!/bin/bash

jobspec=$(flux python samenuma.py)

JOBID1=$(echo $jobspec | flux job submit)
echo "Submitted $JOBID1"
JOBID2=$(echo $jobspec | flux job submit)
echo "Submitted $JOBID2"

for id in $JOBID1 $JOBID2; do
    echo "Output for $id:"
    flux job attach $id
    echo "Concretized R for $id:"
    flux job info $id R | jq .
done

samenuma.py is:

#!/usr/bin/flux python

import flux
from flux.job import JobspecV1

def insert_numanode(resource_list):
    for idx, resource in enumerate(resource_list):
        try:
            if resource["type"] == "slot":
                resource_list[idx] = {
                    "type": "numanode",
                    "count": 1,
                    "with": [resource],
                }
                return
            else:
                insert_numanode(resource["with"])
        except (IndexError, KeyError):
            return

jobspec = JobspecV1.from_command(["bash", "-c", "printenv CUDA_VISIBLE_DEVICES; sleep 3"],
                                 num_tasks=1,
                                 cores_per_task=2,
                                 gpus_per_task=2,
                                 num_nodes=1)
insert_numanode(jobspec.jobspec["resources"])
jobspec.duration = 4
jobspec.setattr_shell_option("gpu-affinity", "per-task")
print(jobspec.dumps())

And the output of the whole thing is:

Submitted ƒgU3t7R
Submitted ƒhkBGCf
Output for ƒgU3t7R:
2-3
Concretized R for ƒgU3t7R:
{
  "version": 1,
  "execution": {
    "R_lite": [
      {
        "rank": "0",
        "node": "lassen1",
        "children": {
          "core": "38-39",
          "gpu": "2-3"
        }
      }
    ],
    "starttime": 1600825654,
    "expiration": 1600825665
  }
}
Output for ƒhkBGCf:
0-1
Concretized R for ƒhkBGCf:
{
  "version": 1,
  "execution": {
    "R_lite": [
      {
        "rank": "0",
        "node": "lassen1",
        "children": {
          "core": "18-19",
          "gpu": "0-1"
        }
      }
    ],
    "starttime": 1600825654,
    "expiration": 1600825665
  }
}

Without the numanode being included in the jobspec, the second job would have been assigned cores 36-37.