Skip to content

Commit

Permalink
Fixed select_node_type error (#614)
Browse files Browse the repository at this point in the history
## Changes
<!-- Summary of your changes that are easy to understand -->
Calling `ws.clusters.select_node_type(local_disk=True,
min_memory_gb=16)` results in `AttributeError: NotAvailableInRegion`.
This happened because the enum for `CloudProviderNodeStatus` changed.
Updated the enum correctly

## Tests
<!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->

- [x] `make test` run locally
- [x] `make fmt` applied
  • Loading branch information
nkvuong authored Apr 12, 2024
1 parent a93140d commit 855a5e0
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions databricks/sdk/mixins/compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ def select_spark_version(self,
:param beta: bool
:param latest: bool
:param ml: bool
:param genomics: bool
:param gpu: bool
:param scala: str
:param spark_version: str
Expand All @@ -100,7 +101,7 @@ def select_spark_version(self,
for version in sv.versions:
if "-scala" + scala not in version.key:
continue
matches = ((not "apache-spark-" in version.key) and (("-ml-" in version.key) == ml)
matches = (("apache-spark-" not in version.key) and (("-ml-" in version.key) == ml)
and (("-hls-" in version.key) == genomics) and (("-gpu-" in version.key) == gpu)
and (("-photon-" in version.key) == photon)
and (("-aarch64-" in version.key) == graviton) and (("Beta" in version.name) == beta))
Expand Down Expand Up @@ -137,7 +138,7 @@ def _should_node_be_skipped(nt: compute.NodeType) -> bool:
return False
val = compute.CloudProviderNodeStatus
for st in nt.node_info.status:
if st in (val.NotAvailableInRegion, val.NotEnabledOnSubscription):
if st in (val.NOT_AVAILABLE_IN_REGION, val.NOT_ENABLED_ON_SUBSCRIPTION):
return True
return False

Expand Down Expand Up @@ -168,6 +169,8 @@ def select_node_type(self,
:param photon_driver_capable: bool
:param graviton: bool
:param is_io_cache_enabled: bool
:param support_port_forwarding: bool
:param fleet: bool
:returns: `node_type` compatible string
"""
Expand Down

0 comments on commit 855a5e0

Please sign in to comment.