aind-data-transfer-service 1.16.0__py3-none-any.whl → 1.17.2__py3-none-any.whl

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/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/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}"

aind_data_transfer_service/server.py

@@ -7,10 +7,9 @@ import os
 import re
 from asyncio import gather, sleep
 from pathlib import PurePosixPath
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import boto3
-import requests
 from aind_data_transfer_models import (
     __version__ as aind_data_transfer_models_version,
 )
@@ -96,12 +95,13 @@ logger = get_logger(log_configs=LoggingConfigs())
 project_names_url = os.getenv("AIND_METADATA_SERVICE_PROJECT_NAMES_URL")
 
 
-def get_project_names() -> List[str]:
+async def get_project_names() -> List[str]:
     """Get a list of project_names"""
     # TODO: Cache response for 5 minutes
-    response = requests.get(project_names_url)
-    response.raise_for_status()
-    project_names = response.json()["data"]
+    async with AsyncClient() as async_client:
+        response = await async_client.get(project_names_url)
+        response.raise_for_status()
+        project_names = response.json()["data"]
     return project_names
 
 
@@ -156,6 +156,7 @@ def get_parameter_infos(version: Optional[str] = None) -> List[JobParamInfo]:
                     job_type=match.group("job_type"),
                     task_id=match.group("task_id"),
                     modality=match.group("modality"),
+                    version=version,
                 )
                 params.append(param_info)
             else:
@@ -173,6 +174,19 @@ def get_parameter_value(param_name: str) -> dict:
     return param_value
 
 
+def put_parameter_value(param_name: str, param_value: dict) -> Any:
+    """Set a parameter value in AWS param store based on parameter name"""
+    param_value_str = json.dumps(param_value)
+    ssm_client = boto3.client("ssm")
+    result = ssm_client.put_parameter(
+        Name=param_name,
+        Value=param_value_str,
+        Type="String",
+        Overwrite=True,
+    )
+    return result
+
+
 async def get_airflow_jobs(
     params: AirflowDagRunsRequestParameters, get_confs: bool = False
 ) -> tuple[int, Union[List[JobStatus], List[dict]]]:
@@ -267,7 +281,7 @@ async def validate_csv(request: Request):
             )
             context = {
                 "job_types": get_job_types("v2"),
-                "project_names": get_project_names(),
+                "project_names": await get_project_names(),
                 "current_jobs": current_jobs,
             }
             for row in csv_reader:
@@ -361,7 +375,7 @@ async def validate_json_v2(request: Request):
         _, current_jobs = await get_airflow_jobs(params=params, get_confs=True)
         context = {
             "job_types": get_job_types("v2"),
-            "project_names": get_project_names(),
+            "project_names": await get_project_names(),
             "current_jobs": current_jobs,
         }
         with validation_context_v2(context):
@@ -417,7 +431,7 @@ async def validate_json(request: Request):
     logger.info("Received request to validate json")
     content = await request.json()
     try:
-        project_names = get_project_names()
+        project_names = await get_project_names()
         with validation_context({"project_names": project_names}):
             validated_model = SubmitJobRequest.model_validate_json(
                 json.dumps(content)
@@ -477,7 +491,7 @@ async def submit_jobs_v2(request: Request):
         _, current_jobs = await get_airflow_jobs(params=params, get_confs=True)
         context = {
             "job_types": get_job_types("v2"),
-            "project_names": get_project_names(),
+            "project_names": await get_project_names(),
             "current_jobs": current_jobs,
         }
         with validation_context_v2(context):
@@ -485,7 +499,6 @@ async def submit_jobs_v2(request: Request):
         full_content = json.loads(
             model.model_dump_json(warnings=False, exclude_none=True)
         )
-        # TODO: Replace with httpx async client
         logger.info(
             f"Valid request detected. Sending list of jobs. "
             f"dag_id: {model.dag_id}"
@@ -497,19 +510,23 @@ async def submit_jobs_v2(request: Request):
                 f"{job_index} of {total_jobs}."
             )
 
-        response = requests.post(
-            url=os.getenv("AIND_AIRFLOW_SERVICE_URL"),
+        async with AsyncClient(
             auth=(
                 os.getenv("AIND_AIRFLOW_SERVICE_USER"),
                 os.getenv("AIND_AIRFLOW_SERVICE_PASSWORD"),
-            ),
-            json={"conf": full_content},
-        )
+            )
+        ) as async_client:
+            response = await async_client.post(
+                url=os.getenv("AIND_AIRFLOW_SERVICE_URL"),
+                json={"conf": full_content},
+            )
+            status_code = response.status_code
+            response_json = response.json()
         return JSONResponse(
-            status_code=response.status_code,
+            status_code=status_code,
             content={
                 "message": "Submitted request to airflow",
-                "data": {"responses": [response.json()], "errors": []},
+                "data": {"responses": [response_json], "errors": []},
             },
         )
     except ValidationError as e:
@@ -537,13 +554,12 @@ async def submit_jobs(request: Request):
     logger.info("Received request to submit jobs")
     content = await request.json()
     try:
-        project_names = get_project_names()
+        project_names = await get_project_names()
         with validation_context({"project_names": project_names}):
             model = SubmitJobRequest.model_validate_json(json.dumps(content))
         full_content = json.loads(
             model.model_dump_json(warnings=False, exclude_none=True)
         )
-        # TODO: Replace with httpx async client
         logger.info(
             f"Valid request detected. Sending list of jobs. "
             f"Job Type: {model.job_type}"
@@ -555,19 +571,23 @@ async def submit_jobs(request: Request):
                 f"{job_index} of {total_jobs}."
             )
 
-        response = requests.post(
-            url=os.getenv("AIND_AIRFLOW_SERVICE_URL"),
+        async with AsyncClient(
             auth=(
                 os.getenv("AIND_AIRFLOW_SERVICE_USER"),
                 os.getenv("AIND_AIRFLOW_SERVICE_PASSWORD"),
-            ),
-            json={"conf": full_content},
-        )
+            )
+        ) as async_client:
+            response = await async_client.post(
+                url=os.getenv("AIND_AIRFLOW_SERVICE_URL"),
+                json={"conf": full_content},
+            )
+            status_code = response.status_code
+            response_json = response.json()
         return JSONResponse(
-            status_code=response.status_code,
+            status_code=status_code,
             content={
                 "message": "Submitted request to airflow",
-                "data": {"responses": [response.json()], "errors": []},
+                "data": {"responses": [response_json], "errors": []},
             },
         )
 
@@ -631,7 +651,7 @@ async def submit_basic_jobs(request: Request):
         for hpc_job in hpc_jobs:
             try:
                 job_def = hpc_job.job_definition
-                response = hpc_client.submit_job(job_def)
+                response = await hpc_client.submit_job(job_def)
                 response_json = response.json()
                 responses.append(response_json)
                 # Add pause to stagger job requests to the hpc
@@ -755,7 +775,7 @@ async def submit_hpc_jobs(request: Request):  # noqa: C901
             hpc_job_def = hpc_job[0]
             try:
                 script = hpc_job[1]
-                response = hpc_client.submit_hpc_job(
+                response = await hpc_client.submit_hpc_job(
                     job=hpc_job_def, script=script
                 )
                 response_json = response.json()
@@ -828,20 +848,23 @@ async def get_tasks_list(request: Request):
             request.query_params
         )
         params_dict = json.loads(params.model_dump_json())
-        response_tasks = requests.get(
-            url=(
-                f"{url}/{params.dag_id}/dagRuns/{params.dag_run_id}/"
-                "taskInstances"
-            ),
+        async with AsyncClient(
             auth=(
                 os.getenv("AIND_AIRFLOW_SERVICE_USER"),
                 os.getenv("AIND_AIRFLOW_SERVICE_PASSWORD"),
-            ),
-        )
-        status_code = response_tasks.status_code
-        if response_tasks.status_code == 200:
+            )
+        ) as async_client:
+            response_tasks = await async_client.get(
+                url=(
+                    f"{url}/{params.dag_id}/dagRuns/{params.dag_run_id}/"
+                    "taskInstances"
+                ),
+            )
+            status_code = response_tasks.status_code
+            response_json = response_tasks.json()
+        if status_code == 200:
             task_instances = AirflowTaskInstancesResponse.model_validate_json(
-                json.dumps(response_tasks.json())
+                json.dumps(response_json)
             )
             job_tasks_list = sorted(
                 [
@@ -862,7 +885,7 @@ async def get_tasks_list(request: Request):
             message = "Error retrieving job tasks list from airflow"
             data = {
                 "params": params_dict,
-                "errors": [response_tasks.json()],
+                "errors": [response_json],
             }
     except ValidationError as e:
         logger.warning(f"There was a validation error process task_list: {e}")
@@ -892,27 +915,29 @@ async def get_task_logs(request: Request):
         )
         params_dict = json.loads(params.model_dump_json())
         params_full = dict(params)
-        response_logs = requests.get(
-            url=(
-                f"{url}/{params.dag_id}/dagRuns/{params.dag_run_id}"
-                f"/taskInstances/{params.task_id}/logs/{params.try_number}"
-            ),
+        async with AsyncClient(
             auth=(
                 os.getenv("AIND_AIRFLOW_SERVICE_USER"),
                 os.getenv("AIND_AIRFLOW_SERVICE_PASSWORD"),
-            ),
-            params=params_dict,
-        )
-        status_code = response_logs.status_code
-        if response_logs.status_code == 200:
-            message = "Retrieved task logs from airflow"
-            data = {"params": params_full, "logs": response_logs.text}
-        else:
-            message = "Error retrieving task logs from airflow"
-            data = {
-                "params": params_full,
-                "errors": [response_logs.json()],
-            }
+            )
+        ) as async_client:
+            response_logs = await async_client.get(
+                url=(
+                    f"{url}/{params.dag_id}/dagRuns/{params.dag_run_id}"
+                    f"/taskInstances/{params.task_id}/logs/{params.try_number}"
+                ),
+                params=params_dict,
+            )
+            status_code = response_logs.status_code
+            if status_code == 200:
+                message = "Retrieved task logs from airflow"
+                data = {"params": params_full, "logs": response_logs.text}
+            else:
+                message = "Error retrieving task logs from airflow"
+                data = {
+                    "params": params_full,
+                    "errors": [response_logs.json()],
+                }
     except ValidationError as e:
         logger.warning(f"Error validating request parameters: {e}")
         status_code = 406
@@ -1001,16 +1026,20 @@ async def jobs(request: Request):
 
 async def job_params(request: Request):
     """Get Job Parameters page"""
+    user = request.session.get("user")
     return templates.TemplateResponse(
         request=request,
         name="job_params.html",
         context=(
             {
+                "user_signed_in": user is not None,
                 "project_names_url": os.getenv(
                     "AIND_METADATA_SERVICE_PROJECT_NAMES_URL"
                 ),
                 "versions": ["v1", "v2"],
                 "default_version": "v1",
+                "modalities": JobParamInfo._MODALITIES_LIST,
+                "modality_tasks": JobParamInfo._MODALITY_TASKS,
             }
         ),
     )
@@ -1074,7 +1103,10 @@ def get_parameter_v2(request: Request):
     # path params are auto validated
     job_type = request.path_params.get("job_type")
     task_id = request.path_params.get("task_id")
-    param_name = JobParamInfo.get_parameter_name(job_type, task_id, "v2")
+    modality = request.path_params.get("modality")
+    param_name = JobParamInfo.get_parameter_name(
+        job_type=job_type, task_id=task_id, modality=modality, version="v2"
+    )
     try:
         param_value = get_parameter_value(param_name)
         return JSONResponse(
@@ -1095,12 +1127,79 @@ def get_parameter_v2(request: Request):
         )
 
 
+async def put_parameter(request: Request):
+    """Set v1/v2 parameter in AWS param store based on job_type and task_id"""
+    # User must be signed in
+    user = request.session.get("user")
+    if not user:
+        return JSONResponse(
+            content={
+                "message": "User not authenticated",
+                "data": {"error": "User not authenticated"},
+            },
+            status_code=401,
+        )
+    try:
+        # path params
+        param_info = JobParamInfo(
+            name=None,
+            last_modified=None,
+            job_type=request.path_params.get("job_type"),
+            task_id=request.path_params.get("task_id"),
+            modality=request.path_params.get("modality"),
+            version=request.path_params.get("version"),
+        )
+        param_name = JobParamInfo.get_parameter_name(
+            job_type=param_info.job_type,
+            task_id=param_info.task_id,
+            modality=param_info.modality,
+            version=param_info.version,
+        )
+        # update param store
+        logger.info(
+            f"Received request from {user} to set parameter {param_name}"
+        )
+        param_value = await request.json()
+        logger.info(f"Setting parameter {param_name} to {param_value}")
+        result = put_parameter_value(
+            param_name=param_name, param_value=param_value
+        )
+        logger.info(result)
+        return JSONResponse(
+            content={
+                "message": f"Set parameter for {param_name}",
+                "data": param_value,
+            },
+            status_code=200,
+        )
+    except ValidationError as error:
+        return JSONResponse(
+            content={
+                "message": "Invalid parameter",
+                "data": {"errors": json.loads(error.json())},
+            },
+            status_code=400,
+        )
+    except Exception as e:
+        logger.exception(f"Error setting parameter {param_name}: {e}")
+        return JSONResponse(
+            content={
+                "message": f"Error setting parameter {param_name}",
+                "data": {"error": f"{e.__class__.__name__}{e.args}"},
+            },
+            status_code=500,
+        )
+
+
 def get_parameter(request: Request):
     """Get parameter from AWS parameter store based on job_type and task_id"""
     # path params are auto validated
     job_type = request.path_params.get("job_type")
     task_id = request.path_params.get("task_id")
-    param_name = JobParamInfo.get_parameter_name(job_type, task_id)
+    modality = request.path_params.get("modality")
+    param_name = JobParamInfo.get_parameter_name(
+        job_type=job_type, task_id=task_id, modality=modality
+    )
     try:
         param_value = get_parameter_value(param_name)
         return JSONResponse(
@@ -1124,8 +1223,6 @@ def get_parameter(request: Request):
 async def admin(request: Request):
     """Get admin page if authenticated, else redirect to login."""
     user = request.session.get("user")
-    if os.getenv("ENV_NAME") == "local":
-        user = {"name": "local user"}
     if user:
         return templates.TemplateResponse(
             request=request,
@@ -1143,6 +1240,9 @@ async def admin(request: Request):
 
 async def login(request: Request):
     """Redirect to Azure login page"""
+    if os.getenv("ENV_NAME") == "local":
+        request.session["user"] = {"name": "local user"}
+        return RedirectResponse(url="/admin")
     oauth = set_oauth()
     redirect_uri = request.url_for("auth")
     response = await oauth.azure.authorize_redirect(request, redirect_uri)
@@ -1193,7 +1293,13 @@ routes = [
     Route("/api/v1/get_task_logs", endpoint=get_task_logs, methods=["GET"]),
     Route("/api/v1/parameters", endpoint=list_parameters, methods=["GET"]),
     Route(
-        "/api/v1/parameters/job_types/{job_type:str}/tasks/{task_id:path}",
+        "/api/v1/parameters/job_types/{job_type:str}/tasks/{task_id:str}",
+        endpoint=get_parameter,
+        methods=["GET"],
+    ),
+    Route(
+        "/api/v1/parameters/job_types/{job_type:str}/tasks/{task_id:str}"
+        "/{modality:str}",
         endpoint=get_parameter,
         methods=["GET"],
     ),
@@ -1204,10 +1310,28 @@ routes = [
     Route("/api/v2/submit_jobs", endpoint=submit_jobs_v2, methods=["POST"]),
     Route("/api/v2/parameters", endpoint=list_parameters_v2, methods=["GET"]),
     Route(
-        "/api/v2/parameters/job_types/{job_type:str}/tasks/{task_id:path}",
+        "/api/v2/parameters/job_types/{job_type:str}/tasks/{task_id:str}",
+        endpoint=get_parameter_v2,
+        methods=["GET"],
+    ),
+    Route(
+        "/api/v2/parameters/job_types/{job_type:str}/tasks/{task_id:str}"
+        "/{modality:str}",
         endpoint=get_parameter_v2,
         methods=["GET"],
     ),
+    Route(
+        "/api/{version:str}/parameters/job_types/{job_type:str}"
+        "/tasks/{task_id:str}",
+        endpoint=put_parameter,
+        methods=["PUT"],
+    ),
+    Route(
+        "/api/{version:str}/parameters/job_types/{job_type:str}"
+        "/tasks/{task_id:str}/{modality:str}",
+        endpoint=put_parameter,
+        methods=["PUT"],
+    ),
     Route("/jobs", endpoint=jobs, methods=["GET"]),
     Route("/job_tasks_table", endpoint=job_tasks_table, methods=["GET"]),
     Route("/task_logs", endpoint=task_logs, methods=["GET"]),

{aind_data_transfer_service-1.16.0.dist-info → aind_data_transfer_service-1.17.2.dist-info}/METADATA

@@ -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"