-
Notifications
You must be signed in to change notification settings - Fork 134
fix: Respect cc_model_filenames in Python backend auto-fill #489
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
4ef4eba
de6bf52
38b404e
6881a86
eba4837
9631b7a
a18162d
6c42f72
c5ec7ac
6037019
0fc4a6d
135f58e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,5 @@ | ||
| /build | ||
| /build_test | ||
| /.vscode | ||
| *.so | ||
| *__pycache__/ |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,178 @@ | ||
| // Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| // | ||
| // Redistribution and use in source and binary forms, with or without | ||
| // modification, are permitted provided that the following conditions | ||
| // are met: | ||
| // * Redistributions of source code must retain the above copyright | ||
| // notice, this list of conditions and the following disclaimer. | ||
| // * Redistributions in binary form must reproduce the above copyright | ||
| // notice, this list of conditions and the following disclaimer in the | ||
| // documentation and/or other materials provided with the distribution. | ||
| // * Neither the name of NVIDIA CORPORATION nor the names of its | ||
| // contributors may be used to endorse or promote products derived | ||
| // from this software without specific prior written permission. | ||
| // | ||
| // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY | ||
| // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE | ||
| // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR | ||
| // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR | ||
| // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, | ||
| // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, | ||
| // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR | ||
| // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY | ||
| // OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT | ||
| // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE | ||
| // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. | ||
|
|
||
| #include "model_config_utils.h" | ||
|
|
||
| #include <sys/stat.h> | ||
|
|
||
| #include <fstream> | ||
| #include <string> | ||
|
|
||
| #include "constants.h" | ||
| #include "filesystem/api.h" | ||
| #include "gtest/gtest.h" | ||
|
|
||
| namespace tc = triton::core; | ||
|
|
||
| namespace { | ||
|
|
||
| // Helper to create a temporary model directory with a version subdirectory | ||
| // and an optional file inside it. | ||
| class TempModelDir { | ||
| public: | ||
| TempModelDir() | ||
| { | ||
| auto status = | ||
| tc::MakeTemporaryDirectory(tc::FileSystemType::LOCAL, &root_path_); | ||
| EXPECT_TRUE(status.IsOk()) << status.AsString(); | ||
| } | ||
|
|
||
| ~TempModelDir() | ||
| { | ||
| // Best-effort cleanup | ||
| std::string cmd = "rm -rf " + root_path_; | ||
| (void)system(cmd.c_str()); | ||
| } | ||
|
|
||
| // Create version subdir (e.g., "1") and optionally place a file in it. | ||
| void AddVersionWithFile( | ||
| const std::string& version, const std::string& filename) | ||
| { | ||
| std::string version_dir = tc::JoinPath({root_path_, version}); | ||
| mkdir(version_dir.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH); | ||
| if (!filename.empty()) { | ||
| std::ofstream f(tc::JoinPath({version_dir, filename})); | ||
| f << "# placeholder"; | ||
| } | ||
|
Comment on lines
+64
to
+69
|
||
| } | ||
|
|
||
| const std::string& Path() const { return root_path_; } | ||
|
|
||
| private: | ||
| std::string root_path_; | ||
| }; | ||
|
|
||
| class AutoCompleteBackendFieldsTest : public ::testing::Test {}; | ||
|
|
||
| // When backend is "python" and default_model_filename is empty and | ||
| // cc_model_filenames is empty, default_model_filename should be set to | ||
| // "model.py". | ||
| TEST_F(AutoCompleteBackendFieldsTest, PythonBackendSetsDefaultFilename) | ||
| { | ||
| TempModelDir dir; | ||
| dir.AddVersionWithFile("1", "model.py"); | ||
|
|
||
| inference::ModelConfig config; | ||
| config.set_backend("python"); | ||
| // default_model_filename and cc_model_filenames are both empty | ||
|
|
||
| auto status = | ||
| tc::AutoCompleteBackendFields("test_model", dir.Path(), &config); | ||
| ASSERT_TRUE(status.IsOk()) << status.AsString(); | ||
| EXPECT_EQ(config.default_model_filename(), "model.py"); | ||
| } | ||
|
|
||
| // When backend is "python" and default_model_filename is empty but | ||
| // cc_model_filenames is populated, default_model_filename should NOT be | ||
| // auto-filled to "model.py". | ||
| TEST_F( | ||
| AutoCompleteBackendFieldsTest, | ||
| PythonBackendSkipsDefaultFilenameWhenCcModelFilenamesSet) | ||
| { | ||
| TempModelDir dir; | ||
| dir.AddVersionWithFile("1", "custom_model.py"); | ||
|
|
||
| inference::ModelConfig config; | ||
| config.set_backend("python"); | ||
| (*config.mutable_cc_model_filenames())["gpu"] = "custom_model.py"; | ||
| // default_model_filename is empty, cc_model_filenames is set | ||
|
|
||
| auto status = | ||
| tc::AutoCompleteBackendFields("test_model", dir.Path(), &config); | ||
| ASSERT_TRUE(status.IsOk()) << status.AsString(); | ||
| EXPECT_EQ(config.default_model_filename(), "") | ||
| << "default_model_filename should remain empty when cc_model_filenames " | ||
| "is set"; | ||
|
Comment on lines
+116
to
+118
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In this case, which python file backend should load if not "gpu"? |
||
| } | ||
|
|
||
| // When backend is "python" and default_model_filename is already set, | ||
| // it should be preserved regardless of cc_model_filenames. | ||
|
nightflight-dk marked this conversation as resolved.
|
||
| TEST_F( | ||
| AutoCompleteBackendFieldsTest, | ||
| PythonBackendPreservesExplicitDefaultFilename) | ||
| { | ||
| TempModelDir dir; | ||
| dir.AddVersionWithFile("1", "my_model.py"); | ||
|
|
||
| inference::ModelConfig config; | ||
| config.set_backend("python"); | ||
| config.set_default_model_filename("my_model.py"); | ||
|
|
||
| auto status = | ||
| tc::AutoCompleteBackendFields("test_model", dir.Path(), &config); | ||
| ASSERT_TRUE(status.IsOk()) << status.AsString(); | ||
| EXPECT_EQ(config.default_model_filename(), "my_model.py"); | ||
| } | ||
|
|
||
| // When backend is "python" and both default_model_filename and | ||
| // cc_model_filenames are set, both should be preserved as-is. | ||
| TEST_F( | ||
| AutoCompleteBackendFieldsTest, | ||
| PythonBackendPreservesBothDefaultAndCcModelFilenames) | ||
| { | ||
| TempModelDir dir; | ||
| dir.AddVersionWithFile("1", "my_model.py"); | ||
|
|
||
| inference::ModelConfig config; | ||
| config.set_backend("python"); | ||
| config.set_default_model_filename("my_model.py"); | ||
| (*config.mutable_cc_model_filenames())["gpu"] = "gpu_model.py"; | ||
|
|
||
| auto status = | ||
| tc::AutoCompleteBackendFields("test_model", dir.Path(), &config); | ||
| ASSERT_TRUE(status.IsOk()) << status.AsString(); | ||
| EXPECT_EQ(config.default_model_filename(), "my_model.py"); | ||
| EXPECT_EQ(config.cc_model_filenames().at("gpu"), "gpu_model.py"); | ||
| } | ||
|
|
||
| // When backend is empty but version dir contains model.py, backend should be | ||
| // auto-detected as "python" and default_model_filename set to "model.py". | ||
| TEST_F(AutoCompleteBackendFieldsTest, AutoDetectPythonBackendFromModelFile) | ||
| { | ||
| TempModelDir dir; | ||
| dir.AddVersionWithFile("1", "model.py"); | ||
|
|
||
| inference::ModelConfig config; | ||
| // backend, platform, default_model_filename all empty | ||
|
|
||
| auto status = | ||
| tc::AutoCompleteBackendFields("test_model", dir.Path(), &config); | ||
| ASSERT_TRUE(status.IsOk()) << status.AsString(); | ||
| EXPECT_EQ(config.backend(), "python"); | ||
| EXPECT_EQ(config.default_model_filename(), "model.py"); | ||
| } | ||
|
|
||
| } // namespace | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
TempModelDir::~TempModelDir uses
system("rm -rf ...")for cleanup. This is non-portable (breaks on Windows), can behave incorrectly if the path contains spaces/shell metacharacters, andsystemmay not be declared here (no include), causing a compile failure in C++ builds. Prefer using the existing filesystem helper (e.g.,triton::core::DeletePath(root_path_)), and optionally assert/expect the returned Status in teardown/cleanup.