Add docker backend

Signed-off-by: Brian Gallagher <[email protected]>

Author: briangallagher

kubeflow/trainer/__init__.py

@@ -15,6 +15,7 @@
 
 # Import the Kubeflow Trainer client.
 from kubeflow.trainer.api.trainer_client import TrainerClient  # noqa: F401
+from kubeflow.trainer.backends.docker.types import LocalDockerBackendConfig
 
 # import backends and its associated configs
 from kubeflow.trainer.backends.kubernetes.types import KubernetesBackendConfig
@@ -58,5 +59,6 @@
     "TrainerClient",
     "TrainerType",
     "LocalProcessBackendConfig",
+    "LocalDockerBackendConfig",
     "KubernetesBackendConfig",
 ]

kubeflow/trainer/api/trainer_client.py

@@ -16,6 +16,8 @@
 import logging
 from typing import Optional, Union
 
+from kubeflow.trainer.backends.docker.backend import LocalDockerBackend
+from kubeflow.trainer.backends.docker.types import LocalDockerBackendConfig
 from kubeflow.trainer.backends.kubernetes.backend import KubernetesBackend
 from kubeflow.trainer.backends.kubernetes.types import KubernetesBackendConfig
 from kubeflow.trainer.backends.localprocess.backend import (
@@ -31,7 +33,11 @@
 class TrainerClient:
     def __init__(
         self,
-        backend_config: Union[KubernetesBackendConfig, LocalProcessBackendConfig] = None,
+        backend_config: Union[
+            KubernetesBackendConfig,
+            LocalProcessBackendConfig,
+            LocalDockerBackendConfig,
+        ] = None,
     ):
         """Initialize a Kubeflow Trainer client.
 
@@ -52,6 +58,8 @@ def __init__(
             self.backend = KubernetesBackend(backend_config)
         elif isinstance(backend_config, LocalProcessBackendConfig):
             self.backend = LocalProcessBackend(backend_config)
+        elif isinstance(backend_config, LocalDockerBackendConfig):
+            self.backend = LocalDockerBackend(backend_config)
         else:
             raise ValueError(f"Invalid backend config '{backend_config}'")