Change type for --runtime_env in ray in-fw deployment script (#505)

Signed-off-by: Abhishree <[email protected]>

Author: athitten

scripts/deploy/nlp/deploy_ray_inframework.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import argparse
+import json
 import logging
 import multiprocessing
 
@@ -26,6 +27,14 @@ def get_available_cpus():
     return multiprocessing.cpu_count()
 
 
+def json_type(string):
+    """Parse JSON string into a dictionary."""
+    try:
+        return json.loads(string)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"Invalid JSON: {e}")
+
+
 def parse_args():
     """Parse command-line arguments for the Ray deployment script."""
     parser = argparse.ArgumentParser(description="Deploy a Megatron model using Ray")
@@ -187,9 +196,9 @@ def parse_args():
     )
     parser.add_argument(
         "--runtime_env",
-        type=dict,
+        type=json_type,
         default={},
-        help="Runtime environment for the deployment",
+        help="Runtime environment for the deployment (JSON string)",
     )
     return parser.parse_args()
 

tests/unit_tests/deploy/test_deploy_ray.py

@@ -13,11 +13,17 @@
 # limitations under the License.
 
 
+import argparse
 import unittest
 from unittest.mock import MagicMock, patch
 
 from nemo_deploy.deploy_ray import DeployRay
 
+# Import the functions from the deploy script
+from scripts.deploy.nlp.deploy_ray_inframework import (
+    json_type,
+)
+
 
 class TestDeployRay(unittest.TestCase):
     def setUp(self):
@@ -297,5 +303,42 @@ def test_stop_with_errors(self, mock_logger, mock_serve, mock_ray):
         mock_logger.warning.assert_any_call("Error during ray.shutdown(): Ray shutdown error")
 
 
+class TestDeployRayInFrameworkScriptJsonType(unittest.TestCase):
+    """Test suite for deploy_ray_inframework.py script's json_type function."""
+
+    def test_json_type_valid_json(self):
+        """Test json_type with valid JSON strings."""
+        # Test valid dictionary
+        result = json_type('{"key": "value", "number": 42}')
+        self.assertEqual(result, {"key": "value", "number": 42})
+
+        # Test valid list
+        result = json_type("[1, 2, 3]")
+        self.assertEqual(result, [1, 2, 3])
+
+        # Test nested structure
+        result = json_type('{"pip": ["numpy", "pandas"], "env_vars": {"PATH": "/usr/bin"}}')
+        expected = {"pip": ["numpy", "pandas"], "env_vars": {"PATH": "/usr/bin"}}
+        self.assertEqual(result, expected)
+
+    def test_json_type_invalid_json(self):
+        """Test json_type with invalid JSON strings."""
+        with self.assertRaises(argparse.ArgumentTypeError) as context:
+            json_type("not a valid json")
+        self.assertIn("Invalid JSON", str(context.exception))
+
+        with self.assertRaises(argparse.ArgumentTypeError) as context:
+            json_type('{"incomplete": ')
+        self.assertIn("Invalid JSON", str(context.exception))
+
+    def test_json_type_empty_json(self):
+        """Test json_type with empty JSON objects."""
+        result = json_type("{}")
+        self.assertEqual(result, {})
+
+        result = json_type("[]")
+        self.assertEqual(result, [])
+
+
 if __name__ == "__main__":
     unittest.main()