[RAPTOR-13911] refactor to remove distutils usage (#1528)

Author: yakov-g

custom_model_runner/datarobot_drum/drum/args_parser.py

@@ -6,7 +6,6 @@
 """
 import argparse
 import os
-import subprocess
 import sys
 from urllib.parse import urlparse
 

custom_model_runner/datarobot_drum/drum/common.py

@@ -8,8 +8,8 @@
 import logging
 import os
 import sys
+import trafaret as t
 from contextvars import ContextVar
-from distutils.util import strtobool
 from urllib.parse import urlparse, urlunparse
 
 from contextlib import contextmanager
@@ -126,7 +126,7 @@ def to_bool(value):
         return False
     if isinstance(value, bool):
         return value
-    return strtobool(value)
+    return t.ToBool().check(value)
 
 
 FIT_METADATA_FILENAME = "fit_runtime_data.json"

custom_model_runner/datarobot_drum/drum/data_marshalling.py

@@ -5,7 +5,7 @@
 Released under the terms of DataRobot Tool and Utility Agreement.
 """
 import logging
-from distutils.util import strtobool
+import trafaret as t
 from typing import Any, List, Optional, Union
 
 import numpy as np
@@ -81,9 +81,9 @@ def _standardize(label):
     except ValueError:
         pass
 
-    # Maybe if its a boolean we can make it floaty anyways
+    # Maybe if it's a boolean we can make it floaty anyways
     try:
-        return float(strtobool(label))
+        return float(t.ToBool().check(label))
     except ValueError:
         pass
 

custom_model_runner/datarobot_drum/drum/drum.py

@@ -6,7 +6,6 @@
 """
 import contextlib
 import copy
-import glob
 import json
 import logging
 import os
@@ -17,7 +16,6 @@
 import sys
 import tempfile
 import time
-from distutils.dir_util import copy_tree
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import Callable, Optional
@@ -1211,35 +1209,74 @@ def _generate_runtime_report_file(self, fit_mem_usage: float, pred_mem_usage: fl
         json.dump(report_information, open(output_path, "w"))
 
 
-def output_in_code_dir(code_dir, output_dir):
-    """Does the code directory house the output directory?"""
-    code_abs_path = os.path.abspath(code_dir)
-    output_abs_path = os.path.abspath(output_dir)
-    return os.path.commonpath([code_dir, output_abs_path]) == code_abs_path
+def _output_in_code_dir(code_dir: Path, output_dir: Path) -> bool:
+    """Return True if output_dir is inside code_dir."""
+    try:
+        output_dir.resolve().relative_to(code_dir.resolve())
+        return True
+    except ValueError:
+        return False
 
 
-def create_custom_inference_model_folder(code_dir, output_dir):
-    readme = """
-    This folder was generated by the DRUM tool. It provides functionality for making
-    predictions using the model trained by DRUM
+def _copy_tree(src_dir: Path, dst_dir: Path) -> set[Path]:
     """
-    files_in_output = set(glob.glob(output_dir + "/**"))
-    if output_in_code_dir(code_dir, output_dir):
-        # since the output directory is in the code directory use a tempdir to copy into first and
-        # cleanup files and prevent errors related to copying the output into itself.
-        with tempfile.TemporaryDirectory() as tempdir:
-            copy_tree(code_dir, tempdir)
-            # remove the temporary version of the target dir
-            shutil.rmtree(os.path.join(tempdir, os.path.relpath(output_dir, code_dir)))
-            shutil.rmtree(os.path.join(tempdir, "__pycache__"), ignore_errors=True)
-            copied_files = set(copy_tree(tempdir, output_dir))
+    Recursively copy contents of src_dir into dst_dir.
+    Returns a set of all copied file paths.
+    """
+    copied_files = set()
+    dst_dir.mkdir(parents=True, exist_ok=True)
+
+    for item in src_dir.iterdir():
+        dst_item = dst_dir / item.name
+        if item.is_dir():
+            shutil.copytree(item, dst_item, dirs_exist_ok=True)
+            copied_files.update(p for p in dst_item.rglob("*") if p.is_file())
+        else:
+            shutil.copy2(item, dst_item)
+            copied_files.add(dst_item)
+
+    return copied_files
+
+
+def create_custom_inference_model_folder(code_dir: str, output_dir: str) -> None:
+    """
+    Prepares a model inference folder by copying code_dir into output_dir,
+    avoiding recursive self-copying if output_dir is inside code_dir.
+    """
+    code_path = Path(code_dir).resolve()
+    output_path = Path(output_dir).resolve()
+
+    readme_content = (
+        "This folder was generated by the DRUM tool. It provides functionality for making\n"
+        "predictions using the model trained by DRUM\n"
+    )
+
+    existing_files = set(p for p in output_path.rglob("*") if p.is_file())
+
+    if _output_in_code_dir(code_path, output_path):
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            _copy_tree(code_path, tmp_path)
+
+            # Remove the output subdir inside the copied tree
+            rel_output = output_path.relative_to(code_path)
+            shutil.rmtree(tmp_path / rel_output, ignore_errors=True)
+
+            # Clean up __pycache__
+            shutil.rmtree(tmp_path / "__pycache__", ignore_errors=True)
+
+            copied_files = _copy_tree(tmp_path, output_path)
     else:
-        copied_files = set(copy_tree(code_dir, output_dir))
-        shutil.rmtree(os.path.join(output_dir, "__pycache__"), ignore_errors=True)
-    with open(os.path.join(output_dir, "README.md"), "w") as fp:
-        fp.write(readme)
-    if files_in_output & copied_files:
-        print("Files were overwritten: {}".format(files_in_output & copied_files))
+        copied_files = _copy_tree(code_path, output_path)
+        shutil.rmtree(output_path / "__pycache__", ignore_errors=True)
+
+    # Add README
+    (output_path / "README.md").write_text(readme_content)
+
+    # Check overwritten files
+    overwritten = existing_files & copied_files
+    if overwritten:
+        print(f"Files were overwritten: {sorted(overwritten)}")
 
 
 def _get_default_numeric_param_value(param_config: Dict, cast_to_int: bool) -> Union[int, float]:

tests/unit/datarobot_drum/drum/test_drum.py

@@ -8,6 +8,7 @@
 
 import os
 import sys
+import tempfile
 from copy import deepcopy
 from pathlib import Path
 from tempfile import TemporaryDirectory, NamedTemporaryFile
@@ -23,7 +24,7 @@
 from datarobot_drum.drum.drum import (
     CMRunner,
     create_custom_inference_model_folder,
-    output_in_code_dir,
+    _output_in_code_dir,
 )
 from datarobot_drum.drum.enum import (
     RunLanguage,
@@ -498,12 +499,49 @@ def test_output_dir_copy(self):
             assert not Path(out_dir, "__pycache__").exists()
             assert not Path(out_dir, "out").exists()
 
-    def test_output_in_code_dir(self):
-        code_dir = "/test/code/is/here"
-        output_other = "/test/not/code"
-        output_code_dir = "/test/code/is/here/output"
-        assert not output_in_code_dir(code_dir, output_other)
-        assert output_in_code_dir(code_dir, output_code_dir)
+
+class TestUtilityFunctionsOutputInCommonDir:
+    def test_output_inside_code_dir(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            code_dir = Path(tmp)
+            output_dir = code_dir / "subdir"
+            output_dir.mkdir()
+            assert _output_in_code_dir(code_dir, output_dir) is True
+
+    def test_output_is_code_dir(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp)
+            assert _output_in_code_dir(path, path) is True
+
+    def test_output_outside_code_dir(self):
+        with tempfile.TemporaryDirectory() as tmp1, tempfile.TemporaryDirectory() as tmp2:
+            code_dir = Path(tmp1)
+            output_dir = Path(tmp2)
+            assert _output_in_code_dir(code_dir, output_dir) is False
+
+    def test_output_sibling_directory(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            parent = Path(tmp)
+            code_dir = parent / "code"
+            output_dir = parent / "output"
+            code_dir.mkdir()
+            output_dir.mkdir()
+            assert _output_in_code_dir(code_dir, output_dir) is False
+
+    def test_relative_path_handling(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            base = Path(tmp)
+            code_dir = base / "code"
+            output_dir = code_dir / "out"
+            output_dir.mkdir(parents=True)
+
+            # Use relative paths
+            cwd = Path.cwd()
+            try:
+                os.chdir(base)
+                assert _output_in_code_dir(Path("code"), Path("code/out")) is True
+            finally:
+                os.chdir(cwd)
 
 
 class TestRuntimeParametersDockerCommand: