Rr 139 error creating pool

cfa/cloudops/_cloudclient.py

@@ -354,7 +354,7 @@ def create_pool(
                 raise ValueError(
                     "monitoring_script_url is required when enabling node monitoring"
                 )
-
+            else:
                 start_task_command = rf"""/bin/bash -c 'set -euo pipefail mkdir
                                        -p /mnt/batch/tasks/startup/wd/node-metrics chmod +x ./start-metrics.sh
                                         nohup ./start-metrics.sh {monitoring_interval_seconds} \

cfa/cloudops/blob.py

@@ -6,10 +6,11 @@
 import os
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
+from typing import Any
 
 import anyio
-from azure.batch import models
 from azure.identity import ManagedIdentityCredential
+from azure.mgmt.batch import models
 from azure.storage.blob import (
     BlobImmutabilityPolicyMode,
     BlobServiceClient,
@@ -290,9 +291,7 @@ def download_from_storage_container(
 def get_node_mount_config(
     storage_containers: str | list[str],
     account_names: str | list[str],
-    identity_references: (
-        models.ComputeNodeIdentityReference | list[models.ComputeNodeIdentityReference]
-    ),
+    identity_references: Any,
     shared_relative_mount_path: str = "",
     mount_names: list[str] = None,
     blobfuse_options: str | list[str] = "",
@@ -334,7 +333,7 @@ def get_node_mount_config(
             storage_containers and isn't exactly 1.
 
     Example:
-        >>> from azure.batch import models
+        >>> from azure.mgmt.batch import models
         >>> identity_ref = models.ComputeNodeIdentityReference(
         ...     resource_id="/subscriptions/.../resourceGroups/.../providers/..."
         ... )
@@ -431,10 +430,33 @@ def get_node_mount_config(
         blob_str = " -o direct_io"
         logger.debug("Caching disabled - adding direct_io option")
 
+    def _to_mgmt_identity_reference(
+        identity_reference: Any,
+    ) -> models.ComputeNodeIdentityReference:
+        if isinstance(identity_reference, models.ComputeNodeIdentityReference):
+            return identity_reference
+
+        resource_id = None
+        if isinstance(identity_reference, str):
+            resource_id = identity_reference
+        elif isinstance(identity_reference, dict):
+            resource_id = identity_reference.get("resource_id")
+        else:
+            resource_id = getattr(identity_reference, "resource_id", None)
+
+        if not resource_id:
+            raise TypeError(
+                "Each identity reference must be a ComputeNodeIdentityReference, "
+                "a dict containing 'resource_id', or a resource_id string."
+            )
+
+        return models.ComputeNodeIdentityReference(resource_id=resource_id)
+
     mount_configs = []
     for account_name, container_name, relative_mount_path, identity_reference in zip(
         account_names, storage_containers, relative_mount_paths, identity_references
     ):
+        mgmt_identity_reference = _to_mgmt_identity_reference(identity_reference)
         logger.debug(
             f"Creating mount config: container '{container_name}' from account '{account_name}' -> '{relative_mount_path}'"
         )
@@ -446,7 +468,7 @@ def get_node_mount_config(
                     container_name=container_name,
                     relative_mount_path=relative_mount_path,
                     blobfuse_options=blobfuse_options + blob_str,
-                    identity_reference=identity_reference,
+                    identity_reference=mgmt_identity_reference,
                     **kwargs,
                 )
             )

changelog.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 The versioning pattern is `major.minor.patch`.
 
 ---
+## v0.3.15
+- set an upper limit for azure-mgmt-batch package to undo its breaking changes
+
 ## v0.3.14
 - fixed defect in `batch_helpers.get_rel_mnt_path` function for getting optional mount configuration in a pool
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cfa.cloudops"
-version = "0.3.14"
+version = "0.3.15"
 description = "Cloud storage, batch, functions, MLOps assistance"
 authors = [
     {name = "Ryan Raasch", email = "xng3@cdc.gov"}
@@ -13,7 +13,7 @@ dependencies = [
     "pandas>=2.3.2",
     "polars>=1.33.1",
     "azure-core>=1.35.1",
-    "azure-mgmt-batch>=19.0.0",
+    "azure-mgmt-batch>=18.0.0,<19.0.0",
     "azure-mgmt-web>=10.1.0",
     "azure-storage-blob>=12.26.0",
     "azure-identity>=1.25.0",

tests/test_blob.py

@@ -4,6 +4,7 @@
 import anyio
 import pytest
 from azure.batch import models
+from azure.mgmt.batch import models as mgmt_models
 from shared_fixtures import FAKE_BLOBS, MockLogger
 
 from cfa.cloudops.blob import (
@@ -80,6 +81,38 @@ def test_get_node_mount_config_success(mock_compute_node):
     )
     assert mounts
     assert len(mounts) == 2
+    assert isinstance(
+        mounts[0].azure_blob_file_system_configuration.identity_reference,
+        mgmt_models.ComputeNodeIdentityReference,
+    )
+
+
+def test_get_node_mount_config_success_identity_dict():
+    mounts = get_node_mount_config(
+        storage_containers=["mock-container-1", "mock-container-2"],
+        account_names=["mock-account-1"],
+        identity_references={"resource_id": "mock-resource-id"},
+    )
+    assert mounts
+    assert len(mounts) == 2
+    assert isinstance(
+        mounts[0].azure_blob_file_system_configuration.identity_reference,
+        mgmt_models.ComputeNodeIdentityReference,
+    )
+
+
+def test_get_node_mount_config_success_identity_str():
+    mounts = get_node_mount_config(
+        storage_containers=["mock-container-1", "mock-container-2"],
+        account_names=["mock-account-1"],
+        identity_references="mock-resource-id",
+    )
+    assert mounts
+    assert len(mounts) == 2
+    assert isinstance(
+        mounts[0].azure_blob_file_system_configuration.identity_reference,
+        mgmt_models.ComputeNodeIdentityReference,
+    )
 
 
 def test_get_node_mount_config_success_alternate(mock_compute_node):