diff --git a/singlestoredb/fusion/handlers/models.py b/singlestoredb/fusion/handlers/models.py new file mode 100644 index 00000000..f696b23c --- /dev/null +++ b/singlestoredb/fusion/handlers/models.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +import os +from typing import Any +from typing import Dict +from typing import Optional + +from ..handler import SQLHandler +from ..result import FusionSQLResult +from .files import ShowFilesHandler +from .utils import get_file_space + + +class ShowModelsHandler(ShowFilesHandler): + """ + SHOW MODELS + [ at_path ] [ ] + [ ] + [ ] [ recursive ] [ extended ]; + + # File path to list + at_path = AT '' + + # Should the listing be recursive? + recursive = RECURSIVE + + # Should extended attributes be shown? + extended = EXTENDED + + Description + ----------- + Displays the list of models in models space. + + Arguments + --------- + * ````: A path in the models space. + * ````: A pattern similar to SQL LIKE clause. + Uses ``%`` as the wildcard character. + + Remarks + ------- + * Use the ``LIKE`` clause to specify a pattern and return only the + files that match the specified pattern. + * The ``LIMIT`` clause limits the number of results to the + specified number. + * Use the ``ORDER BY`` clause to sort the results by the specified + key. By default, the results are sorted in the ascending order. + * The ``AT PATH`` clause specifies the path in the models + space to list the files from. + * To return more information about the files, use the ``EXTENDED`` + clause. + + Examples + -------- + The following command lists the models:: + + SHOW MODELS; + + The following command lists the models with additional information:: + + SHOW MODELS EXTENDED; + + See Also + -------- + * ``UPLOAD MODEL model_name FROM path`` + * ``DOWNLOAD MODEL model_name`` + + + """ # noqa: E501 + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + params['file_location'] = 'MODELS' + + return super().run(params) + + +ShowModelsHandler.register(overwrite=True) + + +class UploadModelHandler(SQLHandler): + """ + UPLOAD MODEL model_name + FROM local_path [ overwrite ]; + + # Model Name + model_name = '' + + # Path to local file or directory + local_path = '' + + # Should an existing file be overwritten? + overwrite = OVERWRITE + + Description + ----------- + Uploads a file or folder to models space. + + Arguments + --------- + * ````: Model name. + * ````: The path to the file or folder to upload in the local + directory. + + Remarks + ------- + * If the ``OVERWRITE`` clause is specified, any existing file at the + specified path in the models space is overwritten. + + Examples + -------- + The following command uploads a file to models space and overwrite any + existing files at the specified path:: + + UPLOAD MODEL model_name + FROM 'llama3/' OVERWRITE; + + See Also + -------- + * ``DOWNLOAD MODEL model_name`` + + """ # noqa: E501 + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + params['file_location'] = 'MODELS' + + model_name = params['model_name'] + local_path = params['local_path'] + + file_space = get_file_space(params) + + if os.path.isdir(local_path): + file_space.upload_folder( + local_path=local_path, + path=os.path.join(model_name, ''), + overwrite=params['overwrite'], + ) + else: + file_space.upload_file( + local_path=local_path, + path=os.path.join(model_name, local_path), + overwrite=params['overwrite'], + ) + + return None + + +UploadModelHandler.register(overwrite=True) + + +class DownloadModelHandler(SQLHandler): + """ + DOWNLOAD MODEL model_name + [ local_path ] + [ overwrite ]; + + # Model Name + model_name = '' + + # Path to local directory + local_path = TO '' + + # Should an existing directory be overwritten? + overwrite = OVERWRITE + + Description + ----------- + Download a model from models space. + + Arguments + --------- + * ````: Model name to download in models space. + * ````: Specifies the path in the local directory + where the model is downloaded. + + Remarks + ------- + * If the ``OVERWRITE`` clause is specified, any existing file or folder at + the download location is overwritten. + * If ```` is not specified, the model is downloaded to the current location. + + Examples + -------- + The following command displays the contents of the file on the + standard output:: + + DOWNLOAD MODEL llama3; + + The following command downloads a model to a specific location and + overwrites any existing models folder with the name ``local_llama3`` on the local storage:: + + DOWNLOAD MODEL llama3 + TO 'local_llama3' OVERWRITE; + + See Also + -------- + * ``UPLOAD MODEL model_name FROM local_path`` + + """ # noqa: E501 + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + params['file_location'] = 'MODELS' + + file_space = get_file_space(params) + + model_name = params['model_name'] + file_space.download_folder( + path=os.path.join(model_name, ''), + local_path=params['local_path'] or model_name, + overwrite=params['overwrite'], + ) + + return None + + +DownloadModelHandler.register(overwrite=True) + + +class DropModelsHandler(SQLHandler): + """ + DROP MODEL model_name; + + # Model Name + model_name = '' + + Description + ----------- + Deletes a model from models space. + + Arguments + --------- + * ````: Model name to delete in models space. + + Example + -------- + The following commands deletes a model from a model space:: + + DROP MODEL llama3; + + """ # noqa: E501 + + def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: + params['file_location'] = 'MODELS' + path = os.path.join(params['model_name'], '') + + file_space = get_file_space(params) + file_space.removedirs(path=path) + + return None + + +DropModelsHandler.register(overwrite=True) diff --git a/singlestoredb/fusion/handlers/utils.py b/singlestoredb/fusion/handlers/utils.py index a3df6a74..85c55cab 100644 --- a/singlestoredb/fusion/handlers/utils.py +++ b/singlestoredb/fusion/handlers/utils.py @@ -11,6 +11,7 @@ from ...management.files import FilesManager from ...management.files import FileSpace from ...management.files import manage_files +from ...management.files import MODELS_SPACE from ...management.files import PERSONAL_SPACE from ...management.files import SHARED_SPACE from ...management.workspace import StarterWorkspace @@ -296,15 +297,14 @@ def get_file_space(params: Dict[str, Any]) -> FileSpace: file_location = params.get('file_location') if file_location: file_location_lower_case = file_location.lower() - if ( - file_location_lower_case != PERSONAL_SPACE and - file_location_lower_case != SHARED_SPACE - ): - raise ValueError(f'invalid file location: {file_location}') if file_location_lower_case == PERSONAL_SPACE: return manager.personal_space elif file_location_lower_case == SHARED_SPACE: return manager.shared_space + elif file_location_lower_case == MODELS_SPACE: + return manager.models_space + else: + raise ValueError(f'invalid file location: {file_location}') raise KeyError('no file space was specified') diff --git a/singlestoredb/management/files.py b/singlestoredb/management/files.py index 4680cc01..de4d5342 100644 --- a/singlestoredb/management/files.py +++ b/singlestoredb/management/files.py @@ -3,6 +3,7 @@ from __future__ import annotations import datetime +import glob import io import os import re @@ -23,9 +24,9 @@ from .utils import to_datetime from .utils import vars_to_str - PERSONAL_SPACE = 'personal' SHARED_SPACE = 'shared' +MODELS_SPACE = 'models' class FilesObject(object): @@ -35,8 +36,8 @@ class FilesObject(object): It can belong to either a workspace stage or personal/shared space. This object is not instantiated directly. It is used in the results - of various operations in ``WorkspaceGroup.stage``, ``FilesManager.personal_space`` - and ``FilesManager.shared_space`` methods. + of various operations in ``WorkspaceGroup.stage``, ``FilesManager.personal_space``, + ``FilesManager.shared_space`` and ``FilesManager.models_space`` methods. """ @@ -513,6 +514,11 @@ def shared_space(self) -> FileSpace: """Return the shared file space.""" return FileSpace(SHARED_SPACE, self) + @property + def models_space(self) -> FileSpace: + """Return the models file space.""" + return FileSpace(MODELS_SPACE, self) + def manage_files( access_token: Optional[str] = None, @@ -551,7 +557,8 @@ class FileSpace(FileLocation): FileSpace manager. This object is not instantiated directly. - It is returned by ``FilesManager.personal_space`` or ``FilesManager.shared_space``. + It is returned by ``FilesManager.personal_space``, ``FilesManager.shared_space`` + or ``FileManger.models_space``. """ @@ -687,10 +694,36 @@ def upload_folder( ignore all '*.pyc' files in the directory tree """ - raise ManagementError( - msg='Operation not supported: directories are currently not allowed ' - 'in Files API', - ) + if not os.path.isdir(local_path): + raise NotADirectoryError(f'local path is not a directory: {local_path}') + + if not path: + path = local_path + + ignore_files = set() + if ignore: + if isinstance(ignore, list): + for item in ignore: + ignore_files.update(glob.glob(str(item), recursive=recursive)) + else: + ignore_files.update(glob.glob(str(ignore), recursive=recursive)) + + for dir_path, _, files in os.walk(str(local_path)): + for fname in files: + if ignore_files and fname in ignore_files: + continue + + local_file_path = os.path.join(dir_path, fname) + remote_path = os.path.join( + path, + local_file_path.lstrip(str(local_path)), + ) + self.upload_file( + local_path=local_file_path, + path=remote_path, + overwrite=overwrite, + ) + return self.info(path) def _upload( self, @@ -875,15 +908,30 @@ def is_file(self, path: PathLike) -> bool: return False raise - def _list_root_dir(self) -> List[str]: + def _listdir(self, path: PathLike, *, recursive: bool = False) -> List[str]: """ - Return the names of files in the root directory. + Return the names of files in a directory. + Parameters ---------- + path : Path or str + Path to the folder + recursive : bool, optional + Should folders be listed recursively? + """ res = self._manager._get( - f'files/fs/{self._location}', + f'files/fs/{self._location}/{path}', ).json() + + if recursive: + out = [] + for item in res['content'] or []: + out.append(item['path']) + if item['type'] == 'directory': + out.extend(self._listdir(item['path'], recursive=recursive)) + return out + return [x['path'] for x in res['content'] or []] def listdir( @@ -905,13 +953,17 @@ def listdir( List[str] """ - if path == '' or path == '/': - return self._list_root_dir() + path = re.sub(r'^(\./|/)+', r'', str(path)) + path = re.sub(r'/+$', r'', path) + '/' - raise ManagementError( - msg='Operation not supported: directories are currently not allowed ' - 'in Files API', - ) + if not self.is_dir(path): + raise NotADirectoryError(f'path is not a directory: {path}') + + out = self._listdir(path, recursive=recursive) + if path != '/': + path_n = len(path.split('/')) - 1 + out = ['/'.join(x.split('/')[path_n:]) for x in out] + return out def download_file( self, @@ -973,17 +1025,28 @@ def download_folder( Parameters ---------- path : Path or str - Path to the file + Directory path local_path : Path or str Path to local directory target location overwrite : bool, optional Should an existing directory / files be overwritten if they exist? """ - raise ManagementError( - msg='Operation not supported: directories are currently not allowed ' - 'in Files API', - ) + + if local_path is not None and not overwrite and os.path.exists(local_path): + raise OSError('target path already exists; use overwrite=True to replace') + + if not self.is_dir(path): + raise NotADirectoryError(f'path is not a directory: {path}') + + files = self.listdir(path, recursive=True) + for f in files: + remote_path = os.path.join(path, f) + if self.is_dir(remote_path): + continue + target = os.path.normpath(os.path.join(local_path, f)) + os.makedirs(os.path.dirname(target), exist_ok=True) + self.download_file(remote_path, target, overwrite=overwrite) def remove(self, path: PathLike) -> None: """ @@ -1010,10 +1073,10 @@ def removedirs(self, path: PathLike) -> None: Path to the file location """ - raise ManagementError( - msg='Operation not supported: directories are currently not allowed ' - 'in Files API', - ) + if not self.is_dir(path): + raise NotADirectoryError('path is not a directory') + + self._manager._delete(f'files/fs/{self._location}/{path}') def rmdir(self, path: PathLike) -> None: """