Merge branch 'master' into fw-and-version-bug

Author: aviruthen

sagemaker-train/src/sagemaker/train/remote_function/runtime_environment/runtime_environment_manager.py

@@ -94,6 +94,50 @@ def from_dependency_file_path(dependency_file_path):
 class RuntimeEnvironmentManager:
     """Runtime Environment Manager class to manage runtime environment."""
 
+    def _validate_path(self, path: str) -> str:
+        """Validate and sanitize file path to prevent path traversal attacks.
+        
+        Args:
+            path (str): The file path to validate
+            
+        Returns:
+            str: The validated absolute path
+            
+        Raises:
+            ValueError: If the path is invalid or contains suspicious patterns
+        """
+        if not path:
+            raise ValueError("Path cannot be empty")
+        
+        # Get absolute path to prevent path traversal
+        abs_path = os.path.abspath(path)
+        
+        # Check for null bytes (common in path traversal attacks)
+        if '\x00' in path:
+            raise ValueError(f"Invalid path contains null byte: {path}")
+        
+        return abs_path
+
+    def _validate_env_name(self, env_name: str) -> None:
+        """Validate conda environment name to prevent command injection.
+        
+        Args:
+            env_name (str): The environment name to validate
+            
+        Raises:
+            ValueError: If the environment name contains invalid characters
+        """
+        if not env_name:
+            raise ValueError("Environment name cannot be empty")
+        
+        # Allow only alphanumeric, underscore, and hyphen
+        import re
+        if not re.match(r'^[a-zA-Z0-9_-]+$', env_name):
+            raise ValueError(
+                f"Invalid environment name '{env_name}'. "
+                "Only alphanumeric characters, underscores, and hyphens are allowed."
+            )
+
     def snapshot(self, dependencies: str = None) -> str:
         """Creates snapshot of the user's environment
 
@@ -252,39 +296,50 @@ def _is_file_exists(self, dependencies):
 
     def _install_requirements_txt(self, local_path, python_executable):
         """Install requirements.txt file"""
-        cmd = f"{python_executable} -m pip install -r {local_path} -U"
-        logger.info("Running command: '%s' in the dir: '%s' ", cmd, os.getcwd())
+        # Validate path to prevent command injection
+        validated_path = self._validate_path(local_path)
+        cmd = [python_executable, "-m", "pip", "install", "-r", validated_path, "-U"]
+        logger.info("Running command: '%s' in the dir: '%s' ", " ".join(cmd), os.getcwd())
         _run_shell_cmd(cmd)
-        logger.info("Command %s ran successfully", cmd)
+        logger.info("Command %s ran successfully", " ".join(cmd))
 
     def _create_conda_env(self, env_name, local_path):
         """Create conda env using conda yml file"""
+        # Validate inputs to prevent command injection
+        self._validate_env_name(env_name)
+        validated_path = self._validate_path(local_path)
 
-        cmd = f"{self._get_conda_exe()} env create -n {env_name} --file {local_path}"
-        logger.info("Creating conda environment %s using: %s.", env_name, cmd)
+        cmd = [self._get_conda_exe(), "env", "create", "-n", env_name, "--file", validated_path]
+        logger.info("Creating conda environment %s using: %s.", env_name, " ".join(cmd))
         _run_shell_cmd(cmd)
         logger.info("Conda environment %s created successfully.", env_name)
 
     def _install_req_txt_in_conda_env(self, env_name, local_path):
         """Install requirements.txt in the given conda environment"""
+        # Validate inputs to prevent command injection
+        self._validate_env_name(env_name)
+        validated_path = self._validate_path(local_path)
 
-        cmd = f"{self._get_conda_exe()} run -n {env_name} pip install -r {local_path} -U"
-        logger.info("Activating conda env and installing requirements: %s", cmd)
+        cmd = [self._get_conda_exe(), "run", "-n", env_name, "pip", "install", "-r", validated_path, "-U"]
+        logger.info("Activating conda env and installing requirements: %s", " ".join(cmd))
         _run_shell_cmd(cmd)
         logger.info("Requirements installed successfully in conda env %s", env_name)
 
     def _update_conda_env(self, env_name, local_path):
         """Update conda env using conda yml file"""
+        # Validate inputs to prevent command injection
+        self._validate_env_name(env_name)
+        validated_path = self._validate_path(local_path)
 
-        cmd = f"{self._get_conda_exe()} env update -n {env_name} --file {local_path}"
-        logger.info("Updating conda env: %s", cmd)
+        cmd = [self._get_conda_exe(), "env", "update", "-n", env_name, "--file", validated_path]
+        logger.info("Updating conda env: %s", " ".join(cmd))
         _run_shell_cmd(cmd)
         logger.info("Conda env %s updated succesfully", env_name)
 
     def _export_conda_env_from_prefix(self, prefix, local_path):
         """Export the conda env to a conda yml file"""
 
-        cmd = f"{self._get_conda_exe()} env export -p {prefix} --no-builds > {local_path}"
+        cmd = [self._get_conda_exe(), "env", "export", "-p", prefix, "--no-builds", ">", local_path]
         logger.info("Exporting conda environment: %s", cmd)
         _run_shell_cmd(cmd)
         logger.info("Conda environment %s exported successfully", prefix)
@@ -402,19 +457,26 @@ def _run_pre_execution_command_script(script_path: str):
     return return_code, error_logs
 
 
-def _run_shell_cmd(cmd: str):
+def _run_shell_cmd(cmd: list):
     """This method runs a given shell command using subprocess
 
-    Raises RuntimeEnvironmentError if the command fails
+    Args:
+        cmd (list): Command and arguments as a list (e.g., ['pip', 'install', '-r', 'requirements.txt'])
+
+    Raises:
+        RuntimeEnvironmentError: If the command fails
+        ValueError: If cmd is not a list
     """
+    if not isinstance(cmd, list):
+        raise ValueError("Command must be a list of arguments for security reasons")
 
-    process = subprocess.Popen((cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False)
 
     _log_output(process)
     error_logs = _log_error(process)
     return_code = process.wait()
     if return_code:
-        error_message = f"Encountered error while running command '{cmd}'. Reason: {error_logs}"
+        error_message = f"Encountered error while running command '{' '.join(cmd)}'. Reason: {error_logs}"
         raise RuntimeEnvironmentError(error_message)
 
 

sagemaker-train/tests/unit/ai_registry/test_dataset_domain_id.py

@@ -11,27 +11,58 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Unit tests for domain-id tagging in DataSet."""
+import json
+import tempfile
+import os
 import pytest
 from unittest.mock import Mock, patch, MagicMock
 from sagemaker.ai_registry.dataset import DataSet
 from sagemaker.ai_registry.dataset_utils import CustomizationTechnique
 
 
+# Sample RLVR format dataset (GSM8K style)
+SAMPLE_DATASET = {
+    "data_source": "openai/gsm8k",
+    "prompt": [{"content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? Let's think step by step and output the final answer after \"####\".", "role": "user"}],
+    "ability": "math",
+    "reward_model": {"ground_truth": "72", "style": "rule"},
+    "extra_info": {"answer": "Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72", "index": 0, "question": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?", "split": "train"}
+}
+
+
+@pytest.fixture
+def sample_dataset_file():
+    """Create a temporary JSONL file with sample dataset."""
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f:
+        json.dump(SAMPLE_DATASET, f)
+        temp_path = f.name
+    
+    yield temp_path
+    
+    # Cleanup
+    if os.path.exists(temp_path):
+        os.remove(temp_path)
+
+
 class TestDataSetDomainId:
     """Test domain-id is added to SearchKeywords when available."""
 
     @patch('sagemaker.core.helper.session_helper.Session')
     @patch('sagemaker.ai_registry.dataset._get_current_domain_id')
     @patch('sagemaker.ai_registry.dataset.AIRHub')
-    @patch('sagemaker.ai_registry.dataset.validate_dataset')
+    @patch('sagemaker.train.defaults.TrainDefaults.get_sagemaker_session')
+    @patch('sagemaker.train.defaults.TrainDefaults.get_role')
     def test_domain_id_added_when_available(
-        self, mock_validate, mock_air_hub, mock_get_domain_id, mock_session
+        self, mock_get_role, mock_get_session, mock_air_hub, mock_get_domain_id, mock_session, sample_dataset_file
     ):
         """Test that domain-id is added to tags when available."""
         # Setup mocks
         mock_domain_id = "d-test123456"
         mock_get_domain_id.return_value = mock_domain_id
-        mock_session.return_value = Mock()
+        mock_session_instance = Mock()
+        mock_session.return_value = mock_session_instance
+        mock_get_session.return_value = mock_session_instance
+        mock_get_role.return_value = "arn:aws:iam::123456789012:role/test-role"
 
         # Mock AIRHub methods
         mock_air_hub.upload_to_s3 = Mock()
@@ -46,11 +77,11 @@ def test_domain_id_added_when_available(
             'HubContentDocument': '{"DatasetS3Bucket": "bucket", "DatasetS3Prefix": "prefix"}'
         })
 
-        # Create dataset
+        # Create dataset with real file
         with patch('sagemaker.ai_registry.dataset.DataSet.wait'):
             dataset = DataSet.create(
                 name="test-dataset",
-                source="test-data.jsonl",
+                source=sample_dataset_file,
                 customization_technique=CustomizationTechnique.SFT
             )
 
@@ -67,14 +98,18 @@ def test_domain_id_added_when_available(
     @patch('sagemaker.core.helper.session_helper.Session')
     @patch('sagemaker.ai_registry.dataset._get_current_domain_id')
     @patch('sagemaker.ai_registry.dataset.AIRHub')
-    @patch('sagemaker.ai_registry.dataset.validate_dataset')
+    @patch('sagemaker.train.defaults.TrainDefaults.get_sagemaker_session')
+    @patch('sagemaker.train.defaults.TrainDefaults.get_role')
     def test_domain_id_not_added_when_unavailable(
-        self, mock_validate, mock_air_hub, mock_get_domain_id, mock_session
+        self, mock_get_role, mock_get_session, mock_air_hub, mock_get_domain_id, mock_session, sample_dataset_file
     ):
         """Test that domain-id is not added when unavailable (non-Studio)."""
         # Setup mocks - domain_id returns None
         mock_get_domain_id.return_value = None
-        mock_session.return_value = Mock()
+        mock_session_instance = Mock()
+        mock_session.return_value = mock_session_instance
+        mock_get_session.return_value = mock_session_instance
+        mock_get_role.return_value = "arn:aws:iam::123456789012:role/test-role"
 
         # Mock AIRHub methods
         mock_air_hub.upload_to_s3 = Mock()
@@ -89,11 +124,11 @@ def test_domain_id_not_added_when_unavailable(
             'HubContentDocument': '{"DatasetS3Bucket": "bucket", "DatasetS3Prefix": "prefix"}'
         })
 
-        # Create dataset
+        # Create dataset with real file
         with patch('sagemaker.ai_registry.dataset.DataSet.wait'):
             dataset = DataSet.create(
                 name="test-dataset",
-                source="test-data.jsonl",
+                source=sample_dataset_file,
                 customization_technique=CustomizationTechnique.SFT
             )
 
@@ -110,14 +145,19 @@ def test_domain_id_not_added_when_unavailable(
     @patch('sagemaker.core.helper.session_helper.Session')
     @patch('sagemaker.ai_registry.dataset._get_current_domain_id')
     @patch('sagemaker.ai_registry.dataset.AIRHub')
+    @patch('sagemaker.train.defaults.TrainDefaults.get_sagemaker_session')
+    @patch('sagemaker.train.defaults.TrainDefaults.get_role')
     def test_domain_id_added_without_customization_technique(
-        self, mock_air_hub, mock_get_domain_id, mock_session
+        self, mock_get_role, mock_get_session, mock_air_hub, mock_get_domain_id, mock_session, sample_dataset_file
     ):
         """Test that domain-id is added even without customization_technique."""
         # Setup mocks
         mock_domain_id = "d-test789"
         mock_get_domain_id.return_value = mock_domain_id
-        mock_session.return_value = Mock()
+        mock_session_instance = Mock()
+        mock_session.return_value = mock_session_instance
+        mock_get_session.return_value = mock_session_instance
+        mock_get_role.return_value = "arn:aws:iam::123456789012:role/test-role"
 
         # Mock AIRHub methods
         mock_air_hub.upload_to_s3 = Mock()
@@ -132,11 +172,11 @@ def test_domain_id_added_without_customization_technique(
             'HubContentDocument': '{"DatasetS3Bucket": "bucket", "DatasetS3Prefix": "prefix"}'
         })
 
-        # Create dataset WITHOUT customization_technique
+        # Create dataset WITHOUT customization_technique using real file
         with patch('sagemaker.ai_registry.dataset.DataSet.wait'):
             dataset = DataSet.create(
                 name="test-dataset",
-                source="test-data.jsonl"
+                source=sample_dataset_file
                 # No customization_technique
             )
 

sagemaker-train/tests/unit/train/remote_function/test_runtime_environment_manager.py

@@ -490,7 +490,7 @@ def test_runs_command_successfully(self, mock_popen, mock_log_output, mock_log_e
         mock_popen.return_value = mock_process
         mock_log_error.return_value = ""
 
-        _run_shell_cmd("echo test")
+        _run_shell_cmd(["echo", "test"])
 
         mock_popen.assert_called_once()
 
@@ -505,7 +505,7 @@ def test_runs_command_raises_error_on_failure(self, mock_popen, mock_log_output,
         mock_log_error.return_value = "Error message"
 
         with pytest.raises(RuntimeEnvironmentError):
-            _run_shell_cmd("false")
+            _run_shell_cmd(["false"])
 
 
 class TestLogOutput: