Refactor the common modules between main package and elasticdl_client. (#2100)

* Remove duplicated k8s_resource and k8s_volume

* Remove api and client module, and duplicated method in args.py from main package

* Remove the entry point in elasticdl main package and rename elasticdl_client entry to elasticdl.

* Add dependency for elasticdl_client in main package.

* Make k8s_client in main package inherit from k8s_client from client package.

* Revert setup.py and setup_client.py

* Update the duplicated module to be imported from main package to elasticdl_client package.

* Fix the inheritance issue.

* Revert the changes in build_and_test.sh

* Keep JobType in elasticdl main package.

* Set default value for image_name argument.

* Rename the console entry of elasticdl_client package to elasticdl.

* Use python ps for evaluate and predict job in integration test.

* Revert the change that use python ps in predict/evaluate integration test.

Author: brightcoder01

elasticdl/python/common/args.py

@@ -55,12 +55,6 @@ class MetricsDictKey(object):
     LABEL = "label"
 
 
-class DistributionStrategy(object):
-    LOCAL = "Local"
-    PARAMETER_SERVER = "ParameterServerStrategy"
-    ALLREDUCE = "AllreduceStrategy"
-
-
 class SaveModelConfig(object):
     SAVED_MODEL_PATH = "saved_model_path"
 
@@ -86,8 +80,3 @@ class ReaderType(object):
     CSV_READER = "CSV"
     ODPS_READER = "ODPS"
     RECORDIO_READER = "RecordIO"
-
-
-class BashCommandTemplate(object):
-    REDIRECTION = " 2>&1 | tee {}"
-    SET_PIPEFAIL = "set -o pipefail;"

elasticdl/python/common/constants.py

@@ -11,32 +11,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import threading
 import time
-import traceback
 
-import yaml
-from kubernetes import client, config, watch
-from kubernetes.client import V1EnvVar, V1EnvVarSource, V1ObjectFieldSelector
+from kubernetes import client, watch
 
-from elasticdl.python.common.k8s_resource import parse as parse_resource
-from elasticdl.python.common.k8s_volume import parse_volume_and_mount
 from elasticdl.python.common.log_utils import default_logger as logger
-from elasticdl.python.common.model_utils import load_module
+from elasticdl_client.common.k8s_client import (
+    ELASTICDL_APP_NAME,
+    ELASTICDL_JOB_KEY,
+    ELASTICDL_REPLICA_INDEX_KEY,
+    ELASTICDL_REPLICA_TYPE_KEY,
+)
+from elasticdl_client.common.k8s_client import Client as BaseClient
+from elasticdl_client.common.k8s_client import append_pod_ip_to_env
 
-ELASTICDL_APP_NAME = "elasticdl"
-ELASTICDL_JOB_KEY = "elasticdl-job-name"
-ELASTICDL_REPLICA_TYPE_KEY = "elasticdl-replica-type"
-ELASTICDL_REPLICA_INDEX_KEY = "elasticdl-replica-index"
 _PS_SERVICE_PORT = 2222
 _WORKER_SERVICE_PORT = 3333
 _FTLIB_GOSSIP_CONTAINER_PORT = 7946
-_FTLIB_SSH_CONTAINER_PORT = 22
-
-
-def get_master_pod_name(job_name):
-    return "elasticdl-%s-master" % job_name
 
 
 def get_worker_pod_name(job_name, worker_id):
@@ -47,21 +39,7 @@ def get_ps_pod_name(job_name, ps_id):
     return "elasticdl-%s-ps-%s" % (job_name, str(ps_id))
 
 
-def append_pod_ip_to_env(env):
-    pod_ip_var = V1EnvVar(
-        name="MY_POD_IP",
-        value_from=V1EnvVarSource(
-            field_ref=V1ObjectFieldSelector(field_path="status.podIP")
-        ),
-    )
-    if env:
-        env.append(pod_ip_var)
-    else:
-        env = [pod_ip_var]
-    return env
-
-
-class Client(object):
+class Client(BaseClient):
     def __init__(
         self,
         *,
@@ -88,37 +66,18 @@ def __init__(
                 running in a K8S environment, it loads the incluster config,
                 if not, it loads the kube config file.
         """
-        try:
-            if (
-                os.getenv("KUBERNETES_SERVICE_HOST")
-                and not force_use_kube_config_file
-            ):
-                # We are running inside a k8s cluster
-                config.load_incluster_config()
-                logger.info("Load the incluster config.")
-            else:
-                # Use user's kube config
-                config.load_kube_config()
-                logger.info("Load the kube config file.")
-        except Exception as ex:
-            traceback.print_exc()
-            raise Exception(
-                "Failed to load configuration for Kubernetes:\n%s" % str(ex)
-            )
-
-        self.client = client.CoreV1Api()
-        self.namespace = namespace
-        self.job_name = job_name
-        self._image_name = image_name
+        super().__init__(
+            image_name=image_name,
+            namespace=namespace,
+            job_name=job_name,
+            cluster_spec=cluster_spec,
+            force_use_kube_config_file=force_use_kube_config_file,
+        )
         self._event_cb = event_callback
         if self._event_cb:
             threading.Thread(
                 target=self._watch, name="event_watcher", daemon=True
             ).start()
-        self.cluster = None
-        if cluster_spec:
-            cluster_spec_module = load_module(cluster_spec)
-            self.cluster = cluster_spec_module.cluster
 
     def _watch(self):
         while True:
@@ -139,9 +98,6 @@ def _watch(self):
     def _get_service_address(self, service_name, port):
         return "%s.%s.svc:%d" % (service_name, self.namespace, port)
 
-    def get_master_pod_name(self):
-        return get_master_pod_name(self.job_name)
-
     def get_worker_pod_name(self, worker_id):
         return get_worker_pod_name(self.job_name, worker_id)
 
@@ -164,16 +120,6 @@ def get_ps_service_address(self, ps_id):
             self.get_ps_service_name(ps_id), _PS_SERVICE_PORT
         )
 
-    def patch_labels_to_pod(self, pod_name, labels_dict):
-        body = {"metadata": {"labels": labels_dict}}
-        try:
-            return self.client.patch_namespaced_pod(
-                name=pod_name, namespace=self.namespace, body=body
-            )
-        except client.api_client.ApiException as e:
-            logger.warning("Exception when patching labels to pod: %s\n" % e)
-            return None
-
     def get_master_pod(self):
         return self.get_pod(self.get_master_pod_name())
 
@@ -216,129 +162,6 @@ def get_worker_service(self, worker_id):
             logger.warning("Exception when reading worker service: %s\n" % e)
             return None
 
-    @staticmethod
-    def create_owner_reference(owner_pod):
-        owner_ref = (
-            [
-                client.V1OwnerReference(
-                    api_version="v1",
-                    block_owner_deletion=True,
-                    kind="Pod",
-                    name=owner_pod.metadata.name,
-                    uid=owner_pod.metadata.uid,
-                )
-            ]
-            if owner_pod
-            else None
-        )
-        return owner_ref
-
-    def create_pod(self, **kargs):
-        # Container
-        pod_resource_requests = kargs["resource_requests"]
-        pod_resource_limits = kargs["resource_limits"]
-        pod_resource_limits = (
-            pod_resource_limits
-            if pod_resource_limits
-            else pod_resource_requests
-        )
-        ports = (
-            [
-                client.V1ContainerPort(
-                    container_port=_FTLIB_GOSSIP_CONTAINER_PORT, name="gossip"
-                ),
-            ]
-            if "expose_ports" in kargs and kargs["expose_ports"]
-            else None
-        )
-        container = client.V1Container(
-            name=kargs["pod_name"],
-            image=kargs["image_name"],
-            command=kargs["command"],
-            resources=client.V1ResourceRequirements(
-                requests=parse_resource(pod_resource_requests),
-                limits=parse_resource(pod_resource_limits),
-            ),
-            args=kargs["container_args"],
-            image_pull_policy=kargs["image_pull_policy"],
-            env=kargs["env"],
-            ports=ports,
-        )
-
-        # Pod
-        spec = client.V1PodSpec(
-            containers=[container],
-            restart_policy=kargs["restart_policy"],
-            priority_class_name=kargs["pod_priority"],
-            termination_grace_period_seconds=kargs.get(
-                "termination_period", None
-            ),
-        )
-
-        # Mount data path
-        if kargs["volume"]:
-            volumes, volume_mounts = parse_volume_and_mount(
-                kargs["volume"], kargs["pod_name"]
-            )
-            spec.volumes = volumes
-            container.volume_mounts = volume_mounts
-
-        pod = client.V1Pod(
-            spec=spec,
-            metadata=client.V1ObjectMeta(
-                name=kargs["pod_name"],
-                labels=self._get_common_labels(),
-                owner_references=self.create_owner_reference(
-                    kargs["owner_pod"]
-                ),
-                namespace=self.namespace,
-            ),
-        )
-        if self.cluster:
-            pod = self.cluster.with_pod(pod)
-
-        return pod
-
-    def create_master(self, **kargs):
-        pod = self._create_master_pod_obj(**kargs)
-        self.client.create_namespaced_pod(self.namespace, pod)
-        logger.info("Master launched.")
-
-    def dump_master_yaml(self, **kargs):
-        pod = self._create_master_pod_obj(**kargs)
-        pod_dict = self.client.api_client.sanitize_for_serialization(pod)
-        with open(kargs["yaml"], "w") as f:
-            yaml.safe_dump(pod_dict, f, default_flow_style=False)
-
-    def _create_master_pod_obj(self, **kargs):
-        env = []
-        if "envs" in kargs:
-            for key in kargs["envs"]:
-                env.append(V1EnvVar(name=key, value=kargs["envs"][key]))
-        env = append_pod_ip_to_env(env)
-
-        pod = self.create_pod(
-            pod_name=self.get_master_pod_name(),
-            job_name=self.job_name,
-            image_name=self._image_name,
-            command=["/bin/bash"],
-            resource_requests=kargs["resource_requests"],
-            resource_limits=kargs["resource_limits"],
-            container_args=kargs["args"],
-            pod_priority=kargs["pod_priority"],
-            image_pull_policy=kargs["image_pull_policy"],
-            restart_policy=kargs["restart_policy"],
-            volume=kargs["volume"],
-            owner_pod=None,
-            env=env,
-        )
-        # Add replica type and index
-        pod.metadata.labels[ELASTICDL_REPLICA_TYPE_KEY] = "master"
-        pod.metadata.labels[ELASTICDL_REPLICA_INDEX_KEY] = "0"
-        pod.api_version = "v1"
-        pod.kind = "Pod"
-        return pod
-
     def _create_ps_worker_pod(self, pod_name, type_key, index_key, **kargs):
         # Find that master pod that will be used as the owner reference
         # for the ps or worker pod.
@@ -380,10 +203,6 @@ def create_ps(self, **kargs):
             pod_name, "ps", kargs["ps_id"], **kargs
         )
 
-    def delete_master(self):
-        logger.info("pod name is %s" % self.get_master_pod_name())
-        self.delete_pod(self.get_master_pod_name())
-
     def delete_worker(self, worker_id):
         self.delete_pod(self.get_worker_pod_name(worker_id))
 
@@ -489,12 +308,6 @@ def _create_service(self, **kargs):
             service = self.cluster.with_service(service)
         return self.client.create_namespaced_service(self.namespace, service)
 
-    def _get_common_labels(self):
-        """Labels that should be attached to all k8s objects belong to
-           current job.
-        """
-        return {"app": ELASTICDL_APP_NAME, ELASTICDL_JOB_KEY: self.job_name}
-
     def get_master_log(self):
         return self.get_pod_log(self.get_master_pod_name())