본문으로 건너뛰기

[ultralytics] [Ultralytics] NDJSON 변환 최적화: 보안과 성능을 동시에 잡는 설계 전략

PR 링크: ultralytics/ultralytics#25196 상태: Merged | 변경: +38 / -50

들어가며

소프트웨어 엔지니어링에서 보안과 성능은 종종 트레이드오프(Trade-off) 관계에 놓입니다. 최근 Ultralytics 레포지토리에 반영된 PR은 NDJSON(Newline Delimited JSON) 데이터셋을 YOLO 포맷으로 변환하는 과정에서 발생했던 심각한 보안 취약점(Path Traversal)을 해결함과 동시에, 이전 방식보다 최대 7.8배 빠른 성능을 달성한 흥미로운 사례를 보여줍니다.

기존의 해결책(#25188)은 보안을 위해 모든 레코드를 사전에 전수 검사하는 'Prepass' 방식을 도입했으나, 이는 대규모 데이터셋에서 상당한 성능 저하를 야기했습니다. 이번 PR은 이러한 비효율적인 사전 검사 단계를 제거하고, 파일 시스템 구조 자체를 안전하게 재설계함으로써 두 마리 토끼를 모두 잡았습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

1. 사전 검증 루프 제거 및 인덱스 기반 파일명 생성

가장 큰 변화는 사용자로부터 입력받은 파일명(file)을 더 이상 신뢰하지 않는다는 점입니다. 이전에는 파일명이 안전한지 확인하기 위해 복잡한 유효성 검사를 수행했지만, 이제는 레코드의 인덱스를 파일명으로 사용하여 충돌과 보안 위협을 원천 차단합니다.

Before (기존의 비효율적인 전수 검사):

# 모든 레코드를 미리 돌며 경로의 안전성을 검사함 (성능 저하의 원인)
for i, record in enumerate(image_records, start=1):
    split, file = record.get("split"), record.get("file")
    if not isinstance(split, str) or split not in {"train", "val", "test"}:
        raise ValueError(f"Unsafe NDJSON split in record {i}: {split!r}")
    if not isinstance(file, str) or PureWindowsPath(file).name != file or file.rstrip(" .") in {"", ".", ".."}:
        raise ValueError(f"Unsafe NDJSON file path in record {i}: {file!r}")

After (해시 스캔 루프 내 통합 및 인덱스 기반 명명):

# 해시를 생성하는 단일 루프 내에서 검증과 파일명 재할당을 동시에 수행
for i, r in enumerate(lines):
    if i:
        split, source_name = r.get("split"), r.get("file")
        if split not in {"train", "val", "test"}:
            raise ValueError(f"Invalid NDJSON split: {split!r}")
        # 레코드 인덱스(i)를 사용하여 충돌 없는 안전한 파일명 생성
        suffix = source_name.rsplit(".", 1)[-1]
        r["file"] = f"{i}.{suffix}" if suffix.isalnum() and len(suffix) <= 10 else f"{i}.jpg"

이 변경을 통해 PureWindowsPath와 같은 무거운 객체 할당이 제거되었고, 데이터셋의 크기가 커질수록 기하급수적으로 늘어나던 오버헤드를 줄였습니다.

2. Classification 디렉토리의 익명화 및 메타데이터 분리

분류(Classification) 작업에서 클래스 이름에 슬래시(/)나 백슬래시(\)가 포함될 경우, 디렉토리 구조가 깨지거나 의도치 않은 경로에 파일이 생성될 위험이 있었습니다. 이를 해결하기 위해 실제 디렉토리는 숫자로 구성하고, 원본 이름은 별도의 YAML 파일에 저장하는 방식을 채택했습니다.

After (ultralytics/data/converter.py):

# 클래스 이름을 직접 디렉토리명으로 쓰지 않고 6자리 숫자로 변환
class_dirs = {class_id: f"{i:06d}" for i, class_id in enumerate(sorted(classification_ids))}
classification_names = {i: class_names.get(class_id, str(class_id)) for i, class_id in enumerate(class_dirs)}

# ... 중략 ...

if is_classification:
    # 안전한 경로 구조를 유지하면서 원본 이름을 복원하기 위한 메타데이터 저장
    YAML.save(dataset_dir / ".ndjson.yaml", {"names": classification_names})

3. 데이터셋 로드 시 이름 복원 로직 (ultralytics/data/utils.py)

디렉토리 이름을 숫자로 바꿨더라도 사용자는 원래의 클래스 이름을 보고 싶어 합니다. 이를 위해 check_cls_dataset 함수에서 .ndjson.yaml 파일이 존재할 경우 이를 로드하여 이름을 복원하도록 개선되었습니다.

After:

names = dict(enumerate(sorted(names)))
if (ndjson_names := data_dir / ".ndjson.yaml").is_file():
    names = YAML.load(ndjson_names)["names"]
    if len(names) != nc:
        raise ValueError(f"NDJSON class names length {len(names)} does not match directory count {nc}")

왜 이게 좋은가?

1. 압도적인 성능 향상

PR 설명에 기재된 벤치마크 결과에 따르면, 100만 개의 레코드를 처리할 때 Detection 작업 기준으로 기존 방식(#25188)은 1.809초가 걸렸으나, 이번 최적화 이후 0.233초로 단축되었습니다. 이는 약 7.8배의 성능 향상입니다. 불필요한 Path 객체 생성과 중복 루프를 제거한 결과입니다.

2. 보안 강화 (Path Traversal 방지)

리뷰어의 피드백 중에는 Windows 환경에서 C:escaped.jpg와 같은 경로가 입력될 경우 데이터셋 디렉토리를 탈출할 수 있다는 날카로운 지적이 있었습니다. 이번 PR은 아예 소스 파일명을 무시하고 1.jpg, 2.jpg와 같이 시스템이 제어하는 파일명을 강제함으로써 이러한 공격 벡터를 완전히 차단했습니다.

3. 플랫폼 이식성

파일 시스템마다 허용하는 특수문자가 다르기 때문에 발생하는 문제를 '숫자형 디렉토리 + 메타데이터' 구조로 해결했습니다. 이는 Windows, Linux, macOS 간의 데이터셋 이동 시 발생할 수 있는 잠재적 버그를 예방합니다.

마치며

이번 최적화의 핵심 교훈은 "신뢰할 수 없는 입력값(User Input)을 검증하려 애쓰기보다, 시스템이 제어 가능한 값으로 대체하는 것이 보안과 성능 모두에 유리할 수 있다"는 점입니다. 또한, 대규모 데이터를 다룰 때는 단 한 번의 추가 루프나 객체 할당도 성능에 치명적일 수 있음을 다시 한번 상기시켜 줍니다.

시니어 엔지니어로서 이러한 변경사항은 코드의 간결함을 유지하면서도 견고함을 높이는 아주 모범적인 리팩토링 사례라고 평가합니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글