Skip to content

Commit 0392688

Browse files
#4448 - Fail early if engine file path is unwritable before ONNX parsing
Signed-off-by: Paul Pacheco <154653205+PierreMarieCurie@users.noreply.github.com>
1 parent 2114dc7 commit 0392688

File tree

3 files changed

+37
-0
lines changed

3 files changed

+37
-0
lines changed

samples/common/sampleOptions.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1386,6 +1386,10 @@ void BuildOptions::parse(Arguments& arguments)
13861386

13871387
if (getAndDelOption(arguments, "--saveEngine", engine))
13881388
{
1389+
if (!canWriteFile(engine))
1390+
{
1391+
throw std::invalid_argument(std::string("Cannot write engine file to path: ") + engine);
1392+
}
13891393
save = true;
13901394
}
13911395
if (load && save)

samples/common/sampleUtils.cpp

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,36 @@ void loadFromFile(std::string const& fileName, char* dst, size_t size)
105105
}
106106
}
107107

108+
// Check if the file at the given path can be written to.
109+
bool canWriteFile(const std::string& path)
110+
{
111+
// Verify that the target directory exists
112+
namespace fs = std::filesystem;
113+
fs::path p(path);
114+
fs::path dir = p.has_parent_path() ? p.parent_path() : fs::current_path();
115+
if (!fs::exists(dir) || !fs::is_directory(dir))
116+
{
117+
return false;
118+
}
119+
120+
// Try creating and writing to a temporary file in the directory
121+
const fs::path tempFilePath = dir / ".writetest.tmp";
122+
std::ofstream test(tempFilePath.string(), std::ios::out | std::ios::trunc);
123+
if (!test.is_open())
124+
{
125+
return false;
126+
}
127+
test << "test";
128+
const bool ok = test.good();
129+
test.close();
130+
131+
// Clean up the temporary file without throwing on failure
132+
std::error_code ec;
133+
fs::remove(tempFilePath, ec);
134+
135+
return ok;
136+
}
137+
108138
std::vector<std::string> splitToStringVec(std::string const& s, char separator, int64_t maxSplit)
109139
{
110140
std::vector<std::string> splitted;

samples/common/sampleUtils.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#define TRT_SAMPLE_UTILS_H
2020

2121
#include <fstream>
22+
#include <filesystem>
2223
#include <iostream>
2324
#include <memory>
2425
#include <numeric>
@@ -75,6 +76,8 @@ void dumpInt4Buffer(void const* buffer, std::string const& separator, std::ostre
7576

7677
void loadFromFile(std::string const& fileName, char* dst, size_t size);
7778

79+
bool canWriteFile(const std::string& path);
80+
7881
std::vector<std::string> splitToStringVec(std::string const& option, char separator, int64_t maxSplit = -1);
7982

8083
bool broadcastIOFormats(std::vector<IOFormat> const& formats, size_t nbBindings, bool isInput = true);

0 commit comments

Comments
 (0)