feat(scripts): add pre_start_hook (#877)

Dockerfile

@@ -7,6 +7,7 @@ RUN apt-get -y update \
     && apt-get -y install libgmp-dev \
     && apt-get -y install libmpfr-dev \
     && apt-get -y install libmpc-dev \
+    && apt-get -y install krb5-user \
     && rm -rf /var/lib/apt/lists/* 
 
 RUN pip install --upgrade pip \

deploy/scripts/data_join/run_data_join_follower_master.sh

@@ -16,6 +16,8 @@
 
 set -ex
 
+source /app/deploy/scripts/pre_start_hook.sh || true
+
 echo "Reset the ROLE=$ROLE as follower"
 export ROLE=follower
 data_join_master_cmd=/app/deploy/scripts/data_join/run_data_join_master.sh

deploy/scripts/data_join/run_data_join_leader_master.sh

@@ -16,6 +16,8 @@
 
 set -ex
 
+source /app/deploy/scripts/pre_start_hook.sh || true
+
 echo "Reset the ROLE=$ROLE as leader"
 export ROLE=leader
 data_join_master_cmd=/app/deploy/scripts/data_join/run_data_join_master.sh

deploy/scripts/data_join/run_data_join_master.sh

@@ -18,8 +18,10 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 source /app/deploy/scripts/env_to_args.sh
 
+
 kvstore_type=$(normalize_env_to_args '--kvstore_type' $KVSTORE_TYPE)
 
 python -m fedlearner.data_join.cmd.prepare_launch_data_join_cli \

deploy/scripts/data_join/run_data_join_worker.sh

@@ -18,8 +18,10 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 source /app/deploy/scripts/env_to_args.sh
 
+
 MASTER_POD_NAMES=`python -c 'import json, os; print(json.loads(os.environ["CLUSTER_SPEC"])["clusterSpec"]["Master"][0])'`
 
 raw_data_iter=$(normalize_env_to_args "--raw_data_iter" $RAW_DATA_ITER)

deploy/scripts/data_join/run_psi_data_join_follower_worker.sh

@@ -16,6 +16,8 @@
 
 set -ex
 
+source /app/deploy/scripts/pre_start_hook.sh || true
+
 WORKER_REPLICAS=`python -c 'import json, os; print(len(json.loads(os.environ["CLUSTER_SPEC"])["clusterSpec"]["Worker"]))'`
 
 echo "${WORKER_REPLICAS:?Need to set WORKER_REPLICAS non-empty}"

deploy/scripts/data_join/run_psi_data_join_follower_worker_v2.sh

@@ -16,6 +16,8 @@
 
 set -ex
 
+source /app/deploy/scripts/pre_start_hook.sh || true
+
 psi_preprocessor_cmd=/app/deploy/scripts/rsa_psi/run_psi_preprocessor.sh
 data_join_worker_cmd=/app/deploy/scripts/data_join/run_data_join_worker.sh
 

deploy/scripts/data_join/run_psi_data_join_leader_worker.sh

@@ -16,6 +16,8 @@
 
 set -ex
 
+source /app/deploy/scripts/pre_start_hook.sh || true
+
 WORKER_REPLICAS=`python -c 'import json, os; print(len(json.loads(os.environ["CLUSTER_SPEC"])["clusterSpec"]["Worker"]))'`
 echo "${WORKER_REPLICAS:?Need to set WORKER_REPLICAS non-empty}"
 echo "${INDEX:?Need to set INDEX non-empty}"

deploy/scripts/data_join/run_psi_data_join_leader_worker_v2.sh

@@ -16,6 +16,8 @@
 
 set -ex
 
+source /app/deploy/scripts/pre_start_hook.sh || true
+
 psi_signer_cmd=/app/deploy/scripts/rsa_psi/run_rsa_psi_signer.sh
 psi_preprocessor_cmd=/app/deploy/scripts/rsa_psi/run_psi_preprocessor.sh 
 data_join_worker_cmd=/app/deploy/scripts/data_join/run_data_join_worker.sh

deploy/scripts/data_portal/run_data_portal_master.sh

@@ -18,6 +18,7 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 source /app/deploy/scripts/env_to_args.sh
 
 input_file_wildcard=$(normalize_env_to_args "--input_file_wildcard" $FILE_WILDCARD)

deploy/scripts/data_portal/run_data_portal_worker.sh

@@ -18,6 +18,7 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 source /app/deploy/scripts/env_to_args.sh 
 
 MASTER_POD_NAMES=`python -c 'import json, os; print(json.loads(os.environ["CLUSTER_SPEC"])["clusterSpec"]["Master"][0])'`

deploy/scripts/pre_start_hook.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python -c 'from fedlearner.common.hooks import pre_start_hook; pre_start_hook()'

deploy/scripts/trainer/run_merge_scores.sh

@@ -18,6 +18,7 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 
 python example/tree_model/merge_scores.py \
     --worker-rank=$WORKER_RANK \

deploy/scripts/trainer/run_trainer_master.sh

@@ -18,6 +18,7 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 source /app/deploy/scripts/env_to_args.sh
 
 mode=$(normalize_env_to_args "--mode" "$MODE")

deploy/scripts/trainer/run_trainer_ps.sh

@@ -18,5 +18,6 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 
 python -m fedlearner.trainer.parameter_server $POD_IP:50051 

deploy/scripts/trainer/run_trainer_worker.sh

@@ -20,6 +20,7 @@ export CUDA_VISIBLE_DEVICES=
 export MODEL_NAME=${APPLICATION_ID}
 
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 source /app/deploy/scripts/env_to_args.sh
 
 # When the WORKER_GROUPS is "2,4", this script would update the WORKER_RANK

deploy/scripts/trainer/run_tree_worker.sh

@@ -18,6 +18,7 @@ set -ex
 
 export CUDA_VISIBLE_DEVICES=
 source /app/deploy/scripts/hdfs_common.sh || true
+source /app/deploy/scripts/pre_start_hook.sh || true
 source /app/deploy/scripts/env_to_args.sh
 
 NUM_WORKERS=`python -c 'import json, os; print(len(json.loads(os.environ["CLUSTER_SPEC"])["clusterSpec"]["Worker"]))'`

fedlearner/common/hooks.py

@@ -0,0 +1,35 @@
+# Copyright 2021 The FedLearner Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# coding: utf-8
+import os
+import importlib
+from typing import Any
+
+
+def parse_and_call_fn(module_fn_path: str) -> Any:
+    if module_fn_path.find(':') == -1:
+        raise RuntimeError(f'Invalid module_fn_path: {module_fn_path}')
+
+    module_path, func_name = module_fn_path.split(':')
+    module = importlib.import_module(module_path)
+    # Dynamically run the function
+    return getattr(module, func_name)()
+
+
+def pre_start_hook() -> Any:
+    before_hook_path = os.getenv('PRE_START_HOOK', None)
+    if before_hook_path:
+        return parse_and_call_fn(before_hook_path)
+    return None

test/common/for_test_hooks.py

@@ -0,0 +1,17 @@
+# Copyright 2021 The FedLearner Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# coding: utf-8
+def test():
+    return 1

test/common/test_hooks.py

@@ -0,0 +1,27 @@
+# Copyright 2021 The FedLearner Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# coding: utf-8
+import unittest
+from fedlearner.common.hooks import parse_and_call_fn
+
+class HooksTest(unittest.TestCase):
+    def test_parse_and_call_fn(self):
+        with self.assertRaises((RuntimeError)):
+            parse_and_call_fn('something_you_cannot_understand')
+        self.assertEqual(parse_and_call_fn('test.common.for_test_hooks:test'), 1)
+
+
+if __name__ == '__main__':
+    unittest.main()