This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 63000a77ba Added priority to Azure Container Instances (#40616)
63000a77ba is described below
commit 63000a77ba4ab70b50950f506cee6964e10c88a3
Author: Liam Alsopp <[email protected]>
AuthorDate: Tue Jul 9 11:54:44 2024 +0100
Added priority to Azure Container Instances (#40616)
---------
Co-authored-by: Liam Alsopp <[email protected]>
---
.../azure/operators/container_instances.py | 11 +++++++
.../azure/operators/test_container_instances.py | 37 ++++++++++++++++++++++
2 files changed, 48 insertions(+)
diff --git a/airflow/providers/microsoft/azure/operators/container_instances.py
b/airflow/providers/microsoft/azure/operators/container_instances.py
index 3d55776b5a..605c0e1608 100644
--- a/airflow/providers/microsoft/azure/operators/container_instances.py
+++ b/airflow/providers/microsoft/azure/operators/container_instances.py
@@ -94,6 +94,7 @@ class AzureContainerInstancesOperator(BaseOperator):
:param subnet_ids: The subnet resource IDs for a container group
:param dns_config: The DNS configuration for a container group.
:param diagnostics: Container group diagnostic information (Log Analytics).
+ :param priority: Container group priority, Possible values include:
'Regular', 'Spot'
**Example**::
@@ -129,6 +130,7 @@ class AzureContainerInstancesOperator(BaseOperator):
"workspaceKey": "workspaceKey",
}
},
+ priority="Regular",
command=["/bin/echo", "world"],
task_id="start_container",
)
@@ -163,6 +165,7 @@ class AzureContainerInstancesOperator(BaseOperator):
subnet_ids: list[ContainerGroupSubnetId] | None = None,
dns_config: DnsConfiguration | None = None,
diagnostics: ContainerGroupDiagnostics | None = None,
+ priority: str | None = "Regular",
**kwargs,
) -> None:
super().__init__(**kwargs)
@@ -203,6 +206,13 @@ class AzureContainerInstancesOperator(BaseOperator):
self.subnet_ids = subnet_ids
self.dns_config = dns_config
self.diagnostics = diagnostics
+ self.priority = priority
+ if self.priority not in ["Regular", "Spot"]:
+ raise AirflowException(
+ "Invalid value for the priority argument. "
+ "Please set 'Regular' or 'Spot' as the priority. "
+ f"Found `{self.priority}`."
+ )
def execute(self, context: Context) -> int:
# Check name again in case it was templated.
@@ -278,6 +288,7 @@ class AzureContainerInstancesOperator(BaseOperator):
subnet_ids=self.subnet_ids,
dns_config=self.dns_config,
diagnostics=self.diagnostics,
+ priority=self.priority,
)
self._ci_hook.create_or_update(self.resource_group, self.name,
container_group)
diff --git
a/tests/providers/microsoft/azure/operators/test_container_instances.py
b/tests/providers/microsoft/azure/operators/test_container_instances.py
index 0a153e2fdd..41d884e863 100644
--- a/tests/providers/microsoft/azure/operators/test_container_instances.py
+++ b/tests/providers/microsoft/azure/operators/test_container_instances.py
@@ -460,3 +460,40 @@ class TestACIOperator:
(_, _, called_cg), _ = aci_mock.return_value.create_or_update.call_args
assert called_cg.diagnostics == diagnostics
+
+
@mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook")
+ def test_execute_with_spot_discount(self, aci_mock):
+ expected_cg = make_mock_container(state="Terminated", exit_code=0,
detail_status="test")
+ aci_mock.return_value.get_state.return_value = expected_cg
+
+ aci_mock.return_value.exists.return_value = False
+
+ aci = AzureContainerInstancesOperator(
+ ci_conn_id=None,
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ priority="Spot",
+ )
+ aci.execute(None)
+
+ assert aci_mock.return_value.create_or_update.call_count == 1
+ (called_rg, called_cn, called_cg), _ =
aci_mock.return_value.create_or_update.call_args
+
+ assert called_rg == "resource-group"
+ assert called_cn == "container-name"
+
+ assert called_cg.location == "region"
+ assert called_cg.image_registry_credentials is None
+ assert called_cg.restart_policy == "Never"
+ assert called_cg.os_type == "Linux"
+ assert called_cg.priority == "Spot"
+
+ called_cg_container = called_cg.containers[0]
+ assert called_cg_container.name == "container-name"
+ assert called_cg_container.image == "container-image"
+
+ assert aci_mock.return_value.delete.call_count == 1
