component_integration_tests: use compute_world instead of dummy_app (#710)

Author: d4l3k

dev-requirements.txt

@@ -13,7 +13,7 @@ hydra-core
 ipython
 kfp==1.8.9
 mlflow-skinny
-moto==4.1.3
+moto==4.1.6
 pyre-extensions
 pyre-check
 pytest

scripts/component_integration_tests.py

@@ -166,6 +166,12 @@ def _mock_aws_batch() -> None:
     This sets up a mock AWS batch backend that uses Docker to execute the jobs
     locally.
     """
+    # setup the docker network so DNS works correctly
+    from torchx.schedulers.docker_scheduler import ensure_network, NETWORK
+
+    ensure_network()
+    os.environ.setdefault("MOTO_DOCKER_NETWORK_NAME", NETWORK)
+
     from moto import mock_batch, mock_ec2, mock_ecs, mock_iam, mock_logs
 
     mock_batch().__enter__()

torchx/components/integration_tests/component_provider.py

@@ -37,12 +37,15 @@ def tearDown(self) -> None:
 class DDPComponentProvider(ComponentProvider):
     def get_app_def(self) -> AppDef:
         return dist_components.ddp(
-            script="torchx/components/integration_tests/test/dummy_app.py",
+            m="torchx.examples.apps.compute_world_size.main",
             name="ddp-trainer",
             image=self._image,
             cpu=1,
             j="2x2",
             max_retries=3,
+            env={
+                "LOGLEVEL": "INFO",
+            },
         )
 
 

torchx/schedulers/docker_scheduler.py

@@ -43,6 +43,7 @@
 
 
 if TYPE_CHECKING:
+    from docker import DockerClient
     from docker.models.containers import Container
 
 log: logging.Logger = logging.getLogger(__name__)
@@ -94,6 +95,30 @@ def has_docker() -> bool:
         return False
 
 
+def ensure_network(client: Optional["DockerClient"] = None) -> None:
+    """
+    This creates the torchx docker network. Multi-process safe.
+    """
+    import filelock
+    from docker.errors import APIError
+
+    if client is None:
+        import docker
+
+        client = docker.from_env()
+
+    lock_path = os.path.join(tempfile.gettempdir(), "torchx_docker_network_lock")
+
+    # Docker networks.create check_duplicate has a race condition so we need
+    # to do client side locking to ensure only one network is created.
+    with filelock.FileLock(lock_path, timeout=10):
+        try:
+            client.networks.create(name=NETWORK, driver="bridge", check_duplicate=True)
+        except APIError as e:
+            if "already exists" not in str(e):
+                raise
+
+
 class DockerOpts(TypedDict, total=False):
     copy_env: Optional[List[str]]
 
@@ -145,24 +170,6 @@ class DockerScheduler(DockerWorkspaceMixin, Scheduler[DockerOpts]):
     def __init__(self, session_name: str) -> None:
         super().__init__("docker", session_name)
 
-    def _ensure_network(self) -> None:
-        import filelock
-        from docker.errors import APIError
-
-        client = self._docker_client
-        lock_path = os.path.join(tempfile.gettempdir(), "torchx_docker_network_lock")
-
-        # Docker networks.create check_duplicate has a race condition so we need
-        # to do client side locking to ensure only one network is created.
-        with filelock.FileLock(lock_path, timeout=10):
-            try:
-                client.networks.create(
-                    name=NETWORK, driver="bridge", check_duplicate=True
-                )
-            except APIError as e:
-                if "already exists" not in str(e):
-                    raise
-
     def schedule(self, dryrun_info: AppDryRunInfo[DockerJob]) -> str:
         client = self._docker_client
 
@@ -180,7 +187,7 @@ def schedule(self, dryrun_info: AppDryRunInfo[DockerJob]) -> str:
             except Exception as e:
                 log.warning(f"failed to pull image {image}, falling back to local: {e}")
 
-        self._ensure_network()
+        ensure_network(self._docker_client)
 
         for container in req.containers:
             client.containers.run(