viettelcloud-aiplatform 0.3.0__py3-none-any.whl

viettelcloud/aiplatform/trainer/backends/localprocess/backend_test.py

@@ -0,0 +1,501 @@
+# Copyright 2025 The Kubeflow Authors.
+#
+# Licensed 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.
+
+"""
+Unit tests for the LocalProcessBackend class in the Kubeflow Trainer SDK.
+"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+
+from viettelcloud.aiplatform.trainer.backends.localprocess.backend import LocalProcessBackend
+from viettelcloud.aiplatform.trainer.backends.localprocess.constants import LOCAL_RUNTIME_IMAGE
+from viettelcloud.aiplatform.trainer.backends.localprocess.types import (
+    LocalProcessBackendConfig,
+    LocalRuntimeTrainer,
+)
+from viettelcloud.aiplatform.trainer.constants import constants
+from viettelcloud.aiplatform.trainer.options import (
+    Annotations,
+    Labels,
+    Name,
+    PodTemplateOverride,
+    PodTemplateOverrides,
+)
+from viettelcloud.aiplatform.trainer.test.common import FAILED, SUCCESS, TestCase
+from viettelcloud.aiplatform.trainer.types import types
+
+# Test constants
+TORCH_RUNTIME = constants.DEFAULT_TRAINING_RUNTIME
+BASIC_TRAIN_JOB_NAME = "test-job"
+
+
+def dummy_training_function():
+    """Dummy training function for testing."""
+    print("Training started")
+    return {"loss": 0.5, "accuracy": 0.95}
+
+
+@pytest.fixture
+def local_backend():
+    """Create LocalProcessBackend for testing."""
+    cfg = LocalProcessBackendConfig()
+    backend = LocalProcessBackend(cfg)
+    yield backend
+    # Cleanup: Clear jobs to prevent test pollution
+    backend._LocalProcessBackend__local_jobs.clear()
+
+
+@pytest.fixture
+def mock_train_environment():
+    """Mock the training environment to avoid actual subprocess execution."""
+    with (
+        patch(
+            "viettelcloud.aiplatform.trainer.backends.localprocess.job.LocalJob.start"
+        ) as mock_start,
+        patch(
+            "viettelcloud.aiplatform.trainer.backends.localprocess.utils.get_local_runtime_trainer"
+        ) as mock_get_trainer,
+        patch(
+            "viettelcloud.aiplatform.trainer.backends.localprocess.utils.get_local_train_job_script"
+        ) as mock_get_script,
+        patch("tempfile.mkdtemp") as mock_mkdtemp,
+    ):
+        # Setup mock return values
+        mock_mkdtemp.return_value = "/tmp/test-venv"
+        mock_get_script.return_value = ["/bin/bash", "-c", "echo 'training'"]
+
+        mock_trainer = LocalRuntimeTrainer(
+            trainer_type=types.TrainerType.CUSTOM_TRAINER,
+            framework="torch",
+            num_nodes=1,
+            device_count="1",
+            device="cpu",
+            packages=["torch"],
+            image=LOCAL_RUNTIME_IMAGE,
+        )
+        mock_trainer.set_command = Mock()
+        mock_get_trainer.return_value = mock_trainer
+
+        yield {
+            "start": mock_start,
+            "get_trainer": mock_get_trainer,
+            "get_script": mock_get_script,
+            "mkdtemp": mock_mkdtemp,
+        }
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="list_all_local_runtimes",
+            expected_status=SUCCESS,
+            config={},
+        ),
+    ],
+)
+def test_list_runtimes(local_backend, test_case):
+    """Test LocalProcessBackend.list_runtimes()."""
+    runtimes = local_backend.list_runtimes()
+    assert len(runtimes) > 0
+    assert all(isinstance(rt, types.Runtime) for rt in runtimes)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="get_existing_runtime",
+            expected_status=SUCCESS,
+            config={"runtime_name": TORCH_RUNTIME},
+        ),
+        TestCase(
+            name="get_nonexistent_runtime",
+            expected_status=FAILED,
+            config={"runtime_name": "nonexistent-runtime"},
+            expected_error=ValueError,
+        ),
+    ],
+)
+def test_get_runtime(local_backend, test_case):
+    """Test LocalProcessBackend.get_runtime()."""
+    runtime_name = test_case.config.get("runtime_name")
+
+    if test_case.expected_status == FAILED:
+        with pytest.raises(test_case.expected_error):
+            local_backend.get_runtime(runtime_name)
+    else:
+        runtime = local_backend.get_runtime(runtime_name)
+        assert runtime is not None
+        assert runtime.name == runtime_name
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="get_packages_for_existing_runtime",
+            expected_status=SUCCESS,
+            config={
+                "runtime": types.Runtime(
+                    name=TORCH_RUNTIME,
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+            },
+        ),
+        TestCase(
+            name="get_packages_for_nonexistent_runtime",
+            expected_status=FAILED,
+            config={
+                "runtime": types.Runtime(
+                    name="nonexistent-runtime",
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+            },
+            expected_error=ValueError,
+        ),
+    ],
+)
+def test_get_runtime_packages(local_backend, test_case):
+    """Test LocalProcessBackend.get_runtime_packages()."""
+    runtime = test_case.config.get("runtime")
+
+    if test_case.expected_status == FAILED:
+        with pytest.raises(test_case.expected_error):
+            local_backend.get_runtime_packages(runtime)
+    else:
+        packages = local_backend.get_runtime_packages(runtime)
+        assert packages is not None
+        assert isinstance(packages, list)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="train with basic custom trainer - no options",
+            expected_status=SUCCESS,
+            config={
+                "runtime": types.Runtime(
+                    name=TORCH_RUNTIME,
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+                "trainer": types.CustomTrainer(
+                    func=dummy_training_function,
+                    packages_to_install=["numpy", "torch"],
+                ),
+                "options": [],
+            },
+        ),
+        TestCase(
+            name="train with custom trainer and environment variables",
+            expected_status=SUCCESS,
+            config={
+                "runtime": types.Runtime(
+                    name=TORCH_RUNTIME,
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+                "trainer": types.CustomTrainer(
+                    func=dummy_training_function,
+                    packages_to_install=["torch"],
+                    env={"CUDA_VISIBLE_DEVICES": "0", "OMP_NUM_THREADS": "4"},
+                ),
+                "options": [],
+            },
+        ),
+        TestCase(
+            name="train rejects kubernetes labels option",
+            expected_status=FAILED,
+            config={
+                "runtime": types.Runtime(
+                    name=TORCH_RUNTIME,
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+                "trainer": types.CustomTrainer(
+                    func=dummy_training_function,
+                ),
+                "options": [Labels({"app": "test"})],
+            },
+            expected_error=ValueError,
+        ),
+        TestCase(
+            name="train rejects kubernetes annotations option",
+            expected_status=FAILED,
+            config={
+                "runtime": types.Runtime(
+                    name=TORCH_RUNTIME,
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+                "trainer": types.CustomTrainer(
+                    func=dummy_training_function,
+                ),
+                "options": [Annotations({"description": "test"})],
+            },
+            expected_error=ValueError,
+        ),
+        TestCase(
+            name="train rejects pod template overrides option",
+            expected_status=FAILED,
+            config={
+                "runtime": types.Runtime(
+                    name=TORCH_RUNTIME,
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+                "trainer": types.CustomTrainer(
+                    func=dummy_training_function,
+                ),
+                "options": [
+                    PodTemplateOverrides(
+                        PodTemplateOverride(
+                            target_jobs=["node"],
+                        )
+                    )
+                ],
+            },
+            expected_error=ValueError,
+        ),
+        TestCase(
+            name="train fails without runtime",
+            expected_status=FAILED,
+            config={
+                "runtime": None,
+                "trainer": types.CustomTrainer(
+                    func=dummy_training_function,
+                ),
+                "options": [],
+            },
+            expected_error=ValueError,
+        ),
+        TestCase(
+            name="train fails without custom trainer",
+            expected_status=FAILED,
+            config={
+                "runtime": types.Runtime(
+                    name=TORCH_RUNTIME,
+                    trainer=types.RuntimeTrainer(
+                        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+                        framework="torch",
+                        num_nodes=1,
+                        image=LOCAL_RUNTIME_IMAGE,
+                    ),
+                ),
+                "trainer": None,
+            },
+            expected_error=ValueError,
+        ),
+    ],
+)
+def test_train(local_backend, mock_train_environment, test_case):
+    """Test LocalProcessBackend.train() with success and failure cases."""
+    runtime = test_case.config.get("runtime")
+    trainer = test_case.config.get("trainer")
+    options = test_case.config.get("options", [])
+
+    mocks = mock_train_environment
+
+    if test_case.expected_status == FAILED:
+        with pytest.raises(test_case.expected_error) as exc_info:
+            local_backend.train(
+                runtime=runtime,
+                trainer=trainer,
+                options=options,
+            )
+
+        # Verify specific error messages
+        error_msg = str(exc_info.value)
+        if "rejects kubernetes" in test_case.name:
+            assert "not compatible with" in error_msg
+        elif "without runtime" in test_case.name:
+            assert "Runtime must be provided" in error_msg
+        elif "without custom trainer" in test_case.name:
+            assert "CustomTrainer must be set" in error_msg
+    else:
+        train_job_name = local_backend.train(
+            runtime=runtime,
+            trainer=trainer,
+            options=options,
+        )
+
+        assert train_job_name is not None
+        assert len(train_job_name) > 0
+        mocks["start"].assert_called_once()
+        mocks["get_trainer"].assert_called_once()
+        mocks["get_script"].assert_called_once()
+
+        # Verify job is tracked
+        jobs = local_backend.list_jobs(runtime=runtime)
+        assert any(job.name == train_job_name for job in jobs)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="get_nonexistent_job",
+            expected_status=FAILED,
+            config={"job_name": "nonexistent-job"},
+            expected_error=ValueError,
+        ),
+    ],
+)
+def test_get_job(local_backend, test_case):
+    """Test LocalProcessBackend.get_job()."""
+    job_name = test_case.config.get("job_name")
+
+    if test_case.expected_status == FAILED:
+        with pytest.raises(test_case.expected_error):
+            local_backend.get_job(job_name)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="list_jobs_empty",
+            expected_status=SUCCESS,
+            config={"runtime": None},
+        ),
+    ],
+)
+def test_list_jobs(local_backend, test_case):
+    """Test LocalProcessBackend.list_jobs()."""
+    runtime = test_case.config.get("runtime")
+    jobs = local_backend.list_jobs(runtime=runtime)
+    assert isinstance(jobs, list)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="get_logs_nonexistent_job",
+            expected_status=FAILED,
+            config={"job_name": "nonexistent-job", "step": "train"},
+            expected_error=ValueError,
+        ),
+    ],
+)
+def test_get_job_logs(local_backend, test_case):
+    """Test LocalProcessBackend.get_job_logs()."""
+    job_name = test_case.config.get("job_name")
+    step = test_case.config.get("step", "train")
+
+    if test_case.expected_status == FAILED:
+        with pytest.raises(test_case.expected_error):
+            list(local_backend.get_job_logs(job_name, step=step))
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="wait_for_nonexistent_job",
+            expected_status=FAILED,
+            config={"job_name": "nonexistent-job"},
+            expected_error=ValueError,
+        ),
+    ],
+)
+def test_wait_for_job_status(local_backend, test_case):
+    """Test LocalProcessBackend.wait_for_job_status()."""
+    job_name = test_case.config.get("job_name")
+
+    if test_case.expected_status == FAILED:
+        with pytest.raises(test_case.expected_error):
+            local_backend.wait_for_job_status(job_name)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="delete_nonexistent_job",
+            expected_status=FAILED,
+            config={"job_name": "nonexistent-job"},
+            expected_error=ValueError,
+        ),
+    ],
+)
+def test_delete_job(local_backend, test_case):
+    """Test LocalProcessBackend.delete_job()."""
+    job_name = test_case.config.get("job_name")
+
+    if test_case.expected_status == FAILED:
+        with pytest.raises(test_case.expected_error):
+            local_backend.delete_job(job_name)
+
+
+def test_name_option_sets_job_name(local_backend, mock_train_environment):
+    """Test that Name option sets the custom job name."""
+    custom_name = "my-custom-job-name"
+
+    def dummy_func():
+        pass
+
+    runtime = types.Runtime(
+        name=TORCH_RUNTIME,
+        trainer=types.RuntimeTrainer(
+            trainer_type=types.TrainerType.CUSTOM_TRAINER,
+            framework="torch",
+            image=LOCAL_RUNTIME_IMAGE,
+        ),
+    )
+
+    trainer = types.CustomTrainer(func=dummy_func)
+    options = [Name(name=custom_name)]
+
+    job_name = local_backend.train(
+        runtime=runtime,
+        trainer=trainer,
+        options=options,
+    )
+
+    assert job_name == custom_name

viettelcloud/aiplatform/trainer/backends/localprocess/constants.py

@@ -0,0 +1,90 @@
+# Copyright 2025 The Kubeflow Authors.
+#
+# Licensed 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.
+
+import re
+import textwrap
+
+import viettelcloud.aiplatform.common.constants as common_constants
+from viettelcloud.aiplatform.trainer.backends.localprocess import types
+from viettelcloud.aiplatform.trainer.constants import constants
+from viettelcloud.aiplatform.trainer.types import types as base_types
+
+TORCH_FRAMEWORK_TYPE = "torch"
+
+# Image name for the local runtime.
+LOCAL_RUNTIME_IMAGE = "local"
+
+local_runtimes = [
+    base_types.Runtime(
+        name=constants.DEFAULT_TRAINING_RUNTIME,
+        trainer=types.LocalRuntimeTrainer(
+            trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
+            framework=TORCH_FRAMEWORK_TYPE,
+            num_nodes=1,
+            device_count=common_constants.UNKNOWN,
+            device=common_constants.UNKNOWN,
+            packages=["torch"],
+            image=LOCAL_RUNTIME_IMAGE,
+        ),
+    )
+]
+
+
+# Create venv script
+
+
+# The exec script to embed training function into container command.
+DEPENDENCIES_SCRIPT = textwrap.dedent(
+    """
+        PIP_DISABLE_PIP_VERSION_CHECK=1 pip install $QUIET \
+    --no-warn-script-location $PIP_INDEX $PACKAGE_STR
+    """
+)
+
+# activate virtualenv, then run the entrypoint from the virtualenv bin
+LOCAL_EXEC_ENTRYPOINT = textwrap.dedent(
+    """
+    $ENTRYPOINT "$FUNC_FILE" "$PARAMETERS"
+    """
+)
+
+TORCH_COMMAND = "torchrun"
+
+# default command, will run from within the virtualenv
+DEFAULT_COMMAND = "python"
+
+# remove virtualenv after training is completed.
+LOCAL_EXEC_JOB_CLEANUP_SCRIPT = textwrap.dedent(
+    """
+    rm -rf $PYENV_LOCATION
+    """
+)
+
+
+LOCAL_EXEC_JOB_TEMPLATE = textwrap.dedent(
+    """
+    set -e
+    $OS_PYTHON_BIN -m venv --without-pip $PYENV_LOCATION
+    echo "Operating inside $PYENV_LOCATION"
+    source $PYENV_LOCATION/bin/activate
+    $PYENV_LOCATION/bin/python -m ensurepip --upgrade --default-pip
+    $DEPENDENCIES_SCRIPT
+    $ENTRYPOINT
+    $CLEANUP_SCRIPT
+    """
+)
+
+LOCAL_EXEC_FILENAME = "train_{}.py"
+
+PYTHON_PACKAGE_NAME_RE = re.compile(r"^\s*([A-Za-z0-9][A-Za-z0-9._-]*)")