Skip to content

Commit

Permalink
Initial commit of upgrade to Slurm REST API v0.0.40
Browse files Browse the repository at this point in the history
  • Loading branch information
isikhar committed Mar 27, 2024
1 parent 22fcd05 commit 7abc692
Show file tree
Hide file tree
Showing 5 changed files with 3,453 additions and 602 deletions.
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ jmx = "zocalo.configuration.plugin_jmx:JMX"
logging = "zocalo.configuration.plugin_logging:Logging"
rabbitmqapi = "zocalo.configuration.plugin_rabbitmqapi:RabbitAPI"
slurm = "zocalo.configuration.plugin_slurm:Slurm"
iris = "zocalo.configuration.plugin_slurm:Slurm"
smtp = "zocalo.configuration.plugin_smtp:SMTP"
storage = "zocalo.configuration.plugin_storage:Storage"

Expand Down
31 changes: 17 additions & 14 deletions src/zocalo/util/slurm/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

import requests

import zocalo.configuration
from zocalo.configuration import Configuration

from . import models

Expand All @@ -15,7 +15,7 @@ class SlurmRestApi:
def __init__(
self,
url: str,
version: str = "v0.0.36",
version: str = "v0.0.40",
user_name: str | None = None,
user_token: str | pathlib.Path | None = None,
):
Expand All @@ -35,12 +35,13 @@ def __init__(
self.session.headers["X-SLURM-USER-TOKEN"] = self.user_token

@classmethod
def from_zocalo_configuration(cls, zc: zocalo.configuration.Configuration):
def from_zocalo_configuration(cls, zc: Configuration, cluster: str = "slurm"):
cluster_config = getattr(zc, cluster)
return cls(
url=zc.slurm["url"],
version=zc.slurm["api_version"],
user_name=zc.slurm.get("user"),
user_token=zc.slurm.get("user_token"),
url=cluster_config["url"],
version=cluster_config["api_version"],
user_name=cluster_config.get("user"),
user_token=cluster_config.get("user_token"),
)

def get(
Expand Down Expand Up @@ -93,19 +94,21 @@ def delete(
response.raise_for_status()
return response

def get_jobs(self) -> models.JobsResponse:
def get_jobs(self) -> models.OpenapiJobInfoResp:
endpoint = f"slurm/{self.version}/jobs"
response = self.get(endpoint)
return models.JobsResponse(**response.json())
return models.OpenapiJobInfoResp(**response.json())

def get_job_info(self, job_id: int) -> models.JobsResponse:
def get_job_info(self, job_id: int) -> models.JobInfo:
endpoint = f"slurm/{self.version}/job/{job_id}"
response = self.get(endpoint)
return models.JobsResponse(**response.json())
job_info_resp = models.OpenapiJobInfoResp(**response.json())
jobinfo = next(iter(dict(job_info_resp.jobs).get("__root__", [])))
return jobinfo

def submit_job(
self, job_submission: models.JobSubmission
) -> models.JobSubmissionResponse:
self, job_submission: models.JobSubmitReq
) -> models.JobSubmitResponseMsg:
endpoint = f"slurm/{self.version}/job/submit"
response = self.post(endpoint, json=job_submission.dict(exclude_defaults=True))
return models.JobSubmissionResponse(**response.json())
return models.JobSubmitResponseMsg(**response.json())
Loading

0 comments on commit 7abc692

Please sign in to comment.