chore: import upstream snapshot with attribution

tests/python/support/test_ccache.py

@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test support.cc with ccache"""
+
+import os
+import shutil
+import tempfile
+
+import pytest
+
+import tvm
+from tvm.support.cc import _is_linux_like, _is_windows_like, create_executable, create_shared
+
+
+def _src_gen(text):
+    return """
+#include <iostream>
+
+int main() {
+    std::cout << "text";
+    return 0;
+}""".replace("text", text)
+
+
+def _compile(f_create, text, output):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        src_path = os.path.join(temp_dir, "src.cpp")
+        with open(src_path, "w", encoding="utf-8") as file:
+            file.write(_src_gen(text))
+        log_path = os.path.join(temp_dir, "log.txt")
+        ccache_env = {
+            "CCACHE_COMPILERCHECK": "content",
+            "CCACHE_LOGFILE": log_path,
+        }
+        f_create(output, ["src.cpp"], ["-c"], cwd=temp_dir, ccache_env=ccache_env)
+        with open(log_path, encoding="utf-8") as file:
+            log = file.read()
+        return log
+
+
+@pytest.mark.skipif(shutil.which("ccache") is None, reason="ccache not installed")
+def test_shared():
+    if _is_linux_like():
+        _ = _compile(create_shared, "shared", "main.o")
+        log = _compile(create_shared, "shared", "main.o")
+        assert "Succeeded getting cached result" in log
+    elif _is_windows_like():
+        _ = _compile(create_shared, "shared", "main.obj")
+        log = _compile(create_shared, "shared", "main.obj")
+        assert "Succeeded getting cached result" in log
+
+
+@pytest.mark.skipif(shutil.which("ccache") is None, reason="ccache not installed")
+def test_executable():
+    if _is_linux_like():
+        _ = _compile(create_executable, "executable", "main")
+        log = _compile(create_executable, "executable", "main")
+        assert "Succeeded getting cached result" in log
+    elif _is_windows_like():
+        _ = _compile(create_executable, "executable", "main.exe")
+        log = _compile(create_executable, "executable", "main.exe")
+        assert "Succeeded getting cached result" in log
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/support/test_popen_pool.py

@@ -0,0 +1,197 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F401, F841
+"""Test PopenPoolExecutor."""
+
+import os
+import time
+
+import psutil
+import pytest
+
+pytest.importorskip("cloudpickle")  # PopenWorker.send requires cloudpickle
+
+from tvm.support.popen_pool import PopenPoolExecutor, PopenWorker
+from tvm.testing import (
+    identity_after,
+    terminate_self,
+)
+from tvm.testing.popen_pool import (
+    after_initializer,
+    call_cpp_ffi,
+    call_cpp_py_ffi,
+    call_py_ffi,
+    initializer,
+    register_ffi,
+    timeout_job,
+)
+
+
+def test_popen_worker():
+    proc = PopenWorker()
+
+    with pytest.raises(TimeoutError):
+        proc.send(identity_after, [1, 100], timeout=0.01)
+        proc.recv()
+
+    with pytest.raises(ChildProcessError):
+        proc.send(terminate_self)
+        proc.recv()
+
+    proc.send(identity_after, [2, 0])
+    assert proc.recv() == 2
+
+    proc.send(identity_after, [4, 0.0001])
+    assert proc.recv() == 4
+
+
+def test_popen_worker_reuses():
+    proc = PopenWorker(maximum_uses=None)
+
+    proc.send(os.getpid)
+    initial_pid = proc.recv()
+
+    proc.send(os.getpid)
+    assert proc.recv() == initial_pid
+
+
+def test_popen_worker_recycles():
+    proc = PopenWorker(maximum_uses=2)
+
+    proc.send(os.getpid)
+    initial_pid = proc.recv()
+    assert psutil.pid_exists(initial_pid)
+
+    proc.send(os.getpid)
+    assert proc.recv() == initial_pid
+    assert psutil.pid_exists(initial_pid)
+
+    proc.send(os.getpid)
+    assert proc.recv() != initial_pid
+    assert not psutil.pid_exists(initial_pid)
+
+
+def test_popen_pool_executor():
+    import tvm_ffi
+
+    import tvm
+
+    pool = PopenPoolExecutor(max_workers=2, timeout=0.015)
+    value1 = pool.submit(identity_after, 1, 100)
+    value2 = pool.submit(terminate_self)
+    value3 = pool.submit(identity_after, 3, 0)
+    value4 = pool.submit(tvm_ffi.core.String, "xyz")
+
+    with pytest.raises(TimeoutError):
+        value1.result()
+
+    with pytest.raises(ChildProcessError):
+        value2.result()
+
+    assert value3.result() == 3
+    value = value4.result()
+    assert value == "xyz"
+
+    pool = PopenPoolExecutor(max_workers=4, timeout=None)
+    values = pool.map_with_error_catching(lambda x: x, range(100))
+
+    for idx, val in enumerate(values):
+        assert val.value == idx
+
+
+def test_popen_initializer():
+    initargs = [1, 2, 3]
+    proc = PopenWorker(initializer=initializer, initargs=initargs)
+    proc.send(after_initializer)
+    test_global_state_1, test_global_state_2, test_global_state_3 = proc.recv()
+    assert test_global_state_1 == initargs[0]
+    assert test_global_state_2 == initargs[1]
+    assert test_global_state_3 == initargs[2]
+
+
+def test_popen_worker_recycles_with_initializer():
+    initargs = [1, 2, 3]
+    proc = PopenWorker(initializer=initializer, initargs=initargs, maximum_uses=3)
+
+    proc.send(os.getpid)
+    initial_pid = proc.recv()
+
+    proc.send(after_initializer)
+    assert list(proc.recv()) == initargs
+
+    proc.send(os.getpid)
+    assert proc.recv() == initial_pid
+
+    # The process should be recycled with this send.
+    proc.send(os.getpid)
+    assert proc.recv() != initial_pid
+
+    # But the initializer should've run this time as well.
+    proc.send(after_initializer)
+    assert list(proc.recv()) == initargs
+
+
+def test_popen_ffi():
+    proc = PopenWorker(register_ffi)
+
+    # call python function via ffi
+    initargs = [0]
+    proc.send(call_py_ffi, initargs)
+    assert proc.recv() == initargs[0]
+
+    # call cpp function (testing.echo) via ffi
+    initargs = [1]
+    proc.send(call_cpp_ffi, initargs)
+    assert proc.recv() == initargs[0]
+
+    # call python function from ffi registry via cross-language dispatch
+    initargs = [2]
+    proc.send(call_cpp_py_ffi, initargs)
+    assert proc.recv() == initargs[0]
+
+
+def test_popen_pool_executor_timeout():
+    timeout = 0.5
+
+    pool = PopenPoolExecutor(timeout=timeout)
+
+    f1 = pool.submit(timeout_job, timeout)
+    while not f1.done():
+        pass
+    try:
+        res = f1.result()
+    except Exception as ex:
+        assert isinstance(ex, TimeoutError)
+
+
+def test_popen_pool_executor_recycles():
+    pool = PopenPoolExecutor(max_workers=1, timeout=None, maximum_process_uses=2)
+
+    initial_pid = pool.submit(os.getpid).result()
+    assert initial_pid == pool.submit(os.getpid).result()
+    assert initial_pid != pool.submit(os.getpid).result()
+
+
+if __name__ == "__main__":
+    test_popen_worker()
+    test_popen_worker_recycles()
+    test_popen_pool_executor()
+    test_popen_initializer()
+    test_popen_worker_recycles_with_initializer()
+    test_popen_ffi()
+    test_popen_pool_executor_timeout()
+    test_popen_pool_executor_recycles()

tests/python/support/test_util.py

@@ -0,0 +1,90 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for functions in tvm/python/tvm/contrib/util.py."""
+
+import datetime
+import os
+import shutil
+
+from tvm.support import utils
+
+
+def validate_debug_dir_path(temp_dir, expected_basename):
+    """Validate the dir path of debugging"""
+    dirname, basename = os.path.split(temp_dir.temp_dir)
+    assert basename == expected_basename, f"unexpected basename: {basename}"
+
+    parent_dir = os.path.basename(dirname)
+    create_time = datetime.datetime.strptime(parent_dir.split("___", 1)[0], "%Y-%m-%dT%H-%M-%S")
+    assert abs(datetime.datetime.now() - create_time) < datetime.timedelta(seconds=60)
+
+
+def test_tempdir():
+    """Tests for temporary dir"""
+    assert utils.TempDirectory._KEEP_FOR_DEBUG is False, "don't submit with KEEP_FOR_DEBUG == True"
+
+    temp_dir = utils.tempdir()
+    assert os.path.exists(temp_dir.temp_dir)
+
+    old_debug_mode = utils.TempDirectory._KEEP_FOR_DEBUG
+    old_tempdirs = utils.TempDirectory.TEMPDIRS
+    try:
+        for temp_dir_number in range(0, 3):
+            with utils.TempDirectory.set_keep_for_debug():
+                debug_temp_dir = utils.tempdir()
+                try:
+                    validate_debug_dir_path(debug_temp_dir, "0000" + str(temp_dir_number))
+                finally:
+                    shutil.rmtree(debug_temp_dir.temp_dir)
+
+        with utils.TempDirectory.set_keep_for_debug():
+            # Create 2 temp_dir within the same session.
+            debug_temp_dir = utils.tempdir()
+            try:
+                validate_debug_dir_path(debug_temp_dir, "00003")
+            finally:
+                shutil.rmtree(debug_temp_dir.temp_dir)
+
+            debug_temp_dir = utils.tempdir()
+            try:
+                validate_debug_dir_path(debug_temp_dir, "00004")
+            finally:
+                shutil.rmtree(debug_temp_dir.temp_dir)
+
+            with utils.TempDirectory.set_keep_for_debug(False):
+                debug_temp_dir = utils.tempdir()  # This one should get deleted.
+
+                # Simulate atexit hook
+                utils.TempDirectory.remove_tempdirs()
+
+                # Calling twice should be a no-op.
+                utils.TempDirectory.remove_tempdirs()
+
+                # Creating a new TempDirectory should fail now
+                try:
+                    utils.tempdir()
+                    assert False, "creation should fail"
+                except utils.DirectoryCreatedPastAtExit:
+                    pass
+
+    finally:
+        utils.TempDirectory.DEBUG_MODE = old_debug_mode
+        utils.TempDirectory.TEMPDIRS = old_tempdirs
+
+
+if __name__ == "__main__":
+    test_tempdir()