From dabb69b2224b5b66b04040957bd61985a5a64969 Mon Sep 17 00:00:00 2001
From: Camille Scott <cswel@ucdavis.edu>
Date: Wed, 18 Sep 2024 17:50:10 -0700
Subject: [PATCH] Fix slurm sacctmgr account add

---
 cheeto/database.py | 24 +++++++++++++++++++++---
 cheeto/slurm.py    | 11 ++++++-----
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/cheeto/database.py b/cheeto/database.py
index fee7947..caf2ec3 100644
--- a/cheeto/database.py
+++ b/cheeto/database.py
@@ -539,6 +539,21 @@ class SiteGroup(BaseDocument):
     ldap_synced = BooleanField(default=False)
     iam_synced = BooleanField(default=False)
 
+    meta = {
+        'indexes': [
+            {
+                'fields': ('groupname', 'sitename'),
+                'unique': True
+            },
+            {
+                'fields': ['_members']
+            },
+            {
+                'fields': ['_slurmers']
+            }
+        ]
+    }
+
     @property
     def gid(self):
         return self.parent.gid
@@ -1513,7 +1528,9 @@ def group_to_puppet(group: SiteGroup):
          'group': s.group,
          'autofs': {
             'nas': s.host,
-            'path': str(s.host_path)},
+            'path': str(s.host_path),
+            'options': ','.join(sorted(set(s.mount_options) - {'fstype=nfs'}, reverse=True))
+         },
          'zfs': get_puppet_zfs(s),
          'globus': s.globus
         } for s in storages
@@ -1630,12 +1647,13 @@ def site_write_to_puppet(args: argparse.Namespace):
 @subcommand('sync-old-puppet',
             add_site_args_req,
             lambda parser: parser.add_argument('repo', type=Path),
+            lambda parser: parser.add_argument('--base-branch', default='main'),
             lambda parser: parser.add_argument('--push-merge', default=False, action='store_true'))
 def site_sync_old_puppet(args: argparse.Namespace):
     connect_to_database(args.config.mongo)
 
     site = Site.objects.get(sitename=args.site)
-    repo = GitRepo(args.repo)
+    repo = GitRepo(args.repo, base_branch=args.base_branch)
     prefix = (args.repo / 'domains' / site.fqdn).absolute()
     yaml_path = prefix / 'merged' / 'all.yaml'
     puppet_map = site_to_puppet(args.site)
@@ -1652,7 +1670,7 @@ def site_sync_old_puppet(args: argparse.Namespace):
             with keyfile.open('w') as fp:
                 for key in user.parent.ssh_key:
                     print(key, file=fp)
-        add(args.repo)
+        add(args.repo.absolute())
 
 
 @subcommand('to-ldap',
diff --git a/cheeto/slurm.py b/cheeto/slurm.py
index 93ec7bd..bb86410 100644
--- a/cheeto/slurm.py
+++ b/cheeto/slurm.py
@@ -91,11 +91,12 @@ def __init__(self, sacctmgr_path: Optional[str] = None,
         self.show = self.cmd.bake('show', '-P')
 
     def add_account(self, account_name: str,
-                          max_user_jobs: int = -1,
-                          max_group_jobs: int = -1) -> sh.Command:
-        args = ['account', account_name]
-        args.append(f'MaxJobs={max_user_jobs}')
-        args.append(f'MaxSubmit={max_group_jobs}')
+                          account: SlurmAccount) -> sh.Command:
+        args = ['account', account_name,
+                f'MaxJobs={account.max_user_jobs}',
+                f'GrpJobs={account.max_group_jobs}',
+                f'MaxWall={account.max_job_length}',
+                f'MaxSubmit={account.max_submit_jobs}']
         return self.add.bake(*args)
 
     def modify_account(self, account_name: str,