Add `archive_model_in_azure` fn for deleted repositories (#1)

- Add `archive_model_in_azure` fn for deleted repositories (a0e77b2a3664491d50fb850d5c4d255e567735af)
- Add Azure-related requirements (fd0ba9751f257887fef0613fdc5cc571fa23825b)


Co-authored-by: Alvaro Bartolome <[email protected]>

main.py

@@ -6,6 +6,9 @@ import json
 import logging
 import time
 
+from azure.ai.ml import MLClient
+from azure.identity import DefaultAzureCredential
+
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -18,6 +21,17 @@ CATALOG_WEBHOOKS_FLEET_DATASET = os.environ.get("CATALOG_WEBHOOKS_FLEET_DATASET"
 SIMSHIP_WEBHOOK = os.environ.get("SIMSHIP_WEBHOOK")
 SIMSHIP_WEBHOOK_SECRET = os.environ.get("SIMSHIP_WEBHOOK_SECRET")
 
+# Initialize Azure ML Client
+AZURE_SUBSCRIPTION_ID = os.getenv('AZURE_SUBSCRIPTION_ID')
+AZURE_RESOURCE_GROUP = os.getenv('AZURE_RESOURCE_GROUP')
+
+client = MLClient(
+    credential=DefaultAzureCredential(),
+    subscription_id=AZURE_SUBSCRIPTION_ID,
+    resource_group_name=AZURE_RESOURCE_GROUP,
+    registry_name="HuggingFace",
+)
+
 def send_slack_message(payload: dict, max_retries: int = 3, timeout: int = 10, retry_delay: int = 2):
     """
     Send a message to Slack using webhook URL with timeout and retry mechanism
@@ -122,6 +136,11 @@ def process_model_catalog_webhook(data: dict):
         )
         send_slack_message({"text": message, "type": "", "link": ""})
 
+        try:
+            archive_model_in_azure(model_id=repo.get("name", ""))
+        except Exception as e:
+            logger.error(f"Model {repo.get('name','')} couldn't be archived with {e}")
+
     # other events
     else:
         pass
@@ -167,6 +186,24 @@ def process_simship_webhook(data: dict):
     else:
         pass
 
+def archive_model_in_azure(model_id: str) -> None:
+    name = model_id.replace("/", "-").replace("_", "-").lower()
+    models = list(client.models.list(name=name))
+    if len(models) < 1:
+        msg = f"Model {name} not available in Azure ML, then no need to delete anything."
+        logger.error(msg)
+        raise ValueError(msg)
+
+    for model in models:
+        try:
+            client.models.archive(name=model.name, version=model.version)  # type: ignore
+            logger.info(f"Archived {model.name=} {model.version=}")
+        except Exception as e:
+            msg = f"Failed when archiving model with {name=} with exception {e}"
+            logger.error(msg)
+            raise RuntimeError(msg)
+
+    logger.info(f"Model {name} already archived on Azure ML!")
 
 app = FastAPI()
 

requirements.txt

@@ -1,4 +1,7 @@
 fastapi
 uvicorn
 requests
-datasets
+datasets
+azure-ai-ml
+azureml
+azureml-core