aind-data-transfer-service 1.17.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.17.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/server.py

@@ -10,7 +10,6 @@ from pathlib import PurePosixPath
 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
 
 
@@ -281,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:
@@ -375,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):
@@ -431,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)
@@ -491,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):
@@ -499,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}"
@@ -511,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:
@@ -551,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}"
@@ -569,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": []},
             },
         )
 
@@ -645,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
@@ -769,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()
@@ -842,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(
                 [
@@ -876,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}")
@@ -906,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

{aind_data_transfer_service-1.17.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.17.0
+Version: 1.17.2
 Summary: Service that handles requests to upload data to the cloud
 Author: Allen Institute for Neural Dynamics
 License: MIT

aind_data_transfer_service-1.17.2.dist-info/RECORD

@@ -0,0 +1,18 @@
+aind_data_transfer_service/__init__.py,sha256=cCMYl-lxHhXlEMAHdJ7INgHPd0O6ETEp9wsqf6DT8TQ,272
+aind_data_transfer_service/log_handler.py,sha256=c7a-gLmZeRpeCUBwCz6XsTszWXQeQdR7eKZtas4llXM,1700
+aind_data_transfer_service/server.py,sha256=CQdOkX86qn8RQm36NqS5b0aUyNCjvaEbTIObaza6RY0,48910
+aind_data_transfer_service/configs/__init__.py,sha256=9W5GTuso9Is1B9X16RXcdb_GxasZvj6qDzOBDv0AbTc,36
+aind_data_transfer_service/configs/csv_handler.py,sha256=hCdfAYZW_49-l1rbua5On2Tw2ks674Z-MgB_NJlIkU4,5746
+aind_data_transfer_service/configs/job_configs.py,sha256=T-h5N6lyY9xTZ_xg_5FxkyYuMdagApbE6xalxFQ-bqA,18848
+aind_data_transfer_service/configs/job_upload_template.py,sha256=aC5m1uD_YcpbggFQ-yZ7ZJSUUGX1yQqQLF3SwmljrLk,5127
+aind_data_transfer_service/hpc/__init__.py,sha256=YNc68YNlmXwKIPFMIViz_K4XzVVHkLPEBOFyO5DKMKI,53
+aind_data_transfer_service/hpc/client.py,sha256=-WwNjpV2FAJMiXT3Y_JCllSbXPyvqrkPxpNWE0BAWzY,5084
+aind_data_transfer_service/hpc/models.py,sha256=-7HhV16s_MUyKPy0x0FGIbnq8DPL2qJAzJO5G7003AE,16184
+aind_data_transfer_service/models/__init__.py,sha256=Meym73bEZ9nQr4QoeyhQmV3nRTYtd_4kWKPNygsBfJg,25
+aind_data_transfer_service/models/core.py,sha256=uXtPUqjxKalg-sE8MxaJr11w_T_KKBRBSJuUgwoMZlQ,10135
+aind_data_transfer_service/models/internal.py,sha256=tWO1yRMu9hHLv0mt7QhOwTPWGbKApWmCf1wwajx5qjI,10681
+aind_data_transfer_service-1.17.2.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
+aind_data_transfer_service-1.17.2.dist-info/METADATA,sha256=DotVNzRhHX0woD5qMefPipnYAja4EbXje8luub37g9I,2478
+aind_data_transfer_service-1.17.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aind_data_transfer_service-1.17.2.dist-info/top_level.txt,sha256=XmxH0q27Jholj2-VYh-6WMrh9Lw6kkuCX_fdsj3SaFE,27
+aind_data_transfer_service-1.17.2.dist-info/RECORD,,

aind_data_transfer_service/templates/admin.html

@@ -1,45 +0,0 @@
-<!DOCTYPE html>
-<html>
-<head>
-  <meta charset="UTF-8">
-  <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/css/bootstrap.min.css" rel="stylesheet">
-  <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.3/dist/js/bootstrap.bundle.min.js"></script>
-  <title>{% block title %} {% endblock %} AIND Data Transfer Service Admin</title>
-  <style>
-    body {
-      margin: 20px;
-      font-family: arial, sans-serif;
-    }
-    nav {
-      height: 40px;
-    }
-  </style>
-</head>
-<body>
-  <nav>
-    <a href="/">Submit Jobs</a> |
-    <a href="/jobs">Job Status</a> |
-    <a href="/job_params">Job Parameters</a> |
-    <a title="Download job template as .xslx" href="/api/job_upload_template" download>Job Submit Template</a> |
-    <a title="List of project names" href="{{ project_names_url }}" target="_blank">Project Names</a> |
-    <a title="For more information click here" href="https://aind-data-transfer-service.readthedocs.io"
-      target="_blank">Help</a> |
-    <a href="/admin">Admin</a>
-    <a href="/logout" class="float-end">Log out</a>
-  </nav>
-  <div>
-    <h3>Admin</h3>
-    <div>Hello {{user_name}}, welcome to the admin page</div>
-    <div>Email: {{user_email}}</div>
-    <h4 class="mt-4">Job Type Parameters</h4>
-     <div class="mb-2">
-      <div class="alert alert-info mb-0 p-4" role="alert">
-        To manage job type parameters, go to the <a href="/job_params" class="alert-link">Job Parameters</a> page.
-        Select an existing parameter, or click the
-        <button type="button" class="btn btn-success btn-sm" disabled><i class="bi bi-plus-circle"></i> Add New Parameter</button>
-        button.
-      </div>
-    </div>
-  </div>
-</body>
-</html>