Throttling concurrent Azure Automation Runbooks

I don’t take any credit for this post. It’s mostly stuff I learnt from Google searches, and a script I put together for my use from other places. I should have credited them, but all this was put together over time and today I had to do it again and that’s when I went through my past snippets etc. and created something in a formal way.

Anyways, here’s the issue. I have a bunch of Azure Automation Runbooks that run on a schedule. Thing is, sometimes a Runbook takes longer than expected and so while it is running another one launches. This leads to both of them slowing down or conflicting with each others tasks and soon I have a whole bunch of Runbooks launching and making a mess of things. I need a way of ensuring only one Runbook is active at a time.

There’s a variable within Runbooks called PSPrivateMetadata. It doesn’t have anything except the JobId. You can extract this via $PSPrivateMetadata.JobId.Guid. Then you must search all running jobs to find the Automation Account name, Runbook name, and Resource Group name of this Runbook; and using that find what other jobs are running. And then decide what to do – wait for quit. Simple really. Wish the variable had the various names so I can save some trouble finding them… but oh well.

Here’s what I came up with:

# This is a Function I created (from various Google results) to throttle a Runbook.
# It will either wait or quit the runbook. 
function Throttle-AzRunbook {
    param(
        [switch]$quitRatherThanWait,
        [int]$numberOfInstances = 1
    )

    # Connect to Azure. With a Managed Identity in this case as that's what I use. 
    # It's like I am already connectes but I can't assume that within this function. 
    # Must connect to Azure before running Get-AzAutomationJob or Get-AzResource
    try {
        # From https://docs.microsoft.com/en-us/azure/automation/enable-managed-identity-for-automation#authenticate-access-with-system-assigned-managed-identity
        # Ensures you do not inherit an AzContext in your runbook
        Disable-AzContextAutosave -Scope Process | Out-Null

        # Connect to Azure with system-assigned managed identity
        Connect-AzAccount -Identity | Out-Null

    } catch {
        Write-Error "Runbook could not connect to Azure: $($_.Exception.Message)"
        exit
    }

    # Get the Job ID from PSPrivateMetadata. That's the only thing it contains!
    $automationJobId = $PSPrivateMetadata.JobId.Guid

    # Get all Runbooks in the current subscription
    $allAutomationAccounts = Get-AzResource -ResourceType Microsoft.Automation/automationAccounts

    $automationAccountName = $null
    $resourceGroupName = $null
    $runbookName = $null

    foreach ($automationAccount in $allAutomationAccounts) {
        $runbookJob = Get-AzAutomationJob -AutomationAccountName $automationAccount.Name `
                                    -ResourceGroupName $automationAccount.ResourceGroupName `
                                    -Id $automationJobId `
                                    -ErrorAction SilentlyContinue

        if (!([string]::IsNullOrEmpty($runbookJob))) {
            $automationAccountName = $runbookJob.AutomationAccountName
            $resourceGroupName = $runbookJob.ResourceGroupName
            $runbookName = $runbookJob.RunbookName
        }
    }

    # At this point I'll have the Automation Account Name, Runbook Name, Job ID and Resource Group Name, 
    # Find all other active jobs of this Runbook.

    $allActiveJobs = Get-AzAutomationJob -AutomationAccountName $automationAccountName `
                                    -ResourceGroupName $resourceGroupName `
                                    -RunbookName $runbookName | 
                    Where-Object { ($_.Status -eq "Running") -or ($_.Status -eq "Starting") -or ($_.Status -eq "Queued")}  

    if ($quitRatherThanWait.IsPresent -and $allActiveJobs.Count -gt $numberOfInstances) {
        Write-Output "Exiting as another job is already running"
        exit

    } else {
        $oldestJob = $AllActiveJobs | Sort-Object -Property CreationTime  | Select-Object -First 1

        # If this job is not the oldest created job we will wait until the existing jobs complete or the number of jobs is less than numberOfInstances
        while (($AutomationJobID -ne $oldestJob.JobId) -and ($allActiveJobs.Count -ge $numberOfInstances)) {
            Write-Output "Waiting as there are currently running $($allActiveJobs.Count) active jobs for this runbook already. Sleeping 30 seconds..."
            Write-Output "Oldest Job is $($oldestJob.JobId)"
        
            Start-Sleep -Seconds 30
        
            $allActiveJobs = Get-AzAutomationJob -AutomationAccountName $automationAccountName `
                                        -ResourceGroupName $resourceGroupName `
                                        -RunbookName $runbookName | 
                        Where-Object { ($_.Status -eq "Running") -or ($_.Status -eq "Starting") -or ($_.Status -eq "Queued")}  
        
            $oldestJob = $allActiveJobs | Sort-Object -Property CreationTime | Select-Object -First 1
        } 
        
        Write-Output "Job can continue..."
    }
}

# This is a Function I created (from various Google results) to throttle a Runbook.

# It will either wait or quit the runbook.

function Throttle-AzRunbook {

param(

[switch]$quitRatherThanWait,

[int]$numberOfInstances = 1

)

# Connect to Azure. With a Managed Identity in this case as that's what I use.

# It's like I am already connectes but I can't assume that within this function.

# Must connect to Azure before running Get-AzAutomationJob or Get-AzResource

try {

# From https://docs.microsoft.com/en-us/azure/automation/enable-managed-identity-for-automation#authenticate-access-with-system-assigned-managed-identity

# Ensures you do not inherit an AzContext in your runbook

Disable-AzContextAutosave -Scope Process | Out-Null

# Connect to Azure with system-assigned managed identity

Connect-AzAccount -Identity | Out-Null

} catch {

Write-Error "Runbook could not connect to Azure: $($_.Exception.Message)"

exit

}

# Get the Job ID from PSPrivateMetadata. That's the only thing it contains!

$automationJobId = $PSPrivateMetadata.JobId.Guid

# Get all Runbooks in the current subscription

$allAutomationAccounts = Get-AzResource -ResourceType Microsoft.Automation/automationAccounts

$automationAccountName = $null

$resourceGroupName = $null

$runbookName = $null

foreach ($automationAccount in $allAutomationAccounts) {

$runbookJob = Get-AzAutomationJob -AutomationAccountName $automationAccount.Name `

-ResourceGroupName $automationAccount.ResourceGroupName `

-Id $automationJobId `

-ErrorAction SilentlyContinue

if (!([string]::IsNullOrEmpty($runbookJob))) {

$automationAccountName = $runbookJob.AutomationAccountName

$resourceGroupName = $runbookJob.ResourceGroupName

$runbookName = $runbookJob.RunbookName

}

# At this point I'll have the Automation Account Name, Runbook Name, Job ID and Resource Group Name,

# Find all other active jobs of this Runbook.

$allActiveJobs = Get-AzAutomationJob -AutomationAccountName $automationAccountName `

-ResourceGroupName $resourceGroupName `

-RunbookName $runbookName |

Where-Object { ($_.Status -eq "Running") -or ($_.Status -eq "Starting") -or ($_.Status -eq "Queued")}

if ($quitRatherThanWait.IsPresent -and $allActiveJobs.Count -gt $numberOfInstances) {

Write-Output "Exiting as another job is already running"

exit

} else {

$oldestJob = $AllActiveJobs | Sort-Object -Property CreationTime | Select-Object -First 1

# If this job is not the oldest created job we will wait until the existing jobs complete or the number of jobs is less than numberOfInstances

while (($AutomationJobID -ne $oldestJob.JobId) -and ($allActiveJobs.Count -ge $numberOfInstances)) {

Write-Output "Waiting as there are currently running $($allActiveJobs.Count) active jobs for this runbook already. Sleeping 30 seconds..."

Write-Output "Oldest Job is $($oldestJob.JobId)"

Start-Sleep -Seconds 30

$allActiveJobs = Get-AzAutomationJob -AutomationAccountName $automationAccountName `

-ResourceGroupName $resourceGroupName `

-RunbookName $runbookName |

Where-Object { ($_.Status -eq "Running") -or ($_.Status -eq "Starting") -or ($_.Status -eq "Queued")}

$oldestJob = $allActiveJobs | Sort-Object -Property CreationTime | Select-Object -First 1

}

Write-Output "Job can continue..."

}

Will update this in case of any bugs later.

I put a line like this before my main script in the Runbook:

# === [ Main Script ] ===
# Quit if another instance is active
Throttle-AzRunbook -quitRatherThanWait

# === [ Main Script ] ===

# Quit if another instance is active

Throttle-AzRunbook -quitRatherThanWait

And that’s it.