From 7abcaadec3a4579337fd9af4e3a88604a04edf8d Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 13 Oct 2026 13:01:15 -0700 Subject: [PATCH] fix(core): place repack dependencies in code root Copy directory and file dependencies alongside inference code so SageMaker inference toolkits can import them without extending sys.path. Preserve the existing copytree collision error and copy2 file overwrite behavior. Fixes #3491 --- .../src/sagemaker/core/common_utils.py | 7 +- .../tests/unit/test_common_utils.py | 116 ++++++++++++++---- 2 files changed, 96 insertions(+), 27 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index e8704ef600..2de18493e5 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/sagemaker-core/src/sagemaker/core/common_utils.py @@ -744,13 +744,10 @@ def _create_or_update_code_dir( for dependency in dependencies: # Validate dependency path for security _validate_dependency_path(dependency) - lib_dir = os.path.join(code_dir, "lib") if os.path.isdir(dependency): - shutil.copytree(dependency, os.path.join(lib_dir, os.path.basename(dependency))) + shutil.copytree(dependency, os.path.join(code_dir, os.path.basename(dependency))) else: - if not os.path.exists(lib_dir): - os.mkdir(lib_dir) - shutil.copy2(dependency, lib_dir) + shutil.copy2(dependency, code_dir) def _extract_model(model_uri, sagemaker_session, tmp): diff --git a/sagemaker-core/tests/unit/test_common_utils.py b/sagemaker-core/tests/unit/test_common_utils.py index 58863335bb..1461dd9d64 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/sagemaker-core/tests/unit/test_common_utils.py @@ -2660,37 +2660,109 @@ def test_create_or_update_code_dir_with_source_directory(self): assert os.path.exists(code_dir) assert os.path.exists(os.path.join(code_dir, "app.py")) - def test_create_or_update_code_dir_with_dependencies(self): - """Test code directory creation with dependencies.""" + def test_create_or_update_code_dir_copies_directory_dependency_to_code_root(self, tmp_path): + """Test directory dependencies are importable from the code root.""" from sagemaker.core.common_utils import _create_or_update_code_dir - with tempfile.TemporaryDirectory() as tmpdir: - model_dir = os.path.join(tmpdir, "model") - os.makedirs(model_dir) + model_dir = tmp_path / "model" + model_dir.mkdir() + inference_script = tmp_path / "inference.py" + inference_script.write_text("# inference code") + dependency = tmp_path / "my_lib" + dependency.mkdir() + (dependency / "helper.py").write_text("# helper code") + + _create_or_update_code_dir( + str(model_dir), str(inference_script), None, [str(dependency)], None, str(tmp_path) + ) - inference_script = os.path.join(tmpdir, "inference.py") - with open(inference_script, "w") as f: - f.write("# inference code") + assert (model_dir / "code" / "my_lib" / "helper.py").exists() + assert not (model_dir / "code" / "lib").exists() + + def test_create_or_update_code_dir_copies_file_dependency_to_code_root(self, tmp_path): + """Test single-file dependencies are importable from the code root.""" + from sagemaker.core.common_utils import _create_or_update_code_dir + + model_dir = tmp_path / "model" + model_dir.mkdir() + inference_script = tmp_path / "inference.py" + inference_script.write_text("# inference code") + dependency = tmp_path / "helper.py" + dependency.write_text("# helper code") + + _create_or_update_code_dir( + str(model_dir), str(inference_script), None, [str(dependency)], None, str(tmp_path) + ) + + assert (model_dir / "code" / "helper.py").exists() + assert not (model_dir / "code" / "lib").exists() + + def test_create_or_update_code_dir_dependency_coexists_with_source_directory(self, tmp_path): + """Test dependencies and source-directory files share the code root.""" + from sagemaker.core.common_utils import _create_or_update_code_dir + + model_dir = tmp_path / "model" + model_dir.mkdir() + source_dir = tmp_path / "source" + source_dir.mkdir() + (source_dir / "inference.py").write_text("# inference code") + dependency = tmp_path / "my_lib" + dependency.mkdir() + (dependency / "helper.py").write_text("# helper code") + + _create_or_update_code_dir( + str(model_dir), "inference.py", str(source_dir), [str(dependency)], None, str(tmp_path) + ) - dep_dir = os.path.join(tmpdir, "my_lib") - os.makedirs(dep_dir) - with open(os.path.join(dep_dir, "helper.py"), "w") as f: - f.write("# helper code") + assert (model_dir / "code" / "inference.py").exists() + assert (model_dir / "code" / "my_lib" / "helper.py").exists() + + def test_create_or_update_code_dir_rejects_directory_dependency_collision(self, tmp_path): + """Test directory dependencies do not silently merge with source code.""" + from sagemaker.core.common_utils import _create_or_update_code_dir - # Should create code directory with dependencies + model_dir = tmp_path / "model" + model_dir.mkdir() + source_dir = tmp_path / "source" + (source_dir / "my_lib").mkdir(parents=True) + (source_dir / "my_lib" / "source.py").write_text("# source code") + dependency_parent = tmp_path / "dependency" + dependency = dependency_parent / "my_lib" + dependency.mkdir(parents=True) + (dependency / "helper.py").write_text("# helper code") + + with pytest.raises(FileExistsError): _create_or_update_code_dir( - model_dir, - inference_script, - None, - [dep_dir], + str(model_dir), + "inference.py", + str(source_dir), + [str(dependency)], None, - tmpdir, + str(tmp_path), ) - code_dir = os.path.join(model_dir, "code") - lib_dir = os.path.join(code_dir, "lib") - assert os.path.exists(lib_dir) - assert os.path.exists(os.path.join(lib_dir, "my_lib", "helper.py")) + assert (model_dir / "code" / "my_lib" / "source.py").exists() + assert not (model_dir / "code" / "my_lib" / "helper.py").exists() + + def test_create_or_update_code_dir_file_dependency_overwrites_collision(self, tmp_path): + """Test a file dependency preserves copy2 overwrite semantics.""" + from sagemaker.core.common_utils import _create_or_update_code_dir + + model_dir = tmp_path / "model" + model_dir.mkdir() + source_dir = tmp_path / "source" + source_dir.mkdir() + (source_dir / "helper.py").write_text("source") + dependency_parent = tmp_path / "dependency" + dependency_parent.mkdir() + dependency = dependency_parent / "helper.py" + dependency.write_text("dependency") + + _create_or_update_code_dir( + str(model_dir), "inference.py", str(source_dir), [str(dependency)], None, str(tmp_path) + ) + + assert (model_dir / "code" / "helper.py").read_text() == "dependency" def test_create_or_update_code_dir_rejects_sensitive_paths(self): """Test that code_dir validation rejects sensitive system paths."""