vitor-aignosi
dac4f1718f
QTZPOC-20: Implement model import workflow for .sientia bundles Introduces the import_model workflow and its associated activities to enable importing .sientia bundles into MLflow and MongoDB. This includes: * New core workflow and activity logic for handling bundle download, decryption, validation, MLflow provisioning (experiments, runs, registered versions), and MongoDB model document creation. * Comprehensive E2E and unit test suites, including a dedicated test harness and a manual run script for development. * Updated .env.example and README.md with new configurations and usage instructions. * Ensures correct MLflow model source resolution for legacy loader compatibility and robust MongoDB document management. * Refactors existing utility scripts into a scripts/ directory.
114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
import argparse
|
|
import os
|
|
from typing import Any
|
|
|
|
import yaml # type: ignore
|
|
from pathspec import PathSpec
|
|
|
|
r"""
|
|
Usage:
|
|
python .\encrypt.py path_to_dir output_file --ignore ignore_file --chunk-size 100000
|
|
"""
|
|
|
|
|
|
def load_ignore_patterns(ignore_file, include_library):
|
|
# Ensure the .gitignore file exists
|
|
if not os.path.exists(ignore_file):
|
|
raise FileNotFoundError(f'Ignore file not found at {ignore_file}')
|
|
|
|
# Load and parse the .gitignore patterns
|
|
with open(ignore_file) as file:
|
|
patterns = file.readlines()
|
|
if not include_library:
|
|
patterns.append('**/deploy/library/')
|
|
|
|
spec = PathSpec.from_lines('gitwildmatch', patterns)
|
|
return spec
|
|
|
|
|
|
def is_ignored(file_path, spec):
|
|
"""Check if a file should be ignored based on the ignore patterns."""
|
|
return spec.match_file(file_path) if spec else False
|
|
|
|
|
|
def encode_file_tree_to_yaml(directory, ignore_file, include_library):
|
|
"""Encode the file tree into a single YAML file."""
|
|
ignore_patterns = load_ignore_patterns(ignore_file, include_library) if ignore_file else None
|
|
file_tree: dict[str, Any] = {}
|
|
|
|
for root, dirs, files in os.walk(directory):
|
|
# Skip ignored directories
|
|
dirs[:] = [d for d in dirs if not is_ignored(os.path.join(root, d), ignore_patterns)]
|
|
|
|
for file in files:
|
|
file_path = os.path.join(root, file)
|
|
|
|
# Skip ignored files
|
|
if is_ignored(file_path, ignore_patterns):
|
|
continue
|
|
|
|
# Read file content
|
|
try:
|
|
with open(file_path, encoding='utf-8') as f:
|
|
content = f.read()
|
|
except Exception as e:
|
|
print(f'Error reading file {file_path}: {e}')
|
|
raise
|
|
|
|
# Create nested dictionary structure
|
|
path_parts = os.path.relpath(file_path, directory).split(os.sep)
|
|
current_level = file_tree
|
|
|
|
# all except the last part (the file name)
|
|
for part in path_parts[:-1]:
|
|
current_level = current_level.setdefault(part, {})
|
|
|
|
# Add the file and its content
|
|
current_level[path_parts[-1]] = content
|
|
return yaml.dump(file_tree, default_flow_style=False)
|
|
|
|
|
|
def chunk_and_write_file_tree_to_yaml(yaml_content, output_file, chunk_size=None):
|
|
"""Chunk the YAML content and write it to the output file."""
|
|
|
|
chunks = (
|
|
[yaml_content]
|
|
if chunk_size is None
|
|
else [yaml_content[i : i + chunk_size] for i in range(0, len(yaml_content), chunk_size)]
|
|
)
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
chunk_file = f'{output_file}_{i}.yaml'
|
|
# Write the file tree to the output YAML file
|
|
with open(chunk_file, 'w', encoding='utf-8') as yaml_file:
|
|
yaml_file.write(chunk)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Encrypts file tree to yaml file')
|
|
parser.add_argument('input_directory', help='Directory to encode')
|
|
parser.add_argument('output_yaml_file', help='Output YAML file')
|
|
parser.add_argument('--ignore', default=None, help='Path to the ignore file')
|
|
parser.add_argument(
|
|
'--chunk-size', type=int, default=None, help='Chunk size for the output YAML file'
|
|
)
|
|
parser.add_argument(
|
|
'--library', type=bool, default=False, help='Incude the library in the output YAML file'
|
|
)
|
|
|
|
# Parse arguments
|
|
args = parser.parse_args()
|
|
|
|
# Example usage
|
|
directory_to_encode = args.input_directory
|
|
ignore_file_path = args.ignore
|
|
output_yaml_file = args.output_yaml_file
|
|
include_library = args.library
|
|
|
|
content = encode_file_tree_to_yaml(directory_to_encode, ignore_file_path, include_library)
|
|
chunk_and_write_file_tree_to_yaml(content, output_yaml_file, args.chunk_size)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|