Skip to content

Commit

Permalink
Added E2E test to run Python wheels on interactive cluster created in…
Browse files Browse the repository at this point in the history
… bundle (#1864)

## Changes
Added E2E test to run python wheels on interactive cluster created in
bundle.

We had a gap in testing wheel on all purpose clusters, so this PR
addresses the gap
  • Loading branch information
andrewnester authored Nov 1, 2024
1 parent f3bf33d commit 71cf426
Show file tree
Hide file tree
Showing 6 changed files with 102 additions and 4 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
{
"properties": {
"project_name": {
"type": "string",
"default": "my_test_code",
"description": "Unique name for this project"
},
"spark_version": {
"type": "string",
"description": "Spark version used for job cluster"
},
"node_type_id": {
"type": "string",
"description": "Node type id for job cluster"
},
"unique_id": {
"type": "string",
"description": "Unique ID for job name"
},
"instance_pool_id": {
"type": "string",
"description": "Instance pool id for job cluster"
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
bundle:
name: wheel-task

workspace:
root_path: "~/.bundle/{{.unique_id}}"

resources:
clusters:
test_cluster:
cluster_name: "test-cluster-{{.unique_id}}"
spark_version: "{{.spark_version}}"
node_type_id: "{{.node_type_id}}"
num_workers: 1
data_security_mode: USER_ISOLATION

jobs:
some_other_job:
name: "[${bundle.target}] Test Wheel Job {{.unique_id}}"
tasks:
- task_key: TestTask
existing_cluster_id: "${resources.clusters.test_cluster.cluster_id}"
python_wheel_task:
package_name: my_test_code
entry_point: run
parameters:
- "one"
- "two"
libraries:
- whl: ./dist/*.whl
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from setuptools import setup, find_packages

import {{.project_name}}

setup(
name="{{.project_name}}",
version={{.project_name}}.__version__,
author={{.project_name}}.__author__,
url="https://databricks.com",
author_email="[email protected]",
description="my example wheel",
packages=find_packages(include=["{{.project_name}}"]),
entry_points={"group1": "run={{.project_name}}.__main__:main"},
install_requires=["setuptools"],
)
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
__version__ = "0.0.1"
__author__ = "Databricks"
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
"""
The entry point of the Python Wheel
"""

import sys


def main():
# This method will print the provided arguments
print("Hello from my func")
print("Got arguments:")
print(sys.argv)


if __name__ == "__main__":
main()
19 changes: 15 additions & 4 deletions internal/bundle/python_wheel_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,18 @@ import (

"github.com/databricks/cli/internal"
"github.com/databricks/cli/internal/acc"
"github.com/databricks/cli/internal/testutil"
"github.com/databricks/cli/libs/env"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
)

func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bool) {
func runPythonWheelTest(t *testing.T, templateName string, sparkVersion string, pythonWheelWrapper bool) {
ctx, _ := acc.WorkspaceTest(t)

nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
instancePoolId := env.Get(ctx, "TEST_INSTANCE_POOL_ID")
bundleRoot, err := initTestTemplate(t, ctx, "python_wheel_task", map[string]any{
bundleRoot, err := initTestTemplate(t, ctx, templateName, map[string]any{
"node_type_id": nodeTypeId,
"unique_id": uuid.New().String(),
"spark_version": sparkVersion,
Expand Down Expand Up @@ -45,9 +46,19 @@ func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bo
}

func TestAccPythonWheelTaskDeployAndRunWithoutWrapper(t *testing.T) {
runPythonWheelTest(t, "13.3.x-snapshot-scala2.12", false)
runPythonWheelTest(t, "python_wheel_task", "13.3.x-snapshot-scala2.12", false)
}

func TestAccPythonWheelTaskDeployAndRunWithWrapper(t *testing.T) {
runPythonWheelTest(t, "12.2.x-scala2.12", true)
runPythonWheelTest(t, "python_wheel_task", "12.2.x-scala2.12", true)
}

func TestAccPythonWheelTaskDeployAndRunOnInteractiveCluster(t *testing.T) {
_, wt := acc.WorkspaceTest(t)

if testutil.IsAWSCloud(wt.T) {
t.Skip("Skipping test for AWS cloud because it is not permitted to create clusters")
}

runPythonWheelTest(t, "python_wheel_task_with_cluster", defaultSparkVersion, false)
}

0 comments on commit 71cf426

Please sign in to comment.