Create hf_artifact_repo.py

hf_artifact_repo.py

@@ -0,0 +1,189 @@
+# hf_artifact_repo.py
+import os
+import tempfile
+import shutil
+from pathlib import Path
+from typing import List
+from huggingface_hub import HfApi, hf_hub_download, list_repo_files
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.entities import FileInfo
+
+
+class HuggingFaceArtifactRepository(ArtifactRepository):
+    """
+    MLflow Artifact Repository that uses HuggingFace Hub as storage
+    """
+    
+    def __init__(self, artifact_uri, hf_token=None, **kwargs):
+        super().__init__(artifact_uri)
+        self.hf_token = hf_token or os.getenv("HF_TOKEN")
+        self.api = HfApi()
+        
+        if artifact_uri.startswith("hf://"):
+            parts = artifact_uri.replace("hf://", "").split("/")
+            raw_type = parts[0] if len(parts) > 0 else "models"
+            type_mapping = {
+                "models": "model",
+                "datasets": "dataset",
+                "spaces": "space",
+                "model": "model",
+                "dataset": "dataset",
+                "space": "space"
+            }
+            self.repo_type = type_mapping.get(raw_type, None)
+            self.repo_id = f"{parts[1]}/{parts[2]}" if len(parts) > 2 else ""
+            self.base_path = "/".join(parts[3:]) if len(parts) > 3 else ""
+        else:
+            raise ValueError(f"Invalid HF URI: {artifact_uri}")
+        
+        print(f"🔧 Initialized HF Artifact Repository:")
+        print(f"   - Repo: {self.repo_id}")
+        print(f"   - Type: {self.repo_type}")
+        print(f"   - Base Path: {self.base_path}")
+    
+    def _get_full_path_in_repo(self, artifact_path):
+        if self.base_path and artifact_path:
+            return f"{self.base_path}/{artifact_path}".strip("/")
+        elif self.base_path:
+            return self.base_path.strip("/")
+        elif artifact_path:
+            return artifact_path.strip("/")
+        return ""
+    
+    def log_artifact(self, local_file, artifact_path=None):
+        artifact_path = artifact_path or ""
+        full_path = self._get_full_path_in_repo(artifact_path)
+        file_name = os.path.basename(local_file)
+        path_in_repo = f"{full_path}/{file_name}".strip("/") if full_path else file_name
+        
+        self.api.upload_file(
+            path_or_fileobj=local_file,
+            path_in_repo=path_in_repo,
+            repo_id=self.repo_id,
+            repo_type=self.repo_type,
+            token=self.hf_token
+        )
+    
+    def log_artifacts(self, local_dir, artifact_path=None):
+        artifact_path = artifact_path or ""
+        full_path = self._get_full_path_in_repo(artifact_path)
+        
+        for root, dirs, files in os.walk(local_dir):
+            for file in files:
+                local_file_path = os.path.join(root, file)
+                relative_path = os.path.relpath(local_file_path, local_dir)
+                
+                if full_path:
+                    path_in_repo = f"{full_path}/{relative_path}".replace("\\", "/")
+                else:
+                    path_in_repo = relative_path.replace("\\", "/")
+                
+                self.api.upload_file(
+                    path_or_fileobj=local_file_path,
+                    path_in_repo=path_in_repo,
+                    repo_id=self.repo_id,
+                    repo_type=self.repo_type,
+                    token=self.hf_token
+                )
+    
+    def list_artifacts(self, path=None):
+        try:
+            full_path = self._get_full_path_in_repo(path or "")
+            all_files = list_repo_files(
+                repo_id=self.repo_id,
+                repo_type=self.repo_type,
+                token=self.hf_token
+            )
+            
+            artifacts = []
+            for file_path in all_files:
+                if full_path:
+                    if file_path.startswith(full_path + "/") or file_path == full_path:
+                        relative_path = file_path[len(full_path):].lstrip("/")
+                        if relative_path:
+                            artifacts.append(FileInfo(relative_path, False, 0))
+                else:
+                    artifacts.append(FileInfo(file_path, False, 0))
+            
+            return artifacts
+        except Exception as e:
+            print(f"⚠️ Error listing artifacts: {e}")
+            return []
+    
+    def _download_file(self, remote_file_path, local_path):
+        full_path = self._get_full_path_in_repo(remote_file_path)
+        
+        downloaded_path = hf_hub_download(
+            repo_id=self.repo_id,
+            filename=full_path,
+            repo_type=self.repo_type,
+            token=self.hf_token,
+            local_dir=os.path.dirname(local_path)
+        )
+        
+        if downloaded_path != local_path:
+            shutil.copy2(downloaded_path, local_path)
+    
+    def download_artifacts(self, artifact_path, dst_path=None):
+        if dst_path is None:
+            dst_path = tempfile.mkdtemp()
+        
+        artifact_path = artifact_path or ""
+        full_path = self._get_full_path_in_repo(artifact_path)
+        
+        try:
+            all_files = list_repo_files(
+                repo_id=self.repo_id,
+                repo_type=self.repo_type,
+                token=self.hf_token
+            )
+            
+            for file_path in all_files:
+                if full_path:
+                    if file_path.startswith(full_path + "/") or file_path == full_path:
+                        relative_path = file_path[len(full_path):].lstrip("/")
+                        if relative_path:
+                            local_file_path = os.path.join(dst_path, relative_path)
+                            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+                            
+                            hf_hub_download(
+                                repo_id=self.repo_id,
+                                filename=file_path,
+                                repo_type=self.repo_type,
+                                token=self.hf_token,
+                                local_dir=dst_path
+                            )
+                else:
+                    local_file_path = os.path.join(dst_path, file_path)
+                    os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+                    
+                    hf_hub_download(
+                        repo_id=self.repo_id,
+                        filename=file_path,
+                        repo_type=self.repo_type,
+                        token=self.hf_token,
+                        local_dir=dst_path
+                    )
+            
+            return dst_path
+        except Exception as e:
+            print(f"❌ Error downloading artifacts: {e}")
+            raise
+
+
+# ======================
+# Auto-register when imported
+# ======================
+def _register_plugin():
+    try:
+        from mlflow.store.artifact import artifact_repository_registry
+        artifact_repository_registry._artifact_repository_registry.register(
+            "hf", HuggingFaceArtifactRepository
+        )
+        print("✅ HuggingFace Artifact Repository registered!")
+        return True
+    except Exception as e:
+        print(f"⚠️ Could not auto-register: {e}")
+        return False
+
+_register_plugin()