aind-data-transfer-service 1.16.0__tar.gz → 1.17.2__tar.gz

{aind_data_transfer_service-1.16.0/src/aind_data_transfer_service.egg-info → aind_data_transfer_service-1.17.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aind-data-transfer-service
-Version: 1.16.0
+Version: 1.17.2
 Summary: Service that handles requests to upload data to the cloud
 Author: Allen Institute for Neural Dynamics
 License: MIT
@@ -25,7 +25,7 @@ Requires-Dist: furo; extra == "docs"
 Provides-Extra: server
 Requires-Dist: aind-data-schema<2.0,>=1.0.0; extra == "server"
 Requires-Dist: aind-data-transfer-models==0.17.0; extra == "server"
-Requires-Dist: aind-metadata-mapper==0.23.0; extra == "server"
+Requires-Dist: aind-metadata-mapper>=0.23.0; extra == "server"
 Requires-Dist: boto3; extra == "server"
 Requires-Dist: boto3-stubs[ssm]; extra == "server"
 Requires-Dist: fastapi>=0.115.13; extra == "server"

{aind_data_transfer_service-1.16.0 → aind_data_transfer_service-1.17.2}/pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools", "setuptools-scm"]
+requires = ["setuptools"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -41,7 +41,7 @@ docs = [
 server = [
     'aind-data-schema>=1.0.0,<2.0',
     'aind-data-transfer-models==0.17.0',
-    'aind-metadata-mapper==0.23.0',
+    'aind-metadata-mapper>=0.23.0',
     'boto3',
     'boto3-stubs[ssm]',
     'fastapi>=0.115.13', 

{aind_data_transfer_service-1.16.0 → aind_data_transfer_service-1.17.2}/src/aind_data_transfer_service/__init__.py

@@ -1,7 +1,7 @@
 """Init package"""
 import os
 
-__version__ = "1.16.0"
+__version__ = "1.17.2"
 
 # Global constants
 OPEN_DATA_BUCKET_NAME = os.getenv("OPEN_DATA_BUCKET_NAME", "open")

{aind_data_transfer_service-1.16.0 → aind_data_transfer_service-1.17.2}/src/aind_data_transfer_service/hpc/client.py

@@ -3,10 +3,9 @@
 import json
 from typing import List, Optional, Union
 
-import requests
+from httpx import AsyncClient, Response
 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
 
@@ -75,37 +74,41 @@ class HpcClient:
             "X-SLURM-USER-TOKEN": self.configs.hpc_token.get_secret_value(),
         }
 
-    def get_node_status(self) -> Response:
+    async def get_node_status(self) -> Response:
         """Get status of nodes"""
-        response = requests.get(
-            url=self._node_status_url, headers=self.__headers
-        )
+        async with AsyncClient() as async_client:
+            response = await async_client.get(
+                url=self._node_status_url, headers=self.__headers
+            )
         return response
 
-    def get_job_status(self, job_id: Union[str, int]) -> Response:
+    async 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,
-        )
+        async with AsyncClient() as async_client:
+            response = await async_client.get(
+                url=self._job_status_url + "/" + str(job_id),
+                headers=self.__headers,
+            )
         return response
 
-    def get_jobs(self) -> Response:
+    async def get_jobs(self) -> Response:
         """Get status of job"""
-        response = requests.get(
-            url=self._jobs_url,
-            headers=self.__headers,
-        )
+        async with AsyncClient() as async_client:
+            response = await async_client.get(
+                url=self._jobs_url,
+                headers=self.__headers,
+            )
         return response
 
-    def submit_job(self, job_def: dict) -> Response:
+    async 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
-        )
+        async with AsyncClient() as async_client:
+            response = await async_client.post(
+                url=self._job_submit_url, json=job_def, headers=self.__headers
+            )
         return response
 
-    def submit_hpc_job(
+    async def submit_hpc_job(
         self,
         script: str,
         job: Optional[HpcJobSubmitSettings] = None,
@@ -144,8 +147,8 @@ class HpcClient:
                 ],
                 "script": script,
             }
-
-        response = requests.post(
-            url=self._job_submit_url, json=job_def, headers=self.__headers
-        )
+        async with AsyncClient() as async_client:
+            response = await async_client.post(
+                url=self._job_submit_url, json=job_def, headers=self.__headers
+            )
         return response

{aind_data_transfer_service-1.16.0 → aind_data_transfer_service-1.17.2}/src/aind_data_transfer_service/models/internal.py

@@ -3,8 +3,9 @@
 import ast
 import os
 from datetime import datetime, timedelta, timezone
-from typing import List, Optional, Union
+from typing import ClassVar, List, Optional, Union
 
+from aind_data_schema_models.modalities import Modality
 from mypy_boto3_ssm.type_defs import ParameterMetadataTypeDef
 from pydantic import AwareDatetime, BaseModel, Field, field_validator
 from starlette.datastructures import QueryParams
@@ -223,11 +224,30 @@ class JobTasks(BaseModel):
 class JobParamInfo(BaseModel):
     """Model for job parameter info from AWS Parameter Store"""
 
+    _MODALITIES_LIST: ClassVar[list[str]] = list(
+        Modality.abbreviation_map.keys()
+    )
+    _MODALITY_TASKS: ClassVar[list[str]] = [
+        "modality_transformation_settings",
+        "codeocean_pipeline_settings",
+    ]
+
     name: Optional[str]
     last_modified: Optional[datetime]
-    job_type: str
-    task_id: str
+    job_type: str = Field(..., pattern=r"^[^\s/]+$")
+    task_id: str = Field(..., pattern=r"^[^\s/]+$")
     modality: Optional[str]
+    version: Optional[str] = Field(..., pattern=r"^(v1|v2)?$")
+
+    @field_validator("modality", mode="after")
+    def validate_modality(cls, v):
+        """Check that modality is one of aind-data-schema modalities"""
+        if v is not None and v not in JobParamInfo._MODALITIES_LIST:
+            raise ValueError(
+                "Invalid modality: modality must be one of "
+                f"{JobParamInfo._MODALITIES_LIST}"
+            )
+        return v
 
     @classmethod
     def from_aws_describe_parameter(
@@ -236,6 +256,7 @@ class JobParamInfo(BaseModel):
         job_type: str,
         task_id: str,
         modality: Optional[str],
+        version: Optional[str],
     ):
         """Map the parameter to the model"""
         return cls(
@@ -244,13 +265,14 @@ class JobParamInfo(BaseModel):
             job_type=job_type,
             task_id=task_id,
             modality=modality,
+            version=version,
         )
 
     @staticmethod
     def get_parameter_prefix(version: Optional[str] = None) -> str:
         """Get the prefix for job_type parameters"""
         prefix = os.getenv("AIND_AIRFLOW_PARAM_PREFIX")
-        if version is None:
+        if version is None or version == "v1":
             return prefix
         return f"{prefix}/{version}"
 
@@ -262,16 +284,19 @@ class JobParamInfo(BaseModel):
             "(?P<job_type>[^/]+)/tasks/(?P<task_id>[^/]+)"
             "(?:/(?P<modality>[^/]+))?"
         )
-        if version is None:
+        if version is None or version == "v1":
             return f"{prefix}/{regex}"
         return f"{prefix}/{version}/{regex}"
 
     @staticmethod
     def get_parameter_name(
-        job_type: str, task_id: str, version: Optional[str] = None
+        job_type: str,
+        task_id: str,
+        modality: Optional[str],
+        version: Optional[str] = None,
     ) -> str:
         """Create the parameter name from job_type and task_id"""
-        prefix = os.getenv("AIND_AIRFLOW_PARAM_PREFIX")
-        if version is None:
-            return f"{prefix}/{job_type}/tasks/{task_id}"
-        return f"{prefix}/{version}/{job_type}/tasks/{task_id}"
+        prefix = JobParamInfo.get_parameter_prefix(version)
+        if modality:
+            return f"{prefix}/{job_type}/tasks/{task_id}/{modality}"
+        return f"{prefix}/{job_type}/tasks/{task_id}"