Update test_jobs_pytorch_ddp.py

Author: qiuosier

tests/unitary/default_setup/jobs/test_jobs_pytorch_ddp.py

@@ -1,16 +1,20 @@
 import json
+import os
+import sys
 import unittest
 import zipfile
-from ads.jobs import PyTorchDistributedRuntime, DataScienceJob
+from unittest import mock
+from ads.jobs import PyTorchDistributedRuntime, DataScienceJob, DataScienceJobRun
 from ads.jobs.builders.infrastructure.dsc_job_runtime import (
     PyTorchDistributedRuntimeHandler as Handler,
 )
 from ads.jobs.builders.runtimes.pytorch_runtime import (
     PyTorchDistributedArtifact,
     GitPythonArtifact,
 )
-from ads.opctl.distributed.common import cluster_config_helper as Cluster
-from ads.jobs.templates import driver_utils as Driver
+from ads.opctl.distributed.common import cluster_config_helper as cluster
+from ads.jobs.templates import driver_utils as utils
+from ads.jobs.templates import driver_pytorch as driver
 
 
 class PyTorchRuntimeHandlerTest(unittest.TestCase):
@@ -26,6 +30,7 @@ class PyTorchRuntimeHandlerTest(unittest.TestCase):
     )
 
     def init_runtime(self):
+        """Initializes a PyTorchDistributedRuntime for testing."""
         return (
             PyTorchDistributedRuntime()
             .with_replica(self.REPLICAS)
@@ -43,6 +48,7 @@ def init_runtime(self):
         )
 
     def test_translate_artifact(self):
+        """Tests preparing ADS driver scripts in job artifacts."""
         artifact = Handler(DataScienceJob())._translate_artifact(self.init_runtime())
         self.assertIsInstance(artifact, PyTorchDistributedArtifact)
         self.assertEqual(
@@ -65,6 +71,7 @@ def test_translate_artifact(self):
             self.assertIn(GitPythonArtifact.CONST_DRIVER_SCRIPT, file_list)
 
     def test_translate_env(self):
+        """Tests setting up environment variables"""
         envs = Handler(DataScienceJob())._translate_env(self.init_runtime())
         self.assertIsInstance(envs, dict)
         self.assertEqual(envs[Handler.CONST_WORKER_COUNT], str(self.REPLICAS - 1))
@@ -73,12 +80,87 @@ def test_translate_env(self):
             PyTorchDistributedArtifact.CONST_DRIVER_SCRIPT,
         )
         self.assertEqual(envs[Handler.CONST_COMMAND], self.TORCHRUN_CMD)
-        self.assertEqual(envs[Cluster.OCI__RUNTIME_URI], self.TEST_REPO)
-        self.assertEqual(envs[Cluster.OCI__RUNTIME_GIT_COMMIT], self.TEST_COMMIT)
-        self.assertEqual(envs[Driver.CONST_ENV_PIP_PKG], self.PIP_PKG)
-        self.assertEqual(envs[Driver.CONST_ENV_PIP_REQ], self.PIP_REQ)
+        self.assertEqual(envs[cluster.OCI__RUNTIME_URI], self.TEST_REPO)
+        self.assertEqual(envs[cluster.OCI__RUNTIME_GIT_COMMIT], self.TEST_COMMIT)
+        self.assertEqual(envs[utils.CONST_ENV_PIP_PKG], self.PIP_PKG)
+        self.assertEqual(envs[utils.CONST_ENV_PIP_REQ], self.PIP_REQ)
         self.assertEqual(
-            envs[Driver.CONST_ENV_INPUT_MAPPINGS],
+            envs[utils.CONST_ENV_INPUT_MAPPINGS],
             json.dumps({self.INPUT_SRC: self.INPUT_DST}),
         )
         self.assertNotIn(Handler.CONST_DEEPSPEED, envs)
+
+
+class PyTorchRunnerTest(unittest.TestCase):
+    TEST_IP = "10.0.0.1"
+    TEST_HOST_IP = "10.0.0.100"
+    TEST_HOST_OCID = "ocid_host"
+    TEST_NODE_OCID = "ocid_node"
+
+    def init_torch_runner(self):
+        with mock.patch(
+            "ads.jobs.templates.driver_pytorch.TorchRunner.build_c_library"
+        ), mock.patch("socket.gethostbyname") as GetHostIP, mock.patch(
+            "ads.jobs.DataScienceJobRun.from_ocid"
+        ) as GetJobRun:
+            GetHostIP.return_value = self.TEST_IP
+            GetJobRun.return_value = DataScienceJobRun(id="ocid.abcdefghijk")
+            return driver.TorchRunner()
+
+    @mock.patch.dict(os.environ, {driver.CONST_ENV_HOST_JOB_RUN_OCID: TEST_HOST_OCID})
+    def test_init_torch_runner_at_node(self):
+        runner = self.init_torch_runner()
+        self.assertEqual(runner.host_ocid, self.TEST_HOST_OCID)
+        self.assertEqual(runner.host_ip, None)
+
+    @mock.patch.dict(os.environ, {driver.CONST_ENV_JOB_RUN_OCID: TEST_NODE_OCID})
+    def test_init_torch_runner_at_host(self):
+        runner = self.init_torch_runner()
+        self.assertEqual(runner.host_ocid, self.TEST_NODE_OCID)
+        self.assertEqual(runner.host_ip, self.TEST_IP)
+
+    @mock.patch.dict(os.environ, {driver.CONST_ENV_HOST_JOB_RUN_OCID: TEST_HOST_OCID})
+    def test_wait_for_host_ip(self):
+        with mock.patch("ads.jobs.DataScienceJobRun.logs") as get_logs:
+            get_logs.return_value = [
+                {"message": f"{driver.LOG_PREFIX_HOST_IP} {self.TEST_HOST_IP}"}
+            ]
+            runner = self.init_torch_runner()
+            self.assertEqual(runner.host_ip, None)
+            runner.wait_for_host_ip_address()
+            self.assertEqual(runner.host_ip, self.TEST_HOST_IP)
+
+    @mock.patch.dict(
+        os.environ, {driver.CONST_ENV_LAUNCH_CMD: "torchrun train.py --data abc"}
+    )
+    def test_launch_cmd(self):
+        runner = self.init_torch_runner()
+        self.assertTrue(runner.launch_cmd_contains("data"))
+        self.assertFalse(runner.launch_cmd_contains("data1"))
+        self.assertEqual(
+            runner.prepare_cmd(prefix="A=1"), "A=1 torchrun train.py --data abc"
+        )
+
+    @mock.patch.dict(os.environ, {Handler.CONST_CODE_ENTRYPOINT: "train.py"})
+    @mock.patch.object(sys, "argv", ["python", "hello", "--data", "abc"])
+    def test_prepare_cmd_with_entrypoint_args(self):
+        runner = self.init_torch_runner()
+        self.assertEqual(
+            runner.prepare_cmd(launch_args=["--key", "val"], prefix="A=1"),
+            "A=1 torchrun --key val train.py hello --data abc",
+        )
+
+
+class LazyEvaluateTest(unittest.TestCase):
+    def test_lazy_evaluation(self):
+        def func(a, b):
+            return a + b
+
+        def func_with_error():
+            raise ValueError()
+
+        lazy_val = driver.LazyEvaluate(func, 1, 1)
+        self.assertEqual(str(lazy_val), "2")
+
+        lazy_val = driver.LazyEvaluate(func_with_error)
+        self.assertEqual(str(lazy_val), "ERROR: ")