# Copyright 2023 Alibaba, Inc. or its affiliates. # # 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 # # https://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 typing import Any, Dict, List, Optional from ..api.image import ImageLabel from ..common.logging import get_logger from ..common.utils import to_semantic_version from ..estimator import Estimator from ..session import Session logger = get_logger(__name__) class ModelScopeEstimator(Estimator): """Handle training of custom ModelScope model. The ModelScope Estimator is optimized to run a ModelScope training script in the PAI Training Service with a specific image. Example:: est = ModelScopeEstimator( source_dir="./train/src/", command="python train.py", modelscope_version = 'latest', instance_type="ecs.c6.xlarge", ) est.fit() print(est.model_data()) """ def __init__( self, command: str, source_dir: Optional[str] = None, git_config: Optional[Dict[str, str]] = None, image_uri: Optional[str] = None, modelscope_version: Optional[str] = None, hyperparameters: Optional[Dict[str, Any]] = None, base_job_name: Optional[str] = None, checkpoints_path: Optional[str] = None, output_path: Optional[str] = None, instance_type: Optional[str] = None, instance_count: int = 1, session: Optional[Session] = None, **kwargs, ): """Initialize a ModelScope Estimator. Args: command (str): The command used to run the training job. source_dir (str, optional): The local source code directory used in the training job. The directory will be packaged and uploaded to an OSS bucket, then downloaded to the `/ml/usercode` directory in the training job container. If there is a `requirements.txt` file in the source code directory, the corresponding dependencies will be installed before the training script runs. If 'git_config' is provided, 'source_dir' should be a relative location to a directory in the Git repo. With the following GitHub repo directory structure: .. code:: |----- README.md |----- src |----- train.py |----- test.py if you need 'src' directory as the source code directory, you can assign source_dir='./src/'. git_config (Dict[str, str]): Git configuration used to clone the repo. Including ``repo``, ``branch``, ``commit``, ``username``, ``password`` and ``token``. The ``repo`` is required. All other fields are optional. ``repo`` specifies the Git repository. If you don't provide ``branch``, the default value 'master' is used. If you don't provide ``commit``, the latest commit in the specified branch is used. ``username``, ``password`` and ``token`` are for authentication purpose. For example, the following config: .. code:: python git_config = { 'repo': 'https://github.com/modelscope/modelscope.git', 'branch': 'master', 'commit': '9bfc4a9d83c4beaf8378d0a186261ffc1cd9f960' } results in cloning the git repo specified in 'repo', then checking out the 'master' branch, and checking out the specified commit. image_uri (str, optional): If specified, the estimator will use this image in the training job, instead of selecting the appropriate PAI official image based on modelscope_version. It can be an image provided by PAI or a user customized image. To view the images provided by PAI, please refer to the document: https://help.aliyun.com/document_detail/202834.htm. If ``modelscope_version`` is ``None``, then ``image_uri`` is required. If also ``None``, then a ``ValueError`` will be raised. modelscope_version (str, optional): Modelscope version you want to use for executing your model training code. Defaults to ``None``. Required unless ``image_uri`` is provided. hyperparameters (dict, optional): A dictionary that represents the hyperparameters used in the training job. The hyperparameters will be stored in the `/ml/input/config/hyperparameters.json` as a JSON dictionary in the training container. base_job_name (str, optional): The base name used to generate the training job name. checkpoints_path (str, optional): An OSS URI that stores the checkpoint of the training job. If provided, the OSS URI will be mounted to the directory `/ml/output/checkpoints/`. output_path (str, optional): An OSS URI to store the outputs of the training jobs. If not provided, an OSS URI will be generated using the default OSS bucket in the session. When the `estimator.fit` method is called, a specific OSS URI under the output_path for each channel is generated and mounted to the training container. A completed training container directory structure example:: /ml |-- usercode // User source code directory. | |-- requirements.txt | `-- train.py |-- input // TrainingJob input | `-- config | |-- hyperparameters.json // Hyperparameters in JSON | | // dictionary format for the | | // TrainingJob | | | `-- data // TrainingJob input channels | | // `/ml/input/data/` is a input | | // channel, and the directory | | // name is the channel name. | | // Each directory under the | |-- test-data | | `-- test.csv | `-- train-data | `-- train.csv `-- output // TrainingJob output channels. | // Each directory under the | // `/ml/output/` is an output | // channel, and the directory | // name is the channel name. `-- model `-- checkpoints instance_type (str): The machine instance type used to run the training job. To view the supported machine instance types, please refer to the document: https://help.aliyun.com/document_detail/171758.htm#section-55y-4tq-84y. If the instance_type is "local", the training job is executed locally using docker. instance_count (int): The number of machines used to run the training job. session (:class:`pai.session.Session`, optional): A pai session object manages interactions with PAI REST API. **kwargs: Additional kwargs passed to the :class:`~pai.estimator.Estimator` constructor. .. tip:: You can find additional parameters for initializing this class at :class:`~pai.estimator.Estimator`. """ self._validate_image_uri( image_uri=image_uri, modelscope_version=modelscope_version ) self.image_uri = image_uri self.modelscope_version = modelscope_version super(ModelScopeEstimator, self).__init__( image_uri=self.image_uri, command=command, source_dir=source_dir, git_config=git_config, hyperparameters=hyperparameters, base_job_name=base_job_name, checkpoints_path=checkpoints_path, output_path=output_path, instance_type=instance_type, instance_count=instance_count, session=session, **kwargs, ) # Check image_uri and modelscope_version self.training_image_uri() def _validate_image_uri(self, image_uri: str, modelscope_version: str) -> None: """Check if image_uri or modelscope_version arguments are specified.""" if not image_uri and not modelscope_version: raise ValueError( "modelscope_version, and image_uri are both None. " "Specify either modelscope_version or image_uri." ) def training_image_uri(self) -> str: """Return the Docker image to use for training. The :meth:`~pai.estimator.Estimator.fit` method, which does the model training, calls this method to find the image to use for model training. Returns: str: The URI of the Docker image. """ if self.image_uri: return self.image_uri labels = [ ImageLabel.OFFICIAL_LABEL, ImageLabel.DSW_LABEL, ImageLabel.DEVICE_TYPE_GPU, ] # Filter images by ModelScope version if self.modelscope_version == "latest": latest_version = self._get_latest_ms_version_for_training() labels.append(ImageLabel.framework_version("ModelScope", latest_version)) else: labels.append( ImageLabel.framework_version("ModelScope", self.modelscope_version) ) resp = self.session.image_api.list( labels=labels, workspace_id=0, verbose=True, ) if resp.total_count == 0: raise ValueError( "No official image found for modelscope version:" f" {self.modelscope_version}. Currently supported versions are:" f" {self._get_supported_ms_versions_for_training()}" ) image = resp.items[0]["ImageUri"] return image def _get_supported_ms_versions_for_training(self) -> List[str]: """Return the list of supported ModelScope versions for training.""" label_keys = "system.framework.ModelScope" label_filter = [ ImageLabel.OFFICIAL_LABEL, ImageLabel.DSW_LABEL, ImageLabel.DEVICE_TYPE_GPU, ImageLabel.framework_version("ModelScope", "*"), ] list_image_labels = self.session.image_api.list_labels( label_keys=label_keys, label_filter=label_filter, workspace_id=0, ) res = [] for label in list_image_labels: if label["Value"] not in res: res.append(label["Value"]) res.sort(key=lambda x: to_semantic_version(x)) return res def _get_latest_ms_version_for_training(self) -> str: """Return the latest ModelScope version for training.""" res = self._get_supported_ms_versions_for_training() return max( res, key=lambda x: to_semantic_version(x), )