Add functionality to export slurm config bucket as controller NFS

Author: mr0re1

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/setup_network_storage.py

@@ -22,6 +22,7 @@
 import stat
 import time
 import logging
+import uuid
 
 import shutil
 from pathlib import Path
@@ -35,41 +36,56 @@
 
 log = logging.getLogger()
 
-def mounts_by_local(mounts):
+def mounts_by_local(mounts: list[NSMount]) -> dict[str, NSMount]:
     """convert list of mounts to dict of mounts, local_mount as key"""
-    return {str(Path(m.local_mount).resolve()): m for m in mounts}
+    return {str(m.local_mount.resolve()): m for m in mounts}
 
 
-def _get_default_mounts(lkp: util.Lookup) -> List[NSDict]:
+def _get_default_mounts(lkp: util.Lookup) -> list[NSMount]:
     if lkp.cfg.disable_default_mounts:
         return []
     return [
-        NSDict(
-                server_ip= "$controller",
-                remote_mount= path,
-                local_mount= path,
-                fs_type= "nfs",
-                mount_options= "defaults,hard,intr",
+        NSMount(
+                server_ip=lkp.controller_mount_server_ip(),
+                remote_mount=path,
+                local_mount=path,
+                fs_type="nfs",
+                mount_options="defaults,hard,intr",
         )
         for path in (
             dirs.home,
             dirs.apps,
         )
     ]
 
-def resolve_network_storage(nodeset=None) -> List[NSMount]:
+def get_slurm_bucket_mount() -> NSMount:
+    bucket, path = util._get_bucket_and_common_prefix()
+    return  NSMount(
+        fs_type="gcsfuse",
+        server_ip="",
+        remote_mount=Path(bucket),
+        local_mount=dirs.slurm_bucket_mount,
+        mount_options=f"defaults,_netdev,implicit_dirs,only_dir={path}",
+    )
+
+def resolve_network_storage() -> List[NSMount]:
     """Combine appropriate network_storage fields to a single list"""
     lkp = lookup()
 
     # create dict of mounts, local_mount: mount_info
     mounts = mounts_by_local(_get_default_mounts(lkp))
 
+    if lkp.is_controller and util.should_mount_slurm_bucket():
+        mounts.update(mounts_by_local([get_slurm_bucket_mount()]))
+
     # On non-controller instances, entries in network_storage could overwrite
     # default exports from the controller. Be careful, of course
-    mounts.update(mounts_by_local(lkp.cfg.network_storage))
+    common = [lkp.normalize_ns_mount(m) for m in lkp.cfg.network_storage]
+    mounts.update(mounts_by_local(common))
 
     if lkp.is_login_node:
-        login_ns = lkp.cfg.login_groups[util.instance_login_group()].network_storage
+        login_group = lkp.cfg.login_groups[util.instance_login_group()] 
+        login_ns = [lkp.normalize_ns_mount(m) for m in login_group.network_storage]
         mounts.update(mounts_by_local(login_ns))
 
     if lkp.instance_role == "compute":
@@ -78,74 +94,46 @@ def resolve_network_storage(nodeset=None) -> List[NSMount]:
         except Exception:
             pass # external nodename, skip lookup
         else:
-            mounts.update(mounts_by_local(nodeset.network_storage))
-
-    return [lkp.normalize_ns_mount(mnt) for mnt in mounts.values()]
-
+            nodeset_ns = [lkp.normalize_ns_mount(m) for m in nodeset.network_storage]
+            mounts.update(mounts_by_local(nodeset_ns))
 
-def separate_external_internal_mounts(mounts):
-    """separate into cluster-external and internal mounts"""
+    return list(mounts.values())
 
-    def internal_mount(mount):
-        # NOTE: Valid Lustre server_ip can take the form of '<IP>@tcp'
-        server_ip = mount.server_ip.split("@")[0]
-        mount_addr = util.host_lookup(server_ip)
-        return mount_addr == lookup().control_host_addr
-
-    return separate(internal_mount, mounts)
 
+def is_controller_mount(mount) -> bool:
+    # NOTE: Valid Lustre server_ip can take the form of '<IP>@tcp'
+    server_ip = mount.server_ip.split("@")[0]
+    mount_addr = util.host_lookup(server_ip)
+    return mount_addr == lookup().control_host_addr
 
 def setup_network_storage():
     """prepare network fs mounts and add them to fstab"""
     log.info("Set up network storage")
-    # filter mounts into two dicts, cluster-internal and external mounts
-
+    
     all_mounts = resolve_network_storage()
-    ext_mounts, int_mounts = separate_external_internal_mounts(all_mounts)
-
     if lookup().is_controller:
-        mounts = ext_mounts
+        mounts, _ = separate(is_controller_mount, all_mounts)
     else:
-        mounts = ext_mounts + int_mounts
+        mounts = all_mounts
 
     # Determine fstab entries and write them out
     fstab_entries = []
     for mount in mounts:
-        local_mount = Path(mount.local_mount)
-        remote_mount = mount.remote_mount
+        local_mount = mount.local_mount
         fs_type = mount.fs_type
         server_ip = mount.server_ip or ""
+        src = mount.remote_mount if fs_type == "gcsfuse" else f"{server_ip}:{mount.remote_mount}"
+        
+        log.info(f"Setting up mount ({fs_type}) {src} to {local_mount}")
         util.mkdirp(local_mount)
 
-        log.info(
-            "Setting up mount ({}) {}{} to {}".format(
-                fs_type,
-                server_ip + ":" if fs_type != "gcsfuse" else "",
-                remote_mount,
-                local_mount,
-            )
-        )
-
         mount_options = mount.mount_options.split(",") if mount.mount_options else []
-        if not mount_options or "_netdev" not in mount_options:
+        if "_netdev" not in mount_options:
             mount_options += ["_netdev"]
-
-        if fs_type == "gcsfuse":
-            fstab_entries.append(
-                "{0}   {1}     {2}     {3}     0 0".format(
-                    remote_mount, local_mount, fs_type, ",".join(mount_options)
-                )
-            )
-        else:
-            fstab_entries.append(
-                "{0}:{1}    {2}     {3}      {4}  0 0".format(
-                    server_ip,
-                    remote_mount,
-                    local_mount,
-                    fs_type,
-                    ",".join(mount_options),
-                )
-            )
+        options_line = ",".join(mount_options)
+        
+        
+        fstab_entries.append(f"{src}   {local_mount}     {fs_type}     {options_line}     0 0")
 
     fstab = Path("/etc/fstab")
     if not Path(fstab.with_suffix(".bak")).is_file():
@@ -157,16 +145,16 @@ def setup_network_storage():
             f.write(entry)
             f.write("\n")
 
-    mount_fstab(mounts_by_local(mounts), log)
+    mount_fstab(mounts, log)
     if lookup().cfg.enable_slurm_auth:
       slurm_key_mount_handler()
     else:
       munge_mount_handler()
 
 
-def mount_fstab(mounts, log):
+def mount_fstab(mounts: list[NSMount], log):
     """Wait on each mount, then make sure all fstab is mounted"""
-    def mount_path(path):
+    def mount_path(path: Path):
         log.info(f"Waiting for '{path}' to be mounted...")
         try:
             run(f"mount {path}", timeout=120)
@@ -184,8 +172,8 @@ def mount_path(path):
     with Executors.thread_pool().with_timeout(MAX_MOUNT_TIMEOUT).with_retry(
         retry_policy=retry_policy
     ) as exe:
-        for path in mounts:
-            future = exe.submit(mount_path, path)
+        for m in mounts:
+            future = exe.submit(mount_path, m.local_mount)
             future_list.append(future)
 
         # Iterate over futures, checking for exceptions
@@ -306,36 +294,34 @@ def setup_nfs_exports():
     lkp = util.lookup()
     assert lkp.is_controller
 
-
     # The controller only needs to set up exports for cluster-internal mounts
-    # switch the key to remote mount path since that is what needs exporting
-    mounts = resolve_network_storage()
-
-    if lkp.cfg.enable_slurm_auth:
-      mounts.append(lkp.slurm_key_mount)
-    else:
-      mounts.append(lkp.munge_mount)
-
-    # controller mounts
-    _, con_mounts = separate_external_internal_mounts(mounts)
-    con_mounts = {m.remote_mount: m for m in con_mounts}
-    for nodeset in lkp.cfg.nodeset.values():
-        # get internal mounts for each nodeset by calling
-        # resolve_network_storage as from a node in each nodeset
-        ns_mounts = resolve_network_storage(nodeset=nodeset)
-        _, int_mounts = separate_external_internal_mounts(ns_mounts)
-        con_mounts.update({m.remote_mount: m for m in int_mounts})
+    exported_mounts = [m for m in resolve_network_storage() if is_controller_mount(m)]
+
+    # key by remote mount path since that is what needs exporting
+    to_export = {m.remote_mount: "*(rw,no_subtree_check,no_root_squash)" for m in exported_mounts}
+
+    key_mount = lkp.slurm_key_mount if lkp.cfg.enable_slurm_auth else lkp.munge_mount
+    if is_controller_mount(key_mount):
+        # Export key mount as read-only
+        to_export[key_mount.remote_mount] = "*(ro,no_subtree_check,no_root_squash)"
+    
+    if util.should_mount_slurm_bucket():
+        mnt = get_slurm_bucket_mount()
+        # FSID is required for virtual filesystem that is not based on a device
+        # Also export it as read-only
+        fsid=str(uuid.uuid4())
+        to_export[mnt.local_mount] = f"*(ro,no_subtree_check,no_root_squash,fsid={fsid})"
 
     # export path if corresponding selector boolean is True
-    exports = []
-    for path in con_mounts:
+    lines = []
+    for path,options in to_export.items():
         util.mkdirp(Path(path))
         run(rf"sed -i '\#{path}#d' /etc/exports", timeout=30)
-        exports.append(f"{path}  *(rw,no_subtree_check,no_root_squash)")
+        lines.append(f"{path}  {options}")
 
     exportsd = Path("/etc/exports.d")
     util.mkdirp(exportsd)
     with (exportsd / "slurm.exports").open("w") as f:
         f.write("\n")
-        f.write("\n".join(exports))
+        f.write("\n".join(lines))
     run("exportfs -a", timeout=30)

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py

@@ -14,8 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from importlib import metadata
-from typing import Iterable, List, Tuple, Optional, Any, Dict, Sequence, Type, Callable, Literal
+from typing import Iterable, List, Tuple, Optional, Any, Dict, Sequence, Type, Callable, Union
 import argparse
 import base64
 from dataclasses import dataclass, field
@@ -34,7 +33,6 @@
 import socket
 import subprocess
 import sys
-import tempfile
 from enum import Enum
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -58,8 +56,7 @@
 
 import google.api_core.exceptions as gExceptions
 
-from requests import get as get_url
-from requests.exceptions import RequestException
+import requests as requests_lib
 
 import yaml
 from addict import Dict as NSDict # type: ignore
@@ -93,6 +90,7 @@ def mkdirp(path: Path) -> None:
     munge = Path("/etc/munge"),
     secdisk = Path("/mnt/disks/sec"),
     log = Path("/var/log/slurm"),
+    slurm_bucket_mount = Path("/slurm/bucket"),
 )
 
 slurmdirs = NSDict(
@@ -428,6 +426,13 @@ def map_with_futures(func, seq):
                 res = e
             yield res
 
+def should_mount_slurm_bucket() -> bool:
+    try:
+        return instance_metadata("attributes/slurm_bucket_mount").lower() == "true"
+    except MetadataNotFoundError:
+        return False
+
+
 def _get_bucket_and_common_prefix() -> Tuple[str, str]:
     uri = instance_metadata("attributes/slurm_bucket_path")
     return parse_bucket_uri(uri)
@@ -1066,18 +1071,20 @@ def backoff_delay(start, timeout=None, ratio=None, count: int = 0):
 
 ROOT_URL = "http://metadata.google.internal/computeMetadata/v1"
 
+class MetadataNotFoundError(Exception):
+    pass
 
 def get_metadata(path, root=ROOT_URL):
     """Get metadata relative to metadata/computeMetadata/v1"""
     HEADERS = {"Metadata-Flavor": "Google"}
     url = f"{root}/{path}"
     try:
-        resp = get_url(url, headers=HEADERS)
+        resp = requests_lib.get(url, headers=HEADERS)
         resp.raise_for_status()
         return resp.text
-    except RequestException:
-        log.debug(f"metadata not found ({url})")
-        raise Exception(f"failed to get_metadata from {url}")
+    except requests_lib.exceptions.HTTPError:
+        log.exception(f"metadata not found ({url})")
+        raise MetadataNotFoundError(f"failed to get_metadata from {url}")
 
 
 @lru_cache(maxsize=None)
@@ -1925,10 +1932,16 @@ def job(self, job_id: int) -> Optional[Job]:
     def etc_dir(self) -> Path:
         return Path(self.cfg.output_dir or slurmdirs.etc)
 
-    def normalize_ns_mount(self, ns: Dict[str, str]) -> NSMount:
+    def controller_mount_server_ip(self) -> str:
+        return self.control_addr or self.control_host
+
+    def normalize_ns_mount(self, ns: Union[dict, NSMount]) -> NSMount:
+        if isinstance(ns, NSMount):
+            return ns
+
         server_ip = ns.get("server_ip") or "$controller"
         if server_ip == "$controller":
-            server_ip = self.control_addr or self.control_host
+            server_ip = self.controller_mount_server_ip()
 
         return NSMount(
             server_ip=server_ip,
@@ -1943,30 +1956,30 @@ def munge_mount(self) -> NSMount:
         if self.cfg.munge_mount:
             mnt = self.cfg.munge_mount
             mnt.local_mount = mnt.local_mount or "/mnt/munge"
+            return self.normalize_ns_mount(mnt)
         else:
-            mnt = NSDict(
-                server_ip="$controller",
-                local_mount="/mnt/munge",
+            return NSMount(
+                server_ip=self.controller_mount_server_ip(),
+                local_mount=Path("/mnt/munge"),
                 remote_mount=dirs.munge,
                 fs_type="nfs",
                 mount_options="defaults,hard,intr,_netdev",
             )
-        return self.normalize_ns_mount(mnt)
 
     @property
     def slurm_key_mount(self) -> NSMount:
         if self.cfg.slurm_key_mount:
             mnt = self.cfg.slurm_key_mount
             mnt.local_mount = mnt.local_mount or slurmdirs.key_distribution
+            return self.normalize_ns_mount(mnt)
         else:
-            mnt = NSDict(
-                server_ip="$controller",
+            return NSMount(
+                server_ip=self.controller_mount_server_ip(),
                 local_mount=slurmdirs.key_distribution,
                 remote_mount=slurmdirs.key_distribution,
                 fs_type="nfs",
                 mount_options="defaults,hard,intr,_netdev",
             )
-        return self.normalize_ns_mount(mnt)
 
     def is_flex_node(self, node: str) -> bool:
         try: