Merge branch 'main' into issue-81-local-setup

astronomer · Nov 27, 2024 · bde4aee · bde4aee
2 parents 5f1c4e0 + 1abf239
commit bde4aee
Show file tree

Hide file tree

Showing 27 changed files with 188 additions and 166 deletions.
diff --git a/.astro-registry.yaml b/.astro-registry.yaml
@@ -4,15 +4,15 @@ display-name: Ray
 docs-url: https://github.com/astronomer/astro-provider-ray/blob/main/README.md
 
 hooks:
-  - module: ray_provider.hooks.ray.RayHook
+  - module: ray_provider.hooks.RayHook
 
 decorators:
-  - module: ray_provider.decorators.ray.ray
+  - module: ray_provider.decorators.ray
 
 operators:
-  - module: ray_provider.operators.ray.SetupRayCluster
-  - module: ray_provider.operators.ray.SubmitRayJob
-  - module: ray_provider.operators.ray.DeleteRayCluster
+  - module: ray_provider.operators.SetupRayCluster
+  - module: ray_provider.operators.SubmitRayJob
+  - module: ray_provider.operators.DeleteRayCluster
 
 triggers:
-  - module: ray_provider.triggers.ray.RayJobTrigger
+  - module: ray_provider.triggers.RayJobTrigger
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -1,6 +1,28 @@
 CHANGELOG
 =========
 
+0.3.0 (2024-11-29)
+------------------
+
+**Breaking changes**
+
+In order to improve the development and troubleshooting DAGs created with this provider, we introduced breaking changes
+to the folder structure. It was flattened and the import paths to existing decorators, hooks, operators and trigger
+changed, as documented in the table below:
+
++-----------+---------------------------------------------+-----------------------------------------+
+| Type      | Previous import path                        | Current import path                     |
++===========+=============================================+=========================================+
+| Decorator | ray_provider.decorators.ray.ray             | ray_provider.decorators.ray             |
+| Hook      | ray_provider.hooks.ray.RayHook              | ray_provider.hooks.RayHook              |
+| Operator  | ray_provider.operators.ray.DeleteRayCluster | ray_provider.operators.DeleteRayCluster |
+| Operator  | ray_provider.operators.ray.SetupRayCluster  | ray_provider.operators.SetupRayCluster  |
+| Operator  | ray_provider.operators.ray.SubmitRayJob     | ray_provider.operators.SubmitRayJob     |
+| Trigger   | ray_provider.triggers.ray.RayJobTrigger     | ray_provider.triggers.RayJobTrigger     |
++-----------+---------------------------------------------+-----------------------------------------+
+
+
+
 0.2.1 (2024-09-04)
 ------------------
 

diff --git a/dev/dags/ray_single_operator.py b/dev/dags/ray_single_operator.py
@@ -3,7 +3,7 @@
 
 from airflow import DAG
 
-from ray_provider.operators.ray import SubmitRayJob
+from ray_provider.operators import SubmitRayJob
 
 CONN_ID = "ray_conn"
 RAY_SPEC = Path(__file__).parent / "scripts/ray.yaml"

diff --git a/dev/dags/ray_taskflow_example.py b/dev/dags/ray_taskflow_example.py
@@ -4,7 +4,7 @@
 
 from airflow.decorators import dag, task
 
-from ray_provider.decorators.ray import ray
+from ray_provider.decorators import ray
 
 CONN_ID = "ray_conn"
 RAY_SPEC_FILENAME = os.getenv("RAY_SPEC_FILENAME", "ray.yaml")

diff --git a/dev/dags/ray_taskflow_example_existing_cluster.py b/dev/dags/ray_taskflow_example_existing_cluster.py
@@ -3,7 +3,7 @@
 
 from airflow.decorators import dag, task
 
-from ray_provider.decorators.ray import ray
+from ray_provider.decorators import ray
 
 CONN_ID = "ray_conn"
 FOLDER_PATH = Path(__file__).parent / "ray_scripts"

diff --git a/dev/dags/setup-teardown.py b/dev/dags/setup-teardown.py
@@ -3,7 +3,7 @@
 
 from airflow import DAG
 
-from ray_provider.operators.ray import DeleteRayCluster, SetupRayCluster, SubmitRayJob
+from ray_provider.operators import DeleteRayCluster, SetupRayCluster, SubmitRayJob
 
 CONN_ID = "ray_conn"
 RAY_SPEC = Path(__file__).parent / "scripts/ray.yaml"

diff --git a/docs/api/ray_provider.decorators.rst b/docs/api/ray_provider.decorators.rst
@@ -2,7 +2,7 @@
 Decorators
 ----------
 
-.. automodule:: ray_provider.decorators.ray
+.. automodule:: ray_provider.decorators
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/docs/api/ray_provider.hooks.rst b/docs/api/ray_provider.hooks.rst
@@ -1,7 +1,7 @@
 Hook
 -----
 
-.. automodule:: ray_provider.hooks.ray
+.. automodule:: ray_provider.hooks
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/docs/api/ray_provider.operators.rst b/docs/api/ray_provider.operators.rst
@@ -1,7 +1,7 @@
 Operators
 ---------
 
-.. automodule:: ray_provider.operators.ray
+.. automodule:: ray_provider.operators
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/docs/api/ray_provider.triggers.rst b/docs/api/ray_provider.triggers.rst
@@ -1,7 +1,7 @@
 Trigger
 --------
 
-.. automodule:: ray_provider.triggers.ray
+.. automodule:: ray_provider.triggers
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/ray_provider/__init__.py b/ray_provider/__init__.py
@@ -11,6 +11,6 @@ def get_provider_info() -> dict[str, Any]:
         "package-name": "astro-provider-ray",  # Required
         "name": "Ray",  # Required
         "description": "An integration between airflow and ray",  # Required
-        "connection-types": [{"connection-type": "ray", "hook-class-name": "ray_provider.hooks.ray.RayHook"}],
+        "connection-types": [{"connection-type": "ray", "hook-class-name": "ray_provider.hooks.RayHook"}],
         "versions": [__version__],  # Required
     }
diff --git a/ray_provider/decorators/ray.py → ray_provider/decorators.py b/ray_provider/decorators/ray.py → ray_provider/decorators.py
@@ -12,7 +12,7 @@
 from airflow.exceptions import AirflowException
 from airflow.utils.context import Context
 
-from ray_provider.operators.ray import SubmitRayJob
+from ray_provider.operators import SubmitRayJob
 
 
 class _RayDecoratedOperator(DecoratedOperator, SubmitRayJob):

diff --git a/ray_provider/decorators/__init__.py b/ray_provider/decorators/__init__.py
diff --git a/ray_provider/hooks/ray.py → ray_provider/hooks.py b/ray_provider/hooks/ray.py → ray_provider/hooks.py
diff --git a/ray_provider/hooks/__init__.py b/ray_provider/hooks/__init__.py
diff --git a/ray_provider/operators/ray.py → ray_provider/operators.py b/ray_provider/operators/ray.py → ray_provider/operators.py
@@ -10,8 +10,8 @@
 from airflow.utils.context import Context
 from ray.job_submission import JobStatus
 
-from ray_provider.hooks.ray import RayHook
-from ray_provider.triggers.ray import RayJobTrigger
+from ray_provider.hooks import RayHook
+from ray_provider.triggers import RayJobTrigger
 
 
 class SetupRayCluster(BaseOperator):

diff --git a/ray_provider/operators/__init__.py b/ray_provider/operators/__init__.py
diff --git a/ray_provider/triggers/ray.py → ray_provider/triggers.py b/ray_provider/triggers/ray.py → ray_provider/triggers.py
@@ -7,7 +7,7 @@
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 from ray.job_submission import JobStatus
 
-from ray_provider.hooks.ray import RayHook
+from ray_provider.hooks import RayHook
 
 
 class RayJobTrigger(BaseTrigger):
@@ -51,7 +51,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
         :return: A tuple containing the fully qualified class name and a dictionary of its parameters.
         """
         return (
-            "ray_provider.triggers.ray.RayJobTrigger",
+            "ray_provider.triggers.RayJobTrigger",
             {
                 "job_id": self.job_id,
                 "conn_id": self.conn_id,

diff --git a/ray_provider/triggers/__init__.py b/ray_provider/triggers/__init__.py
diff --git a/tests/decorators/__init__.py b/tests/decorators/__init__.py
diff --git a/tests/hooks/__init__.py b/tests/hooks/__init__.py
diff --git a/tests/operators/__init__.py b/tests/operators/__init__.py
diff --git a/tests/decorators/test_ray_decorators.py → tests/test_decorators.py b/tests/decorators/test_ray_decorators.py → tests/test_decorators.py
@@ -5,7 +5,7 @@
 from airflow.exceptions import AirflowException
 from airflow.utils.context import Context
 
-from ray_provider.decorators.ray import _RayDecoratedOperator, ray
+from ray_provider.decorators import _RayDecoratedOperator, ray
 
 
 class TestRayDecoratedOperator:
@@ -81,7 +81,7 @@ def dummy_callable():
             _RayDecoratedOperator(task_id="test_task", config=config, python_callable=dummy_callable)
 
     @patch.object(_RayDecoratedOperator, "_extract_function_body")
-    @patch("ray_provider.decorators.ray.SubmitRayJob.execute")
+    @patch("ray_provider.decorators.SubmitRayJob.execute")
     def test_execute_decorated_function(self, mock_super_execute, mock_extract_function_body):
         config = {
             "runtime_env": {"pip": ["ray"]},
@@ -101,7 +101,7 @@ def dummy_callable():
         assert operator.entrypoint == "python script.py"
         assert "working_dir" in operator.runtime_env
 
-    @patch("ray_provider.decorators.ray.SubmitRayJob.execute")
+    @patch("ray_provider.decorators.SubmitRayJob.execute")
     def test_execute_with_entrypoint(self, mock_super_execute):
         config = {
             "entrypoint": "python my_script.py",
@@ -119,7 +119,7 @@ def dummy_callable():
         assert result == "success"
         assert operator.entrypoint == "python my_script.py"
 
-    @patch("ray_provider.decorators.ray.SubmitRayJob.execute")
+    @patch("ray_provider.decorators.SubmitRayJob.execute")
     def test_execute_failure(self, mock_super_execute):
         config = {}