Merge pull request #176 from AllenNeuralDynamics/release-v0.21.0

Release v0.21.0

Author: helen-m-lin

pyproject.toml

@@ -23,8 +23,8 @@ dependencies = [
     "pydantic>=2.10",
     "dask==2023.5.0",
     "aind-data-schema==1.4.0",
-    "codeocean==0.3.0",
-    "aind-data-access-api>=1.6.0",
+    "codeocean>=0.10.0",
+    "aind-data-access-api>=1.9.0",
 ]
 
 [project.optional-dependencies]

src/aind_data_asset_indexer/__init__.py

@@ -1,3 +1,3 @@
 """Package"""
 
-__version__ = "0.20.2"
+__version__ = "0.21.0"

src/aind_data_asset_indexer/codeocean_bucket_indexer.py

@@ -2,28 +2,24 @@
 DocDB."""
 
 import argparse
-import json
 import logging
 import os
 import sys
 import warnings
 from typing import List, Optional
 
-import boto3
 import dask.bag as dask_bag
 import requests
 from aind_data_access_api.document_db import MetadataDbClient
-from aind_data_access_api.utils import get_s3_bucket_and_prefix, paginate_docdb
+from aind_data_access_api.utils import paginate_docdb
 from aind_data_schema.core.metadata import ExternalPlatforms
 from codeocean import CodeOcean
 from codeocean.data_asset import DataAssetSearchOrigin, DataAssetSearchParams
-from mypy_boto3_s3 import S3Client
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
 
 from aind_data_asset_indexer.models import CodeOceanIndexBucketJobSettings
 from aind_data_asset_indexer.utils import (
-    build_metadata_record_from_prefix,
     get_all_processed_codeocean_asset_records,
 )
 
@@ -264,48 +260,31 @@ def _process_codeocean_record(
         self,
         codeocean_record: dict,
         docdb_client: MetadataDbClient,
-        s3_client: S3Client,
     ):
         """
         Processes a code ocean record. It's assumed that the check to verify
         the record is not in DocDB is done upstream.
-        1) Using the s3 location in the codeocean record, build metadata file.
-        2) Save metadata record to DocDB if no issue
+        1) Uses the codeocean record info to register the asset to DocDB.
 
         Parameters
         ----------
         codeocean_record : dict
         docdb_client : MetadataDbClient
-        s3_client : S3Client
 
         """
-        location = codeocean_record["location"]
-        created = codeocean_record["created"]
-        external_links = codeocean_record["external_links"]
         name = codeocean_record["name"]
-        url_parts = get_s3_bucket_and_prefix(location)
-        bucket = url_parts["bucket"]
-        s3_prefix = url_parts["prefix"]
-        new_metadata_contents = build_metadata_record_from_prefix(
-            bucket=bucket,
-            prefix=s3_prefix,
-            s3_client=s3_client,
-            optional_name=name,
-            optional_created=created,
-            optional_external_links=external_links,
+        location = codeocean_record["location"]
+        co_asset_id = codeocean_record["co_asset_id"]
+        co_computation_id = codeocean_record["co_computation_id"]
+
+        logging.info(f"Uploading metadata record for: {location}")
+        register_response = docdb_client.register_co_result(
+            s3_location=location,
+            name=name,
+            co_asset_id=co_asset_id,
+            co_computation_id=co_computation_id,
         )
-        if new_metadata_contents is not None:
-            logging.info(f"Uploading metadata record for: {location}")
-            # noinspection PyTypeChecker
-            json_contents = json.loads(new_metadata_contents)
-            response = docdb_client.insert_one_docdb_record(
-                record=json_contents
-            )
-            logging.debug(response.json())
-        else:
-            logging.warning(
-                f"Unable to build metadata record for: {location}!"
-            )
+        logging.info(register_response)
 
     def _dask_task_to_process_record_list(self, record_list: List[dict]):
         """
@@ -318,15 +297,13 @@ def _dask_task_to_process_record_list(self, record_list: List[dict]):
         record_list : List[dict]
 
         """
-        # create clients here since dask doesn't serialize them
-        s3_client = boto3.client("s3")
+        # create client here since dask doesn't serialize them
         with self._create_docdb_client() as doc_db_client:
             for record in record_list:
                 try:
                     self._process_codeocean_record(
                         codeocean_record=record,
                         docdb_client=doc_db_client,
-                        s3_client=s3_client,
                     )
                 except requests.HTTPError as e:
                     logging.error(
@@ -337,7 +314,6 @@ def _dask_task_to_process_record_list(self, record_list: List[dict]):
                     logging.error(
                         f'Error processing {record.get("location")}: {repr(e)}'
                     )
-        s3_client.close()
 
     def _process_codeocean_records(self, records: List[dict]):
         """

src/aind_data_asset_indexer/utils.py

@@ -4,7 +4,7 @@
 import json
 import logging
 import re
-from datetime import datetime, timezone
+from datetime import datetime
 from json.decoder import JSONDecodeError
 from typing import Dict, Iterator, List, Optional
 from urllib.parse import urlparse
@@ -13,7 +13,6 @@
 from aind_data_schema.core.data_description import DataLevel, DataRegex
 from aind_data_schema.core.metadata import CORE_FILES as CORE_SCHEMAS
 from aind_data_schema.core.metadata import (
-    ExternalPlatforms,
     Metadata,
     create_metadata_json,
 )
@@ -798,12 +797,13 @@ def get_all_processed_codeocean_asset_records(
     Returns
     -------
     Dict[str, dict]
-      {data_asset_location:
-      {"name": data_asset_name,
-      "location": data_asset_location,
-      "created": data_asset_created,
-      "external_links": {"Code Ocean": [data_asset_id]}
-      }
+      {
+        data_asset_location: {
+            "name": data_asset_name,
+            "location": data_asset_location,
+            "co_asset_id": data_asset_id,
+            "co_computation_id": data_asset_computation_id,
+        }
       }
 
     """
@@ -826,24 +826,33 @@ def get_all_processed_codeocean_asset_records(
         for data_asset_info in iter_response:
             data_asset_id = data_asset_info.id
             data_asset_name = data_asset_info.name
-            created_timestamp = data_asset_info.created
-            created_datetime = datetime.fromtimestamp(
-                created_timestamp, tz=timezone.utc
-            )
+            # Get source computation id from provenance
+            if (
+                data_asset_info.provenance is not None
+                and data_asset_info.provenance.computation is not None
+            ):
+                data_asset_computation_id = (
+                    data_asset_info.provenance.computation
+                )
+            else:
+                data_asset_computation_id = None
+                logging.warning(
+                    f"Data asset {data_asset_id}, {data_asset_name} does not"
+                    "have computation provenance!"
+                )
             # Results hosted externally have a source_bucket field
             is_external = data_asset_info.source_bucket is not None
             if (
                 not is_external
                 and data_asset_info.state == DataAssetState.Ready
+                and data_asset_computation_id is not None
             ):
                 location = f"s3://{co_data_asset_bucket}/{data_asset_id}"
                 extracted_info[location] = {
                     "name": data_asset_name,
                     "location": location,
-                    "created": created_datetime,
-                    "external_links": {
-                        ExternalPlatforms.CODEOCEAN.value: [data_asset_id]
-                    },
+                    "co_asset_id": data_asset_id,
+                    "co_computation_id": data_asset_computation_id,
                 }
         # Occasionally, there are duplicate items returned. This is one
         # way to remove the duplicates.

tests/resources/utils/example_search_co_assets.json

@@ -16,6 +16,7 @@
          "provenance": {
             "capsule": "8f8f888f-8888-88f8-fff8-888888ff8f8f",
             "commit": "3333333333333f33be33ae33fb33333333e33333",
+            "computation": "77a777aa-a77a-7a77-a7aa-77777a777aaa",
             "data_assets": [
                "aaa3a3aa-a33a-3aa3-a33a-aa3aaa333a3a"
             ],
@@ -47,6 +48,7 @@
          "provenance": {
             "capsule": "8f8f888f-8888-88f8-fff8-888888ff8f8f",
             "commit": "3333333333333f33be33ae33fb33333333e33333",
+            "computation": "2ff2222f-ff22-2f22-2222-ff22fff22f22",
             "data_assets": [
                "d11dd111-111d-1d11-d111-d1d11111dd11"
             ],
@@ -61,6 +63,61 @@
             "ecephys"
          ],
          "type": "result"
+      },
+      {
+         "created": 1720045200,
+         "custom_metadata": {
+            "data level": "derived",
+            "experiment type": "ecephys",
+            "subject id": "734567"
+         },
+         "description": "",
+         "files": 8921,
+         "id": "888888ee-88ee-8e88-888e-8e88e8888888",
+         "last_used": 1720045470,
+         "mount": "ecephys_734567_2024-07-03_09-15-30_sorted_2024-07-04_08-00-00",
+         "name": "ecephys_734567_2024-07-03_09-15-30_sorted_2024-07-04_08-00-00",
+         "provenance": null,
+         "size": 17234567890,
+         "state": "ready",
+         "tags": [
+            "derived",
+            "734567",
+            "ecephys"
+         ],
+         "type": "result"
+      },
+      {
+         "created": 1719435600,
+         "custom_metadata": {
+            "data level": "derived",
+            "experiment type": "ecephys",
+            "subject id": "723456"
+         },
+         "description": "",
+         "files": 7543,
+         "id": "999999dd-99dd-9d99-999d-9d99d9999999",
+         "last_used": 1719435870,
+         "mount": "ecephys_723456_2024-06-26_14-30-15_sorted_2024-06-27_10-00-00",
+         "name": "ecephys_723456_2024-06-26_14-30-15_sorted_2024-06-27_10-00-00",
+         "provenance": {
+            "capsule": "8f8f888f-8888-88f8-fff8-888888ff8f8f",
+            "commit": "3333333333333f33be33ae33fb33333333e33333",
+            "data_assets": [
+               "e22ee222-222e-2e22-e222-e2e22222ee22"
+            ],
+            "computation": null,
+            "docker_image": "",
+            "run_script": "code/run"
+         },
+         "size": 15847923456,
+         "state": "ready",
+         "tags": [
+            "derived",
+            "723456",
+            "ecephys"
+         ],
+         "type": "result"
       }
    ]
 }