mantatech-sdk 0.5b0.dev65__py3-none-any.whl

manta/apis/module.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import io
+import os
+import shutil
+import struct
+import tempfile
+import zipfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+from manta_common.build.common.tasks import MaximumType
+from manta_common.build.common.tasks import Module as ModuleProto
+from manta_common.build.common.tasks import Scheduling as SchedulingProto
+from manta_common.build.core.user_services import ModuleRequest
+
+# Import here to avoid circular imports
+
+
+def zip_module(python_program: Union[str, Path]) -> bytes:
+    """
+    Zip multiple files into an executable python module
+
+    Parameters
+    ----------
+    python_program : Union[str, Path]
+        Python program location
+
+    Returns
+    -------
+    bytes
+        Zipped binary in memory
+    """
+    buffer = io.BytesIO()
+    python_program = Path(python_program)
+
+    def nopycache(path: Path):
+        return "pycache" not in str(path)
+
+    with tempfile.TemporaryDirectory() as folder:
+        folder = Path(folder)
+        copy = shutil.copytree if python_program.is_dir() else shutil.copyfile
+        copy(python_program, folder / python_program.name)
+
+        # Create deterministic ZIP archive by manually setting fixed timestamps
+        # Fixed timestamp for deterministic ZIP creation (2023-01-01 00:00:00)
+        FIXED_TIMESTAMP = (2023, 1, 1, 0, 0, 0)
+
+        with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
+            # Add all files from the source folder
+            for root, dirs, files in os.walk(folder):
+                # Filter out pycache directories
+                dirs[:] = [d for d in dirs if "pycache" not in d]
+
+                for file in files:
+                    file_path = Path(root) / file
+                    if nopycache(file_path):
+                        # Calculate relative path for archive
+                        arc_path = file_path.relative_to(folder)
+
+                        # Create ZipInfo with fixed timestamp
+                        zip_info = zipfile.ZipInfo(str(arc_path))
+                        zip_info.date_time = FIXED_TIMESTAMP
+                        zip_info.compress_type = zipfile.ZIP_DEFLATED
+
+                        # Read file content and add to archive
+                        with open(file_path, "rb") as src_file:
+                            zf.writestr(zip_info, src_file.read())
+
+            # Add the __main__.py file for zipapp functionality
+            main_content = f"""# -*- coding: utf-8 -*-
+import {python_program.stem}
+{python_program.stem}.main()
+"""
+            main_zip_info = zipfile.ZipInfo("__main__.py")
+            main_zip_info.date_time = FIXED_TIMESTAMP
+            main_zip_info.compress_type = zipfile.ZIP_DEFLATED
+            zf.writestr(main_zip_info, main_content.encode("utf-8"))
+
+    return buffer.getvalue()
+
+
+@dataclass
+class Module:
+    """
+    Class which holds information about a `Module`
+
+    It helps the user to indicate which Nodes must executed a :class:`Module <manta.module.Module>`.
+
+    During the execution of a :class:`Swarm <manta.swarm.Swarm>`, the Manager selects nodes
+    depending on attributes set in this class.
+
+    For instance, :code:`Module(method="any", excepted_ids=["node_28"], maximum=0.8)`
+    indicates that the module will be executed on 80% of nodes excepted :code:`"node_28"`.
+
+    Attributes
+    ----------
+    python_program : Union[str, Path]
+        python program code
+    image : str
+        image of the container in which the program will be executed
+    alias : Optional[str]
+        Alias task name. If specified, it is used to generate to a unique ID
+    network : Optional[str]
+        It indicates in which network the task will be executed.
+        Network name must exist in :code:`Swarm.networks`.
+    gpu : bool
+        If :code:`True`, the module will be executed on GPU
+    python_files : Dict[str, str]
+        Dictionary containing unzipped python files {filename: content}
+    """
+
+    python_program: Union[str, Path]
+    image: str
+    name: Optional[str] = field(default=None)
+    datasets: Optional[list] = field(default=None)
+    python_files: Dict[str, str] = field(default_factory=dict)
+    module_id: Optional[str] = field(default=None)
+
+    def __post_init__(self):
+        if self.name is None:
+            self.name = Path(self.python_program).stem
+
+    def to_proto(self):
+        """Convert to protobuf object
+
+        Returns
+        -------
+        Object
+            Protobuf object
+        """
+        return ModuleProto(
+            python_program=zip_module(self.python_program),
+            image=self.image,
+            name=self.name or "",
+            datasets=self.datasets or [],
+        )
+
+    def __str__(self):
+        return self.name
+
+    @classmethod
+    def from_proto(cls, proto: Union[ModuleProto, ModuleRequest]) -> "Module":
+        """
+        Create a Module from a protobuf object
+
+        Parameters
+        ----------
+        proto : Union[ModuleProto, ModuleRequest]
+            Protobuf module object
+
+        Returns
+        -------
+        Module
+            Module object with unzipped python_files
+        """
+        # Unzip the python_program bytes
+        python_files: Dict[str, str] = {}
+        if isinstance(proto, ModuleProto):
+            module = proto
+            module_id = ""
+            image = proto.image
+            name = proto.name
+            datasets = list(proto.datasets) if proto.datasets else None
+
+        elif isinstance(proto, ModuleRequest):
+            module = proto.module
+            module_id = proto.module_id
+            image = proto.module.image
+            name = proto.module.name
+            datasets = list(proto.module.datasets) if proto.module.datasets else None
+
+        try:
+            with zipfile.ZipFile(io.BytesIO(module.python_program), "r") as zip_file:
+                for file_name in zip_file.namelist():
+                    # Skip directories and __pycache__ files
+                    if not file_name.endswith("/") and "__pycache__" not in file_name:
+                        with zip_file.open(file_name) as file:
+                            python_files[file_name] = file.read().decode("utf-8")
+        except zipfile.BadZipFile:
+            # If it's not a zip file, treat as raw python code
+            print(
+                f"Warning: {module.python_program} is not a zip file, treating as raw python code"
+            )
+            python_files["main.py"] = module.python_program.decode("utf-8")
+
+        return cls(
+            python_program="",  # Empty since we have python_files
+            image=image,
+            name=name,
+            datasets=datasets,
+            python_files=python_files,
+            module_id=module_id,
+        )
+
+    def print_files(self):
+        """
+        Print the module files in a readable format
+        """
+        if self.python_files:
+            print(f"=== Module: {self.name} ===")
+            print(f"Image: {self.image}")
+            if self.datasets:
+                print(f"Datasets: {self.datasets}")
+            print(f"Files ({len(self.python_files)}):")
+            print("-" * 50)
+
+            for filename, content in self.python_files.items():
+                print(f"\n📁 {filename}")
+                print("-" * len(filename))
+                # Add line numbers for better readability
+                lines = content.split("\n")
+                for i, line in enumerate(lines, 1):
+                    print(f"{i:3d} | {line}")
+                print("-" * 50)
+        else:
+            # Fallback to python_program if python_files is empty
+            print(f"=== Module: {self.name} ===")
+            print(f"Image: {self.image}")
+            if self.datasets:
+                print(f"Datasets: {self.datasets}")
+            print(f"Python program: {self.python_program}")
+
+            # Try to read and display the file if it's a path
+            if isinstance(self.python_program, (str, Path)):
+                try:
+                    path = Path(self.python_program)
+                    if path.exists():
+                        print("-" * 50)
+                        print(f"\n📁 {path.name}")
+                        print("-" * len(path.name))
+                        content = path.read_text(encoding="utf-8")
+                        lines = content.split("\n")
+                        for i, line in enumerate(lines, 1):
+                            print(f"{i:3d} | {line}")
+                        print("-" * 50)
+                except Exception as e:
+                    print(f"Could not read file: {e}")
+
+
+@dataclass
+class Scheduling:
+    """
+    Class which holds scheduling information for a `Module`
+
+    Attributes
+    ----------
+    method : str
+        Currently, "any" or "best_resources" or "tag-<your-tag>"
+        Note with tag, a result, associated to your tag, should
+        be available and sent by previous module
+    fixed : bool
+        if :code:`True`, the chosen nodes must remain the same
+        over the execution of the :class:`Swarm <manta.swarm.Swarm>`
+    excepted_ids : Optional[list]
+        List of ids that must not be chosen for this module
+    specified_ids : Optional[list]
+        List of ids that must be chosen for this module
+    maximum : Optional[Union[int, float]]
+        maximum of ids that can be chosen for this module.
+        If it is a float, it is a percentage between :code:`0.0` and :code:`1.0`
+    gpu : bool
+        If :code:`True`, the module will be executed on GPU
+    datasets : Optional[list]
+        List of datasets required for this module
+    """
+
+    method: str
+    fixed: bool
+    excepted_ids: Optional[list]
+    specified_ids: Optional[list]
+    maximum: Optional[Union[int, float]]
+    gpu: bool
+
+    def __post_init__(self):
+        if self.specified_ids is not None:
+            if self.excepted_ids is not None and not set(self.excepted_ids) & set(
+                self.specified_ids
+            ):
+                raise ValueError("No coherence between specified ids and excepted ids")
+            offset = 0 if self.excepted_ids is None else len(self.excepted_ids)
+            if isinstance(self.maximum, int):
+                # Regulate maximum if maximum is over estimated
+                self.maximum = min(len(self.specified_ids) - offset, self.maximum)
+
+    def to_proto(self):
+        """Convert to protobuf object
+
+        Returns
+        -------
+        Object
+            Protobuf object
+        """
+        # Serialize maximum value
+        if isinstance(self.maximum, float):
+            maximum = struct.pack("f", self.maximum)
+            max_type = MaximumType.FLOAT
+        elif isinstance(self.maximum, int):
+            maximum = struct.pack("i", self.maximum)
+            max_type = MaximumType.INT
+        elif self.maximum is None:
+            maximum = struct.pack("i", 0)
+            max_type = MaximumType.INT
+        else:
+            raise ValueError("Maximum must be a float or an int")
+
+        return SchedulingProto(
+            method=self.method,
+            fixed=self.fixed,
+            excepted_ids=self.excepted_ids or [],
+            specified_ids=self.specified_ids or [],
+            maximum=maximum,
+            max_type=max_type,
+            gpu=self.gpu,
+        )

manta/apis/results.py

@@ -0,0 +1,251 @@
+from functools import reduce
+from operator import ior, itemgetter
+from typing import Any, Dict, List, Tuple, Union
+
+__all__ = ["Results"]
+
+
+def is_int(obj: Any) -> bool:
+    """
+    Check if the :code:`obj` is an integer
+
+    Parameters
+    ----------
+    obj : Any
+        Object input
+
+    Returns
+    -------
+    bool
+        Is an integer
+    """
+    return isinstance(obj, int)
+
+
+class Schema:
+    """
+    This class represents the rows and columns of values from
+    results.
+
+    Parameters
+    ----------
+    values : List[Dict[str, Dict[str, Any]]]
+        List of values from results
+
+    Attributes
+    ----------
+    unordered_rows : Set[str]
+        Node IDs
+    unordered_columns : Set[str]
+        Tags
+    """
+
+    __slots__ = ["unordered_rows", "unordered_columns"]
+
+    def __init__(self, values: List[Dict[str, Dict[str, Any]]]):
+        def ior_dict(dictn, next_item):
+            return ior(dictn, next_item.keys())
+
+        def reduce_ior_dict(dictn, next_item):
+            return reduce(ior_dict, next_item.values(), dictn)
+
+        self.unordered_rows = reduce(ior_dict, values, set())
+        self.unordered_columns = reduce(reduce_ior_dict, values, set())
+
+    def parse(self, values: Dict[str, Dict[str, Any]]) -> List[List[Any]]:
+        """
+        Order results based on the schema values
+
+        Parameters
+        ----------
+        values : Dict[str, Dict[str, Any]]
+            Input values for one iteration
+
+        Returns
+        -------
+        List[List[Any]]
+            Ordered values
+        """
+        return [
+            [values.get(row, {}).get(column, None) for column in self.columns]
+            for row in self.rows
+        ]
+
+    @property
+    def rows(self) -> List[str]:
+        """
+        Convenient to have sorted rows
+
+        Returns
+        -------
+        List[str]
+            Sorted rows
+        """
+        return sorted(self.unordered_rows)
+
+    @property
+    def columns(self) -> List[str]:
+        """
+        Convenient to have sorted columns
+
+        Returns
+        -------
+        List[str]
+            Sorted columns
+        """
+        return sorted(self.unordered_columns)
+
+    def __str__(self):  # pragma: no cover
+        return f"Schema(rows={self.rows}, columns={self.columns})"
+
+    def __repr__(self):  # pragma: no cover
+        return str(self)
+
+
+class Results:
+    """
+    This class organizes the results gotten from the database
+    for ergonomic manipulations
+
+    Parameters
+    ----------
+    values : List[Dict[str, Dict[str, Any]]]
+        Values organized by the class
+
+    Examples
+    --------
+
+    >>> values = [
+    ...     {
+    ...         "node_c": {"accuracy": 0.8, "params": [0.80, 0.84]},
+    ...         "node_a": {"accuracy": 0.7, "params": [0.61, 0.72]},
+    ...         "node_b": {"accuracy": 0.1, "params": [0.23, 0.37]},
+    ...     },
+    ...     {
+    ...         "node_c": {"accuracy": 0.1, "params": [0.53, 0.68]},
+    ...         "node_a": {"accuracy": 0.4, "params": [0.12, 0.90]},
+    ...         "node_b": {"accuracy": 0.9, "params": [0.11, 0.79]},
+    ...     },
+    ... ]
+    >>> results = Results(values)
+    >>> results.schema.rows
+    ["node_a", "node_b", "node_c"]
+    >>> results.schema.columns
+    ["accuracy", "params"]
+    >>> results.iterations
+    [0, 1]
+    >>> results(0)
+    [[0.7, [0.61, 0.72]], [0.1, [0.23, 0.37]], [0.8, [0.8, 0.84]]]
+    >>> results("node_a")
+    [[0.7, [0.61, 0.72]], [0.4, [0.12, 0.9]]]
+    >>> results("accuracy")
+    [[0.7, 0.1, 0.8], [0.4, 0.9, 0.1]]
+    >>> results([0], "accuracy")
+    [[0.7, 0.1, 0.8]]
+    """
+
+    __slots__ = ["schema", "values"]
+
+    def __init__(self, values: List[Dict[str, Dict[str, Any]]]):
+        self.schema = Schema(values)
+        self.values = list(map(self.schema.parse, values))
+
+    @property
+    def iterations(self) -> List[int]:
+        """
+        Return list of iterations
+
+        Returns
+        -------
+        List[int]
+            Iterations
+        """
+        return list(range(len(self.values)))
+
+    def sort_queries(
+        self,
+        indices: Tuple[List[int], List[int], List[int]],
+        query: Union[int, List[int], str],
+    ) -> Tuple[List[int], List[int], List[int]]:
+        """
+        Convert the query into index and add it to its corresponding
+        list of indices
+
+        Parameters
+        ----------
+        indices : Tuple[List[int], List[int], List[int]],
+            Iteration indices, row indices, column indices
+        query : Union[int, List[int], str]
+            Query
+
+        Returns
+        -------
+        Tuple[List[int], List[int], List[int]]
+            Updated iteration indices, row indices, column indices
+
+        Raises
+        ------
+        TypeError
+            When query input is a list which not contains only integers
+        KeyError
+            When query can not be found in keys defined by :code:`self.values`
+        """
+        it_indices, row_indices, column_indices = indices
+        if isinstance(query, int):
+            it_indices.append(query)
+        elif isinstance(query, list):
+            if not all(map(is_int, query)):
+                args = ", ".join(map(repr, query))
+                raise TypeError(
+                    "Only integers in list queries are supported. "
+                    f"Maybe try `results(..., {args})`."
+                )
+            it_indices.extend(query)
+        elif isinstance(query, str):
+            if query in self.schema.unordered_rows:
+                row_indices.append(self.schema.rows.index(query))
+            elif query in self.schema.unordered_columns:
+                column_indices.append(self.schema.columns.index(query))
+            else:
+                raise KeyError(f"Item {query!r} not found in schema.")
+        return (it_indices, row_indices, column_indices)
+
+    def __call__(self, *query: Union[str, int, List[int]]) -> list:
+        """
+        Return a list filtered by queries and ordered by :code:`self.schema`
+
+        Parameters
+        ----------
+        *query : Union[str, int, List[int]]
+            It can be an iteration, a list of iterations, a tag or a node ID
+
+        Returns
+        -------
+        list
+            Filtered list of results
+        """
+        if not query:
+            return self.values
+        it_indices, row_indices, column_indices = reduce(
+            self.sort_queries, query, ([], [], [])
+        )
+        values = itemgetter(*it_indices)(self.values) if it_indices else self.values
+        values = [values] if it_indices and (row_indices or column_indices) else values
+        values = map(itemgetter(*row_indices), values) if row_indices else values
+        values = [values] if row_indices and column_indices else values
+        values = (
+            (list(map(itemgetter(*column_indices), matrix)) for matrix in values)
+            if column_indices
+            else values
+        )
+        return list(values)
+
+    def __str__(self):  # pragma: no cover
+        return (
+            f"Results(len(iterations)={len(self.iterations)},"
+            f" len(nodes)={len(self.schema.unordered_rows)}, "
+            f"len(tags)={len(self.schema.unordered_columns)})"
+        )
+
+    def __repr__(self):  # pragma: no cover
+        return str(self)