Files

343 lines
12 KiB
Ruby
Executable File

#!/usr/bin/env ruby
require "digest"
require "json"
require "optparse"
require "pathname"
require "yaml"
TYPE_BYTES = {
"int32" => 4,
"bfloat16" => 2,
"float32" => 4,
"int8" => 1,
"uint8" => 1,
"PrimitiveType.S32" => 4,
"PrimitiveType.BF16" => 2,
"PrimitiveType.F32" => 4,
"PrimitiveType.S8" => 1,
"PrimitiveType.U8" => 1,
"PrimitiveType.U16" => 2,
"PrimitiveType.S16" => 2,
"PrimitiveType.U32" => 4,
}.freeze
options = {
output: nil,
full_hash: false,
hash_limit: 64 * 1024 * 1024,
}
OptionParser.new do |parser|
parser.banner = "usage: calculet_package_validate.rb [options] MODEL_DIR"
parser.on("-o", "--output FILE", "write JSON report to FILE") { |v| options[:output] = v }
parser.on("--full-hash", "SHA-256 every file, including large parameter blocks") { options[:full_hash] = true }
parser.on("--hash-limit-mib N", Integer, "fast-mode hash limit (default: 64 MiB)") do |v|
options[:hash_limit] = v * 1024 * 1024
end
end.parse!
abort "MODEL_DIR is required" unless ARGV.length == 1
root = Pathname.new(ARGV.first).expand_path
abort "not a directory: #{root}" unless root.directory?
findings = []
add_finding = lambda do |severity, code, message, evidence = nil|
item = { severity: severity, code: code, message: message }
item[:evidence] = evidence if evidence
findings << item
end
safe_yaml = lambda do |path|
YAML.safe_load(path.read, permitted_classes: [], permitted_symbols: [], aliases: true)
rescue Psych::Exception => e
add_finding.call("error", "YAML_PARSE", e.message, path.relative_path_from(root).to_s)
nil
end
relative = lambda do |path|
Pathname.new(path).expand_path.relative_path_from(root).to_s
rescue ArgumentError
path.to_s
end
files = root.glob("**/*", File::FNM_DOTMATCH).select(&:file?)
symlinks = root.glob("**/*", File::FNM_DOTMATCH).select(&:symlink?)
symlinks.each do |link|
begin
resolved = link.realpath
unless resolved.to_s.start_with?(root.realpath.to_s + File::SEPARATOR)
add_finding.call("error", "SYMLINK_ESCAPE", "symlink resolves outside package", relative.call(link))
end
rescue Errno::ENOENT
add_finding.call("error", "BROKEN_SYMLINK", "broken symlink", relative.call(link))
end
end
required_exact = %w[hparam.yaml model_memory_reserved_info.txt param_blk0.bin param_blk1.bin]
required_exact.each do |name|
path = root.join(name)
if !path.file?
add_finding.call("error", "REQUIRED_FILE", "required file is missing", name)
elsif path.size.zero?
add_finding.call("error", "EMPTY_FILE", "required file is empty", name)
end
end
manifest_candidates = root.glob("*.yaml").reject do |path|
path.basename.to_s == "hparam.yaml" || path.basename.to_s.end_with?("_op_io_buf_ptr.yaml")
end
if manifest_candidates.length != 1
add_finding.call("error", "ROOT_MANIFEST_COUNT",
"expected exactly one root model manifest, got #{manifest_candidates.length}")
end
manifest_path = manifest_candidates.first
manifest = manifest_path ? safe_yaml.call(manifest_path) : nil
model_summary = {}
submodels = []
supply_chain = {}
if manifest.is_a?(Hash)
cfg = manifest["one_model_cfg"]
unless cfg.is_a?(Hash)
add_finding.call("error", "MODEL_CFG", "one_model_cfg must be a mapping", relative.call(manifest_path))
cfg = {}
end
cfg.each do |logical_name, item|
unless item.is_a?(Hash)
add_finding.call("error", "SUBMODEL_CFG", "submodel config must be a mapping", logical_name)
next
end
info = item["model_info"] || {}
device = item["device_cfg"] || {}
attrs = item["model_attrs"].to_s
kind = attrs.empty? ? (logical_name.include?("prefill") ? "prefill" : "decode") : attrs
submodels << {
logical_name: logical_name,
kind: kind,
model_path: item["model_path"],
batch: info["batch_num"],
max_seq: info["max_seq_len"],
chips: device["chip_nums"],
calcores_per_chip: device["calcore_num"],
input_shape: item["input_shape"],
flash_attention: item.dig("calcc_env_cfg", "enable_flash_attention"),
one_token: item.dig("calcc_env_cfg", "enable_one_token"),
dynamic_d2d: item.dig("calcc_env_cfg", "dynamic_d2d"),
}
end
kinds = submodels.map { |s| s[:kind] }
%w[prefill decode].each do |kind|
count = kinds.count(kind)
add_finding.call("error", "SUBMODEL_KIND_COUNT", "expected one #{kind}, got #{count}") unless count == 1
end
batches = submodels.map { |s| s[:batch] }.compact.uniq
seqs = submodels.map { |s| s[:max_seq] }.compact.uniq
chips = submodels.map { |s| s[:chips] }.compact.uniq
add_finding.call("error", "BATCH_MISMATCH", "submodel batch values differ: #{batches}") if batches.length > 1
add_finding.call("error", "SEQ_MISMATCH", "submodel max_seq values differ: #{seqs}") if seqs.length > 1
add_finding.call("error", "CHIP_MISMATCH", "submodel chip counts differ: #{chips}") if chips.length > 1
model_summary = {
name: manifest["model_name"],
seq_len: manifest["seq_len"],
batch: batches.first,
max_seq: seqs.first,
chips: chips.first,
}
if model_summary[:seq_len] && model_summary[:max_seq] && model_summary[:seq_len] != model_summary[:max_seq]
add_finding.call("error", "ROOT_SEQ_MISMATCH", "root seq_len differs from submodel max_seq")
end
supply_chain = {
onnx_hash: manifest["onnx_hash"],
compiler_commits: manifest["git_commit_id"],
golden_hash: manifest["golden_hash"],
}
if Array(supply_chain[:onnx_hash]).empty?
add_finding.call("blocked_reproducibility", "ONNX_HASH_MISSING", "ONNX hash is missing")
end
end
memory = { parameter_files: {}, reservations: {}, parameter_maps: [] }
memory_path = root.join("model_memory_reserved_info.txt")
if memory_path.file?
memory_path.each_line.with_index(1) do |line, line_no|
case line
when /\A(dram|sram|sync)_rsvd\s+vaddr\s+(0x[0-9a-fA-F]+)\s+size\s+(\d+)/
memory[:reservations][Regexp.last_match(1)] = {
vaddr: Regexp.last_match(2), bytes: Regexp.last_match(3).to_i,
}
when /\A([^\s]+\.bin)\s+offset\s+(0x[0-9a-fA-F]+)\s+vaddr\s+(0x[0-9a-fA-F]+)\s+size\s+(0x[0-9a-fA-F]+)\s+chip\s+([0-9a-fA-F]+)\s+([0-9a-fA-F]+)\s+([0-9a-fA-F]+)/
file_name = Regexp.last_match(1)
offset = Regexp.last_match(2).to_i(16)
vaddr = Regexp.last_match(3).to_i(16)
size = Regexp.last_match(4).to_i(16)
mask = [Regexp.last_match(5), Regexp.last_match(6), Regexp.last_match(7)].join
target = root.join(file_name)
if !target.file?
add_finding.call("error", "PARAM_FILE_MISSING", "mapped parameter file missing", file_name)
elsif offset > target.size || size > target.size - offset
add_finding.call("error", "PARAM_RANGE", "parameter mapping exceeds file", "#{file_name}:#{line_no}")
end
add_finding.call("error", "CHIP_MASK_ZERO", "parameter map has zero chip mask", "line #{line_no}") if mask.to_i(16).zero?
memory[:parameter_maps] << {
file: file_name, offset: offset, vaddr: vaddr, bytes: size, chip_mask: mask,
}
when /\Adevice_memory_required:\s+(\d+)\s+MB/
memory[:device_memory_required_raw] = { value: Regexp.last_match(1).to_i, unit: "MB" }
add_finding.call("warning_vendor", "DEVICE_MEMORY_UNIT",
"device_memory_required unit/value is not credible for capacity decisions", "line #{line_no}")
end
end
end
%w[param_blk0.bin param_blk1.bin].each do |name|
path = root.join(name)
memory[:parameter_files][name] = path.size if path.file?
end
metadata = []
root.glob("*/submodel_memory_reserved_info.txt").sort.each do |path|
item = { path: relative.call(path), inputs: [], outputs: [], csrs: [], files: [] }
path.each_line.with_index(1) do |line, line_no|
case line
when /\Asmodel_type\s+(\S+)/
item[:smodel_type] = Regexp.last_match(1)
when /\An_batch\s+(\d+)/
item[:batch] = Regexp.last_match(1).to_i
when /\Aicsr\s+(.+?)\s+offset\s+(0x[0-9a-fA-F]+)/
item[:csrs] << { name: Regexp.last_match(1), offset: Regexp.last_match(2) }
when /\A(ibuf|obuf)\s+(.+?)\s+shape\s+\[([^\]]+)\]\s+type\s+(\S+)\s+vaddr\s+(0x[0-9a-fA-F]+)\s+(0x[0-9a-fA-F]+)\s+size\s+(\d+)/
io_kind = Regexp.last_match(1)
dims = Regexp.last_match(3).split(",").map { |v| Integer(v.strip) }
dtype = Regexp.last_match(4)
bytes = Regexp.last_match(7).to_i
tensor = {
name: Regexp.last_match(2), shape: dims, dtype: dtype,
ping_vaddr: Regexp.last_match(5), pong_vaddr: Regexp.last_match(6), bytes: bytes,
}
if TYPE_BYTES.key?(dtype)
expected = dims.inject(1, :*) * TYPE_BYTES.fetch(dtype)
if expected != bytes
add_finding.call("error", "TENSOR_SIZE", "tensor size formula mismatch",
"#{relative.call(path)}:#{line_no} expected=#{expected} actual=#{bytes}")
end
else
add_finding.call("error", "TENSOR_DTYPE", "unknown tensor dtype #{dtype}", "#{relative.call(path)}:#{line_no}")
end
(io_kind == "ibuf" ? item[:inputs] : item[:outputs]) << tensor
when /\A(\S+\.(?:bin|elf|so))\b.*\bchip\s+/
file_name = Regexp.last_match(1)
item[:files] << file_name
declared = path.dirname.join(file_name)
add_finding.call("error", "DECLARED_FILE", "declared submodel file missing", relative.call(declared)) unless declared.file?
end
end
if path.dirname.basename.to_s.include?("prefill") && item[:smodel_type] == "llm_decode"
add_finding.call("warning_vendor", "PREFILL_MODEL_TYPE", "prefill metadata is marked llm_decode", relative.call(path))
end
metadata << item
end
if metadata.length != 2
add_finding.call("error", "SUBMODEL_METADATA_COUNT", "expected two submodel metadata files, got #{metadata.length}")
end
graphs = {}
root.glob("*_op_io_buf_ptr.yaml").sort.each do |path|
graph = safe_yaml.call(path)
next unless graph.is_a?(Hash)
kind = path.basename.to_s.include?("prefill") ? "prefill" : "decode"
histogram = Hash.new(0)
unknown_dtypes = []
graph.each do |op_name, op|
histogram[op_name.sub(/_\d+\z/, "")] += 1
next unless op.is_a?(Hash)
op.each_value do |tensor|
next unless tensor.is_a?(Array) && tensor.length >= 3
dtype = tensor[2].to_s
unknown_dtypes << dtype unless TYPE_BYTES.key?(dtype)
end
end
unknown_dtypes.uniq.each do |dtype|
add_finding.call("error", "GRAPH_DTYPE", "unknown graph dtype #{dtype}", relative.call(path))
end
graphs[kind] = { path: relative.call(path), op_count: graph.length, op_histogram: histogram.sort.to_h }
prefix = path.basename.to_s.sub("_op_io_buf_ptr.yaml", "")
root.glob("#{prefix}.chip[01].profparts").each do |prof|
unknown = prof.each_line.map(&:strip).reject { |name| name.empty? || name == "tail" }.reject { |name| graph.key?(name) }
unless unknown.empty?
add_finding.call("error", "PROFPART_UNKNOWN_OP", "profparts references unknown operators",
"#{relative.call(prof)} first=#{unknown.first} count=#{unknown.length}")
end
end
end
%w[prefill decode].each do |kind|
add_finding.call("error", "GRAPH_MISSING", "#{kind} operator graph is missing") unless graphs.key?(kind)
end
hashes = {}
files.sort.each do |path|
rel = relative.call(path)
if options[:full_hash] || path.size <= options[:hash_limit]
hashes[rel] = { bytes: path.size, sha256: Digest::SHA256.file(path).hexdigest }
else
hashes[rel] = { bytes: path.size, sha256: nil, skipped: "larger_than_fast_mode_limit" }
end
end
severity_order = {
"error" => 0,
"blocked_reproducibility" => 1,
"warning_vendor" => 2,
"warning_service" => 3,
"info" => 4,
}
findings.sort_by! { |f| [severity_order.fetch(f[:severity], 99), f[:code], f[:evidence].to_s] }
verdict = if findings.any? { |f| f[:severity] == "error" }
"fail"
elsif findings.any? { |f| f[:severity] != "info" }
"pass_with_warnings"
else
"pass"
end
report = {
schema: "calculet-package-validation/v1",
package_root: root.to_s,
generated_at_utc: Time.now.utc.strftime("%Y-%m-%dT%H:%M:%SZ"),
model: model_summary,
submodels: submodels,
runtime_metadata: metadata,
memory: memory,
graph: graphs,
supply_chain: supply_chain,
inventory: { file_count: files.length, symlink_count: symlinks.length, files: hashes },
findings: findings,
verdict: verdict,
}
json = JSON.pretty_generate(report) + "\n"
if options[:output]
File.write(options[:output], json)
else
puts json
end
exit(verdict == "fail" ? 1 : 0)