Merge pull request #978 from RS-PYTHON/fix/staging-on-cluster

Fix/staging on cluster

Author: jgaucher-cs

services/common/rs_server_common/__init__.py

@@ -14,5 +14,13 @@
 
 """Main package of commons of rs-server services."""
 
+import os
+
 # Set automatically by running `poetry dynamic-versioning`
 __version__ = "0.0.0"
+
+# Some kind of workaround for boto3 to avoid checksum being added inside
+# the file contents uploaded to the s3 bucket e.g. x-amz-checksum-crc32:xxx
+# See: https://github.com/boto/boto3/issues/4435
+os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"] = "when_required"
+os.environ["AWS_RESPONSE_CHECKSUM_VALIDATION"] = "when_required"

services/staging/rs_server_staging/main.py

@@ -31,7 +31,7 @@
     init_rs_server_config_yaml,
 )
 from rs_server_common.db import Base
-from rs_server_common.settings import env_bool
+from rs_server_common.settings import LOCAL_MODE
 from rs_server_common.utils import opentelemetry
 from rs_server_common.utils.logging import Logging
 from rs_server_common.utils.utils2 import filelock
@@ -207,8 +207,9 @@ async def app_lifespan(fastapi_app: FastAPI):  # pylint: disable=too-many-statem
     # Create jobs table
     process_manager = init_db()
 
+    # In local mode, if the gateway is not defined, create a dask LocalCluster
     cluster = None
-    if env_bool("RSPY_LOCAL_MODE", default=False):
+    if LOCAL_MODE and ("RSPY_DASK_STAGING_CLUSTER_NAME" not in os.environ):
         # Create the LocalCluster only in local mode
         cluster = LocalCluster()
         logger.info("Local Dask cluster created at startup.")
@@ -222,7 +223,7 @@ async def app_lifespan(fastapi_app: FastAPI):  # pylint: disable=too-many-statem
 
     # Shutdown logic (cleanup)
     logger.info("Shutting down the application...")
-    if env_bool("RSPY_LOCAL_MODE", default=False) and cluster:
+    if LOCAL_MODE and cluster:
         cluster.close()
         logger.info("Local Dask cluster shut down.")
 
@@ -325,6 +326,15 @@ async def get_specific_job_result_endpoint(job_id: str = Path(..., title="The ID
         raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail=f"Job with ID {job_id} not found") from error
 
 
+if LOCAL_MODE:
+
+    @router.post("/staging/dask/auth")
+    async def dask_auth(local_dask_username: str, local_dask_password: str):
+        """Set dask cluster authentication, only in local mode."""
+        os.environ["LOCAL_DASK_USERNAME"] = local_dask_username
+        os.environ["LOCAL_DASK_PASSWORD"] = local_dask_password
+
+
 # Configure OpenTelemetry
 opentelemetry.init_traces(app, "rs.server.staging")
 

services/staging/rs_server_staging/processors.py

@@ -24,7 +24,8 @@
 
 import requests
 from dask.distributed import CancelledError, Client, LocalCluster, as_completed
-from dask_gateway import Gateway, JupyterHubAuth
+from dask_gateway import Gateway
+from dask_gateway.auth import BasicAuth, JupyterHubAuth
 from fastapi import HTTPException
 from pygeoapi.process.base import BaseProcessor
 from pygeoapi.process.manager.postgresql import PostgreSQLManager
@@ -36,6 +37,7 @@
     load_external_auth_config_by_domain,
 )
 from rs_server_common.s3_storage_handler.s3_storage_handler import S3StorageHandler
+from rs_server_common.settings import LOCAL_MODE
 from rs_server_common.utils.logging import Logging
 from starlette.datastructures import Headers
 from starlette.requests import Request
@@ -615,7 +617,7 @@ def manage_dask_tasks_results(self, client: Client, catalog_collection: str):
         self.log_job_execution(JobStatus.successful, 100, "Finished")
         self.logger.info("Tasks monitoring finished")
 
-    def dask_cluster_connect(self) -> Client:
+    def dask_cluster_connect(self) -> Client:  # pylint: disable=too-many-branches,too-many-statements
         """Connects a dask cluster scheduler
         Establishes a connection to a Dask cluster, either in a local environment or via a Dask Gateway in
         a Kubernetes cluster. This method checks if the cluster is already created (for local mode) or connects
@@ -672,42 +674,62 @@ def dask_cluster_connect(self) -> Client:
         # If self.cluster is already initialized, it means the application is running in local mode, and
         # the cluster was created when the application started.
         if not self.cluster:
-            # in kubernetes cluster mode, we have to connect to the gateway and get the list of the clusters
+            # Connect to the gateway and get the list of the clusters
             try:
                 # get the name of the cluster
                 cluster_name = os.environ["RSPY_DASK_STAGING_CLUSTER_NAME"]
-                # check the auth type, only jupyterhub type supported for now
-                auth_type = os.environ["DASK_GATEWAY__AUTH__TYPE"]
-                # Handle JupyterHub authentication
-                if auth_type == "jupyterhub":
-                    gateway_auth = JupyterHubAuth(api_token=os.environ["JUPYTERHUB_API_TOKEN"])
+
+                # In local mode, authenticate to the dask cluster with username/password
+                if LOCAL_MODE:
+                    gateway_auth = BasicAuth(
+                        os.environ["LOCAL_DASK_USERNAME"],
+                        os.environ["LOCAL_DASK_PASSWORD"],
+                    )
+
+                # Cluster mode
                 else:
-                    self.logger.error(f"Unsupported authentication type: {auth_type}")
-                    raise RuntimeError(f"Unsupported authentication type: {auth_type}")
+                    # check the auth type, only jupyterhub type supported for now
+                    auth_type = os.environ["DASK_GATEWAY__AUTH__TYPE"]
+                    # Handle JupyterHub authentication
+                    if auth_type == "jupyterhub":
+                        gateway_auth = JupyterHubAuth(api_token=os.environ["JUPYTERHUB_API_TOKEN"])
+                    else:
+                        self.logger.error(f"Unsupported authentication type: {auth_type}")
+                        raise RuntimeError(f"Unsupported authentication type: {auth_type}")
+
                 gateway = Gateway(
                     address=os.environ["DASK_GATEWAY__ADDRESS"],
                     auth=gateway_auth,
                 )
-                clusters = gateway.list_clusters()
-                self.logger.debug(f"The list of clusters: {clusters}")
 
-                # Get the identifier of the cluster whose name is equal to the cluster_name variable
+                # Sort the clusters by newest first
+                clusters = sorted(gateway.list_clusters(), key=lambda cluster: cluster.start_time, reverse=True)
+                self.logger.debug(f"Cluster list for gateway {os.environ['DASK_GATEWAY__ADDRESS']!r}: {clusters}")
+
+                # In local mode, get the first cluster from the gateway.
+                cluster_id = None
+                if LOCAL_MODE:
+                    if clusters:
+                        cluster_id = clusters[0].name
+
+                # In cluster mode, get the identifier of the cluster whose name is equal to the cluster_name variable.
                 # Protection for the case when this cluster does not exit
-                cluster_id = next(
-                    (
-                        cluster.name
-                        for cluster in clusters
-                        if isinstance(cluster.options, dict) and cluster.options.get("cluster_name") == cluster_name
-                    ),
-                    None,
-                )
+                else:
+                    cluster_id = next(
+                        (
+                            cluster.name
+                            for cluster in clusters
+                            if isinstance(cluster.options, dict) and cluster.options.get("cluster_name") == cluster_name
+                        ),
+                        None,
+                    )
 
                 if not cluster_id:
-                    raise IndexError(f"No dask cluster named '{cluster_name}' was found.")
+                    raise IndexError(f"Dask cluster with 'cluster_name'={cluster_name!r} was not found.")
 
                 self.cluster = gateway.connect(cluster_id)
-
                 self.logger.info(f"Successfully connected to the {cluster_name} dask cluster")
+
             except KeyError as e:
                 self.logger.exception(
                     "Failed to retrieve the required connection details for "
@@ -725,6 +747,22 @@ def dask_cluster_connect(self) -> Client:
         # create the client as well
         client = Client(self.cluster)
 
+        # Forward logging from dask workers to the caller
+        client.forward_logging()
+
+        def set_dask_env(host_env: dict):
+            """Pass environment variables to the dask workers."""
+            for name in ["S3_ACCESSKEY", "S3_SECRETKEY", "S3_ENDPOINT", "S3_REGION"]:
+                os.environ[name] = host_env[name]
+
+            # Some kind of workaround for boto3 to avoid checksum being added inside
+            # the file contents uploaded to the s3 bucket e.g. x-amz-checksum-crc32:xxx
+            # See: https://github.com/boto/boto3/issues/4435
+            os.environ["AWS_REQUEST_CHECKSUM_CALCULATION"] = "when_required"
+            os.environ["AWS_RESPONSE_CHECKSUM_VALIDATION"] = "when_required"
+
+        client.run(set_dask_env, os.environ)
+
         # This is a temporary fix for the dask cluster settings which does not create a scheduler by default
         # This code should be removed as soon as this is fixed in the kubernetes cluster
         try:

services/staging/tests/test_rspy_processor.py

@@ -715,6 +715,8 @@ def test_dask_cluster_connect(self, mocker, staging_instance: Staging, cluster_o
                 "RSPY_DASK_STAGING_CLUSTER_NAME": cluster_options["cluster_name"],
             },
         )
+        # Mock the cluster mode
+        mocker.patch("rs_server_staging.processors.LOCAL_MODE", new=False, autospec=False)
         # Mock the logger
         mock_logger = mocker.patch.object(staging_instance, "logger")
         staging_instance.cluster = None
@@ -749,7 +751,9 @@ def test_dask_cluster_connect(self, mocker, staging_instance: Staging, cluster_o
         mock_client.assert_called_once_with(staging_instance.cluster)
 
         # Ensure logging was called as expected
-        mock_logger.debug.assert_any_call(f"The list of clusters: {mock_list_clusters.return_value}")
+        mock_logger.debug.assert_any_call(
+            f"Cluster list for gateway 'gateway-address': {mock_list_clusters.return_value}",
+        )
         mock_logger.info.assert_any_call("Number of running workers: 2")
         mock_logger.debug.assert_any_call(
             f"Dask Client: {client} | Cluster dashboard: {mock_connect.return_value.dashboard_link}",
@@ -768,6 +772,8 @@ def test_dask_cluster_connect_failure_no_cluster_name(self, mocker, staging_inst
                 "RSPY_DASK_STAGING_CLUSTER_NAME": non_existent_cluster,
             },
         )
+        # Mock the cluster mode
+        mocker.patch("rs_server_staging.processors.LOCAL_MODE", new=False, autospec=False)
         # Mock the logger
         mock_logger = mocker.patch.object(staging_instance, "logger")
         staging_instance.cluster = None
@@ -791,7 +797,8 @@ def test_dask_cluster_connect_failure_no_cluster_name(self, mocker, staging_inst
             staging_instance.dask_cluster_connect()
         # Ensure logging was called as expected
         mock_logger.exception.assert_any_call(
-            "Failed to find the specified dask cluster: " f"No dask cluster named '{non_existent_cluster}' was found.",
+            "Failed to find the specified dask cluster: "
+            f"Dask cluster with 'cluster_name'={non_existent_cluster!r} was not found.",
         )
 
     def test_dask_cluster_connect_failure_no_envs(self, mocker, staging_instance: Staging):

services/staging/tests/test_staging.py

@@ -526,12 +526,7 @@ async def test_app_lifespan_gateway_error(
     """Test app_lifespan when there is an error in connecting to the Dask Gateway."""
 
     # Mock environment variables to simulate gateway mode
-    mocker.patch.dict(
-        os.environ,
-        {
-            "RSPY_LOCAL_MODE": "0",
-        },
-    )
+    mocker.patch("rs_server_staging.main.LOCAL_MODE", new=False, autospec=False)
 
     # Mock FastAPI app
     mock_app = FastAPI()