Merge branch 'main' into fix/caching-issue-lab-finish

Author: deep1401

api/pyproject.toml

@@ -31,6 +31,11 @@ dependencies = [
   "watchfiles==1.0.5",
   "werkzeug==3.1.3",
   "google-cloud-storage==2.18.2",
+  "azure-identity==1.25.1",
+  "azure-mgmt-compute==34.1.0",
+  "azure-mgmt-network==30.1.0",
+  "azure-mgmt-resource==24.0.0",
+  "azure-mgmt-authorization==4.0.0",
   "azure-storage-blob==12.28.0",
   "cashews==7.5.0",
   "paramiko==4.0.0",

api/test/api/test_azure_provider_endpoint.py

@@ -0,0 +1,132 @@
+"""Tests for Azure provider creation in team_provider_endpoints."""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from transformerlab.shared.models.models import ProviderType
+
+
+def test_create_azure_provider_injects_resource_group_and_team_id():
+    """create_provider_for_team must inject azure_resource_group and team_id into config for Azure providers."""
+    from transformerlab.services.compute_provider.team_provider_endpoints import create_provider_for_team
+    from pydantic import BaseModel
+
+    class FakeConfig(BaseModel):
+        azure_subscription_id: str = "sub-123"
+        azure_tenant_id: str = "tenant-456"
+        azure_client_id: str = "client-789"
+        azure_client_secret: str = "secret"
+        azure_location: str = "eastus"
+
+        def model_dump(self, **kwargs):
+            return {
+                "azure_subscription_id": self.azure_subscription_id,
+                "azure_tenant_id": self.azure_tenant_id,
+                "azure_client_id": self.azure_client_id,
+                "azure_client_secret": self.azure_client_secret,
+                "azure_location": self.azure_location,
+            }
+
+    captured_config = {}
+
+    async def fake_create_team_provider(session, team_id, name, provider_type, config, created_by_user_id):
+        captured_config.update(config)
+        mock_provider = MagicMock()
+        mock_provider.type = "azure"
+        mock_provider.id = "prov-1"
+        mock_provider.name = name
+        mock_provider.config = config
+        mock_provider.is_default = False
+        mock_provider.disabled = False
+        mock_provider.supported_accelerators = []
+        mock_provider.team_id = team_id
+        mock_provider.created_by_user_id = "user-1"
+        mock_provider.created_at = None
+        mock_provider.updated_at = None
+        return mock_provider
+
+    provider_data = MagicMock()
+    provider_data.type = ProviderType.AZURE
+    provider_data.name = "my-azure"
+    provider_data.config = FakeConfig()
+
+    with (
+        patch(
+            "transformerlab.services.compute_provider.team_provider_endpoints.create_team_provider",
+            side_effect=fake_create_team_provider,
+        ),
+        patch(
+            "transformerlab.services.compute_provider.team_provider_endpoints.list_team_providers",
+            new_callable=AsyncMock,
+            return_value=[],
+        ),
+        patch(
+            "transformerlab.services.compute_provider.team_provider_endpoints.cache.invalidate",
+            new_callable=AsyncMock,
+        ),
+    ):
+        asyncio.run(
+            create_provider_for_team(
+                session=MagicMock(),
+                team_id="team-xyz",
+                user=MagicMock(id="user-1"),
+                provider_data=provider_data,
+                force_refresh=False,
+            )
+        )
+
+    assert captured_config["azure_resource_group"] == "transformerlab-team-xyz"
+    assert captured_config["team_id"] == "team-xyz"
+
+
+def test_azure_is_in_allowed_provider_types():
+    """Azure must be in allowed_provider_types — verifies it no longer raises 400."""
+    from transformerlab.services.compute_provider.team_provider_endpoints import create_provider_for_team
+    from fastapi import HTTPException
+    from pydantic import BaseModel
+
+    class FakeConfig(BaseModel):
+        azure_subscription_id: str = "sub-123"
+
+        def model_dump(self, **kwargs):
+            return {"azure_subscription_id": self.azure_subscription_id}
+
+    provider_data = MagicMock()
+    provider_data.type = ProviderType.AZURE
+    provider_data.name = "my-azure"
+    provider_data.config = FakeConfig()
+
+    did_raise_type_error = False
+    with (
+        patch(
+            "transformerlab.services.compute_provider.team_provider_endpoints.list_team_providers",
+            new_callable=AsyncMock,
+            return_value=[],
+        ),
+        patch(
+            "transformerlab.services.compute_provider.team_provider_endpoints.create_team_provider",
+            new_callable=AsyncMock,
+            side_effect=Exception("stop here"),
+        ),
+        patch(
+            "transformerlab.services.compute_provider.team_provider_endpoints.cache.invalidate",
+            new_callable=AsyncMock,
+        ),
+    ):
+        try:
+            asyncio.run(
+                create_provider_for_team(
+                    session=MagicMock(),
+                    team_id="team-xyz",
+                    user=MagicMock(id="user-1"),
+                    provider_data=provider_data,
+                    force_refresh=False,
+                )
+            )
+        except HTTPException as e:
+            if e.status_code == 400 and "Invalid provider type" in str(e.detail):
+                did_raise_type_error = True
+        except Exception:
+            pass
+
+    assert not did_raise_type_error, "Azure should be in allowed_provider_types"

api/test/api/test_remote_job_status_service.py

@@ -522,6 +522,81 @@ async def fake_update_status(job_id, status, experiment_id=None):
     assert ("status", "STOPPED") in calls
 
 
+@pytest.mark.asyncio
+async def test_check_job_via_provider_azure_unknown_stopping_debounced(monkeypatch):
+    """AZURE STOPPING + UNKNOWN should debounce before STOPPED."""
+    monkeypatch.setattr(remote_job_status_service, "EMPTY_PROVIDER_JOBS_TERMINAL_THRESHOLD", 2)
+    from transformerlab.compute_providers.models import ClusterStatus, ClusterState
+
+    job = _make_job(status="STOPPING")
+    record = _make_provider_record("azure")
+
+    cluster_status = MagicMock(spec=ClusterStatus)
+    cluster_status.state = ClusterState.UNKNOWN
+    cluster_status.status_message = "unknown"
+
+    instance = MagicMock()
+    instance.get_cluster_status = MagicMock(return_value=cluster_status)
+
+    calls = []
+
+    async def fake_update_kv(job_id, key, value, exp_id):
+        calls.append(("kv", key, value))
+
+    monkeypatch.setattr(
+        remote_job_status_service.job_service,
+        "job_update_job_data_insert_key_value",
+        fake_update_kv,
+    )
+    monkeypatch.setattr(
+        remote_job_status_service.job_service,
+        "job_update_status",
+        AsyncMock(side_effect=AssertionError("should not be called on first poll")),
+    )
+
+    result = await _check_job_via_provider(job, "exp-1", record, instance)
+    assert result is False
+    assert ("kv", "azure_stopping_unknown_polls", 1) in calls
+
+    job["job_data"]["azure_stopping_unknown_polls"] = 1
+    monkeypatch.setattr(remote_job_status_service.job_service, "job_update_status", AsyncMock())
+    result = await _check_job_via_provider(job, "exp-1", record, instance)
+    assert result is True
+
+
+@pytest.mark.asyncio
+async def test_check_job_via_provider_azure_resets_unknown_stopping_counter(monkeypatch):
+    """AZURE clears debounce counter once cluster state is no longer UNKNOWN."""
+    from transformerlab.compute_providers.models import ClusterStatus, ClusterState
+
+    job = _make_job(status="STOPPING")
+    job["job_data"]["azure_stopping_unknown_polls"] = 2
+    record = _make_provider_record("azure")
+
+    cluster_status = MagicMock(spec=ClusterStatus)
+    cluster_status.state = ClusterState.DOWN
+    cluster_status.status_message = "ResourceNotFound"
+
+    instance = MagicMock()
+    instance.get_cluster_status = MagicMock(return_value=cluster_status)
+
+    calls = []
+
+    async def fake_update_kv(job_id, key, value, exp_id):
+        calls.append(("kv", key, value))
+
+    monkeypatch.setattr(
+        remote_job_status_service.job_service,
+        "job_update_job_data_insert_key_value",
+        fake_update_kv,
+    )
+    monkeypatch.setattr(remote_job_status_service.job_service, "job_update_status", AsyncMock())
+
+    result = await _check_job_via_provider(job, "exp-1", record, instance)
+    assert result is True
+    assert ("kv", "azure_stopping_unknown_polls", 0) in calls
+
+
 @pytest.mark.asyncio
 async def test_check_job_via_provider_local_still_running(monkeypatch):
     """LOCAL provider cluster still UP → returns False."""