feat: add @huggingface decorator with pluggable auth

- Pluggable auth: HuggingFaceAuthProvider interface and EnvHuggingFaceAuthProvider (HF_TOKEN / HUGGING_FACE_HUB_TOKEN)
- Config METAFLOW_HUGGINGFACE_AUTH_PROVIDER; register hf_auth_provider in plugin system
- @huggingface step decorator: declare models/model_mapping; get token via auth provider; download via huggingface_hub; expose current.huggingface.models[key]

Author: nadeemahmad

metaflow/extension_support/plugins.py

@@ -187,6 +187,7 @@ def resolve_plugins(category, path_only=False):
     "datastore": lambda x: x.TYPE,
     "dataclient": lambda x: x.TYPE,
     "secrets_provider": lambda x: x.TYPE,
+    "hf_auth_provider": lambda x: x.TYPE,
     "gcp_client_provider": lambda x: x.name,
     "deployer_impl_provider": lambda x: x.TYPE,
     "azure_client_provider": lambda x: x.name,

metaflow/metaflow_config.py

@@ -44,6 +44,11 @@
 DEFAULT_SECRETS_BACKEND_TYPE = from_conf("DEFAULT_SECRETS_BACKEND_TYPE")
 DEFAULT_SECRETS_ROLE = from_conf("DEFAULT_SECRETS_ROLE")
 
+# HuggingFace @huggingface decorator: auth provider type (Part 1)
+METAFLOW_HUGGINGFACE_AUTH_PROVIDER = from_conf(
+    "METAFLOW_HUGGINGFACE_AUTH_PROVIDER", "env"
+)
+
 DEFAULT_FROM_DEPLOYMENT_IMPL = from_conf(
     "DEFAULT_FROM_DEPLOYMENT_IMPL", "argo-workflows"
 )
@@ -64,6 +69,7 @@
         "timeout",
         "conda_env_internal",
         "card",
+        "huggingface",
     ],
 )
 

metaflow/metaflow_current.py

@@ -32,6 +32,14 @@ def _raise(ex):
         self.__class__.graph = property(
             fget=lambda self: _raise(RuntimeError("Graph is not available"))
         )
+        self.__class__.huggingface = property(
+            fget=lambda self: _raise(
+                RuntimeError(
+                    "current.huggingface is only available inside a step "
+                    "decorated with @huggingface"
+                )
+            )
+        )
 
     def _set_env(
         self,

metaflow/plugins/__init__.py

@@ -55,6 +55,7 @@
     ("airflow_internal", ".airflow.airflow_decorator.AirflowInternalDecorator"),
     ("pypi", ".pypi.pypi_decorator.PyPIStepDecorator"),
     ("conda", ".pypi.conda_decorator.CondaStepDecorator"),
+    ("huggingface", ".huggingface.huggingface_decorator.HuggingFaceDecorator"),
 ]
 
 # Add new flow decorators here
@@ -139,6 +140,11 @@
 
 FLOW_DECORATORS_DESC += SENSOR_FLOW_DECORATORS
 
+# HuggingFace decorator: auth providers (Part 1)
+HF_AUTH_PROVIDERS_DESC = [
+    ("env", ".huggingface.env_auth_provider.EnvHuggingFaceAuthProvider"),
+]
+
 SECRETS_PROVIDERS_DESC = [
     ("inline", ".secrets.inline_secrets_provider.InlineSecretsProvider"),
     (
@@ -213,6 +219,7 @@ def get_runner_cli_path():
 
 AWS_CLIENT_PROVIDERS = resolve_plugins("aws_client_provider")
 SECRETS_PROVIDERS = resolve_plugins("secrets_provider")
+HF_AUTH_PROVIDERS = resolve_plugins("hf_auth_provider")
 AZURE_CLIENT_PROVIDERS = resolve_plugins("azure_client_provider")
 GCP_CLIENT_PROVIDERS = resolve_plugins("gcp_client_provider")
 

metaflow/plugins/huggingface/__init__.py

@@ -0,0 +1,8 @@
+from .auth import HuggingFaceAuthProvider
+from .huggingface_decorator import HuggingFaceContext, HuggingFaceDecorator
+
+__all__ = [
+    "HuggingFaceAuthProvider",
+    "HuggingFaceContext",
+    "HuggingFaceDecorator",
+]

metaflow/plugins/huggingface/auth.py

@@ -0,0 +1,28 @@
+"""
+Pluggable HuggingFace authentication for the @huggingface decorator.
+
+Implement HuggingFaceAuthProvider and register in HF_AUTH_PROVIDERS_DESC
+to supply tokens from org-specific backends. If no provider is configured,
+the decorator falls back to HF_TOKEN / HUGGING_FACE_HUB_TOKEN environment variables.
+"""
+
+import abc
+from typing import Optional
+
+
+class HuggingFaceAuthProvider(abc.ABC):
+    """
+    Interface for pluggable HuggingFace authentication.
+    Providers are registered via HF_AUTH_PROVIDERS_DESC; the active provider
+    is selected by METAFLOW_HUGGINGFACE_AUTH_PROVIDER (default: env).
+    """
+
+    # Unique identifier for this provider (e.g. "env", "netflix-internal")
+    TYPE = None  # type: Optional[str]
+
+    @abc.abstractmethod
+    def get_token(self) -> Optional[str]:
+        """
+        Return the HuggingFace API token to use for this task, or None if no auth.
+        """
+        pass

metaflow/plugins/huggingface/env_auth_provider.py

@@ -0,0 +1,19 @@
+"""
+Default HuggingFace auth provider: read token from environment variables.
+"""
+
+import os
+from typing import Optional
+
+from .auth import HuggingFaceAuthProvider
+
+
+class EnvHuggingFaceAuthProvider(HuggingFaceAuthProvider):
+    """
+    Supplies HuggingFace token from HF_TOKEN or HUGGING_FACE_HUB_TOKEN.
+    """
+
+    TYPE = "env"
+
+    def get_token(self) -> Optional[str]:
+        return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

metaflow/plugins/huggingface/example_flow.py

@@ -0,0 +1,30 @@
+"""
+Minimal flow to test @huggingface locally.
+
+Run from repo root:
+  PYTHONPATH=. python -m metaflow.plugins.huggingface.example_flow run
+
+Requires: pip install huggingface_hub
+Optional: set HF_TOKEN or HUGGING_FACE_HUB_TOKEN for gated models.
+"""
+
+from metaflow import FlowSpec, step, huggingface
+
+
+class HuggingFaceExampleFlow(FlowSpec):
+    @huggingface(models=["bert-base-uncased"])
+    @step
+    def start(self):
+        from metaflow import current
+
+        self.model_path = current.huggingface.models["bert-base-uncased"]
+        print("Model path:", self.model_path)
+        self.next(self.end)
+
+    @step
+    def end(self):
+        print("Done. Model was at:", self.model_path)
+
+
+if __name__ == "__main__":
+    HuggingFaceExampleFlow()

metaflow/plugins/huggingface/huggingface_decorator.py

@@ -0,0 +1,185 @@
+"""
+@huggingface step decorator: pluggable auth for HuggingFace models (Part 1).
+
+Provides current.huggingface.models[key] -> local path. Supports models=[] and
+model_mapping={alias: repo_id@revision}. Uses huggingface_hub for download.
+"""
+
+import os
+from typing import Dict, List, Optional, Tuple
+
+from metaflow.decorators import StepDecorator
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_current import current
+
+
+# Minimal object exposed as current.huggingface with a .models mapping
+class HuggingFaceContext:
+    """
+    Context object attached to current.huggingface when @huggingface is used.
+    models maps user-facing key (alias or repo_id) to local filesystem path (str).
+    """
+
+    def __init__(self, models: Dict[str, str]):
+        self.models = models
+
+
+def _parse_repo_spec(value: str) -> Tuple[str, str]:
+    """Parse 'repo_id' or 'repo_id@revision' into (repo_id, revision)."""
+    value = (value or "").strip()
+    if not value:
+        raise MetaflowException(
+            "@huggingface: empty model spec; use repo_id or repo_id@revision"
+        )
+    if "@" in value:
+        repo_id, revision = value.rsplit("@", 1)
+        repo_id = repo_id.strip()
+        revision = revision.strip()
+        if not repo_id or not revision:
+            raise MetaflowException(
+                "@huggingface: invalid spec '%s'; use repo_id@revision" % value
+            )
+        return repo_id, revision
+    return value, "main"
+
+
+def _build_spec_map(
+    models: Optional[List[str]], model_mapping: Optional[Dict[str, str]]
+) -> Dict[str, Tuple[str, str]]:
+    """Build key -> (repo_id, revision). Key is alias or repo_id."""
+    spec_map = {}
+    if models:
+        for v in models:
+            if not isinstance(v, str):
+                raise MetaflowException(
+                    "@huggingface: models must be a list of strings, got %s" % type(v)
+                )
+            repo_id, revision = _parse_repo_spec(v)
+            spec_map[repo_id] = (repo_id, revision)
+    if model_mapping:
+        for k, v in model_mapping.items():
+            if not isinstance(k, str) or not isinstance(v, str):
+                raise MetaflowException(
+                    "@huggingface: model_mapping must be dict of str -> str"
+                )
+            repo_id, revision = _parse_repo_spec(v)
+            spec_map[k] = (repo_id, revision)
+    return spec_map
+
+
+def _get_auth_provider():
+    from metaflow.metaflow_config import METAFLOW_HUGGINGFACE_AUTH_PROVIDER
+    from metaflow.plugins import HF_AUTH_PROVIDERS
+
+    provider_type = METAFLOW_HUGGINGFACE_AUTH_PROVIDER or "env"
+    provider_cls = next(
+        (p for p in HF_AUTH_PROVIDERS if getattr(p, "TYPE", None) == provider_type),
+        None,
+    )
+    if provider_cls is None:
+        from metaflow.plugins.huggingface.env_auth_provider import (
+            EnvHuggingFaceAuthProvider,
+        )
+
+        return EnvHuggingFaceAuthProvider()
+    return provider_cls()
+
+
+def _download_model(
+    repo_id: str, revision: str, token: Optional[str], local_dir: str
+) -> str:
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as e:
+        raise MetaflowException(
+            "@huggingface requires the 'huggingface_hub' package. "
+            "Install it with: pip install huggingface_hub. Error: %s" % e
+        ) from e
+    path = snapshot_download(
+        repo_id=repo_id,
+        revision=revision,
+        token=token,
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,
+    )
+    return path
+
+
+class HuggingFaceDecorator(StepDecorator):
+    """
+    Declares HuggingFace models needed for this step. Auth is pluggable;
+    model paths are exposed via current.huggingface.models[key].
+
+    Parameters
+    ----------
+    models : list, optional
+        List of repo ids (and optional revisions), e.g.
+        ["meta-llama/Llama-2-7b", "bert-base-uncased@v1.0"].
+    model_mapping : dict, optional
+        Alias -> repo spec, e.g.
+        {"llama": "meta-llama/Llama-2-7b@main", "bert": "bert-base-uncased"}.
+        Access in step via current.huggingface.models["llama"].
+
+    MF Add To Current
+    -----------------
+    huggingface -> HuggingFaceContext
+        Object with a ``models`` attribute: dict-like mapping from model key
+        (alias or repo_id) to local filesystem path (str). Use
+        current.huggingface.models["key"] to get the path for loading with
+        transformers or other HF APIs.
+    """
+
+    name = "huggingface"
+    defaults = {"models": None, "model_mapping": None}
+
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        models = self.attributes.get("models")
+        model_mapping = self.attributes.get("model_mapping")
+        if not models and not model_mapping:
+            raise MetaflowException(
+                "@huggingface: specify at least one of 'models' or 'model_mapping'"
+            )
+        self._spec_map = _build_spec_map(models, model_mapping)
+        if not self._spec_map:
+            raise MetaflowException(
+                "@huggingface: at least one model or model_mapping entry is required"
+            )
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        token = None
+        try:
+            auth_provider = _get_auth_provider()
+            token = auth_provider.get_token()
+        except Exception as e:
+            raise MetaflowException(
+                "@huggingface: auth provider failed: %s" % e
+            ) from e
+
+        base_dir = os.path.join(current.tempdir or "/tmp", "metaflow_huggingface")
+        os.makedirs(base_dir, exist_ok=True)
+        path_map = {}  # key -> local path
+
+        for key, (repo_id, revision) in self._spec_map.items():
+            task_subdir = os.path.join(
+                base_dir, "%s_%s" % (repo_id.replace("/", "_"), revision)
+            )
+            local_path = _download_model(repo_id, revision, token, task_subdir)
+            path_map[key] = local_path
+
+        ctx = HuggingFaceContext(models=path_map)
+        current._update_env({"huggingface": ctx})

run_huggingface_demo.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+"""
+Demo script for the @huggingface decorator.
+
+Run from repo root:
+  python run_huggingface_demo.py run
+
+Or run a single step (no download):
+  python run_huggingface_demo.py run --no-pylint 2>&1 | head -30
+
+Requires: pip install -e .  and  pip install huggingface_hub
+Optional: HF_TOKEN or HUGGING_FACE_HUB_TOKEN for gated models.
+"""
+from metaflow import FlowSpec, step, huggingface, current
+
+
+class HuggingFaceDemoFlow(FlowSpec):
+    """Single step: download bert-base-uncased via @huggingface and print path."""
+
+    @huggingface(models=["bert-base-uncased"])
+    @step
+    def start(self):
+        path = current.huggingface.models["bert-base-uncased"]
+        print("Model path:", path)
+        self.model_path = path
+        self.next(self.end)
+
+    @step
+    def end(self):
+        print("Done. Model at:", self.model_path)
+
+
+if __name__ == "__main__":
+    HuggingFaceDemoFlow()