deriva-ml 1.17.10__py3-none-any.whl → 1.17.11__py3-none-any.whl

deriva_ml/core/mixins/asset.py

@@ -0,0 +1,384 @@
+"""Asset management mixin for DerivaML.
+
+This module provides the AssetMixin class which handles
+asset table operations including creating, listing, and looking up assets.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Iterable
+
+# Deriva imports - use importlib to avoid shadowing by local 'deriva.py' files
+import importlib
+_ermrest_model = importlib.import_module("deriva.core.ermrest_model")
+Table = _ermrest_model.Table
+
+from deriva_ml.core.definitions import AssetTableDef, ColumnDefinition, MLVocab, RID, VocabularyTerm
+from deriva_ml.core.exceptions import DerivaMLException
+from deriva_ml.schema.annotations import asset_annotation
+
+if TYPE_CHECKING:
+    from deriva_ml.asset.asset import Asset
+    from deriva_ml.execution.execution_record import ExecutionRecord
+    from deriva_ml.model.catalog import DerivaModel
+
+
+class AssetMixin:
+    """Mixin providing asset management operations.
+
+    This mixin requires the host class to have:
+        - model: DerivaModel instance
+        - ml_schema: str - name of the ML schema
+        - domain_schema: str - name of the domain schema
+        - pathBuilder(): method returning catalog path builder
+        - add_term(): method for adding vocabulary terms (from VocabularyMixin)
+        - apply_catalog_annotations(): method to update navbar (from DerivaML base class)
+
+    Methods:
+        create_asset: Create a new asset table
+        list_assets: List contents of an asset table
+    """
+
+    # Type hints for IDE support - actual attributes/methods from host class
+    model: "DerivaModel"
+    ml_schema: str
+    domain_schemas: frozenset[str]
+    default_schema: str | None
+    pathBuilder: Callable[[], Any]
+    add_term: Callable[..., VocabularyTerm]
+    apply_catalog_annotations: Callable[[], None]
+
+    # Note: @validate_call removed because ColumnDefinition is now a dataclass from
+    # deriva.core.typed and Pydantic validation doesn't work well with dataclass fields
+    def create_asset(
+        self,
+        asset_name: str,
+        column_defs: Iterable[ColumnDefinition] | None = None,
+        fkey_defs: Iterable[ColumnDefinition] | None = None,
+        referenced_tables: Iterable[Table] | None = None,
+        comment: str = "",
+        schema: str | None = None,
+        update_navbar: bool = True,
+    ) -> Table:
+        """Creates an asset table.
+
+        Args:
+            asset_name: Name of the asset table.
+            column_defs: Iterable of ColumnDefinition objects to provide additional metadata for asset.
+            fkey_defs: Iterable of ForeignKeyDefinition objects to provide additional metadata for asset.
+            referenced_tables: Iterable of Table objects to which asset should provide foreign-key references to.
+            comment: Description of the asset table. (Default value = '')
+            schema: Schema in which to create the asset table.  Defaults to domain_schema.
+            update_navbar: If True (default), automatically updates the navigation bar to include
+                the new asset table. Set to False during batch asset creation to avoid redundant
+                updates, then call apply_catalog_annotations() once at the end.
+
+        Returns:
+            Table object for the asset table.
+        """
+        # Initialize empty collections if None provided
+        column_defs = column_defs or []
+        fkey_defs = fkey_defs or []
+        referenced_tables = referenced_tables or []
+        schema = schema or self.model._require_default_schema()
+
+        # Add an asset type to vocabulary
+        self.add_term(MLVocab.asset_type, asset_name, description=f"A {asset_name} asset")
+
+        # Create the main asset table
+        # Note: column_defs and fkey_defs should be ColumnDef/ForeignKeyDef objects
+        asset_table = self.model.schemas[schema].create_table(
+            AssetTableDef(
+                schema_name=schema,
+                name=asset_name,
+                columns=list(column_defs),
+                foreign_keys=list(fkey_defs),
+                comment=comment,
+            )
+        )
+
+        # Create an association table between asset and asset type
+        self.model.create_table(
+            Table.define_association(
+                [
+                    (asset_table.name, asset_table),
+                    ("Asset_Type", self.model.name_to_table("Asset_Type")),
+                ]
+            ),
+            schema=schema,
+        )
+
+        # Create references to other tables if specified
+        for t in referenced_tables:
+            asset_table.create_reference(self.model.name_to_table(t))
+
+        # Create an association table for tracking execution
+        atable = self.model.create_table(
+            Table.define_association(
+                [
+                    (asset_name, asset_table),
+                    (
+                        "Execution",
+                        self.model.schemas[self.ml_schema].tables["Execution"],
+                    ),
+                ]
+            ),
+            schema=schema,
+        )
+        atable.create_reference(self.model.name_to_table("Asset_Role"))
+
+        # Add asset annotations
+        asset_annotation(asset_table)
+
+        # Update navbar to include the new asset table
+        if update_navbar:
+            self.apply_catalog_annotations()
+
+        return asset_table
+
+    def list_assets(self, asset_table: Table | str) -> list["Asset"]:
+        """Lists contents of an asset table.
+
+        Returns a list of Asset objects for the specified asset table.
+
+        Args:
+            asset_table: Table or name of the asset table to list assets for.
+
+        Returns:
+            list[Asset]: List of Asset objects for the assets in the table.
+
+        Raises:
+            DerivaMLException: If the table is not an asset table or doesn't exist.
+
+        Example:
+            >>> assets = ml.list_assets("Image")
+            >>> for asset in assets:
+            ...     print(f"{asset.asset_rid}: {asset.filename}")
+        """
+        from deriva_ml.asset.asset import Asset
+
+        # Validate and get asset table reference
+        asset_table_obj = self.model.name_to_table(asset_table)
+        if not self.model.is_asset(asset_table_obj):
+            raise DerivaMLException(f"Table {asset_table_obj.name} is not an asset")
+
+        # Get path builders for asset and type tables
+        pb = self.pathBuilder()
+        asset_path = pb.schemas[asset_table_obj.schema.name].tables[asset_table_obj.name]
+        (
+            asset_type_table,
+            _,
+            _,
+        ) = self.model.find_association(asset_table_obj, MLVocab.asset_type)
+        type_path = pb.schemas[asset_type_table.schema.name].tables[asset_type_table.name]
+
+        # Build a list of Asset objects
+        assets = []
+        for asset_record in asset_path.entities().fetch():
+            # Get associated asset types for each asset
+            asset_types = (
+                type_path.filter(type_path.columns[asset_table_obj.name] == asset_record["RID"])
+                .attributes(type_path.Asset_Type)
+                .fetch()
+            )
+            asset_type_list = [asset_type[MLVocab.asset_type.value] for asset_type in asset_types]
+
+            assets.append(Asset(
+                catalog=self,  # type: ignore[arg-type]
+                asset_rid=asset_record["RID"],
+                asset_table=asset_table_obj.name,
+                filename=asset_record.get("Filename", ""),
+                url=asset_record.get("URL", ""),
+                length=asset_record.get("Length", 0),
+                md5=asset_record.get("MD5", ""),
+                description=asset_record.get("Description", ""),
+                asset_types=asset_type_list,
+            ))
+        return assets
+
+    def list_asset_executions(
+        self, asset_rid: str, asset_role: str | None = None
+    ) -> list["ExecutionRecord"]:
+        """List all executions associated with an asset.
+
+        Given an asset RID, returns a list of executions that created or used
+        the asset, along with the role (Input/Output) in each execution.
+
+        Args:
+            asset_rid: The RID of the asset to look up.
+            asset_role: Optional filter for asset role ('Input' or 'Output').
+                If None, returns all associations.
+
+        Returns:
+            list[ExecutionRecord]: List of ExecutionRecord objects for the
+                executions associated with this asset.
+
+        Raises:
+            DerivaMLException: If the asset RID is not found or not an asset.
+
+        Example:
+            >>> # Find all executions that created this asset
+            >>> executions = ml.list_asset_executions("1-abc123", asset_role="Output")
+            >>> for exe in executions:
+            ...     print(f"Created by execution {exe.execution_rid}")
+
+            >>> # Find all executions that used this asset as input
+            >>> executions = ml.list_asset_executions("1-abc123", asset_role="Input")
+        """
+        # Resolve the RID to find which asset table it belongs to
+        rid_info = self.resolve_rid(asset_rid)  # type: ignore[attr-defined]
+        asset_table = rid_info.table
+
+        if not self.model.is_asset(asset_table):
+            raise DerivaMLException(f"RID {asset_rid} is not an asset (table: {asset_table.name})")
+
+        # Find the association table between this asset table and Execution
+        asset_exe_table, asset_fk, execution_fk = self.model.find_association(asset_table, "Execution")
+
+        # Build the query
+        pb = self.pathBuilder()
+        asset_exe_path = pb.schemas[asset_exe_table.schema.name].tables[asset_exe_table.name]
+
+        # Filter by asset RID
+        query = asset_exe_path.filter(asset_exe_path.columns[asset_fk] == asset_rid)
+
+        # Optionally filter by asset role
+        if asset_role:
+            query = query.filter(asset_exe_path.Asset_Role == asset_role)
+
+        # Convert to ExecutionRecord objects
+        records = list(query.entities().fetch())
+        return [self.lookup_execution(record["Execution"]) for record in records]  # type: ignore[attr-defined]
+
+    def lookup_asset(self, asset_rid: RID) -> "Asset":
+        """Look up an asset by its RID.
+
+        Returns an Asset object for the specified RID. The asset can be from
+        any asset table in the catalog.
+
+        Args:
+            asset_rid: The RID of the asset to look up.
+
+        Returns:
+            Asset object for the specified RID.
+
+        Raises:
+            DerivaMLException: If the RID is not found or is not an asset.
+
+        Example:
+            >>> asset = ml.lookup_asset("3JSE")
+            >>> print(f"File: {asset.filename}, Table: {asset.asset_table}")
+        """
+        from deriva_ml.asset.asset import Asset
+
+        # Resolve the RID to find which table it belongs to
+        rid_info = self.resolve_rid(asset_rid)  # type: ignore[attr-defined]
+        asset_table = rid_info.table
+
+        if not self.model.is_asset(asset_table):
+            raise DerivaMLException(f"RID {asset_rid} is not an asset (table: {asset_table.name})")
+
+        # Query the asset table for this record
+        pb = self.pathBuilder()
+        asset_path = pb.schemas[asset_table.schema.name].tables[asset_table.name]
+
+        records = list(asset_path.filter(asset_path.RID == asset_rid).entities().fetch())
+        if not records:
+            raise DerivaMLException(f"Asset {asset_rid} not found in table {asset_table.name}")
+
+        record = records[0]
+
+        # Get asset types
+        asset_types = []
+        try:
+            type_assoc_table, asset_fk, _ = self.model.find_association(asset_table, "Asset_Type")
+            type_path = pb.schemas[type_assoc_table.schema.name].tables[type_assoc_table.name]
+            types = list(
+                type_path.filter(type_path.columns[asset_fk] == asset_rid)
+                .attributes(type_path.Asset_Type)
+                .fetch()
+            )
+            asset_types = [t["Asset_Type"] for t in types]
+        except Exception:
+            pass  # No type association for this asset table
+
+        return Asset(
+            catalog=self,  # type: ignore[arg-type]
+            asset_rid=asset_rid,
+            asset_table=asset_table.name,
+            filename=record.get("Filename", ""),
+            url=record.get("URL", ""),
+            length=record.get("Length", 0),
+            md5=record.get("MD5", ""),
+            description=record.get("Description", ""),
+            asset_types=asset_types,
+        )
+
+    def list_asset_tables(self) -> list[Table]:
+        """List all asset tables in the catalog.
+
+        Returns:
+            List of Table objects that are asset tables.
+
+        Example:
+            >>> for table in ml.list_asset_tables():
+            ...     print(f"Asset table: {table.name}")
+        """
+        tables = []
+        # Include asset tables from all domain schemas
+        for domain_schema in self.domain_schemas:
+            if domain_schema in self.model.schemas:
+                tables.extend([
+                    t for t in self.model.schemas[domain_schema].tables.values()
+                    if self.model.is_asset(t)
+                ])
+        # Also include ML schema asset tables (like Execution_Asset)
+        tables.extend([
+            t for t in self.model.schemas[self.ml_schema].tables.values()
+            if self.model.is_asset(t)
+        ])
+        return tables
+
+    def find_assets(
+        self,
+        asset_table: Table | str | None = None,
+        asset_type: str | None = None,
+    ) -> Iterable["Asset"]:
+        """Find assets in the catalog.
+
+        Returns an iterable of Asset objects matching the specified criteria.
+        If no criteria are specified, returns all assets from all asset tables.
+
+        Args:
+            asset_table: Optional table or table name to search. If None, searches
+                all asset tables.
+            asset_type: Optional asset type to filter by. Only returns assets
+                with this type.
+
+        Returns:
+            Iterable of Asset objects matching the criteria.
+
+        Example:
+            >>> # Find all assets in the Model table
+            >>> models = list(ml.find_assets(asset_table="Model"))
+
+            >>> # Find all assets with type "Training_Data"
+            >>> training = list(ml.find_assets(asset_type="Training_Data"))
+
+            >>> # Find all assets across all tables
+            >>> all_assets = list(ml.find_assets())
+        """
+        # Determine which tables to search
+        if asset_table is not None:
+            tables = [self.model.name_to_table(asset_table)]
+        else:
+            tables = self.list_asset_tables()
+
+        for table in tables:
+            # Get all assets from this table (now returns Asset objects)
+            for asset in self.list_assets(table):
+                # Filter by asset type if specified
+                if asset_type is not None:
+                    if asset_type not in asset.asset_types:
+                        continue
+                yield asset

deriva_ml/core/mixins/dataset.py

@@ -0,0 +1,237 @@
+"""Dataset management mixin for DerivaML.
+
+This module provides the DatasetMixin class which handles
+dataset operations including finding, creating, looking up,
+deleting, and managing dataset elements.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Iterable
+
+# Deriva imports - use importlib to avoid shadowing by local 'deriva.py' files
+import importlib
+_ermrest_model = importlib.import_module("deriva.core.ermrest_model")
+Table = _ermrest_model.Table
+
+from pydantic import ConfigDict, validate_call
+
+from deriva_ml.core.definitions import RID, MLVocab
+from deriva_ml.core.exceptions import DerivaMLException, DerivaMLTableTypeError
+from deriva_ml.dataset.aux_classes import DatasetSpec
+
+if TYPE_CHECKING:
+    from deriva_ml.dataset.dataset import Dataset
+    from deriva_ml.dataset.dataset_bag import DatasetBag
+    from deriva_ml.model.catalog import DerivaModel
+
+
+class DatasetMixin:
+    """Mixin providing dataset management operations.
+
+    This mixin requires the host class to have:
+        - model: DerivaModel instance
+        - ml_schema: str - name of the ML schema
+        - domain_schema: str - name of the domain schema
+        - s3_bucket: str | None - S3 bucket URL for dataset storage
+        - use_minid: bool - whether to use MINIDs
+        - pathBuilder(): method returning catalog path builder
+        - _dataset_table: property returning the Dataset table
+
+    Methods:
+        find_datasets: List all datasets in the catalog
+        create_dataset: Create a new dataset
+        lookup_dataset: Look up a dataset by RID or spec
+        delete_dataset: Delete a dataset
+        list_dataset_element_types: List types that can be added to datasets
+        add_dataset_element_type: Add a new element type to datasets
+        download_dataset_bag: Download a dataset as a bag
+    """
+
+    # Type hints for IDE support - actual attributes/methods from host class
+    model: "DerivaModel"
+    ml_schema: str
+    domain_schemas: frozenset[str]
+    default_schema: str | None
+    s3_bucket: str | None
+    use_minid: bool
+    pathBuilder: Callable[[], Any]
+
+    @property
+    def _dataset_table(self) -> Table:
+        """Get the Dataset table. Must be provided by host class."""
+        raise NotImplementedError
+
+    def find_datasets(self, deleted: bool = False) -> Iterable["Dataset"]:
+        """List all datasets in the catalog.
+
+        Args:
+            deleted: If True, include datasets that have been marked as deleted.
+
+        Returns:
+            Iterable of Dataset objects.
+
+        Example:
+            >>> datasets = list(ml.find_datasets())
+            >>> for ds in datasets:
+            ...     print(f"{ds.dataset_rid}: {ds.description}")
+        """
+        # Import here to avoid circular imports
+        from deriva_ml.dataset.dataset import Dataset
+
+        # Get datapath to the Dataset table
+        pb = self.pathBuilder()
+        dataset_path = pb.schemas[self._dataset_table.schema.name].tables[self._dataset_table.name]
+
+        if deleted:
+            filtered_path = dataset_path
+        else:
+            filtered_path = dataset_path.filter(
+                (dataset_path.Deleted == False) | (dataset_path.Deleted == None)  # noqa: E711, E712
+            )
+
+        # Create Dataset objects - dataset_types is now a property that fetches from catalog
+        datasets = []
+        for dataset in filtered_path.entities().fetch():
+            datasets.append(
+                Dataset(
+                    self,  # type: ignore[arg-type]
+                    dataset_rid=dataset["RID"],
+                    description=dataset["Description"],
+                )
+            )
+        return datasets
+
+    def lookup_dataset(self, dataset: RID | DatasetSpec, deleted: bool = False) -> "Dataset":
+        """Look up a dataset by RID or DatasetSpec.
+
+        Args:
+            dataset: Dataset RID or DatasetSpec to look up.
+            deleted: If True, include datasets that have been marked as deleted.
+
+        Returns:
+            Dataset: The dataset object for the specified RID.
+
+        Raises:
+            DerivaMLException: If the dataset is not found.
+
+        Example:
+            >>> dataset = ml.lookup_dataset("4HM")
+            >>> print(f"Version: {dataset.current_version}")
+        """
+        if isinstance(dataset, DatasetSpec):
+            dataset_rid = dataset.rid
+        else:
+            dataset_rid = dataset
+
+        try:
+            return [ds for ds in self.find_datasets(deleted=deleted) if ds.dataset_rid == dataset_rid][0]
+        except IndexError:
+            raise DerivaMLException(f"Dataset {dataset_rid} not found.")
+
+    def delete_dataset(self, dataset: "Dataset", recurse: bool = False) -> None:
+        """Delete a dataset from the catalog.
+
+        Args:
+            dataset: The dataset to delete.
+            recurse: If True, delete the dataset along with any nested datasets. (Default value = False)
+        """
+        # Get association table entries for this dataset_table
+        # Delete association table entries
+        dataset_rid = dataset.dataset_rid
+        if not self.model.is_dataset_rid(dataset.dataset_rid):
+            raise DerivaMLException("Dataset_rid is not a dataset.")
+
+        if parents := dataset.list_dataset_parents():
+            raise DerivaMLException(f'Dataset "{dataset}" is in a nested dataset: {parents}.')
+
+        pb = self.pathBuilder()
+        dataset_path = pb.schemas[self._dataset_table.schema.name].tables[self._dataset_table.name]
+
+        # list_dataset_children returns Dataset objects, so extract their RIDs
+        child_rids = [ds.dataset_rid for ds in dataset.list_dataset_children()] if recurse else []
+        rid_list = [dataset_rid] + child_rids
+        dataset_path.update([{"RID": r, "Deleted": True} for r in rid_list])
+
+    def list_dataset_element_types(self) -> Iterable[Table]:
+        """List the types of entities that can be added to a dataset.
+
+        Returns:
+            An iterable of Table objects that can be included as an element of a dataset.
+        """
+
+        def is_domain_or_dataset_table(table: Table) -> bool:
+            return self.model.is_domain_schema(table.schema.name) or table.name == self._dataset_table.name
+
+        return [t for a in self._dataset_table.find_associations() if is_domain_or_dataset_table(t := a.other_fkeys.pop().pk_table)]
+
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def add_dataset_element_type(self, element: str | Table) -> Table:
+        """Makes it possible to add objects from the specified table to a dataset.
+
+        A dataset is a heterogeneous collection of objects, each of which comes from a different table.
+        This routine adds the specified table as a valid element type for datasets.
+
+        Args:
+            element: Name of the table or table object that is to be added to the dataset.
+
+        Returns:
+            The table object that was added to the dataset.
+        """
+        # Import here to avoid circular imports
+        from deriva_ml.dataset.catalog_graph import CatalogGraph
+
+        # Add table to map
+        element_table = self.model.name_to_table(element)
+        atable_def = Table.define_association([self._dataset_table, element_table])
+        try:
+            table = self.model.create_table(atable_def)
+        except ValueError as e:
+            if "already exists" in str(e):
+                table = self.model.name_to_table(atable_def["table_name"])
+            else:
+                raise e
+
+        # self.model = self.catalog.getCatalogModel()
+        annotations = CatalogGraph(self, s3_bucket=self.s3_bucket, use_minid=self.use_minid).generate_dataset_download_annotations()  # type: ignore[arg-type]
+        self._dataset_table.annotations.update(annotations)
+        self.model.model.apply()
+        return table
+
+    def download_dataset_bag(
+        self,
+        dataset: DatasetSpec,
+    ) -> "DatasetBag":
+        """Downloads a dataset to the local filesystem.
+
+        Downloads a dataset specified by DatasetSpec to the local filesystem. If the catalog
+        has s3_bucket configured and use_minid is enabled, the bag will be uploaded to S3
+        and registered with the MINID service.
+
+        Args:
+            dataset: Specification of the dataset to download, including version and materialization options.
+
+        Returns:
+            DatasetBag: Object containing:
+                - path: Local filesystem path to downloaded dataset
+                - rid: Dataset's Resource Identifier
+                - minid: Dataset's Minimal Viable Identifier (if MINID enabled)
+
+        Note:
+            MINID support requires s3_bucket to be configured when creating the DerivaML instance.
+            The catalog's use_minid setting controls whether MINIDs are created.
+
+        Examples:
+            Download with default options:
+                >>> spec = DatasetSpec(rid="1-abc123")
+                >>> bag = ml.download_dataset_bag(dataset=spec)
+                >>> print(f"Downloaded to {bag.path}")
+        """
+        if not self.model.is_dataset_rid(dataset.rid):
+            raise DerivaMLTableTypeError("Dataset", dataset.rid)
+        ds = self.lookup_dataset(dataset)
+        return ds.download_dataset_bag(
+            version=dataset.version,
+            materialize=dataset.materialize,
+            use_minid=self.use_minid,
+        )