Integration test fix & Decorator change task.ray() -> ray.task()

.github/workflows/test.yml

@@ -109,21 +109,24 @@ jobs:
       - name: 'Set GCP Project ID'
         run: gcloud config set project ${{ secrets.PROJECT_ID }}
 
+      - name: Generate unique cluster name
+        id: cluster-name
+        run: echo "name=test-${GITHUB_RUN_ID:0:12}" >> $GITHUB_OUTPUT
+
       - name: Create GKE cluster
         run: |
-          gcloud container clusters create integration-test-cluster \
+          gcloud container clusters create ${{ steps.cluster-name.outputs.name }} \
           --zone us-central1-a \
-          --num-nodes 2 \
-          --machine-type e2-standard-8 \
-          --no-enable-autoupgrade \
-          --no-enable-autorepair \
+          --num-nodes 1 \
+          --machine-type e2-standard-4 \
+          --enable-autoscaling --min-nodes=1 --max-nodes=3 \
           --no-enable-ip-alias \
           --service-account=${{ secrets.SERVICE_ACCOUNT_EMAIL }}
 
       - name: Setup gcloud and get kubeconfig
         run: |
           gcloud components install gke-gcloud-auth-plugin
-          gcloud container clusters get-credentials integration-test-cluster --zone us-central1-a
+          gcloud container clusters get-credentials ${{ steps.cluster-name.outputs.name }} --zone us-central1-a
           kubectl config view --raw > kubeconfig.yaml
           echo "KUBECONFIG=${{ github.workspace }}/kubeconfig.yaml" >> $GITHUB_ENV
           echo "USE_GKE_GCLOUD_AUTH_PLUGIN=True" >> $GITHUB_ENV
@@ -139,4 +142,4 @@ jobs:
       - name: Delete GKE cluster
         if: always()
         run: |
-          gcloud container clusters delete integration-test-cluster --zone us-central1-a --quiet
+          gcloud container clusters delete ${{ steps.cluster-name.outputs.name }} --zone us-central1-a --quiet

example_dags/scripts/ray.yaml

@@ -20,11 +20,11 @@ spec:
           image: rayproject/ray-ml:latest
           resources:
             limits:
-              cpu: 4
-              memory: 8Gi
+              cpu: 1
+              memory: 3Gi
             requests:
-              cpu: 4
-              memory: 8Gi
+              cpu: 1
+              memory: 3Gi
           lifecycle:
             preStop:
               exec:
@@ -40,15 +40,11 @@ spec:
             name: serve
           - containerPort: 8080
             name: metrics
-          - containerPort: 44217
-            name: as-metrics # autoscaler
-          - containerPort: 44227
-            name: dash-metrics # dashboard
   workerGroupSpecs:
   - groupName: small-group
-    replicas: 2
-    minReplicas: 2
-    maxReplicas: 5
+    replicas: 1
+    minReplicas: 1
+    maxReplicas: 2
     rayStartParams:
       block: "true"
     template:
@@ -59,8 +55,8 @@ spec:
           image: rayproject/ray-ml:latest
           resources:
             limits:
-              cpu: 2
-              memory: 4Gi
+              cpu: 1
+              memory: 1Gi
             requests:
-              cpu: 2
-              memory: 4Gi
+              cpu: 1
+              memory: 1Gi

ray_provider/decorators/ray.py

@@ -128,9 +128,9 @@ def _extract_function_name(self) -> str:
         return self.python_callable.__name__
 
 
-class task:
+class ray:
     @staticmethod
-    def ray(
+    def task(
         python_callable: Callable[..., Any] | None = None,
         multiple_outputs: bool | None = None,
         **kwargs: Any,

tests/decorators/test_ray_decorators.py

@@ -5,7 +5,7 @@
 from airflow.exceptions import AirflowException
 from airflow.utils.context import Context
 
-from ray_provider.decorators.ray import _RayDecoratedOperator, task
+from ray_provider.decorators.ray import _RayDecoratedOperator, ray
 from ray_provider.operators.ray import SubmitRayJob
 
 DEFAULT_DATE = "2023-01-01"
@@ -169,14 +169,14 @@ def dummy_callable():
 
 class TestRayTaskDecorator:
     def test_ray_task_decorator(self):
-        @task.ray()
+        @ray.task()
         def dummy_function():
             return "dummy"
 
         assert isinstance(dummy_function, _TaskDecorator)
 
     def test_ray_task_decorator_with_multiple_outputs(self):
-        @task.ray(multiple_outputs=True)
+        @ray.task(multiple_outputs=True)
         def dummy_function():
             return {"key": "value"}
 
@@ -189,7 +189,7 @@ def test_ray_task_decorator_with_config(self):
             "memory": 1024,
         }
 
-        @task.ray(**config)
+        @ray.task(**config)
         def dummy_function():
             return "dummy"