viettelcloud-aiplatform 0.3.0__py3-none-any.whl

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

@@ -0,0 +1,184 @@
+# 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.
+from datetime import datetime
+import logging
+import os
+import subprocess
+import threading
+from typing import Union
+
+from viettelcloud.aiplatform.trainer.constants import constants
+
+logger = logging.getLogger(__name__)
+
+
+class LocalJob(threading.Thread):
+    def __init__(
+        self,
+        name,
+        command: Union[list, tuple[str], str],
+        execution_dir: str = None,
+        env: dict[str, str] = None,
+        dependencies: list = None,
+    ):
+        """Creates a LocalJob.
+
+        Creates a local subprocess with threading to allow users to create background jobs.
+
+        Args:
+            name (str): The name of the job.
+            command (str): The command to run.
+            execution_dir (str): The execution directory.
+            env (Dict[str, str], optional): Environment variables. Defaults to None.
+            dependencies (List[str], optional): List of dependencies. Defaults to None.
+        """
+        super().__init__()
+        self.name = name
+        self.command = command
+        self._stdout = ""
+        self._returncode = None
+        self._success = False
+        self._status = constants.TRAINJOB_CREATED
+        self._lock = threading.Lock()
+        self._process = None
+        self._output_updated = threading.Event()
+        self._cancel_requested = threading.Event()
+        self._start_time = None
+        self._end_time = None
+        self.env = env or {}
+        self.dependencies = dependencies or []
+        self.execution_dir = execution_dir or os.getcwd()
+
+    def run(self):
+        for dep in self.dependencies:
+            dep.join()
+            if not dep.success:
+                with self._lock:
+                    self._stdout = f"Dependency {dep.name} failed. Skipping"
+                return
+
+        current_dir = os.getcwd()
+        try:
+            self._start_time = datetime.now()
+            _c = " ".join(self.command)
+            logger.debug(f"[{self.name}] Started at {self._start_time} with command: \n {_c}")
+
+            # change working directory to venv before executing script
+            os.chdir(self.execution_dir)
+
+            self._process = subprocess.Popen(
+                self.command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                encoding="utf-8",
+                bufsize=1,
+                env=self.env,
+            )
+            # set job status
+            self._status = constants.TRAINJOB_RUNNING
+
+            while True:
+                if self._cancel_requested.is_set():
+                    self._process.terminate()
+                    self._stdout += "[JobCancelled]\n"
+                    self._status = constants.TRAINJOB_FAILED
+                    self._success = False
+                    return
+
+                # Read output line by line (for streaming)
+                output_line = self._process.stdout.readline()
+                with self._lock:
+                    if output_line:
+                        self._stdout += output_line
+                        self._output_updated.set()
+
+                if not output_line and self._process.poll() is not None:
+                    break
+
+            self._process.stdout.close()
+            self._returncode = self._process.wait()
+            self._end_time = datetime.now()
+            self._success = self._process.returncode == 0
+            msg = (
+                f"[{self.name}] Completed with code {self._returncode}"
+                f" in {self._end_time - self._start_time} seconds."
+            )
+            # set status based on success or failure
+            self._status = (
+                constants.TRAINJOB_COMPLETE if self._success else (constants.TRAINJOB_FAILED)
+            )
+            self._stdout += msg
+            logger.debug("Job output: ", self._stdout)
+
+        except Exception as e:
+            with self._lock:
+                self._stdout += f"Exception: {e}\n"
+                self._success = False
+                self._status = constants.TRAINJOB_FAILED
+        finally:
+            os.chdir(current_dir)
+
+    @property
+    def stdout(self):
+        with self._lock:
+            return self._stdout
+
+    @property
+    def success(self):
+        return self._success
+
+    @property
+    def status(self):
+        return self._status
+
+    def cancel(self):
+        self._cancel_requested.set()
+
+    @property
+    def returncode(self):
+        return self._returncode
+
+    def logs(self, follow=False) -> list[str]:
+        if not follow:
+            return self._stdout.splitlines()
+
+        try:
+            for chunk in self.stream_logs():
+                print(chunk, end="", flush=True)  # stream to console live
+        except StopIteration:
+            pass
+
+        return self._stdout.splitlines()
+
+    def stream_logs(self):
+        """Generator that yields new output lines as they come in."""
+        last_index = 0
+        while self.is_alive() or last_index < len(self._stdout):
+            self._output_updated.wait(timeout=1)
+            with self._lock:
+                data = self._stdout
+                new_data = data[last_index:]
+                last_index = len(data)
+                self._output_updated.clear()
+            if new_data:
+                yield new_data
+
+    @property
+    def creation_time(self):
+        return self._start_time
+
+    @property
+    def completion_time(self):
+        return self._end_time

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

@@ -0,0 +1,52 @@
+# 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.
+
+
+from dataclasses import dataclass, field
+from datetime import datetime
+import typing
+from typing import Optional
+
+from pydantic import BaseModel
+
+from viettelcloud.aiplatform.trainer.backends.localprocess.job import LocalJob
+from viettelcloud.aiplatform.trainer.types import types
+
+
+class LocalProcessBackendConfig(BaseModel):
+    cleanup_venv: bool = True
+
+
+@dataclass
+class LocalRuntimeTrainer(types.RuntimeTrainer):
+    packages: list[str] = field(default_factory=list)
+
+
+class LocalBackendStep(BaseModel):
+    step_name: str
+    job: LocalJob
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class LocalBackendJobs(BaseModel):
+    steps: Optional[list[LocalBackendStep]] = []
+    runtime: Optional[types.Runtime] = None
+    name: str
+    created: typing.Optional[datetime] = None
+    completed: typing.Optional[datetime] = None
+
+    class Config:
+        arbitrary_types_allowed = True

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

@@ -0,0 +1,302 @@
+import inspect
+import os
+from pathlib import Path
+import re
+import shutil
+from string import Template
+import textwrap
+from typing import Any, Callable, Optional
+
+from viettelcloud.aiplatform.trainer.backends.localprocess import constants as local_exec_constants
+from viettelcloud.aiplatform.trainer.backends.localprocess.types import LocalRuntimeTrainer
+from viettelcloud.aiplatform.trainer.constants import constants
+from viettelcloud.aiplatform.trainer.types import types
+
+
+def _extract_name(requirement: str) -> str:
+    """
+    Extract the base distribution name from a requirement string without external deps.
+
+    Supports common PEP 508 patterns:
+      - 'package'
+      - 'package[extra1,extra2]'
+      - 'package==1.2.3', 'package>=1.0', 'package~=1.4', etc.
+      - 'package @ https://...'
+      - markers after ';' are irrelevant for name extraction.
+
+    Returns the *raw* (un-normalized) name as it appears.
+    Raises ValueError if a name cannot be parsed.
+    """
+    if requirement is None:
+        raise ValueError("Requirement string cannot be None")
+    s = requirement.strip()
+    if not s:
+        raise ValueError("Empty requirement string")
+
+    m = local_exec_constants.PYTHON_PACKAGE_NAME_RE.match(s)
+    if not m:
+        raise ValueError(f"Could not parse package name from requirement: {requirement!r}")
+    return m.group(1)
+
+
+def _canonicalize_name(name: str) -> str:
+    """
+    PEP 503-style normalization: case-insensitive, and collapse runs of -, _, . into '-'.
+    """
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def get_install_packages(
+    runtime_packages: list[str],
+    trainer_packages: Optional[list[str]] = None,
+) -> list[str]:
+    """
+    Merge two requirement lists into a single list of strings.
+
+    Rules implemented:
+    1) If a package appears in trainer_packages, it overwrites the one in runtime_packages.
+       We keep the *trainer string verbatim* (specifier, markers, extras, spacing).
+    2) Case-insensitive matching of package names (PEP 503-style normalization).
+    3) Output is a list of strings.
+    4) If trainer_packages contains the same dependency multiple times (case-insensitive),
+       raise ValueError.
+    5) If runtime_packages contains duplicates, the last one among *runtime* wins there
+       (no error), but any trainer entry still overwrites it. Runtime packages shouldn't
+        have any duplicates.
+    6) Ordering: keep runtime-only packages in their original order (emitting only their
+       last occurrence), then append all trainer packages in their original order.
+    """
+    if not trainer_packages:
+        return runtime_packages
+
+    # --- Parse + normalize runtime ---
+    runtime_parsed: list[tuple[str, str]] = []  # (orig, canonical_name)
+    last_runtime_index_by_name: dict[str, int] = {}
+
+    for i, orig in enumerate(runtime_packages):
+        raw_name = _extract_name(orig)
+        canon = _canonicalize_name(raw_name)
+        runtime_parsed.append((orig, canon))
+        last_runtime_index_by_name[canon] = i  # last occurrence index wins among runtime
+
+    # --- Parse + validate trainer (detect duplicates) ---
+    trainer_parsed: list[tuple[str, str]] = []
+    seen_trainer: set[str] = set()
+    for orig in trainer_packages:
+        raw_name = _extract_name(orig)
+        canon = _canonicalize_name(raw_name)
+        if canon in seen_trainer:
+            raise ValueError(
+                f"Duplicate dependency in trainer_packages: '{raw_name}' (canonical: '{canon}')"
+            )
+        seen_trainer.add(canon)
+        trainer_parsed.append((orig, canon))
+
+    trainer_names: set[str] = {canon for _, canon in trainer_parsed}
+
+    # --- Build merged list respecting order semantics ---
+    merged: list[str] = []
+
+    # 1) Runtime-only packages (only emit the last occurrence for each name)
+    emitted_runtime_names: set[str] = set()
+    for idx, (orig, canon) in enumerate(runtime_parsed):
+        if canon in trainer_names:
+            continue  # overwritten by trainer
+        if last_runtime_index_by_name[canon] == idx and canon not in emitted_runtime_names:
+            merged.append(orig)
+            emitted_runtime_names.add(canon)
+
+    # 2) Trainer packages (overwrite and preserve trainer's exact strings, original order)
+    for orig, _ in trainer_parsed:
+        merged.append(orig)
+
+    return merged
+
+
+def get_local_runtime_trainer(
+    runtime_name: str,
+    venv_dir: str,
+    framework: str,
+) -> LocalRuntimeTrainer:
+    """
+    Get the LocalRuntimeTrainer object.
+    """
+    local_runtime = next(
+        (rt for rt in local_exec_constants.local_runtimes if rt.name == runtime_name),
+        None,
+    )
+    if not local_runtime:
+        raise ValueError(f"Runtime {runtime_name} not found")
+
+    trainer = LocalRuntimeTrainer(
+        trainer_type=types.TrainerType.CUSTOM_TRAINER,
+        framework=framework,
+        packages=local_runtime.trainer.packages,
+        image=local_exec_constants.LOCAL_RUNTIME_IMAGE,
+    )
+
+    # set command to run from venv
+    venv_bin_dir = str(Path(venv_dir) / "bin")
+    default_cmd = [str(Path(venv_bin_dir) / local_exec_constants.DEFAULT_COMMAND)]
+    # Set the Trainer entrypoint.
+    if framework == local_exec_constants.TORCH_FRAMEWORK_TYPE:
+        _c = [os.path.join(venv_bin_dir, local_exec_constants.TORCH_COMMAND)]
+        trainer.set_command(tuple(_c))
+    else:
+        trainer.set_command(tuple(default_cmd))
+
+    return trainer
+
+
+def get_dependencies_command(
+    runtime_packages: list[str],
+    pip_index_urls: list[str],
+    trainer_packages: list[str],
+    quiet: bool = True,
+) -> str:
+    # resolve runtime dependencies and trainer dependencies.
+    packages = get_install_packages(
+        runtime_packages=runtime_packages,
+        trainer_packages=trainer_packages,
+    )
+
+    options = [f"--index-url {pip_index_urls[0]}"]
+    options.extend(f"--extra-index-url {extra_index_url}" for extra_index_url in pip_index_urls[1:])
+
+    """
+           PIP_DISABLE_PIP_VERSION_CHECK=1 pip install $QUIET $AS_USER \
+       --no-warn-script-location $PIP_INDEX $PACKAGE_STR
+       """
+    mapping = {
+        "QUIET": "--quiet" if quiet else "",
+        "PIP_INDEX": " ".join(options),
+        "PACKAGE_STR": '"{}"'.format('" "'.join(packages)),  # quote deps
+    }
+    t = Template(local_exec_constants.DEPENDENCIES_SCRIPT)
+    result = t.substitute(**mapping)
+    return result
+
+
+def get_command_using_train_func(
+    runtime: types.Runtime,
+    train_func: Callable,
+    train_func_parameters: Optional[dict[str, Any]],
+    venv_dir: str,
+    train_job_name: str,
+) -> str:
+    """
+    Get the Trainer container command from the given training function and parameters.
+    """
+    # Check if the runtime has a Trainer.
+    if not runtime.trainer:
+        raise ValueError(f"Runtime must have a trainer: {runtime}")
+
+    # Check if training function is callable.
+    if not callable(train_func):
+        raise ValueError(
+            f"Training function must be callable, got function type: {type(train_func)}"
+        )
+
+    # Extract the function implementation.
+    func_code = inspect.getsource(train_func)
+
+    # Extract the file name where the function is defined and move it the venv directory.
+    func_file = Path(venv_dir) / local_exec_constants.LOCAL_EXEC_FILENAME.format(train_job_name)
+
+    # Function might be defined in some indented scope (e.g. in another function).
+    # We need to dedent the function code.
+    func_code = textwrap.dedent(func_code)
+
+    # Wrap function code to execute it from the file. For example:
+    # TODO (andreyvelich): Find a better way to run users' scripts.
+    # def train(parameters):
+    #     print('Start Training...')
+    # train({'lr': 0.01})
+    if train_func_parameters is None:
+        func_code = f"{func_code}\n{train_func.__name__}()\n"
+    else:
+        func_code = f"{func_code}\n{train_func.__name__}({train_func_parameters})\n"
+
+    with open(func_file, "w") as f:
+        f.write(func_code)
+    f.close()
+
+    t = Template(local_exec_constants.LOCAL_EXEC_ENTRYPOINT)
+    mapping = {
+        "PARAMETERS": "",  ## Torch Parameters if any
+        "PYENV_LOCATION": venv_dir,
+        "ENTRYPOINT": " ".join(runtime.trainer.command),
+        "FUNC_FILE": func_file,
+    }
+    entrypoint = t.safe_substitute(**mapping)
+
+    return entrypoint
+
+
+def get_cleanup_venv_script(venv_dir: str, cleanup_venv: bool = True) -> str:
+    script = "\n"
+    if not cleanup_venv:
+        return script
+
+    t = Template(local_exec_constants.LOCAL_EXEC_JOB_CLEANUP_SCRIPT)
+    mapping = {
+        "PYENV_LOCATION": venv_dir,
+    }
+    return t.substitute(**mapping)
+
+
+def get_local_train_job_script(
+    train_job_name: str,
+    venv_dir: str,
+    trainer: types.CustomTrainer,
+    runtime: types.Runtime,
+    cleanup_venv: bool = True,
+) -> tuple:
+    # use local-exec train job template
+    t = Template(local_exec_constants.LOCAL_EXEC_JOB_TEMPLATE)
+    # find os python binary to create venv
+    python_bin = shutil.which("python")
+    if not python_bin:
+        python_bin = shutil.which("python3")
+    if not python_bin:
+        raise ValueError("No python executable found")
+
+    # workout if dependencies needs to be installed
+    if isinstance(runtime.trainer, LocalRuntimeTrainer):
+        runtime_trainer: LocalRuntimeTrainer = runtime.trainer
+    else:
+        raise ValueError("Invalid Runtime Trainer type: {type(runtime.trainer)}")
+    dependency_script = "\n"
+    if trainer.packages_to_install:
+        dependency_script = get_dependencies_command(
+            pip_index_urls=(
+                trainer.pip_index_urls
+                if trainer.pip_index_urls
+                else constants.DEFAULT_PIP_INDEX_URLS
+            ),
+            runtime_packages=runtime_trainer.packages,
+            trainer_packages=trainer.packages_to_install,
+            quiet=False,
+        )
+
+    entrypoint = get_command_using_train_func(
+        venv_dir=venv_dir,
+        runtime=runtime,
+        train_func=trainer.func,
+        train_func_parameters=trainer.func_args,
+        train_job_name=train_job_name,
+    )
+
+    cleanup_script = get_cleanup_venv_script(cleanup_venv=cleanup_venv, venv_dir=venv_dir)
+
+    mapping = {
+        "OS_PYTHON_BIN": python_bin,
+        "PYENV_LOCATION": venv_dir,
+        "DEPENDENCIES_SCRIPT": dependency_script,
+        "ENTRYPOINT": entrypoint,
+        "CLEANUP_SCRIPT": cleanup_script,
+    }
+
+    command = t.safe_substitute(**mapping)
+
+    return "bash", "-c", command