aind-data-transfer-service 1.12.0__py3-none-any.whl

aind_data_transfer_service/configs/job_upload_template.py

@@ -0,0 +1,153 @@
+"""Module to configure and create xlsx job upload template"""
+import datetime
+from io import BytesIO
+from typing import Any, Dict, List
+
+from aind_data_schema_models.modalities import Modality
+from aind_data_schema_models.platforms import Platform
+from openpyxl import Workbook
+from openpyxl.styles import Font
+from openpyxl.utils import get_column_letter
+from openpyxl.worksheet.datavalidation import DataValidation
+
+
+# TODO: convert to pydantic model
+class JobUploadTemplate:
+    """Class to configure and create xlsx job upload template"""
+
+    FILE_NAME = "job_upload_template.xlsx"
+    NUM_TEMPLATE_ROWS = 20
+    XLSX_DATETIME_FORMAT = "YYYY-MM-DDTHH:mm:ss"
+    HEADERS = [
+        "project_name",
+        "process_capsule_id",
+        "input_data_mount",
+        "platform",
+        "acq_datetime",
+        "subject_id",
+        "metadata_dir",
+        "modality0",
+        "modality0.source",
+        "modality1",
+        "modality1.source",
+    ]
+    SAMPLE_JOBS = [
+        [
+            "Behavior Platform",
+            "1f999652-00a0-4c4b-99b5-64c2985ad070",
+            "data_mount",
+            Platform.BEHAVIOR.abbreviation,
+            datetime.datetime(2023, 10, 4, 4, 0, 0),
+            "123456",
+            "/allen/aind/stage/fake/metadata_dir",
+            Modality.BEHAVIOR_VIDEOS.abbreviation,
+            "/allen/aind/stage/fake/dir",
+            Modality.BEHAVIOR.abbreviation,
+            "/allen/aind/stage/fake/dir",
+        ],
+        [
+            "Ophys Platform - SLAP2",
+            None,
+            None,
+            Platform.SMARTSPIM.abbreviation,
+            datetime.datetime(2023, 3, 4, 16, 30, 0),
+            "654321",
+            "/allen/aind/stage/fake/Config",
+            Modality.SPIM.abbreviation,
+            "/allen/aind/stage/fake/dir",
+        ],
+        [
+            "Ephys Platform",
+            None,
+            None,
+            Platform.ECEPHYS.abbreviation,
+            datetime.datetime(2023, 1, 30, 19, 1, 0),
+            "654321",
+            None,
+            Modality.ECEPHYS.abbreviation,
+            "/allen/aind/stage/fake/dir",
+            Modality.BEHAVIOR_VIDEOS.abbreviation,
+            "/allen/aind/stage/fake/dir",
+        ],
+    ]
+
+    @property
+    def validators(self) -> List[Dict[str, Any]]:
+        """
+        Returns
+        -------
+        List[Dict[str, Any]]
+          A list of validators for fields that require validation.
+
+        """
+        return [
+            {
+                "name": "platform",
+                "type": "list",
+                "options": list(Platform.abbreviation_map.keys()),
+                "column_indexes": [self.HEADERS.index("platform")],
+            },
+            {
+                "name": "modality",
+                "type": "list",
+                "options": list(Modality.abbreviation_map.keys()),
+                "column_indexes": [
+                    self.HEADERS.index("modality0"),
+                    self.HEADERS.index("modality1"),
+                ],
+            },
+            {
+                "name": "datetime",
+                "type": "date",
+                "column_indexes": [self.HEADERS.index("acq_datetime")],
+            },
+        ]
+
+    @property
+    def excel_sheet_filestream(self) -> BytesIO:
+        """Create job template as xlsx filestream"""
+        xl_io = BytesIO()
+        workbook = Workbook()
+        workbook.iso_dates = True
+        worksheet = workbook.active
+        worksheet.append(self.HEADERS)
+        for job in self.SAMPLE_JOBS:
+            worksheet.append(job)
+        # data validators
+        for validator in self.validators:
+            dv_type = validator["type"]
+            dv_name = validator["name"]
+            dv_params = {
+                "type": dv_type,
+                "promptTitle": dv_name,
+                "error": f"Invalid {dv_name}.",
+                "allow_blank": True,
+                "showErrorMessage": True,
+                "showInputMessage": True,
+            }
+            if dv_type == "list":
+                dv_params["formula1"] = f'"{(",").join(validator["options"])}"'
+                dv_params["prompt"] = f"Select a {dv_name} from the dropdown"
+            elif dv_type == "date":
+                dv_params["prompt"] = "Provide a {} using {}".format(
+                    dv_name, self.XLSX_DATETIME_FORMAT
+                )
+            dv = DataValidation(**dv_params)
+            for i in validator["column_indexes"]:
+                col = get_column_letter(i + 1)
+                col_range = f"{col}2:{col}{self.NUM_TEMPLATE_ROWS}"
+                dv.add(col_range)
+                if dv_type != "date":
+                    continue
+                for (cell,) in worksheet[col_range]:
+                    cell.number_format = self.XLSX_DATETIME_FORMAT
+            worksheet.add_data_validation(dv)
+        # formatting
+        bold = Font(bold=True)
+        for cell in worksheet[1]:
+            cell.font = bold
+            worksheet.column_dimensions[cell.column_letter].auto_size = True
+        # save file
+        workbook.save(xl_io)
+        workbook.close()
+        return xl_io

aind_data_transfer_service/hpc/__init__.py

@@ -0,0 +1 @@
+"""Client to manage connection with slurm cluster"""

aind_data_transfer_service/hpc/client.py

@@ -0,0 +1,151 @@
+"""Module to manage connection with hpc cluster"""
+
+import json
+from typing import List, Optional, Union
+
+import requests
+from pydantic import Field, SecretStr, field_validator
+from pydantic_settings import BaseSettings
+from requests.models import Response
+
+from aind_data_transfer_service.hpc.models import HpcJobSubmitSettings
+
+
+class HpcClientConfigs(BaseSettings):
+    """Configs needed to connect to the hpc cluster"""
+
+    hpc_host: str = Field(...)
+    hpc_port: Optional[int] = Field(default=None)
+    hpc_api_endpoint: Optional[str] = Field(default=None)
+    hpc_username: str = Field(...)
+    hpc_password: SecretStr = Field(...)
+    hpc_token: SecretStr = Field(...)
+
+    @field_validator("hpc_host", "hpc_api_endpoint", mode="before")
+    def _strip_slash(cls, input_str: Optional[str]):
+        """Strips trailing slash from domain."""
+        return None if input_str is None else input_str.strip("/")
+
+    @property
+    def hpc_url(self) -> str:
+        """Construct base url from host, port, and api endpoint"""
+        base_url = f"http://{self.hpc_host}"
+        if self.hpc_port is not None:
+            base_url = base_url + f":{self.hpc_port}"
+        if self.hpc_api_endpoint:
+            base_url = base_url + f"/{self.hpc_api_endpoint}"
+        return base_url
+
+
+class HpcClient:
+    """Class to manage client api"""
+
+    def __init__(self, configs: HpcClientConfigs):
+        """Class constructor"""
+        self.configs = configs
+
+    @property
+    def _job_submit_url(self):
+        """Url for job submission"""
+        return f"{self.configs.hpc_url}/job/submit"
+
+    @property
+    def _node_status_url(self):
+        """Url to check status of nodes"""
+        return f"{self.configs.hpc_url}/nodes"
+
+    @property
+    def _job_status_url(self):
+        """Url to check status of job"""
+        return f"{self.configs.hpc_url}/job"
+
+    @property
+    def _jobs_url(self):
+        """Url to check statuses of all jobs"""
+        return f"{self.configs.hpc_url}/jobs"
+
+    @property
+    def __headers(self):
+        """Headers needed for rest api"""
+        return {
+            "X-SLURM-USER-NAME": self.configs.hpc_username,
+            "X-SLURM-USER-PASSWORD": (
+                self.configs.hpc_password.get_secret_value()
+            ),
+            "X-SLURM-USER-TOKEN": self.configs.hpc_token.get_secret_value(),
+        }
+
+    def get_node_status(self) -> Response:
+        """Get status of nodes"""
+        response = requests.get(
+            url=self._node_status_url, headers=self.__headers
+        )
+        return response
+
+    def get_job_status(self, job_id: Union[str, int]) -> Response:
+        """Get status of job"""
+        response = requests.get(
+            url=self._job_status_url + "/" + str(job_id),
+            headers=self.__headers,
+        )
+        return response
+
+    def get_jobs(self) -> Response:
+        """Get status of job"""
+        response = requests.get(
+            url=self._jobs_url,
+            headers=self.__headers,
+        )
+        return response
+
+    def submit_job(self, job_def: dict) -> Response:
+        """Submit a job defined by job def"""
+        response = requests.post(
+            url=self._job_submit_url, json=job_def, headers=self.__headers
+        )
+        return response
+
+    def submit_hpc_job(
+        self,
+        script: str,
+        job: Optional[HpcJobSubmitSettings] = None,
+        jobs: Optional[List[HpcJobSubmitSettings]] = None,
+    ) -> Response:
+        """
+        Submit a job following the v0.0.36 Slurm rest api job submission guide
+        Parameters
+        ----------
+        script : str
+          Executable script (full contents) to run in batch step
+        job : Optional[HpcJobSubmitSettings]
+          v0.0.36_job_properties (Default is None)
+        jobs : Optional[List[HpcJobSubmitSettings]]
+          List of properties of an HetJob (Default is None)
+
+        Returns
+        -------
+        Response
+
+        """
+        # Assert at least one of job or jobs is defined
+        assert job is not None or jobs is not None
+        # Assert not both job and jobs are defined
+        assert job is None or jobs is None
+        if job is not None:
+            job_def = {
+                "job": json.loads(job.model_dump_json(exclude_none=True)),
+                "script": script,
+            }
+        else:
+            job_def = {
+                "jobs": [
+                    json.loads(j.model_dump_json(exclude_none=True))
+                    for j in jobs
+                ],
+                "script": script,
+            }
+
+        response = requests.post(
+            url=self._job_submit_url, json=job_def, headers=self.__headers
+        )
+        return response