Simplify the project structure and debugging (#93)

This PR flattens the provider folder structure to improve troubleshooting DAGs created with this provider and the provider's developer experience. It contains **breaking changes**, which is fine since the provider has not reached a stable state yet (we're under the 1.0 release). As part of this PR, we are changing 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 | Previously, there were four modules within the `ray_provider` named `ray.py`. Not only was the project structure more nested and complex than it had to be, making the overall project maintainability harder without a benefit, but this caused many problems when troubleshooting the actual ray provider. The following logs illustrate the problem from a debugging perspective: ![Screenshot 2024-11-26 at 15 10 22](https://github.com/user-attachments/assets/1cd87019-af9f-4c36-942e-e75983b2fada) Different lines mentioning `ray.py` referenced different modules. It was time-consuming and frustrating for developers to troubleshoot and identify which `ray.py` the logs referred to: was it the `hooks/ray.py`, `operators/ray.py`, `trigerrer/ray.py`, or `deocrators/ray.py`? This sort of ambiguity significantly delayed the development of #81. We aim to reduce development time by introducing this change.
astronomer · Nov 27, 2024 · 1abf239 · 1abf239
1 parent 9ef6bc3
commit 1abf239
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
@@ -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"
 RAY_SPEC = Path(__file__).parent / "scripts/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_job"
 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 = {}