Nits

FedML-AI · May 9, 2024 · 4534716 · 4534716
1 parent 17c857a
commit 4534716
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/python/fedml/computing/scheduler/comm_utils/job_utils.py b/python/fedml/computing/scheduler/comm_utils/job_utils.py
@@ -66,6 +66,9 @@ class DockerArgs:
     registry: str = ""
     ports: List[int] = field(default_factory=lambda: [2345])
 
+    def __post_init__(self):
+        self.client = JobRunnerUtils.get_docker_client(self)
+
 
 class JobRunnerUtils(Singleton):
     STATIC_RUN_LOCK_KEY_SUFFIX = "STATIC"

diff --git a/python/fedml/computing/scheduler/slave/client_runner.py b/python/fedml/computing/scheduler/slave/client_runner.py
@@ -2,6 +2,7 @@
 import logging
 import multiprocessing
 import sys
+from datetime import datetime
 
 from multiprocessing import Process
 import os
@@ -642,6 +643,7 @@ def execute_train_job_task(self, job_args: JobArgs, unzip_package_path, entry_fi
         container = self.create_docker_container(docker_args=job_args.docker_args,
                                                  unzip_package_path=unzip_package_path,
                                                  image_pull_policy=job_args.image_pull_policy)
+
         try:
             job_executing_commands = JobRunnerUtils.generate_launch_docker_command(docker_args=job_args.docker_args,
                                                                                    run_id=self.run_id,
@@ -674,18 +676,17 @@ def create_docker_container(self, docker_args: DockerArgs,
                                 unzip_package_path: str,
                                 image_pull_policy: str = None):
 
-        docker_client = JobRunnerUtils.get_docker_client(docker_args=docker_args)
         logging.info(f"Start pulling the launch job image {docker_args.image}... "
                      f"with policy {image_pull_policy}")
         try:
             ContainerUtils.get_instance().pull_image_with_policy(image_name=docker_args.image,
-                                                                 client=docker_client,
+                                                                 client=docker_args.client,
                                                                  image_pull_policy=image_pull_policy)
         except Exception as e:
             raise Exception(f"Failed to pull the launch job image {docker_args.image} with Exception {e}")
 
         container_name = JobRunnerUtils.get_run_container_name(self.run_id)
-        JobRunnerUtils.remove_run_container_if_exists(container_name, docker_client)
+        JobRunnerUtils.remove_run_container_if_exists(container_name, docker_args.client)
         device_requests = []
         volumes = []
         binds = {}
@@ -706,11 +707,11 @@ def create_docker_container(self, docker_args: DockerArgs,
         logging.info(f"device_requests: {device_requests}")
 
         try:
-            host_config = docker_client.api.create_host_config(
+            host_config = docker_args.client.api.create_host_config(
                 binds=binds,
                 device_requests=device_requests,
             )
-            container = docker_client.api.create_container(
+            container = docker_args.client.api.create_container(
                 image=docker_args.image,
                 name=container_name,
                 tty=True,
@@ -719,6 +720,7 @@ def create_docker_container(self, docker_args: DockerArgs,
                 volumes=volumes,
                 detach=True  # Run container in detached mode
             )
+            docker_args.client.api.start(container=container.get("Id"))
         except Exception as e:
             logging.error(f"Failed to create docker container with Exception {e}. Traceback: {traceback.format_exc()}")
             raise Exception(f"Failed to create docker container with Exception {e}")