[util] FIx --benchmark flag when value is a file:/// URI.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""A consistent way to interpret a user-specified benchmark from commandline flags."""
+from pathlib import Path
from typing import Optional, Union
from absl import flags
-from compiler_gym.service.proto import Benchmark, File
+from compiler_gym.datasets import Benchmark
flags.DEFINE_string(
"benchmark",
@@ -24,7 +25,9 @@ def benchmark_from_flags() -> Optional[Union[Benchmark, str]]:
"""Returns either the name of the benchmark, or a Benchmark message."""
if FLAGS.benchmark:
if FLAGS.benchmark.startswith("file:///"):
- return Benchmark(uri=FLAGS.benchmark, program=File(uri=FLAGS.benchmark))
+ path = Path(FLAGS.benchmark[len("file:///") :])
+ uri = f"benchmark://user-v0/{path}"
+ return Benchmark.from_file(uri=uri, path=path)
else:
return FLAGS.benchmark
else:
