Spaces:
Running
Running
| import shutil | |
| import os | |
| import json | |
| from pathlib import Path | |
| from datetime import datetime | |
| from mlpipeline.entity import ModelPusherConfig, ModelPusherArtifact | |
| from mlpipeline.logging.logger import get_logger | |
| from mlpipeline.exception import ModelPusherException | |
| import sys | |
| logger = get_logger(__name__) | |
| class ModelPusher: | |
| def __init__(self, config: ModelPusherConfig): | |
| self.config = config | |
| def push_model(self) -> ModelPusherArtifact: | |
| try: | |
| logger.info("Starting model pusher") | |
| os.makedirs(self.config.model_registry_path, exist_ok=True) | |
| os.makedirs(self.config.root_dir, exist_ok=True) | |
| model_source = Path(self.config.model_path) | |
| model_dest = Path(self.config.model_registry_path) / model_source.name | |
| if model_source.is_dir(): | |
| if model_dest.exists(): | |
| shutil.rmtree(model_dest) | |
| shutil.copytree(model_source, model_dest) | |
| else: | |
| shutil.copy2(model_source, model_dest) | |
| logger.info(f"Model pushed to: {model_dest}") | |
| # Save push metadata | |
| metadata = { | |
| "pushed_at": datetime.now().isoformat(), | |
| "source_path": str(model_source), | |
| "destination_path": str(model_dest), | |
| "status": "success" | |
| } | |
| metadata_path = Path(self.config.root_dir) / "push_metadata.json" | |
| with open(metadata_path, 'w') as f: | |
| json.dump(metadata, f, indent=4) | |
| return ModelPusherArtifact( | |
| pushed_model_path=str(model_dest), | |
| is_pushed=True, | |
| message="Model pushed successfully" | |
| ) | |
| except Exception as e: | |
| raise ModelPusherException(str(e), sys) |