nomad-ml-workflows 0.0.6__py3-none-any.whl

nomad_ml_workflows/actions/__init__.py

@@ -0,0 +1,3 @@
+from nomad_ml_workflows.actions.export_entries import export_entries
+
+__all__ = ['export_entries']

nomad_ml_workflows/actions/export_entries/__init__.py

@@ -0,0 +1,53 @@
+from nomad.actions import TaskQueue
+from pydantic import Field
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from nomad.config.models.plugins import ActionEntryPoint
+
+
+class ExportEntriesActionEntryPoint(ActionEntryPoint):
+    search_batch_timeout: int = Field(
+        default=7200,  # 2 hours
+        description='Timeout (in seconds) for each search batch in the Export Entries '
+        'action. Set this accordingly to time out longer searches.',
+    )
+    max_entries_export_limit: int = Field(
+        default=100000,
+        description='Maximum number of entries that can be exported in a single '
+        'Export Entries action.',
+    )
+
+    def load(self):
+        from nomad.actions import Action
+
+        from nomad_ml_workflows.actions.export_entries.activities import (
+            cleanup_artifacts,
+            create_artifact_subdirectory,
+            export_dataset_to_upload,
+            merge_output_files,
+            search,
+        )
+        from nomad_ml_workflows.actions.export_entries.workflows import (
+            ExportEntriesWorkflow,
+        )
+
+        return Action(
+            task_queue=self.task_queue,
+            workflow=ExportEntriesWorkflow,
+            activities=[
+                create_artifact_subdirectory,
+                search,
+                merge_output_files,
+                export_dataset_to_upload,
+                cleanup_artifacts,
+            ],
+        )
+
+
+export_entries = ExportEntriesActionEntryPoint(
+    name='Export Entries Action',
+    description='An action to search entries and export them as a zip file in the '
+    'specified upload.',
+    task_queue=TaskQueue.CPU,
+)

nomad_ml_workflows/actions/export_entries/activities.py

@@ -0,0 +1,209 @@
+import json
+import os
+import shutil
+import zipfile
+from datetime import datetime, timezone
+
+from nomad.actions.manager import action_artifacts_dir, get_upload_files
+from nomad.files import StagingUploadFiles
+from nomad.search import search as nomad_search
+from temporalio import activity
+
+from nomad_ml_workflows.actions.export_entries.models import (
+    CleanupArtifactsInput,
+    CreateArtifactSubdirectoryInput,
+    ExportDatasetInput,
+    MergeOutputFilesInput,
+    SearchInput,
+    SearchOutput,
+)
+from nomad_ml_workflows.actions.export_entries.utils import (
+    merge_files,
+    write_json_file,
+    write_parquet_file,
+)
+
+
+@activity.defn
+async def create_artifact_subdirectory(data: CreateArtifactSubdirectoryInput) -> str:
+    """
+    Creates a subdirectory within the action artifacts directory.
+
+    Args:
+        data (CreateArtifactSubdirectoryInput): Input data for creating subdirectory.
+
+    Returns:
+        str: Path to the created subdirectory.
+    """
+
+    subdir_path = os.path.join(action_artifacts_dir(), data.subdir_name)
+
+    assert not os.path.exists(subdir_path), (
+        f'Artifact subdirectory "{subdir_path}" already exists.'
+    )
+
+    os.makedirs(subdir_path)
+
+    return subdir_path
+
+
+@activity.defn
+async def search(data: SearchInput) -> SearchOutput:
+    """
+    Activity to perform NOMAD search based on the provided input data. The search
+    results are written to a file in the specified format (Parquet or JSON) in the
+    artifacts directory.
+
+    Args:
+        data (SearchInput): Input data for the search activity.
+
+    Returns:
+        SearchOutput: Output data from the search activity.
+    """
+
+    write_dataset_file = {
+        'parquet': write_parquet_file,
+        'json': write_json_file,
+    }.get(data.batch_file_type)
+    if write_dataset_file is None:
+        raise ValueError(f'Unsupported batch file type "{data.batch_file_type}". ')
+
+    start = datetime.now(timezone.utc).isoformat()
+    response = nomad_search(
+        user_id=data.user_id,
+        owner=data.owner,
+        query=data.query,
+        required=data.required,
+        pagination=data.pagination,
+        aggregations={},  # aggregations support can be added later
+    )
+    end = datetime.now(timezone.utc).isoformat()
+
+    # Limit the number of exported entries
+    if len(response.data) > data.max_entries_export_limit:
+        entry_list = response.data[: data.max_entries_export_limit]
+    else:
+        entry_list = response.data
+
+    output = SearchOutput(
+        search_start_time=start,
+        search_end_time=end,
+        num_entries_exported=len(entry_list),
+        num_entries_available=response.pagination.total,
+        pagination_next_page_after_value=response.pagination.next_page_after_value,
+    )
+
+    if len(entry_list) == 0:
+        # skip writing empty files and stop subsequent searches
+        output.pagination_next_page_after_value = None
+    else:
+        write_dataset_file(path=data.output_file_path, data=entry_list)
+
+    return output
+
+
+@activity.defn
+async def merge_output_files(data: MergeOutputFilesInput) -> str | None:
+    """
+    Activity to merge multiple batch files into a single file.
+
+    Args:
+        data (MergeOutputFilesInput): Input data for merging files.
+
+    Returns:
+        str | None: Path of the merged output file, or None if no files were merged.
+    """
+
+    if not data.generated_file_paths:
+        raise ValueError('No generated file paths provided for merging.')
+
+    merged_file_path = os.path.join(
+        data.artifact_subdirectory, 'data.' + data.output_file_type
+    )
+
+    merge_files(data.generated_file_paths, data.output_file_type, merged_file_path)
+
+    return merged_file_path
+
+
+@activity.defn
+async def export_dataset_to_upload(data: ExportDatasetInput) -> str:
+    """
+    Activity to export the generated dataset files as a zip file to the specified
+    upload. A metadata file is also included in the zip.
+
+    Args:
+        data (ExportDatasetInput): Input data for exporting the dataset to the upload.
+    Returns:
+        str: Path to the saved zip file in the upload.
+    """
+
+    def unique_filename(filename: str, upload_files: StagingUploadFiles) -> str:
+        """Generate a unique filename for the upload_files directory."""
+        if not upload_files.raw_path_exists(filename):
+            return filename
+
+        count = 1
+        while True:
+            name, ext = os.path.splitext(filename)
+            _filename = f'{name}({count}){ext}'
+            if not upload_files.raw_path_exists(_filename):
+                return _filename
+            count += 1
+
+    upload_files = get_upload_files(data.upload_id, data.user_id)
+    if not upload_files:
+        raise ValueError(
+            f'Upload with ID {data.upload_id} for user {data.user_id} not found.'
+        )
+
+    # Create a metadata.json file in the artifact subdirectory
+    metadata_dict = {
+        'note': 'This metadata file contains information about the exported dataset '
+        'and the conditions under which it was generated.',
+        'data': data.metadata.model_dump(),
+        'schema': data.metadata.model_json_schema(),
+    }
+    metadata_path = os.path.join(data.artifact_subdirectory, 'metadata.json')
+    with open(metadata_path, 'w', encoding='utf-8') as metafile:
+        json.dump(metadata_dict, metafile, indent=4)
+
+    exportable_filepaths = data.source_paths + [metadata_path]
+    exportable_dir_name = unique_filename(data.exportable_dir_name, upload_files)
+
+    # Create a zip file containing all the source paths and the metadata file
+    if data.zip_output:
+        zipname = exportable_dir_name + '.zip'
+        zippath = os.path.join(data.artifact_subdirectory, zipname)
+        with zipfile.ZipFile(zippath, 'w', compression=zipfile.ZIP_DEFLATED) as zipf:
+            for filepath in exportable_filepaths:
+                arcname = os.path.basename(filepath)
+                zipf.write(filepath, arcname=arcname)
+        # Add zip file to the NOMAD Upload
+        upload_files.add_rawfiles(path=zippath, auto_decompress=False)
+        return zipname
+
+    # If not zipping, copy files to directory named exportable_dir_name
+    exportable_dir_path = os.path.join(data.artifact_subdirectory, exportable_dir_name)
+    os.mkdir(exportable_dir_path)
+    for filepath in exportable_filepaths:
+        temp_path = os.path.join(exportable_dir_path, os.path.basename(filepath))
+        shutil.copy2(filepath, temp_path)
+        # Add directory to the NOMAD Upload
+        upload_files.add_rawfiles(
+            path=exportable_dir_path, target_dir=exportable_dir_name
+        )
+    return exportable_dir_name
+
+
+@activity.defn
+async def cleanup_artifacts(data: CleanupArtifactsInput) -> None:
+    """
+    Activity to clean up the action artifacts directory.
+
+    Args:
+        data (CleanupArtifactsInput): Input data for cleaning up artifacts.
+    """
+
+    if os.path.exists(data.subdir_path):
+        shutil.rmtree(data.subdir_path)

nomad_ml_workflows/actions/export_entries/models.py

@@ -0,0 +1,245 @@
+import json
+from typing import Literal
+
+from nomad.app.v1.models.models import MetadataPagination, MetadataRequired, Query
+from pydantic import BaseModel, Field
+
+OwnerLiteral = Literal['public', 'visible', 'shared', 'user', 'staging']
+BatchFileTypeLiteral = Literal['parquet', 'json']
+OutputFileTypeLiteral = Literal['parquet', 'csv', 'json']
+IndexLiteral = Literal['entries', 'datasets', 'models', 'spaces']
+
+
+class SearchSettings(BaseModel):
+    owner: OwnerLiteral = Field(
+        'visible', description='Owner of the entries to be searched.'
+    )
+    query: str = Field(
+        ...,
+        description="""Query for extracting entries. Should be a valid dictionary
+        string. For example:
+        {
+            'entry_type': 'ELNSample'
+        }""",
+        # TODO: add `ui:widget` though `json_schema_extra` after NOMAD UI supports it
+    )
+    required_include: list[str] = Field(
+        None,
+        description='List of fields to include in the search results. For example: '
+        'results*, data.results*',
+    )
+    required_exclude: list[str] = Field(
+        None,
+        description='List of fields to exclude from the search results. For example: '
+        'results.method.method_name',
+    )
+
+
+class OutputSettings(BaseModel):
+    output_file_type: OutputFileTypeLiteral = Field(
+        'parquet',
+        description='Type of the output file.',
+    )
+    batch_size: int = Field(
+        1000,
+        gt=0,
+        description='Number of entries to be fetched and written per search batch. '
+        'Use smaller batch sizes when exporting large entries to reduce memory usage.',
+    )
+    zip_output: bool = Field(
+        True,
+        description='Whether to create a zip file for the output file(s). Set it '
+        'to true if you want download the dataset for external use. If you want to '
+        'work with the exported data in NOMAD, set it to false. This will export the '
+        'dataset as a directory in the specified project.',
+    )
+
+
+class ExportEntriesUserInput(BaseModel):
+    upload_id: str = Field(
+        ...,
+        description='Unique identifier for the upload associated with the workflow.',
+    )
+    user_id: str = Field(
+        ..., description='Unique identifier for the user who initiated the workflow.'
+    )
+    search_settings: SearchSettings
+    output_settings: OutputSettings
+
+
+class CreateArtifactSubdirectoryInput(BaseModel):
+    subdir_name: str = Field(..., description='Name of the subdirectory to be created.')
+
+
+class SearchInput(BaseModel):
+    user_id: str = Field(..., description='User ID performing the search.')
+    owner: OwnerLiteral = Field(..., description='Owner of the entries to be searched.')
+    query: Query = Field(..., description='Search query parameters.')
+    required: MetadataRequired = Field(
+        ..., description='Required fields for filtering the search results.'
+    )
+    pagination: MetadataPagination = Field(
+        ..., description='Pagination settings for the search results.'
+    )
+    batch_file_type: BatchFileTypeLiteral = Field(
+        ..., description='Type of the output file.'
+    )
+    output_file_path: str = Field(..., description='Path to the generated output file.')
+    max_entries_export_limit: int = Field(
+        ..., description='Maximum number of entries to be exported.'
+    )
+
+    @classmethod
+    def from_user_input(
+        cls,
+        user_input: ExportEntriesUserInput,
+        /,
+        output_file_path: str,
+        max_entries_export_limit: int,
+    ) -> 'SearchInput':
+        """Convert from ExportEntriesUserInput to SearchInput"""
+
+        def _clean_field(field: str) -> str:
+            """
+            Removes trailing whitespaces and inverted commas
+            """
+            return field.strip().strip("'").strip('"')
+
+        query = json.loads(
+            _clean_field(user_input.search_settings.query).replace("'", '"')
+        )
+
+        required = MetadataRequired()
+        if user_input.search_settings.required_include is not None:
+            include = [
+                _clean_field(field)
+                for field in user_input.search_settings.required_include
+            ]
+            required.include = include if include else None
+        if user_input.search_settings.required_exclude:
+            exclude = [
+                _clean_field(field)
+                for field in user_input.search_settings.required_exclude
+            ]
+            required.exclude = exclude if exclude else None
+
+        pagination = MetadataPagination(page_size=user_input.output_settings.batch_size)
+
+        batch_file_type = user_input.output_settings.output_file_type
+        if batch_file_type == 'csv':
+            batch_file_type = 'parquet'  # use parquet batches for csv
+
+        return cls(
+            user_id=user_input.user_id,
+            owner=user_input.search_settings.owner,
+            query=query,
+            required=required,
+            pagination=pagination,
+            batch_file_type=batch_file_type,
+            output_file_path=output_file_path,
+            max_entries_export_limit=max_entries_export_limit,
+        )
+
+
+class SearchOutput(BaseModel):
+    num_entries_exported: int = Field(
+        ..., description='Number of entries exported to the output file.'
+    )
+    num_entries_available: int = Field(
+        ...,
+        description='Total number of entries available for the given search query.',
+    )
+    search_start_time: str = Field(
+        ..., description='Timestamp when the search started.'
+    )
+    search_end_time: str = Field(
+        ..., description='Timestamp when the search completed.'
+    )
+    pagination_next_page_after_value: str | None = Field(
+        None,
+        description='The next_page_after_value from pagination, if more results are '
+        'available.',
+    )
+
+
+class MergeOutputFilesInput(BaseModel):
+    artifact_subdirectory: str = Field(
+        ...,
+        description='Subdirectory where the merged output file will be stored.',
+    )
+    output_file_type: OutputFileTypeLiteral = Field(
+        ...,
+        description='Type of the output file.',
+    )
+    generated_file_paths: list[str] = Field(
+        ...,
+        description='List of the generated file paths to be merged into a single file.',
+    )
+
+
+class ExportDatasetMetadata(BaseModel):
+    num_entries_exported: int = Field(
+        0,
+        description='Total number of entries exported in all the exported dataset '
+        'batches.',
+    )
+    num_entries_available: int = Field(
+        0,
+        description='Total number of entries available for the given search query.',
+    )
+    reached_max_entries_limit: bool = Field(
+        False,
+        description='Indicates whether the export reached the maximum number of '
+        'entries allowed. If true, the exported dataset contains the first N entries '
+        'up to the maximum limit.',
+    )
+    search_start_time: str = Field(
+        '',
+        description='Timestamp when the first search batch started.',
+    )
+    search_end_time: str = Field(
+        '',
+        description='Timestamp when the last search batch completed.',
+    )
+    user_input: ExportEntriesUserInput | None = Field(
+        None, description='Original user input for the export entries workflow.'
+    )
+    error_info: str | None = Field(
+        None,
+        description='Error information if any error occurred during the search and '
+        'merging process.',
+    )
+
+
+class ExportDatasetInput(BaseModel):
+    user_id: str = Field(
+        ..., description='User ID performing the export dataset operation.'
+    )
+    upload_id: str = Field(
+        ..., description='Upload ID associated with the export dataset operation.'
+    )
+    artifact_subdirectory: str = Field(
+        ...,
+        description='Subdirectory where the exported dataset zip file will be stored.',
+    )
+    zip_output: bool = Field(
+        ...,
+        description='Whether to create a zip file for the exported dataset.',
+    )
+    exportable_dir_name: str = Field(
+        ...,
+        description='Name of the directory containing the dataset that will be '
+        'exported.',
+    )
+    source_paths: list[str] = Field(
+        ..., description='List of paths to the source files of the dataset.'
+    )
+    metadata: ExportDatasetMetadata = Field(
+        ..., description='Metadata associated with the exported dataset.'
+    )
+
+
+class CleanupArtifactsInput(BaseModel):
+    subdir_path: str = Field(
+        ..., description='Path to the subdirectory to be cleaned up.'
+    )

nomad_ml_workflows/actions/export_entries/utils.py

@@ -0,0 +1,172 @@
+import json
+
+import json_stream
+from nomad.utils import dict_to_dataframe
+
+try:
+    import pyarrow as pa
+    import pyarrow.csv as pcsv
+    import pyarrow.dataset as ds
+    import pyarrow.parquet as pq
+except ImportError as e:
+    raise ImportError(
+        'pyarrow is required. Install with: pip install nomad-ml-workflows[cpu-action]'
+    ) from e
+
+
+def _is_nested_type(dtype: pa.DataType) -> bool:
+    """Check if a PyArrow type is nested."""
+    return pa.types.is_nested(dtype)
+
+
+def _get_csv_compatible_schema(schema: pa.Schema) -> pa.Schema:
+    """Convert schema to CSV-compatible format by changing nested types to strings."""
+    new_fields = []
+    for field in schema:
+        if _is_nested_type(field.type):
+            new_fields.append(pa.field(field.name, pa.string(), field.nullable))
+        else:
+            new_fields.append(field)
+    return pa.schema(new_fields)
+
+
+def _stringify_nested_columns(batch: pa.RecordBatch) -> pa.RecordBatch:
+    """Convert nested columns (list, struct) in a batch to JSON strings."""
+    new_columns = []
+    for i, column in enumerate(batch.columns):
+        if _is_nested_type(batch.schema.field(i).type):
+            # Convert each element to JSON string
+            stringified = pa.array(
+                [
+                    json.dumps(val.as_py()) if val.as_py() is not None else None
+                    for val in column
+                ],
+                type=pa.string(),
+            )
+            new_columns.append(stringified)
+        else:
+            new_columns.append(column)
+
+    return pa.RecordBatch.from_arrays(
+        new_columns, schema=_get_csv_compatible_schema(batch.schema)
+    )
+
+
+def write_parquet_file(path: str, data: list[dict]):
+    """Writes a list of NOMAD entry dicts to a parquet file.
+
+    Args:
+        path (str): The path where the file will be saved.
+        data (list[dict]): The list of NOMAD entry dicts to be written to the file.
+    """
+    if not path.endswith('parquet'):
+        raise ValueError('Unsupported file type. Please use parquet.')
+
+    df = dict_to_dataframe(data)
+
+    table = pa.Table.from_pandas(df)
+    with pq.ParquetWriter(
+        path,
+        table.schema,
+        compression='snappy',  # snappy for faster write/read for individual files
+        use_dictionary=True,
+    ) as writer:
+        writer.write_table(table)
+
+
+def write_csv_file(path: str, data: list[dict]):
+    """Writes a list of NOMAD entry dicts to a CSV file.
+
+    Args:
+        path (str): The path where the file will be saved.
+        data (list[dict]): The list of NOMAD entry dicts to be written to the file.
+    """
+    if not path.endswith('csv'):
+        raise ValueError('Unsupported file type. Please use csv.')
+
+    df = dict_to_dataframe(data)
+
+    df.to_csv(path, index=False, mode='w', header=True)
+
+
+def write_json_file(path: str, data: list[dict]):
+    """Writes a list of NOMAD entry dicts to a JSON file.
+
+    Args:
+        path (str): The path where the file will be saved.
+        data (list[dict]): The list of NOMAD entry dicts to be written to the file.
+    """
+    if not path.endswith('json'):
+        raise ValueError('Unsupported file type. Please use json.')
+
+    with open(path, 'w') as f:
+        json.dump(data, f, indent=4)
+
+
+def merge_files(
+    input_file_paths: list[str], output_file_type: str, output_file_path: str
+):
+    """Merges multiple Parquet or JSON files into a single file.
+
+    Args:
+        input_file_paths (list[str]): List of file paths to be merged.
+        output_file_type (str): The type of the output file ('parquet', 'csv', or
+            'json').
+        output_file_path (str): Path of the merged output file.
+    """
+    if output_file_type == 'parquet':
+        # Creates a logical dataset from the input files, not loading all data into
+        # memory. Also, unifies the schema across the files.
+        dataset = ds.dataset(input_file_paths, format='parquet')
+
+        # Write the dataset to a single Parquet file in batches
+        with pq.ParquetWriter(
+            output_file_path,
+            dataset.schema,
+            compression='zstd',  # for better compression for merged file
+            compression_level=3,
+            use_dictionary=True,
+        ) as writer:
+            for batch in dataset.to_batches():
+                writer.write_batch(batch)
+
+    elif output_file_type == 'csv':
+        # Creates a logical dataset from the input files, not loading all data into
+        # memory. Also, unifies the schema across the files.
+        # The batch files for `csv` are written in Parquet format for efficiency,
+        # so we read them as Parquet here.
+        dataset = ds.dataset(input_file_paths, format='parquet')
+
+        # PyArrow CSV writer doesn't support nested types (list, struct, etc.)
+        # Convert nested columns to JSON strings
+        csv_schema = _get_csv_compatible_schema(dataset.schema)
+
+        # Write the dataset to a single CSV file in batches
+        with pcsv.CSVWriter(output_file_path, csv_schema) as writer:
+            for batch in dataset.to_batches():
+                csv_batch = _stringify_nested_columns(batch)
+                writer.write_batch(csv_batch)
+
+    elif output_file_type == 'json':
+
+        def _json_stream_files(input_file_paths):
+            """Generator that streams one entry dict at a time from multiple files."""
+            for file_path in input_file_paths:
+                with open(file_path, encoding='utf-8') as f:
+                    data = json_stream.load(f)
+                    yield from data
+
+        # Write a single JSON file by streaming entry dicts and wrapping in a list
+        with open(output_file_path, 'w', encoding='utf-8') as f:
+            f.write('[\n')
+            first_item = True
+            for item in _json_stream_files(input_file_paths):
+                if not first_item:
+                    f.write(',\n')
+                # Convert transient json_stream object to standard Python types
+                json.dump(json_stream.to_standard_types(item), f, indent=4)
+                first_item = False
+            f.write('\n]')
+
+    else:
+        raise ValueError('Unsupported file type. Please use parquet, csv, or json.')

nomad_ml_workflows/actions/export_entries/workflows.py

@@ -0,0 +1,169 @@
+from datetime import timedelta
+
+from temporalio import workflow
+from temporalio.common import RetryPolicy
+from temporalio.exceptions import ApplicationError
+
+with workflow.unsafe.imports_passed_through():
+    from nomad.config import config as nomad_config
+
+    from nomad_ml_workflows.actions.export_entries.activities import (
+        cleanup_artifacts,
+        create_artifact_subdirectory,
+        export_dataset_to_upload,
+        merge_output_files,
+        search,
+    )
+    from nomad_ml_workflows.actions.export_entries.models import (
+        CleanupArtifactsInput,
+        CreateArtifactSubdirectoryInput,
+        ExportDatasetInput,
+        ExportDatasetMetadata,
+        ExportEntriesUserInput,
+        MergeOutputFilesInput,
+        SearchInput,
+    )
+
+
+@workflow.defn
+class ExportEntriesWorkflow:
+    @workflow.run
+    async def run(self, data: ExportEntriesUserInput) -> str:
+        """
+        Workflow to search entries and export them into a datafile in the specified
+        upload.
+
+        Args:
+            data (ExportEntriesUserInput): Input data for the export entries workflow.
+        Returns:
+            str: Path to the saved dataset in the upload's `raw` folder.
+        """
+        retry_policy = RetryPolicy(
+            maximum_attempts=1,
+            initial_interval=timedelta(seconds=10),
+            maximum_interval=timedelta(minutes=1),
+            backoff_coefficient=2.0,
+        )
+        artifact_subdirectory = await workflow.execute_activity(
+            create_artifact_subdirectory,
+            CreateArtifactSubdirectoryInput(subdir_name=workflow.info().workflow_id),
+            start_to_close_timeout=timedelta(minutes=10),
+            retry_policy=retry_policy,
+        )
+        export_dataset_input = ExportDatasetInput(
+            user_id=data.user_id,
+            upload_id=data.upload_id,
+            artifact_subdirectory=artifact_subdirectory,
+            exportable_dir_name='export_entries_error',  # name used in case of error
+            zip_output=data.output_settings.zip_output,
+            source_paths=[],
+            metadata=ExportDatasetMetadata(user_input=data),
+        )
+
+        try:
+            config = nomad_config.get_plugin_entry_point(
+                'nomad_ml_workflows.actions:export_entries'
+            )
+
+            search_counter = 0
+            num_entries_available = 0
+            generated_file_paths = []
+            search_start_times = []
+            search_end_times = []
+            total_num_entries_exported = 0
+            reached_max_entries_limit = False
+            search_input = SearchInput.from_user_input(
+                data,
+                output_file_path='',  # Placeholder, will be set in loop
+                max_entries_export_limit=config.max_entries_export_limit,
+            )
+            while True:
+                search_counter += 1
+                search_input.output_file_path = (
+                    f'{artifact_subdirectory}/{search_counter}.'
+                    f'{search_input.batch_file_type}'
+                )
+                search_output = await workflow.execute_activity(
+                    search,
+                    search_input,
+                    activity_id=f'search-activity-{search_counter}',
+                    start_to_close_timeout=timedelta(
+                        seconds=config.search_batch_timeout
+                    ),
+                    retry_policy=retry_policy,
+                )
+                if search_counter == 1:
+                    # capture the total available entries from the first search output
+                    num_entries_available = search_output.num_entries_available
+                if search_output.num_entries_exported > 0:
+                    # only save paths if the writing files was not skipped
+                    generated_file_paths.append(search_input.output_file_path)
+                search_start_times.append(search_output.search_start_time)
+                search_end_times.append(search_output.search_end_time)
+                total_num_entries_exported += search_output.num_entries_exported
+                # Update pagination for next iteration
+                search_input.pagination.page_after_value = (
+                    search_output.pagination_next_page_after_value
+                )
+                search_input.max_entries_export_limit -= (
+                    search_output.num_entries_exported
+                )
+
+                if search_output.pagination_next_page_after_value is None:
+                    # break if there are no more pages to fetch
+                    break
+                if search_input.max_entries_export_limit <= 0:
+                    # break early if the max entries limit has been reached
+                    reached_max_entries_limit = True
+                    break
+
+            merged_file_path = await workflow.execute_activity(
+                merge_output_files,
+                MergeOutputFilesInput(
+                    artifact_subdirectory=artifact_subdirectory,
+                    output_file_type=data.output_settings.output_file_type,
+                    generated_file_paths=generated_file_paths,
+                ),
+                start_to_close_timeout=timedelta(hours=2),
+                retry_policy=retry_policy,
+            )
+
+            # Prepare export dataset input and metadata
+            export_dataset_input.exportable_dir_name = (
+                'export_entries_' + search_start_times[0].replace(':', '-')
+            )
+            export_dataset_input.source_paths = [merged_file_path]
+            export_dataset_input.metadata = ExportDatasetMetadata(
+                num_entries_exported=total_num_entries_exported,
+                num_entries_available=num_entries_available,
+                reached_max_entries_limit=reached_max_entries_limit,
+                search_start_time=search_start_times[0],
+                search_end_time=search_end_times[-1],
+                user_input=data,
+            )
+
+        except Exception as e:
+            # Capture error info to include in metadata
+            import traceback
+
+            export_dataset_input.metadata.error_info = traceback.format_exc()
+            raise ApplicationError(
+                'Encountered an error during export entries workflow.',
+            ) from e
+
+        finally:
+            saved_dataset_path = await workflow.execute_activity(
+                export_dataset_to_upload,
+                export_dataset_input,
+                start_to_close_timeout=timedelta(hours=2),
+                retry_policy=retry_policy,
+            )
+
+            await workflow.execute_activity(
+                cleanup_artifacts,
+                CleanupArtifactsInput(subdir_path=artifact_subdirectory),
+                start_to_close_timeout=timedelta(hours=2),
+                retry_policy=retry_policy,
+            )
+
+        return saved_dataset_path

nomad_ml_workflows-0.0.6.dist-info/METADATA

@@ -0,0 +1,197 @@
+Metadata-Version: 2.4
+Name: nomad-ml-workflows
+Version: 0.0.6
+Summary: A NOMAD plugin for managing ML workflows.
+Author-email: Sarthak Kapoor <sarthak.kapoor@physik.hu-berlin.de>
+Maintainer-email: Sarthak Kapoor <sarthak.kapoor@physik.hu-berlin.de>
+License: 
+        The MIT License (MIT)
+        
+        Copyright (c) 2025 Sarthak Kapoor
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in
+        all copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+        THE SOFTWARE.
+        
+Project-URL: Repository, https://github.com/FAIRmat-NFDI/nomad-ml-workflows
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: License :: OSI Approved :: MIT License
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: nomad-lab>=1.4.0
+Requires-Dist: json-stream
+Requires-Dist: pydantic
+Requires-Dist: temporalio
+Provides-Extra: dev
+Requires-Dist: nomad-lab[infrastructure]>=1.4.0; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: structlog; extra == "dev"
+Requires-Dist: mkdocs; extra == "dev"
+Requires-Dist: mkdocs-material>=8.1.1; extra == "dev"
+Requires-Dist: pymdown-extensions; extra == "dev"
+Requires-Dist: mkdocs-click; extra == "dev"
+Requires-Dist: pytest-asyncio; extra == "dev"
+Provides-Extra: cpu-action
+Requires-Dist: pyarrow; extra == "cpu-action"
+Dynamic: license-file
+
+# nomad-ml-workflows
+
+A NOMAD plugin for managing ML workflows. Currently, it provides an action to export large number of entries from NOMAD database as tabular data files. Other ML workflow related actions and schemas will be added in future.
+
+## 📦 Installation
+You can install the plugin using pip:
+```sh
+pip install nomad-ml-workflows @ git+https://github.com/FAIRmat-NFDI/nomad-ml-workflows.git
+```
+
+However, to fully utilize the plugin, you need to add it to your NOMAD instance as described [below](#-adding-this-plugin-to-nomad).
+
+## ✨ Features
+
+- Export a large number of NOMAD entries as tabular data files (CSV, Parquet) using NOMAD Actions. Once the action is triggered, it will:
+  - Search entries based on user-defined criteria.
+  - Optionally include or exclude data fields from the entries.
+  - Package the entries into tabular data files like CSV or Parquet (or as JSON)
+  - Export the files to a specified Project (or previously known as Upload) in NOAMD.
+
+  These can then be downloaded from the NOMAD web interface for local use.
+
+## ⚙️ Configuration
+The Export Entries action can be configured using the following parameters in
+the `nomad.yaml` configuration file of your NOMAD Oasis instance:
+
+```yaml
+plugins:
+  entry_points:
+    options:
+      nomad_ml_workflows.actions:export_entries:
+        search_batch_timeout: 7200
+        # Timeout (in seconds) for each search batch in the Export Entries
+        # action. Set this accordingly to time out longer searches.
+        max_entries_export_limit: 100000
+        # Maximum number of entries that can be exported in a single
+        # Export Entries action.
+```
+
+
+## 🚀 Adding this plugin to NOMAD
+
+Currently, NOMAD has two distinct flavors that are relevant depending on your role as an user:
+1. [A NOMAD Oasis](#adding-this-plugin-in-your-nomad-oasis): any user with a NOMAD Oasis instance.
+2. [Local NOMAD installation and the source code of NOMAD](#adding-this-plugin-in-your-local-nomad-installation-and-the-source-code-of-nomad): internal developers.
+
+### Adding this plugin in your NOMAD Oasis
+
+Read the [NOMAD plugin documentation](https://nomad-lab.eu/prod/v1/staging/docs/howto/oasis/plugins_install.html) for all details on how to deploy the plugin on your NOMAD instance.
+
+### Adding this plugin in your local NOMAD installation and the source code of NOMAD
+
+We now recommend using the dedicated [`nomad-distro-dev`](https://github.com/FAIRmat-NFDI/nomad-distro-dev) repository to simplify the process. Please refer to that repository for detailed instructions.
+
+
+## 🛠️ Development
+
+If you want to develop locally this plugin, clone the project and in the plugin folder, create a virtual environment (you can use Python 3.10, 3.11 or 3.12):
+```sh
+git clone https://github.com/FAIRmat-NFDI/nomad-ml-workflows.git
+cd nomad-ml-workflows
+python3.11 -m venv .pyenv
+. .pyenv/bin/activate
+```
+
+Make sure to have `pip` upgraded:
+```sh
+pip install --upgrade pip
+```
+
+We recommend installing `uv` for fast pip installation of the packages:
+```sh
+pip install uv
+```
+
+Install the `nomad-lab` package:
+```sh
+uv pip install -e '.[dev]'
+```
+
+### Run linting and auto-formatting
+
+We use [Ruff](https://docs.astral.sh/ruff/) for linting and formatting the code. Ruff auto-formatting is also a part of the GitHub workflow actions. You can run locally:
+```sh
+ruff check .
+ruff format . --check
+```
+
+### Debugging
+
+For interactive debugging of the tests, use `pytest` with the `--pdb` flag. We recommend using an IDE for debugging, e.g., _VSCode_. If that is the case, add the following snippet to your `.vscode/launch.json`:
+```json
+{
+  "configurations": [
+      {
+        "name": "<descriptive tag>",
+        "type": "debugpy",
+        "request": "launch",
+        "cwd": "${workspaceFolder}",
+        "program": "${workspaceFolder}/.pyenv/bin/pytest",
+        "justMyCode": true,
+        "env": {
+            "_PYTEST_RAISE": "1"
+        },
+        "args": [
+            "-sv",
+            "--pdb",
+            "<path-to-plugin-tests>",
+        ]
+    }
+  ]
+}
+```
+
+where `<path-to-plugin-tests>` must be changed to the local path to the test module to be debugged.
+
+The settings configuration file `.vscode/settings.json` automatically applies the linting and formatting upon saving the modified file.
+
+### Documentation on Github pages
+
+To view the documentation locally, install the related packages using:
+```sh
+uv pip install -r requirements_docs.txt
+```
+
+Run the documentation server:
+```sh
+mkdocs serve
+```
+
+
+## 👥 Main contributors
+| Name | E-mail     |
+|------|------------|
+| Sarthak Kapoor | [sarthak.kapoor@physik.hu-berlin.de](mailto:sarthak.kapoor@physik.hu-berlin.de)
+
+
+## 📄 License
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

nomad_ml_workflows-0.0.6.dist-info/RECORD

@@ -0,0 +1,13 @@
+nomad_ml_workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nomad_ml_workflows/actions/__init__.py,sha256=ooUE9fX622d8JU8ckbr3tQJuZ3oQWvnJFS3CASCINqc,99
+nomad_ml_workflows/actions/export_entries/__init__.py,sha256=bFxKQr8zSQZQcwJwSfJ2bC35orXWh9COZZVOGY5Jy7k,1681
+nomad_ml_workflows/actions/export_entries/activities.py,sha256=ie7uyQMU6I-9u3yuSASN9v2EWU7SxC8xWO_m11uAY44,6943
+nomad_ml_workflows/actions/export_entries/models.py,sha256=YQRrdCbxK9cY0ukrdUxtn8PKvCVYFrrVsde-bDD5glk,8521
+nomad_ml_workflows/actions/export_entries/utils.py,sha256=g2ZbwJZgPEBOcIdXgCYMRGK1XViNYp9414h4ZvUbHRY,6193
+nomad_ml_workflows/actions/export_entries/workflows.py,sha256=T1png3P8JDoi0Yf7X798xFHEao1ePYXQqSUjzDkXoeg,6890
+nomad_ml_workflows-0.0.6.dist-info/licenses/LICENSE,sha256=Wji1LdkrEkb33BQ4YJUjF9EPtOiz74OQW3XO0omcZ3U,1082
+nomad_ml_workflows-0.0.6.dist-info/METADATA,sha256=2sX3U7TM-II4YLu5ctYRRg2tOkudLNXAViKzrDFEqqo,7433
+nomad_ml_workflows-0.0.6.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+nomad_ml_workflows-0.0.6.dist-info/entry_points.txt,sha256=unxYXwft1X1R3FCg9NyANBqevaAZUZNBQJ7We5ez0eA,93
+nomad_ml_workflows-0.0.6.dist-info/top_level.txt,sha256=1NCxJd5BKyy75vikv64WJSwZL9bdISeeHO5FCAeQI_E,19
+nomad_ml_workflows-0.0.6.dist-info/RECORD,,

nomad_ml_workflows-0.0.6.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nomad_ml_workflows-0.0.6.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[nomad.plugin]
+export_entries_action_entry_point = nomad_ml_workflows.actions:export_entries

nomad_ml_workflows-0.0.6.dist-info/licenses/LICENSE

@@ -0,0 +1,22 @@
+
+The MIT License (MIT)
+
+Copyright (c) 2025 Sarthak Kapoor
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

nomad_ml_workflows-0.0.6.dist-info/top_level.txt

@@ -0,0 +1 @@
+nomad_ml_workflows