vdk-huggingface: add new ingest plugin (#2858)

Versatile Data Kit (VDK) plugin for integrating with Huggingface as both
a data source and a target.

This plugin allows you to ingest data payloads into a Huggingface
repository and makes it easier to work with datasets stored in
Huggingface.

projects/vdk-plugins/vdk-huggingface/.plugin-ci.yml

@@ -0,0 +1,22 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+image: "python:3.7"
+
+.build-vdk-huggingface:
+  variables:
+    PLUGIN_NAME: vdk-huggingface
+  extends: .build-plugin
+
+build-py37-vdk-huggingface:
+  extends: .build-vdk-huggingface
+  image: "python:3.7"
+
+build-py311-vdk-huggingface:
+  extends: .build-vdk-huggingface
+  image: "python:3.11"
+
+release-vdk-huggingface:
+  variables:
+    PLUGIN_NAME: vdk-huggingface
+  extends: .release-plugin

projects/vdk-plugins/vdk-huggingface/README.md

@@ -0,0 +1,50 @@
+# Huggingface
+
+Versatile Data Kit (VDK) plugin for integrating with Huggingface as both a data source and a target.
+This plugin allows you to ingest data payloads into a Huggingface repository and makes it easier to work with datasets stored in Huggingface.
+
+
+## Usage
+
+```
+pip install vdk-huggingface
+```
+
+The functionality adds new ingestion method "huggingface" which can be used like that:
+
+```python
+job_input.send_object_for_ingestion(data, method="huggingface")
+```
+
+
+### Configuration
+
+(`vdk config-help` is useful command to browse all config options of your installation of vdk)
+
+| Name                | Description                                                                 | (example)  Value        |
+|---------------------|-----------------------------------------------------------------------------|-------------------------|
+| HUGGINGFACE_TOKEN   | HuggingFace API token for authentication. Get one from HuggingFace Settings | ""                      |
+| HUGGINGFACE_REPO_ID | HuggingFace Dataset repository ID                                           | "username/test-dataset" |
+
+
+
+### Build and testing
+
+```
+pip install -r requirements.txt
+pip install -e .
+pytest
+```
+
+In VDK repo [../build-plugin.sh](https://github.com/vmware/versatile-data-kit/tree/main/projects/vdk-plugins/build-plugin.sh) script can be used also.
+
+
+#### Note about the CICD:
+
+.plugin-ci.yaml is needed only for plugins part of [Versatile Data Kit Plugin repo](https://github.com/vmware/versatile-data-kit/tree/main/projects/vdk-plugins).
+
+The CI/CD is separated in two stages, a build stage and a release stage.
+The build stage is made up of a few jobs, all which inherit from the same
+job configuration and only differ in the Python version they use (3.7, 3.8, 3.9 and 3.10).
+They run according to rules, which are ordered in a way such that changes to a
+plugin's directory trigger the plugin CI, but changes to a different plugin does not.

projects/vdk-plugins/vdk-huggingface/requirements.txt

@@ -0,0 +1,8 @@
+# this file is used to provide testing requirements
+# for requirements (dependencies) needed during and after installation of the plugin see (and update) setup.py install_requires section
+
+datasets
+
+pytest
+vdk-core
+vdk-test-utils

projects/vdk-plugins/vdk-huggingface/setup.py

@@ -0,0 +1,42 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import pathlib
+
+import setuptools
+
+"""
+Builds a package with the help of setuptools in order for this package to be imported in other projects
+"""
+
+__version__ = "0.1.4"
+
+setuptools.setup(
+    name="vdk-huggingface",
+    version=__version__,
+    url="https://github.com/vmware/versatile-data-kit",
+    description="Integrate VDK with Huggingface as both data source and target",
+    long_description=pathlib.Path("README.md").read_text(),
+    long_description_content_type="text/markdown",
+    install_requires=["vdk-core", "huggingface-hub"],
+    package_dir={"": "src"},
+    packages=setuptools.find_namespace_packages(where="src"),
+    # This is the only vdk plugin specifc part
+    # Define entry point called "vdk.plugin.run" with name of plugin and module to act as entry point.
+    entry_points={
+        "vdk.plugin.run": ["vdk-huggingface = vdk.plugin.huggingface.plugin_entry"]
+    },
+    classifiers=[
+        "Development Status :: 2 - Pre-Alpha",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+    ],
+    project_urls={
+        "Documentation": "https://github.com/vmware/versatile-data-kit/tree/main/projects/vdk-plugins/vdk-huggingface",
+        "Source Code": "https://github.com/vmware/versatile-data-kit/tree/main/projects/vdk-plugins/vdk-huggingface",
+        "Bug Tracker": "https://github.com/vmware/versatile-data-kit/issues/new/choose",
+    },
+)

projects/vdk-plugins/vdk-huggingface/src/vdk/plugin/huggingface/ingest.py

@@ -0,0 +1,87 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+import tempfile
+from collections import defaultdict
+from dataclasses import dataclass
+from threading import RLock
+from typing import Dict
+from typing import List
+from typing import Optional
+
+from huggingface_hub import CommitOperationAdd
+from huggingface_hub import HfApi
+from vdk.api.plugin.plugin_input import IIngesterPlugin
+
+
+@dataclass(frozen=True)
+class OpKey:
+    repo_id: str
+    destination_table: str
+
+
+class IngestToHuggingface(IIngesterPlugin):
+    def __init__(self, repo_id: str):
+        self._api = HfApi()
+        self._repo_id = repo_id
+        self._file_handles: Dict[OpKey, any] = {}
+        self._locks: Dict[OpKey, RLock] = defaultdict(RLock)
+
+    def ingest_payload(
+        self,
+        payload: List[dict],
+        destination_table: Optional[str],
+        target: Optional[str] = None,
+        collection_id: Optional[str] = None,
+        metadata: Optional[IIngesterPlugin.IngestionMetadata] = None,
+    ) -> Optional[IIngesterPlugin.IngestionMetadata]:
+        repo_id = target if target is not None else self._repo_id
+
+        op_key = OpKey(repo_id, destination_table)
+        with self._locks[op_key]:
+            if op_key not in self._file_handles:
+                self._check_and_create_repo(repo_id)
+                temp_file = tempfile.NamedTemporaryFile(mode="w+", delete=False)
+                self._file_handles[op_key] = temp_file
+
+            self._append_payload(op_key, payload)
+
+        return None
+
+    def _append_payload(self, op_key: OpKey, payload: List[Dict]):
+        temp_file = self._file_handles[op_key]
+        temp_file.flush()
+        temp_file.seek(0)
+        existing_data_str = temp_file.read()
+
+        existing_data = json.loads(existing_data_str) if existing_data_str else []
+
+        merged_data = existing_data + payload
+        temp_file.seek(0)
+        temp_file.truncate()
+        json.dump(merged_data, temp_file)
+
+    def commit_all(self):
+        for op_key, temp_file in self._file_handles.items():
+            with self._locks[op_key]:
+                temp_file.close()
+                addition = CommitOperationAdd(
+                    path_in_repo=op_key.destination_table,
+                    path_or_fileobj=temp_file.name,
+                )
+                self._api.preupload_lfs_files(
+                    op_key.repo_id, additions=[addition], repo_type="dataset"
+                )
+                self._api.create_commit(
+                    op_key.repo_id,
+                    operations=[addition],
+                    repo_type="dataset",
+                    commit_message="Automatic commit by vdk-huggingface",
+                )
+                os.unlink(temp_file.name)
+        self._file_handles.clear()
+
+    def _check_and_create_repo(self, repo_id: str):
+        if not self._api.repo_exists(repo_id):
+            self._api.create_repo(repo_id=repo_id, exist_ok=True, repo_type="dataset")

projects/vdk-plugins/vdk-huggingface/src/vdk/plugin/huggingface/plugin_entry.py

@@ -0,0 +1,64 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import logging
+from typing import List
+from typing import Optional
+
+import huggingface_hub
+from vdk.api.plugin.hook_markers import hookimpl
+from vdk.api.plugin.plugin_registry import IPluginRegistry
+from vdk.internal.builtin_plugins.run.job_context import JobContext
+from vdk.internal.core.config import ConfigurationBuilder
+from vdk.plugin.huggingface.ingest import IngestToHuggingface
+
+HUGGINGFACE_REPO_ID = "huggingface_repo_id"
+HUGGINGFACE_TOKEN = "huggingface_token"
+
+log = logging.getLogger(__name__)
+
+
+class HuggingfacePlugin:
+    def __init__(self):
+        self._ingester: Optional[IngestToHuggingface] = None
+
+    @hookimpl(tryfirst=True)
+    def vdk_configure(self, config_builder: ConfigurationBuilder):
+        config_builder.add(
+            key=HUGGINGFACE_TOKEN,
+            default_value=None,
+            description="HuggingFace API token for authentication. "
+            "Get one from https://huggingface.co/settings/tokens",
+        )
+        config_builder.add(
+            key=HUGGINGFACE_REPO_ID,
+            default_value="username/test-dataset",
+            description="HuggingFace Dataset repository ID.",
+        )
+
+    @hookimpl
+    def initialize_job(self, context: JobContext) -> None:
+        repo_id = context.core_context.configuration.get_value(HUGGINGFACE_REPO_ID)
+        token = context.core_context.configuration.get_value(HUGGINGFACE_TOKEN)
+
+        if token:
+            log.info("huggingface log in", extra={"huggingface_repo_id": repo_id})
+            huggingface_hub.login(token)
+
+        self._ingester = IngestToHuggingface(repo_id)
+
+        context.ingester.add_ingester_factory_method(
+            "huggingface", lambda: self._ingester
+        )
+
+    @hookimpl(trylast=True)
+    def finalize_job(self, context: JobContext) -> None:
+        if self._ingester:
+            self._ingester.commit_all()
+        # TODO: on exception , this doesn't fail the job ...
+
+
+@hookimpl
+def vdk_start(plugin_registry: IPluginRegistry, command_line_args: List):
+    plugin_registry.load_plugin_with_hooks_impl(
+        HuggingfacePlugin(), "HuggingfacePlugin"
+    )

projects/vdk-plugins/vdk-huggingface/tests/jobs/huggingface-job/10_upload.py

@@ -0,0 +1,23 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+from vdk.api.job_input import IJobInput
+
+
+def run(job_input: IJobInput):
+    payload = {
+        "str_col": "str",
+        "int_col": 2,
+        "bool_col": False,
+        "time": job_input.get_arguments().get("time"),
+    }
+
+    for i in range(10):
+        job_input.send_object_for_ingestion(
+            payload=payload, destination_table="test_vdk_table", method="huggingface"
+        )
+        job_input.send_object_for_ingestion(
+            payload=payload, destination_table="test_vdk_table2", method="huggingface"
+        )
+        job_input.send_object_for_ingestion(
+            payload=payload, destination_table="test_vdk_table3", method="huggingface"
+        )

projects/vdk-plugins/vdk-huggingface/tests/test_huggingface_plugin.py

@@ -0,0 +1,51 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+import pathlib
+import time
+
+import pytest
+from click.testing import Result
+from huggingface_hub import HfApi
+from vdk.plugin.huggingface import plugin_entry
+from vdk.plugin.test_utils.util_funcs import cli_assert_equal
+from vdk.plugin.test_utils.util_funcs import CliEntryBasedTestRunner
+from vdk.plugin.test_utils.util_funcs import jobs_path_from_caller_directory
+
+
+@pytest.mark.skipif(
+    os.environ.get("HUGGINGFACE_TOKEN") is None
+    or os.environ.get("HUGGINGFACE_REPO_ID") is None,
+    reason="Skipping test until HUGGINGFACE_TOKEN and HUGGINGFACE_REPO_ID and are set",
+)
+def test_huggingface_upload():
+    runner = CliEntryBasedTestRunner(plugin_entry)
+
+    time_var = time.time()
+    result: Result = runner.invoke(
+        [
+            "run",
+            jobs_path_from_caller_directory("huggingface-job"),
+            "--arguments",
+            json.dumps({"time": time_var}),
+        ]
+    )
+    cli_assert_equal(0, result)
+
+    api = HfApi()
+    file_path = api.hf_hub_download(
+        repo_id=os.environ["HUGGINGFACE_REPO_ID"],
+        filename="test_vdk_table",
+        repo_type="dataset",
+    )
+    row = [{"str_col": "str", "int_col": 2, "bool_col": False, "time": time_var}]
+    assert json.loads(pathlib.Path(file_path).read_text()) == row * 10
+
+    file_path = api.hf_hub_download(
+        repo_id=os.environ["HUGGINGFACE_REPO_ID"],
+        filename="test_vdk_table2",
+        repo_type="dataset",
+    )
+    row = [{"str_col": "str", "int_col": 2, "bool_col": False, "time": time_var}]
+    assert json.loads(pathlib.Path(file_path).read_text()) == row * 10