diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 096294e..c04fe48 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/example_dags/scripts/ray.yaml b/example_dags/scripts/ray.yaml index 1eb8e70..0c4313f 100644 --- a/example_dags/scripts/ray.yaml +++ b/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 diff --git a/ray_provider/decorators/ray.py b/ray_provider/decorators/ray.py index 5dcc4e7..610e812 100644 --- a/ray_provider/decorators/ray.py +++ b/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, diff --git a/tests/decorators/test_ray_decorators.py b/tests/decorators/test_ray_decorators.py index 545197f..07b7be5 100644 --- a/tests/decorators/test_ray_decorators.py +++ b/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"