22#
33# SPDX-License-Identifier: AGPL-3.0-or-later
44
5- import sys
65import logging
6+ import sys
77
8+ from slurm_rest_api_client import Client
9+ from slurm_rest_api_client .api .slurm import slurm_v0043_delete_jobs
10+ from slurm_rest_api_client .api .slurmdb import (
11+ slurmdb_v0043_delete_account ,
12+ slurmdb_v0043_delete_association ,
13+ slurmdb_v0043_get_accounts ,
14+ slurmdb_v0043_get_associations ,
15+ slurmdb_v0043_post_associations ,
16+ )
17+ from slurm_rest_api_client .models .v0043_account import V0043Account
18+ from slurm_rest_api_client .models .v0043_assoc import V0043Assoc
19+ from slurm_rest_api_client .models .v0043_assoc_max import V0043AssocMax
20+ from slurm_rest_api_client .models .v0043_assoc_max_jobs import V0043AssocMaxJobs
21+ from slurm_rest_api_client .models .v0043_kill_jobs_msg import V0043KillJobsMsg
22+ from slurm_rest_api_client .models .v0043_kill_jobs_resp_job import V0043KillJobsRespJob
23+ from slurm_rest_api_client .models .v0043_openapi_assocs_resp import V0043OpenapiAssocsResp
24+ from slurm_rest_api_client .models .v0043_uint_32_no_val_struct import V0043Uint32NoValStruct
825from tenacity import (
26+ before_sleep_log ,
927 retry ,
10- wait_exponential ,
11- stop_after_attempt ,
1228 retry_if_not_exception_type ,
13- before_sleep_log ,
29+ stop_after_attempt ,
30+ wait_exponential ,
1431)
15- from slurm_rest_api_client import Client
16- from slurm_rest_api_client .models .v0043_account import V0043Account
17- from slurm_rest_api_client .api .slurmdb import slurmdb_v0043_get_accounts
18-
1932
2033logger = logging .getLogger (__name__ )
2134
2235
2336def log_request (request ):
24- logger .info (f"Request event hook: { request .method } { request .url } )
37+ logger .debug (f"Request event hook: { request .method } { request .url } )
2538
2639
2740def log_response (response ):
2841 request = response .request
29- logger .info (f"Response event hook: { request .method } { request .url } { response .status_code } )
42+ logger .debug (f"Response event hook: { request .method } { request .url } { response .status_code } )
3043
3144
3245class SlurmCluster :
3346 def __init__ (self , endpoint , token ):
47+ self .endpoint = endpoint
48+ self .token = token
49+ # for most operations, re-use the root client
3450 self .root_client : Client = Client (
35- base_url = endpoint ,
51+ base_url = self . endpoint ,
3652 headers = {
3753 "X-SLURM-USER-NAME" : "root" ,
38- "X-SLURM-USER-TOKEN" : token ,
54+ "X-SLURM-USER-TOKEN" : self . token ,
3955 },
4056 httpx_args = {"event_hooks" : {"request" : [log_request ], "response" : [log_response ]}},
4157 )
4258
59+ # use as context manager resource to clean up client
60+ def __enter__ (self ):
61+ return self
62+
63+ def __exit__ (self , exc_type , exc_value , traceback ):
64+ self .root_client .get_httpx_client ().close ()
65+ return False
66+
4367 @retry (
4468 wait = wait_exponential (multiplier = 2 , min = 2 , max = 10 ),
4569 stop = stop_after_attempt (10 ),
@@ -48,9 +72,122 @@ def __init__(self, endpoint, token):
4872 retry_error_callback = lambda _ : sys .exit (1 ), # exit if SLURM cannot be reached :(
4973 )
5074 def get_accounts (self ) -> list [V0043Account ]:
51- with self .root_client as client :
52- resp = slurmdb_v0043_get_accounts .sync (client = client , with_associations = str ("true" ))
53- if resp :
54- return resp .accounts
75+ resp = slurmdb_v0043_get_accounts .sync (client = self .root_client , with_associations = str ("true" ))
76+ if resp :
77+ return resp .accounts
78+ else :
79+ raise ConnectionError ("Could not get list of accounts from SLURM endpoint" )
80+
81+ # retry the whole block atomically
82+ @retry (
83+ wait = wait_exponential (multiplier = 2 , min = 2 , max = 10 ),
84+ stop = stop_after_attempt (3 ),
85+ retry = retry_if_not_exception_type (ConnectionError ),
86+ before_sleep = before_sleep_log (logging .getLogger (__name__ ), logging .WARNING ),
87+ )
88+ def delete_association_user_account (self , username : str , account : str , noop : bool = False ) -> None :
89+ if noop :
90+ logging .info (f"noop enabled: skip deleting association between user: { username } { account } )
91+ return
92+
93+ # for some operations, instantiate a user client when possible
94+ # note that this needs the X-SLURM-USER-NAME header upon instantiation
95+ this_user_client : Client = Client (
96+ base_url = self .endpoint ,
97+ headers = {
98+ "X-SLURM-USER-NAME" : f"{ username } ,
99+ "X-SLURM-USER-TOKEN" : self .token ,
100+ },
101+ httpx_args = {"event_hooks" : {"request" : [log_request ], "response" : [log_response ]}},
102+ )
103+
104+ with this_user_client as client :
105+ # start by deleting active jobs for this user
106+ body_job_kill : V0043KillJobsMsg = V0043KillJobsMsg (user_name = username , account = account )
107+ jobs_delete_resp = slurm_v0043_delete_jobs .sync (client = client , body = body_job_kill )
108+ if jobs_delete_resp :
109+ deletion_statuses : list [V0043KillJobsRespJob ] = jobs_delete_resp .status
110+ for status in deletion_statuses :
111+ logging .debug (f"deletion status: { status } )
112+ if jobs_delete_resp .errors :
113+ for error in jobs_delete_resp .errors :
114+ logging .error (f"error deleting job: { error } )
115+ raise RuntimeError (f"Could not delete jobs for user: { username } { account } )
55116 else :
56- raise ConnectionError ("Could not get list of accounts from SLURM endpoint" )
117+ raise ConnectionError (f"Could not delete jobs for user: { username } { account } )
118+
119+ # set maxsubmit to 0 as root
120+ max_submit_assoc : V0043Assoc = V0043Assoc (
121+ user = username ,
122+ account = account ,
123+ max_ = V0043AssocMax (V0043AssocMaxJobs (total = V0043Uint32NoValStruct (number = 0 ))),
124+ )
125+ body_max_submit : V0043OpenapiAssocsResp = V0043OpenapiAssocsResp (associations = [max_submit_assoc ])
126+ maxsubmit_set_resp = slurmdb_v0043_post_associations .sync (client = self .root_client , body = body_max_submit )
127+ if maxsubmit_set_resp :
128+ logger .info ("resp" )
129+ if maxsubmit_set_resp .errors :
130+ for error in maxsubmit_set_resp .errors :
131+ logging .error (f"error posting association: { error } )
132+ raise RuntimeError (f"Could not set maxsubmit=0 for user: { username } { account } )
133+ else :
134+ raise ConnectionError (f"Could not set maxsubmit=0 for user: { username } { account } )
135+
136+ default_assoc_get_resp = slurmdb_v0043_get_associations .sync (
137+ client = self .root_client , filter_to_only_defaults = str (True )
138+ )
139+ if not default_assoc_get_resp :
140+ raise ConnectionError (f"Could not get default association for user: { username } )
141+ if default_assoc_get_resp .errors :
142+ raise RuntimeError (f"Could not get default association for user: { username } )
143+
144+ default_account_for_this_user = default_assoc_get_resp .associations [0 ].account
145+ if default_account_for_this_user == account :
146+ logger .info (
147+ f"User: { username } { account }
148+ )
149+ # Assuming that everyone on the cluster is part of the account "users"!
150+ set_default_account_accoc : V0043Assoc = V0043Assoc (user = username , account = "users" , is_default = True )
151+ default_assoc_set_body : V0043OpenapiAssocsResp = V0043OpenapiAssocsResp (
152+ associations = [set_default_account_accoc ]
153+ )
154+
155+ default_assoc_set_resp = slurmdb_v0043_post_associations .sync (
156+ client = self .root_client , body = default_assoc_set_body
157+ )
158+ if not default_assoc_set_resp :
159+ raise ConnectionError (f"Could not set default association for user: { username } )
160+ if default_assoc_set_resp .errors :
161+ raise RuntimeError (f"Could not set default association for user: { username } )
162+
163+ # finally delete the association!
164+ delete_assoc_resp = slurmdb_v0043_delete_association .sync (
165+ client = self .root_client , user = username , account = account
166+ )
167+ if not delete_assoc_resp :
168+ raise ConnectionError (f"Could not delete association for user: { username } { account } )
169+ if delete_assoc_resp .errors :
170+ raise RuntimeError (f"Could not delete association for user: { username } { account } )
171+
172+ return
173+
174+ @retry (
175+ wait = wait_exponential (multiplier = 2 , min = 2 , max = 10 ),
176+ stop = stop_after_attempt (3 ),
177+ retry = retry_if_not_exception_type (ConnectionError ),
178+ before_sleep = before_sleep_log (logging .getLogger (__name__ ), logging .WARNING ),
179+ )
180+ def delete_slurm_account (self , account : str , noop : bool = False ) -> None :
181+ if noop :
182+ logging .info (f"noop enabled: skip deleting acconut: { account } )
183+ return
184+
185+ account_delete_resp = slurmdb_v0043_delete_account .sync (client = self .root_client , account_name = account )
186+
187+ if not account_delete_resp :
188+ raise ConnectionError (f"Could not delete account: { account } )
189+ if account_delete_resp .errors :
190+ raise RuntimeError (f"Could not delete account: { account } )
191+ for account_removed in account_delete_resp .removed_accounts :
192+ logging .info (f"Removed SLURM account: { account_removed } )
193+ return
0 commit comments