From 809c67b675fac5b8f8fc312cccca667929f94267 Mon Sep 17 00:00:00 2001 From: Andrew Nester Date: Mon, 5 Aug 2024 16:44:23 +0200 Subject: [PATCH] Expand and upload local wheel libraries for all task types (#1649) ## Changes Fixes #1553 ## Tests Added regression test --- bundle/artifacts/whl/autodetect.go | 6 ++--- bundle/artifacts/whl/from_libraries.go | 7 +++++- bundle/libraries/libraries.go | 22 ++++++++----------- bundle/python/warning.go | 2 +- .../.gitignore | 3 +++ .../bundle.yml | 14 ++++++++++++ .../my_test_code/__init__.py | 2 ++ .../my_test_code/__main__.py | 16 ++++++++++++++ .../notebook.py | 3 +++ .../setup.py | 15 +++++++++++++ bundle/tests/python_wheel_test.go | 17 ++++++++++++++ 11 files changed, 89 insertions(+), 18 deletions(-) create mode 100644 bundle/tests/python_wheel/python_wheel_no_artifact_notebook/.gitignore create mode 100644 bundle/tests/python_wheel/python_wheel_no_artifact_notebook/bundle.yml create mode 100644 bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__init__.py create mode 100644 bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__main__.py create mode 100644 bundle/tests/python_wheel/python_wheel_no_artifact_notebook/notebook.py create mode 100644 bundle/tests/python_wheel/python_wheel_no_artifact_notebook/setup.py diff --git a/bundle/artifacts/whl/autodetect.go b/bundle/artifacts/whl/autodetect.go index ee77fff0..1601767f 100644 --- a/bundle/artifacts/whl/autodetect.go +++ b/bundle/artifacts/whl/autodetect.go @@ -27,9 +27,9 @@ func (m *detectPkg) Name() string { } func (m *detectPkg) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { - wheelTasks := libraries.FindAllWheelTasksWithLocalLibraries(b) - if len(wheelTasks) == 0 { - log.Infof(ctx, "No local wheel tasks in databricks.yml config, skipping auto detect") + tasks := libraries.FindTasksWithLocalLibraries(b) + if len(tasks) == 0 { + log.Infof(ctx, "No local tasks in databricks.yml config, skipping auto detect") return nil } log.Infof(ctx, "Detecting Python wheel project...") diff --git a/bundle/artifacts/whl/from_libraries.go b/bundle/artifacts/whl/from_libraries.go index ad321557..79161a82 100644 --- a/bundle/artifacts/whl/from_libraries.go +++ b/bundle/artifacts/whl/from_libraries.go @@ -27,8 +27,13 @@ func (*fromLibraries) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnost return nil } - tasks := libraries.FindAllWheelTasksWithLocalLibraries(b) + tasks := libraries.FindTasksWithLocalLibraries(b) for _, task := range tasks { + // Skip tasks that are not PythonWheelTasks for now, we can later support Jars too + if task.PythonWheelTask == nil { + continue + } + for _, lib := range task.Libraries { matchAndAdd(ctx, lib.Whl, b) } diff --git a/bundle/libraries/libraries.go b/bundle/libraries/libraries.go index 84ead052..72e5bcc6 100644 --- a/bundle/libraries/libraries.go +++ b/bundle/libraries/libraries.go @@ -44,29 +44,25 @@ func isEnvsWithLocalLibraries(envs []jobs.JobEnvironment) bool { return false } -func FindAllWheelTasksWithLocalLibraries(b *bundle.Bundle) []*jobs.Task { +func FindTasksWithLocalLibraries(b *bundle.Bundle) []jobs.Task { tasks := findAllTasks(b) envs := FindAllEnvironments(b) - wheelTasks := make([]*jobs.Task, 0) + allTasks := make([]jobs.Task, 0) for k, jobTasks := range tasks { for i := range jobTasks { - task := &jobTasks[i] - if task.PythonWheelTask == nil { - continue + task := jobTasks[i] + if isTaskWithLocalLibraries(task) { + allTasks = append(allTasks, task) } + } - if isTaskWithLocalLibraries(*task) { - wheelTasks = append(wheelTasks, task) - } - - if envs[k] != nil && isEnvsWithLocalLibraries(envs[k]) { - wheelTasks = append(wheelTasks, task) - } + if envs[k] != nil && isEnvsWithLocalLibraries(envs[k]) { + allTasks = append(allTasks, jobTasks...) } } - return wheelTasks + return allTasks } func isTaskWithLocalLibraries(task jobs.Task) bool { diff --git a/bundle/python/warning.go b/bundle/python/warning.go index 3da88b0d..d53796d7 100644 --- a/bundle/python/warning.go +++ b/bundle/python/warning.go @@ -35,7 +35,7 @@ func isPythonWheelWrapperOn(b *bundle.Bundle) bool { } func hasIncompatibleWheelTasks(ctx context.Context, b *bundle.Bundle) bool { - tasks := libraries.FindAllWheelTasksWithLocalLibraries(b) + tasks := libraries.FindTasksWithLocalLibraries(b) for _, task := range tasks { if task.NewCluster != nil { if lowerThanExpectedVersion(ctx, task.NewCluster.SparkVersion) { diff --git a/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/.gitignore b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/.gitignore new file mode 100644 index 00000000..f03e23bc --- /dev/null +++ b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/.gitignore @@ -0,0 +1,3 @@ +build/ +*.egg-info +.databricks diff --git a/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/bundle.yml b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/bundle.yml new file mode 100644 index 00000000..93e4e691 --- /dev/null +++ b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/bundle.yml @@ -0,0 +1,14 @@ +bundle: + name: python-wheel-notebook + +resources: + jobs: + test_job: + name: "[${bundle.environment}] My Wheel Job" + tasks: + - task_key: TestTask + existing_cluster_id: "0717-aaaaa-bbbbbb" + notebook_task: + notebook_path: "/notebook.py" + libraries: + - whl: ./dist/*.whl diff --git a/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__init__.py b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__init__.py new file mode 100644 index 00000000..909f1f32 --- /dev/null +++ b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__init__.py @@ -0,0 +1,2 @@ +__version__ = "0.0.1" +__author__ = "Databricks" diff --git a/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__main__.py b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__main__.py new file mode 100644 index 00000000..73d045af --- /dev/null +++ b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__main__.py @@ -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() diff --git a/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/notebook.py b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/notebook.py new file mode 100644 index 00000000..24dc150f --- /dev/null +++ b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/notebook.py @@ -0,0 +1,3 @@ +# Databricks notebook source + +print("Hello, World!") diff --git a/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/setup.py b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/setup.py new file mode 100644 index 00000000..7a1317b2 --- /dev/null +++ b/bundle/tests/python_wheel/python_wheel_no_artifact_notebook/setup.py @@ -0,0 +1,15 @@ +from setuptools import setup, find_packages + +import my_test_code + +setup( + name="my_test_code", + version=my_test_code.__version__, + author=my_test_code.__author__, + url="https://databricks.com", + author_email="john.doe@databricks.com", + description="my test wheel", + packages=find_packages(include=["my_test_code"]), + entry_points={"group_1": "run=my_test_code.__main__:main"}, + install_requires=["setuptools"], +) diff --git a/bundle/tests/python_wheel_test.go b/bundle/tests/python_wheel_test.go index 52b3d6e0..53c6764e 100644 --- a/bundle/tests/python_wheel_test.go +++ b/bundle/tests/python_wheel_test.go @@ -45,6 +45,23 @@ func TestPythonWheelBuildAutoDetect(t *testing.T) { require.NoError(t, diags.Error()) } +func TestPythonWheelBuildAutoDetectWithNotebookTask(t *testing.T) { + ctx := context.Background() + b, err := bundle.Load(ctx, "./python_wheel/python_wheel_no_artifact_notebook") + require.NoError(t, err) + + diags := bundle.Apply(ctx, b, bundle.Seq(phases.Load(), phases.Build())) + require.NoError(t, diags.Error()) + + matches, err := filepath.Glob("./python_wheel/python_wheel_no_artifact_notebook/dist/my_test_code-*.whl") + require.NoError(t, err) + require.Equal(t, 1, len(matches)) + + match := libraries.ValidateLocalLibrariesExist() + diags = bundle.Apply(ctx, b, match) + require.NoError(t, diags.Error()) +} + func TestPythonWheelWithDBFSLib(t *testing.T) { ctx := context.Background() b, err := bundle.Load(ctx, "./python_wheel/python_wheel_dbfs_lib")