Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Validate Cluster Access Mode #1551

Merged
merged 15 commits into from
Sep 25, 2024
+41 −4
Merged

Validate Cluster Access Mode #1551

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
0cc5489
validation
Sep 25, 2024
03a1874
invalid
Sep 25, 2024
0fe3ba6
parse
Sep 25, 2024
9589d8c
test
Sep 25, 2024
7130d3d
test
Sep 25, 2024
7be07f2
test
Sep 25, 2024
da5f468
test
Sep 25, 2024
ae43179
test
Sep 25, 2024
2e59cd6
test
Sep 25, 2024
8cf5473
test
Sep 25, 2024
4ab5bc1
works?
Sep 25, 2024
51d8b24
precommit
Sep 25, 2024
1cb5aae
NONE
Sep 25, 2024
8f06e7e
Update llmfoundry/command_utils/data_prep/convert_delta_to_json.py
KuuCi Sep 25, 2024
87edcfd
update test
Sep 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions llmfoundry/command_utils/data_prep/convert_delta_to_json.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

from llmfoundry.utils.exceptions import (
ClusterDoesNotExistError,
ClusterInvalidAccessMode,
FailedToConnectToDatabricksError,
FailedToCreateSQLConnectionError,
InsufficientPermissionsError,
Expand Down Expand Up @@ -546,6 +547,17 @@ def validate_and_get_cluster_info(
if res is None:
raise ClusterDoesNotExistError(cluster_id)

data_security_mode = str(
res.data_security_mode,
).upper()[len('DATASECURITYMODE.'):]

# NONE stands for No Isolation Shared
if data_security_mode == 'NONE':
raise ClusterInvalidAccessMode(
cluster_id=cluster_id,
access_mode=data_security_mode,
)

assert res.spark_version is not None
stripped_runtime = re.sub(
r'[a-zA-Z]',
Expand Down
13 changes: 13 additions & 0 deletions llmfoundry/utils/exceptions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,19 @@ def __init__(self, cluster_id: str) -> None:
super().__init__(message, cluster_id=cluster_id)


class ClusterInvalidAccessMode(NetworkError):
"""Error thrown when the cluster does not exist."""

def __init__(self, cluster_id: str, access_mode: str) -> None:
message = f'Cluster with id {cluster_id} has access mode {access_mode}. ' + \
'please make sure the cluster used has access mode Shared or Single User!'
super().__init__(
message,
cluster_id=cluster_id,
access_mode=access_mode,
)


class FailedToCreateSQLConnectionError(
NetworkError,
):
Expand Down
20 changes: 16 additions & 4 deletions tests/a_scripts/data_prep/test_convert_delta_to_json.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,10 @@ def test_dbconnect_called(
DATABRICKS_TOKEN = 'token'
use_serverless = False

mock_cluster_response = Namespace(spark_version='14.1.0-scala2.12')
mock_cluster_response = Namespace(
spark_version='14.1.0-scala2.12',
data_security_mode='SINGLE_USER',
)
mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response

mock_remote = MagicMock()
Expand Down Expand Up @@ -286,7 +289,10 @@ def test_sqlconnect_called_dbr13(
DATABRICKS_TOKEN = 'token'
use_serverless = False

mock_cluster_response = Namespace(spark_version='13.0.0-scala2.12')
mock_cluster_response = Namespace(
spark_version='13.0.0-scala2.12',
data_security_mode='SINGLE_USER',
)
mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response

fetch_DT(
Expand Down Expand Up @@ -338,7 +344,10 @@ def test_sqlconnect_called_dbr14(
DATABRICKS_TOKEN = 'token'
use_serverless = False

mock_cluster_response = Namespace(spark_version='14.2.0-scala2.12')
mock_cluster_response = Namespace(
spark_version='14.2.0-scala2.12',
data_security_mode='SINGLE_USER',
)
mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response

fetch_DT(
Expand Down Expand Up @@ -390,7 +399,10 @@ def test_sqlconnect_called_https(
DATABRICKS_TOKEN = 'token'
use_serverless = False

mock_cluster_response = Namespace(spark_version='14.2.0-scala2.12')
mock_cluster_response = Namespace(
spark_version='14.2.0-scala2.12',
data_security_mode='SINGLE_USER',
)
mock_workspace_client.return_value.clusters.get.return_value = mock_cluster_response

fetch_DT(
Expand Down
Loading