Coverage for gco/services/slurm

1"""slurmrestd client behind the GCO API's ``/api/v1/slurm`` routes.

3The manifest-processor (the authenticated API ingress) talks to Slurm through

4slurmrestd, the Slurm REST daemon the Slinky chart runs in ``gco-jobs``

5(``gco.slurm_config.SLURM_RESTAPI_NAME``). Authentication is Slurm's own

6``auth/jwt``: the slurm-operator mints a JWT for the API's Slurm identity into

7a Kubernetes Secret through a ``Token`` custom resource

8(post-helm-slurm-api-access.yaml) and refreshes it before it expires; this

9module reads that Secret (a scoped RBAC grant on exactly that one Secret) and

10sends it as ``X-SLURM-USER-TOKEN``. The API never sees the cluster's JWT

11signing key.

13The hop is plain HTTP inside the namespace boundary (the NetworkPolicies in

14post-helm-slurm-network.yaml admit only the manifest-processor pods to the

15REST port); slurmrestd has no TLS listener of its own, and the ALB-facing side

16of the API is the TLS boundary callers see.

18Errors map to the API the same way the cost proxy's do: Slurm not enabled for

19the deployment -> 404 with a remedy, slurmrestd unreachable or the token

20unavailable -> 503, a Slurm-side rejection -> the status slurmrestd meant (404

21for an unknown job, 400 for a bad submission, 502 for anything else).

22"""

24from __future__ import annotations

26import asyncio

27import base64

28import logging

29import os

30import time

31from collections.abc import Mapping

32from typing import Any

34import httpx2

35from fastapi import HTTPException

37from gco.services import tracing

38from gco.slurm_config import (

39 SLURM_API_VERSION,

40 SLURM_NAMESPACE,

41 SLURM_RESTAPI_NAME,

42 SLURM_RESTAPI_PORT,

43)

45logger = logging.getLogger(__name__)

47#: ``SLURM_API_VERSION`` (the slurmrestd plugin version the paths below carry)

48#: lives in :mod:`gco.slurm_config` beside the image tag it moves with; it is

49#: re-exported here because this module is where the version is used.

51#: The Slurm identity the API acts as. ``slurm`` is SlurmUser: it may submit,

52#: inspect and cancel any job, which is what a single authenticated front door

53#: needs; the caller's identity travels in each job's ``comment`` for

54#: accounting attribution.

55SLURM_API_USER = "slurm"

57#: The Token custom resource and the Secret it writes (post-helm-slurm-api-access.yaml).

58SLURM_API_TOKEN_NAME = "slinky-slurm-token-gco-api" # nosec B105 # object name, not a credential

59SLURM_API_TOKEN_SECRET = SLURM_API_TOKEN_NAME # nosec B105 # Kubernetes Secret name

60SLURM_API_TOKEN_KEY = "auth-token" # nosec B105 # Secret data key, not a credential

62DEFAULT_RESTAPI_URL = (

63 f"http://{SLURM_RESTAPI_NAME}.{SLURM_NAMESPACE}.svc.cluster.local:{SLURM_RESTAPI_PORT}"

64)

66#: How long a token read from the Secret is reused before it is re-read. The

67#: operator refreshes the Secret at 80% of the token lifetime, so re-reading

68#: every few minutes (and on any 401) always finds a valid token.

69TOKEN_REFRESH_SECONDS = 300.0

71REQUEST_TIMEOUT_SECONDS = 30.0

73DISABLED_DETAIL = (

74 "Slurm is not enabled in this region. Set helm.slurm.enabled to true in cdk.json "

75 "(or deploy with --enable slurm) and redeploy the regional stack."

76)

77UNAVAILABLE_DETAIL = (

78 "Slurm's REST API is unavailable in this region. Check the slinky-slurm-restapi "

79 "Deployment and the slinky-slurm-controller StatefulSet in gco-jobs."

80)

81TOKEN_UNAVAILABLE_DETAIL = (

82 "The GCO API's Slurm credential is unavailable. Check the Token "

83 f"{SLURM_API_TOKEN_NAME} in gco-jobs and the slurm-operator logs."

84)

85#: slurmrestd answered 404 for ``/slurm/<version>/ping``, an endpoint every

86#: plugin version serves: the deployed image does not ship the version this

87#: API addresses (``{version}`` is filled in from the client).

88API_VERSION_DETAIL = (

89 "Slurm's REST API does not serve OpenAPI plugin version {version}: slurmrestd "

90 "answered 404 for /slurm/{version}/ping. The slurmrestd image in charts.yaml and "

91 "SLURM_API_VERSION in gco/slurm_config.py (or the SLURM_API_VERSION override on the "

92 "API pod) have drifted apart; `slurmrestd -d list` inside the image names the "

93 "versions it ships."

94)

97class SlurmDisabledError(Exception):

98 """Slurm is not deployed in this region (cdk.json helm.slurm.enabled false)."""

100

101class SlurmUnavailableError(Exception):

102 """slurmrestd or the API's token cannot be reached right now.

103

104 ``detail`` is the operator-facing remedy the API answers with when the

105 cause is known precisely; :func:`http_exception_for` falls back to the

106 generic transport / credential remedies otherwise.

107 """

108

109 def __init__(self, message: str, *, detail: str | None = None) -> None:

110 super().__init__(message)

111 self.detail = detail

112

113

114class SlurmApiError(Exception):

115 """slurmrestd answered with an error for the request.

116

117 ``status_code`` is the status the GCO API answers with (see

118 :func:`_map_slurm_status`); ``http_status`` is what slurmrestd itself

119 answered, kept for callers that need the raw transport fact.

120 """

121

122 def __init__(

123 self,

124 status_code: int,

125 errors: list[Any],

126 detail: str,

127 *,

128 http_status: int | None = None,

129 ) -> None:

130 super().__init__(detail)

131 self.status_code = status_code

132 self.errors = errors

133 self.detail = detail

134 self.http_status = http_status

135

136

137def slurm_enabled() -> bool:

138 """Whether the regional stack deployed Slurm (``GCO_SLURM_ENABLED`` on the API pod)."""

139 return os.getenv("GCO_SLURM_ENABLED", "false").strip().lower() == "true"

140

141

142def restapi_base_url() -> str:

143 return os.getenv("SLURM_RESTAPI_URL", DEFAULT_RESTAPI_URL).rstrip("/")

144

145

146def api_version() -> str:

147 return os.getenv("SLURM_API_VERSION", SLURM_API_VERSION).strip() or SLURM_API_VERSION

148

149

150class SecretTokenProvider:

151 """Reads the operator-minted JWT from its Kubernetes Secret, with caching.

152

153 The Secret is read with the manifest-processor's own ServiceAccount, which

154 the post-Helm manifest grants ``get`` on exactly this Secret in

155 ``gco-jobs``. ``invalidate()`` after a 401 forces the next call to re-read.

156 """

157

158 def __init__(

159 self,

160 *,

161 namespace: str = SLURM_NAMESPACE,

162 name: str = SLURM_API_TOKEN_SECRET,

163 key: str = SLURM_API_TOKEN_KEY,

164 refresh_seconds: float = TOKEN_REFRESH_SECONDS,

165 ) -> None:

166 self.namespace = namespace

167 self.name = name

168 self.key = key

169 self.refresh_seconds = refresh_seconds

170 self._token: str | None = None

171 self._read_at = 0.0

172

173 def invalidate(self) -> None:

174 self._token = None

175 self._read_at = 0.0

176

177 def _read_secret(self) -> str:

178 from kubernetes import client as kubernetes_client

179

180 secret = kubernetes_client.CoreV1Api().read_namespaced_secret(self.name, self.namespace)

181 data = getattr(secret, "data", None) or {}

182 encoded = data.get(self.key)

183 if not encoded:

184 raise SlurmUnavailableError(

185 f"Secret {self.namespace}/{self.name} has no {self.key!r} key yet"

186 )

187 token = base64.b64decode(encoded).decode("utf-8").strip()

188 if not token:

189 raise SlurmUnavailableError(f"Secret {self.namespace}/{self.name} holds an empty token")

190 return token

191

192 async def token(self) -> str:

193 now = time.monotonic()

194 if self._token is not None and now - self._read_at < self.refresh_seconds:

195 return self._token

196 try:

197 token = await asyncio.to_thread(self._read_secret)

198 except SlurmUnavailableError:

199 raise

200 except Exception as exc:

201 raise SlurmUnavailableError(

202 f"could not read Secret {self.namespace}/{self.name}: {exc}"

203 ) from exc

204 self._token = token

205 self._read_at = time.monotonic()

206 return token

207

208

209class SlurmRestClient:

210 """A thin, typed wrapper over the slurmrestd endpoints the API exposes."""

211

212 def __init__(

213 self,

214 *,

215 base_url: str | None = None,

216 version: str | None = None,

217 token_provider: SecretTokenProvider | None = None,

218 user: str = SLURM_API_USER,

219 timeout: float = REQUEST_TIMEOUT_SECONDS,

220 ) -> None:

221 self.base_url = (base_url or restapi_base_url()).rstrip("/")

222 self.version = version or api_version()

223 self.token_provider = token_provider or SecretTokenProvider()

224 self.user = user

225 self.timeout = timeout

226

227 # -- transport -----------------------------------------------------------

228

229 async def _send(

230 self, method: str, path: str, *, params: Mapping[str, Any] | None, json: Any

231 ) -> httpx2.Response:

232 token = await self.token_provider.token()

233 headers = {"X-SLURM-USER-NAME": self.user, "X-SLURM-USER-TOKEN": token}

234 url = f"{self.base_url}{path}"

235 try:

236 transport = httpx2.AsyncHTTPTransport(trust_env=False)

237 async with httpx2.AsyncClient(

238 transport=tracing.wrap_async_transport(transport),

239 timeout=self.timeout,

240 trust_env=False,

241 ) as client:

242 return await client.request(

243 method, url, headers=headers, params=dict(params or {}), json=json

244 )

245 except httpx2.HTTPError as exc:

246 logger.warning("slurmrestd unreachable at %s: %s", url, exc)

247 raise SlurmUnavailableError(f"slurmrestd unreachable at {self.base_url}") from exc

248

249 async def request(

250 self,

251 method: str,

252 path: str,

253 *,

254 params: Mapping[str, Any] | None = None,

255 json: Any = None,

256 ) -> dict[str, Any]:

257 """Send one request; retry once with a fresh token on 401; decode the body."""

258 response = await self._send(method, path, params=params, json=json)

259 if response.status_code == 401:

260 self.token_provider.invalidate()

261 response = await self._send(method, path, params=params, json=json)

262 return self._decode(response)

263

264 @staticmethod

265 def _decode(response: httpx2.Response) -> dict[str, Any]:

266 status = response.status_code

267 try:

268 payload = response.json()

269 except ValueError as exc:

270 raise SlurmApiError(

271 502, [], f"slurmrestd returned a non-JSON body (HTTP {status})", http_status=status

272 ) from exc

273 if not isinstance(payload, dict):

274 raise SlurmApiError(

275 502, [], "slurmrestd returned an unexpected JSON shape", http_status=status

276 )

277 errors = payload.get("errors") or []

278 if status >= 400 or errors:

279 raise SlurmApiError(

280 _map_slurm_status(status, errors),

281 errors,

282 _describe_errors(errors),

283 http_status=status,

284 )

285 return payload

286

287 # -- endpoints -------------------------------------------------------------

288

289 def _slurm(self, suffix: str) -> str:

290 return f"/slurm/{self.version}/{suffix}"

291

292 def _slurmdb(self, suffix: str) -> str:

293 return f"/slurmdb/{self.version}/{suffix}"

294

295 async def ping(self) -> dict[str, Any]:

296 """``/ping`` exists in every plugin version, so a 404 here is a version mismatch.

297

298 The status route calls this first; translating the 404 into an

299 unavailable-with-remedy error is what turns a silent image/pin drift

300 into one actionable message instead of a generic "not found".

301 """

302 try:

303 return await self.request("GET", self._slurm("ping"))

304 except SlurmApiError as exc:

305 if exc.http_status == 404:

306 raise SlurmUnavailableError(

307 f"slurmrestd at {self.base_url} does not serve {self.version}",

308 detail=API_VERSION_DETAIL.format(version=self.version),

309 ) from exc

310 raise

311

312 async def diag(self) -> dict[str, Any]:

313 return await self.request("GET", self._slurm("diag"))

314

315 async def partitions(self) -> dict[str, Any]:

316 return await self.request("GET", self._slurm("partitions"))

317

318 async def nodes(self) -> dict[str, Any]:

319 return await self.request("GET", self._slurm("nodes"))

320

321 async def jobs(self) -> dict[str, Any]:

322 return await self.request("GET", self._slurm("jobs"))

323

324 async def job(self, job_id: int) -> dict[str, Any]:

325 return await self.request("GET", self._slurm(f"job/{int(job_id)}"))

326

327 async def submit(self, job: Mapping[str, Any], script: str) -> dict[str, Any]:

328 return await self.request(

329 "POST", self._slurm("job/submit"), json={"job": dict(job), "script": script}

330 )

331

332 async def cancel(self, job_id: int, *, signal: str | None = None) -> dict[str, Any]:

333 params = {"signal": signal} if signal else None

334 return await self.request("DELETE", self._slurm(f"job/{int(job_id)}"), params=params)

335

336 async def accounting_jobs(self, params: Mapping[str, Any]) -> dict[str, Any]:

337 return await self.request("GET", self._slurmdb("jobs"), params=params)

338

339 async def accounting_job(self, job_id: int) -> dict[str, Any]:

340 """One slurmdbd record; an unknown id answers an empty ``jobs`` list, not an error."""

341 return await self.request("GET", self._slurmdb(f"job/{int(job_id)}"))

342

343

344def _describe_errors(errors: list[Any]) -> str:

345 parts: list[str] = []

346 for item in errors:

347 if not isinstance(item, dict):

348 continue

349 description = str(item.get("description") or item.get("error") or "").strip()

350 source = str(item.get("source") or "").strip()

351 number = item.get("error_number")

352 text = description or "Slurm returned an error"

353 if source:

354 text = f"{text} ({source})"

355 if number is not None:

356 text = f"{text} [rc {number}]"

357 parts.append(text)

358 return "; ".join(parts) or "Slurm returned an error"

359

360

361#: Slurm error numbers slurmrestd reports for "no such job" lookups and cancels.

362_SLURM_INVALID_JOB_ID_ERRORS = {2017} # ESLURM_INVALID_JOB_ID

363

364

365def _map_slurm_status(http_status: int, errors: list[Any]) -> int:

366 """Pick the HTTP status the GCO API answers with for a Slurm-side failure."""

367 numbers = {

368 item.get("error_number")

369 for item in errors

370 if isinstance(item, dict) and isinstance(item.get("error_number"), int)

371 }

372 if numbers & _SLURM_INVALID_JOB_ID_ERRORS or http_status == 404:

373 return 404

374 if http_status == 400 or any(

375 "invalid" in str(item.get("description", "")).lower()

376 or "partition" in str(item.get("description", "")).lower()

377 for item in errors

378 if isinstance(item, dict)

379 ):

380 return 400

381 return 502

382

383

384def http_exception_for(error: Exception, *, context: str = "talking to Slurm") -> HTTPException:

385 """Translate a client error into the HTTPException the route raises.

386

387 Anything that is not one of this module's errors is a bug or an

388 unexpected transport failure; it becomes a 500 whose detail names only

389 the operation (``context``), never the exception text.

390 """

391 if isinstance(error, SlurmDisabledError):

392 return HTTPException(status_code=404, detail=DISABLED_DETAIL)

393 if isinstance(error, SlurmUnavailableError):

394 detail = error.detail or (

395 TOKEN_UNAVAILABLE_DETAIL if "Secret" in str(error) else UNAVAILABLE_DETAIL

396 )

397 return HTTPException(status_code=503, detail=detail)

398 if isinstance(error, SlurmApiError):

399 return HTTPException(status_code=error.status_code, detail=error.detail)

400 return HTTPException(status_code=500, detail=f"Internal error while {context}")

401

402

403_client: SlurmRestClient | None = None

404

405

406def get_client() -> SlurmRestClient:

407 """The process-wide client; raises :class:`SlurmDisabledError` when Slurm is off."""

408 global _client

409 if not slurm_enabled():

410 raise SlurmDisabledError(DISABLED_DETAIL)

411 if _client is None:

412 _client = SlurmRestClient()

413 return _client

414

415

416def reset_client() -> None:

417 """Drop the cached client (tests, configuration reloads)."""

418 global _client

419 _client = None

Coverage for gco / services / slurm_client.py: 100.00%

185 statements