psr-factory 5.0.0b1__py3-none-win_amd64.whl → 5.0.0b2__py3-none-win_amd64.whl

psr/cloud/aws.py

@@ -0,0 +1,256 @@
+import os
+import tempfile
+import zipfile
+from typing import Dict, List, Optional
+
+import boto3
+from botocore.exceptions import ClientError
+
+
+def _get_region(url):
+    """Extract the region from the S3 URL."""
+    if url:
+        parts = url.split(".")
+        return parts[0]
+    return None
+
+
+def upload_file_to_s3(
+    s3_client, bucket_name, file_path, object_name=None, extra_args=None
+):
+    """Upload a file to an S3 bucket using a provided S3 client.
+
+    :param s3_client: Initialized S3 client.
+    :param bucket_name: Name of the S3 bucket.
+    :param file_path: Path to the file to upload.
+    :param object_name: S3 object name. If not specified, file_path's basename is used.
+    :param extra_args: A dictionary of extra arguments to pass to S3's upload_file.
+    :return: True if file was uploaded, else False.
+    """
+    if object_name is None:
+        object_name = os.path.basename(file_path)
+
+    try:
+        s3_client.upload_file(file_path, bucket_name, object_name, ExtraArgs=extra_args)
+        return True
+    except ClientError as e:
+        print(f"Error uploading file: {e}")
+        return False
+
+
+def upload_case_to_s3(
+    files: List[str],
+    repository_id: str,
+    cluster_name: str,
+    checksums: Optional[Dict[str, str]] = None,
+    access: Optional[str] = None,
+    secret: Optional[str] = None,
+    session_token: Optional[str] = None,
+    bucket_name: Optional[str] = None,
+    url: Optional[str] = None,
+    zip_compress: bool = False,
+    compress_zip_name: str = None,
+):
+    """Upload files to an S3 bucket."""
+
+    region = _get_region(url)
+
+    if not region or not access or not secret or not session_token or not bucket_name:
+        raise ValueError("Unable to set up AWS connection.")
+
+    s3_client = boto3.client(
+        "s3",
+        aws_access_key_id=access,
+        aws_secret_access_key=secret,
+        aws_session_token=session_token,
+        region_name=region,
+    )
+
+    # Base metadata, common for both zip and individual files
+    base_metadata: Dict[str, str] = {
+        "upload": str(True).lower(),
+        "user-agent": "aws-fsx-lustre",
+        "file-owner": "537",
+        "file-group": "500",
+        "file-permissions": "100777",
+    }
+
+    if zip_compress and not compress_zip_name:
+        compress_zip_name = str(repository_id)
+
+    if zip_compress:
+        # Create a temporary zip file
+        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_zip_file:
+            zip_path = tmp_zip_file.name
+            tmp_zip_file.close()  # Close the file handle so zipfile can open it
+
+        try:
+            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+                for file_path in files:
+                    # Add file to zip, using only the basename inside the zip
+                    zipf.write(file_path, arcname=os.path.basename(file_path))
+
+            # Construct object name for the zip file
+            object_name = f"{repository_id}/uploaded/{compress_zip_name}.zip"
+
+            # For zip files, we use the base metadata without a specific checksum
+            # (as checksums are per-file in the original design)
+            extra_args = {
+                "Metadata": base_metadata.copy()
+            }  # Use a copy to avoid modifying base_metadata
+
+            if not upload_file_to_s3(
+                s3_client, bucket_name, zip_path, object_name, extra_args=extra_args
+            ):
+                raise ValueError(
+                    f"Failed to upload zip file {zip_path} to S3 bucket {bucket_name}."
+                )
+
+        finally:
+            # Clean up the temporary zip file
+            if os.path.exists(zip_path):
+                os.unlink(zip_path)
+
+    else:
+        # Original logic: upload files individually
+        for file_path in files:
+            file_basename = os.path.basename(file_path)
+            object_name = f"{repository_id}/uploaded/{file_basename}"
+
+            current_file_metadata = base_metadata.copy()
+            if checksums:
+                current_file_metadata["checksum"] = checksums.get(file_basename, "")
+
+            extra_args = {"Metadata": current_file_metadata}
+
+            if not upload_file_to_s3(
+                s3_client, bucket_name, file_path, object_name, extra_args=extra_args
+            ):
+                raise ValueError(
+                    f"Failed to upload file {file_path} to S3 bucket {bucket_name}."
+                )
+
+    # Always upload .metadata files if the source 'files' list is provided
+    if files:
+        # Assuming all files in the 'files' list share the same parent directory,
+        # which is the case data directory.
+        data_directory = os.path.dirname(files[0])
+        metadata_dir_local_path = os.path.join(data_directory, ".metadata")
+
+        if os.path.isdir(metadata_dir_local_path):
+            # Iterate through the original list of files to find corresponding metadata files
+            for original_file_path in files:
+                original_file_basename = os.path.basename(original_file_path)
+                local_metadata_file_path = os.path.join(
+                    metadata_dir_local_path, original_file_basename
+                )
+
+                if os.path.isfile(local_metadata_file_path):
+                    # S3 object name for the metadata file (e.g., repository_id/.metadata/original_file_basename)
+                    s3_metadata_object_name = (
+                        f"{repository_id}/.metadata/{original_file_basename}"
+                    )
+                    extra_args = {"Metadata": base_metadata.copy()}
+                    if not upload_file_to_s3(
+                        s3_client,
+                        bucket_name,
+                        local_metadata_file_path,
+                        s3_metadata_object_name,
+                        extra_args=extra_args,
+                    ):
+                        raise ValueError(
+                            f"Failed to upload metadata file {local_metadata_file_path} to S3 bucket {bucket_name}."
+                        )
+
+
+def _download_s3_object(
+    s3_client, bucket_name: str, s3_object_key: str, local_file_path: str
+) -> bool:
+    """
+    Downloads a single object from S3 to a local file path.
+
+    :param s3_client: Initialized S3 client.
+    :param bucket_name: Name of the S3 bucket.
+    :param s3_object_key: The key of the object in S3.
+    :param local_file_path: The local path where the file should be saved.
+    :return: True if download was successful, False otherwise.
+    """
+
+    try:
+        s3_client.download_file(bucket_name, s3_object_key, local_file_path)
+        return True
+    except ClientError as e:
+        print(f"ERROR: Failed to download {s3_object_key} from S3: {e}")
+        return False
+
+
+def download_case_from_s3(
+    repository_id: str,
+    cluster_name: str,  # Kept for consistency with caller, though not used directly in S3 ops
+    access: str,
+    secret: str,
+    session_token: str,
+    bucket_name: str,
+    url: str,  # S3 endpoint URL, used by _get_region
+    output_path: str,
+    file_list: List[str],
+) -> List[str]:
+    """
+    Downloads files from an S3 bucket for a given case repository.
+
+    It iterates through the provided `file_list`, downloads each specified file
+    from the S3 path `{repository_id}/{file_in_list}`, preserving its relative path
+    under `output_path`. It then checks if each downloaded file is gzipped,
+    decompresses it if necessary, and returns a list of basenames of the
+    final downloaded (and potentially decompressed) files.
+
+    :param repository_id: The ID of the repository in S3.
+    :param cluster_name: Name of the cluster (for context, not used in S3 calls).
+    :param access: AWS access key ID.
+    :param secret: AWS secret access key.
+    :param session_token: AWS session token.
+    :param bucket_name: Name of the S3 bucket.
+    :param url: S3 service URL (used to determine region via _get_region).
+    :param output_path: Local directory where files will be downloaded.
+    :param file_list: A list of file names (basenames) to be downloaded.
+    :return: A list of basenames of the downloaded (and decompressed) files.
+    :raises ValueError: If S3 connection parameters are missing or filter is invalid.
+    :raises RuntimeError: If S3 operations fail.
+    """
+    region = _get_region(url)
+    if not all([region, access, secret, session_token, bucket_name]):
+        # TODO: Replace print with proper logging
+        print(
+            "ERROR: Missing S3 connection parameters (region, access, secret, token, or bucket name)."
+        )
+        raise ValueError("Missing S3 connection parameters.")
+
+    s3_client = boto3.client(
+        "s3",
+        aws_access_key_id=access,
+        aws_secret_access_key=secret,
+        aws_session_token=session_token,
+        region_name=region,
+    )
+
+    downloaded_files: List[str] = []
+
+    try:
+        for file_name in file_list:
+            # Construct the full S3 object key
+            s3_object_key = f"{repository_id}/{file_name}"
+
+            local_file_path = os.path.join(output_path, file_name)
+            if _download_s3_object(
+                s3_client, bucket_name, s3_object_key, local_file_path
+            ):
+                downloaded_files.append(os.path.basename(local_file_path))
+
+    except ClientError as e:
+        print(f"ERROR: S3 ClientError during download: {e}")
+        raise RuntimeError(f"Failed to download files from S3: {e}")
+    except Exception as e:
+        print(f"ERROR: An unexpected error occurred during download: {e}")
+        raise RuntimeError(f"An unexpected error occurred during S3 download: {e}")
+
+    return downloaded_files

psr/cloud/cloud.py

@@ -100,7 +100,7 @@ _CONSOLE_REL_PARENT_PATH = r"Oper\Console"
 
 _CONSOLE_APP = r"FakeConsole.exe"
 
-_ALLOWED_PROGRAMS = ["SDDP", "OPTGEN", "PSRIO", "GRAF"]
+_ALLOWED_PROGRAMS = ["SDDP", "OPTGEN", "PSRIO", "GRAF", "MyModel", "GNoMo"]
 
 if os.name == "nt":
     _PSRCLOUD_CREDENTIALS_PATH = os.path.expandvars(
@@ -462,6 +462,7 @@ class Client:
             return selection
 
         program_versions = self.get_program_versions(case.program)
+        case.program_version_name = case.program_version
         case.program_version = validate_selection(
             case.program_version, program_versions, "Version", case.program
         )
@@ -504,6 +505,14 @@ class Client:
                 # Replace partial with complete budget name
                 case.budget = match_list[0]
 
+        # MyModel
+        if case.program == "MyModel":
+            if case.mymodel_program_files is None:
+                raise CloudInputError("MyModel program files not provided")
+
+        if case.program != "MyModel" and case.mymodel_program_files is not None:
+            msg = "Ignoring mymodel_program_files parameter for non MyModel case."
+            warnings.warn(msg)
         return case
 
     def _pre_process_graph(self, path: str, case_id: int) -> None:
@@ -610,6 +619,7 @@ class Client:
             "validacaoModelo": "True",
             "validacaoUsuario": "False",
             "idVersao": case.program_version,
+            "modeloVersao": case.program_version_name,
             "pathModelo": "C:\\PSR",
             "idTipoExecucao": case.execution_type,
             "nomeCaso": case.name,
@@ -618,6 +628,7 @@ class Client:
             "userTag": "(Untagged)",
             "lifecycle": case.repository_duration,
             "versaoInterface": interface_version,
+            "pathPrograma": case.mymodel_program_files,
         }
 
         if case.budget:
@@ -902,26 +913,46 @@ class Client:
         budgets.sort()
         return budgets
 
+    def get_number_of_processes(self, programa_nome):
+        xml = self._get_cloud_versions_xml()
+
+        programa = xml.find(f".//Programa[@nome='{programa_nome}']")
+        if programa is None:
+            raise CloudError(f"Programa '{programa_nome}' não encontrado.")
+
+        cluster = programa.find(f".//Cluster[@nome='{self.cluster['name']}']")
+        if cluster is None:
+            raise CloudError(
+                f"Cluster '{self.cluster['name']}' não encontrado no programa '{programa_nome}'."
+            )
+
+        maximo_processos = cluster.get("maximoProcessos")
+        processos_por_maquina = cluster.get("processosPorMaquina")
+
+        if maximo_processos and processos_por_maquina:
+            maximo_processos = int(maximo_processos)
+            processos_por_maquina = int(processos_por_maquina)
+
+            lista_processos = list(
+                range(
+                    processos_por_maquina, maximo_processos + 1, processos_por_maquina
+                )
+            )
+
+            return lista_processos
+
+        raise CloudError(f"Invalid values for cluster '{self.cluster['name']}'.")
+
     def _make_soap_request(self, service: str, name: str = "", **kwargs) -> ET.Element:
         portal_ws = zeep.Client(self.cluster["url"] + "?WSDL")
         section = str(id(self))
-        password_md5 = _md5sum(self.username + self.__password + section).upper()
-        password_md5 = (
-            password_md5
-            if self.cluster["name"] == "PSR-US"
-            or self.cluster["name"] == "PSR-HOTFIX"
-            or self.cluster["name"] == "PSR-US_OHIO"
-            else self.__password.upper()
-        )
+        password_md5 = self.__password.upper()
         additional_arguments = kwargs.get("additional_arguments", None)
         parameters = {
             "sessao_id": section,
-            "tipo_autenticacao": "portal"
-            if self.cluster["name"] == "PSR-US"
-            or self.cluster["name"] == "PSR-HOTFIX"
-            or self.cluster["name"] == "PSR-US_OHIO"
-            else "bcrypt",
+            "tipo_autenticacao": "bcrypt",
             "idioma": "3",
+            "versao_cliente": self._get_console_version().split("-")[0],
         }
         if additional_arguments:
             parameters.update(additional_arguments)

psr/cloud/data.py

@@ -76,7 +76,9 @@ class Case:
             "Memory per process ratio must be a string",
         )
 
-        self.repository_duration: Optional[Union[str, int]] = kwargs.get("repository_duration", 2)
+        self.repository_duration: Optional[Union[str, int]] = kwargs.get(
+            "repository_duration", 2
+        )
         self._validate_type(
             self.repository_duration,
             (int, str),
@@ -94,6 +96,13 @@ class Case:
         # Save In Cloud
         self.upload_only = kwargs.get("upload_only", False)
 
+        # Model Optional Attributes
+
+        # MyModel
+        self.mymodel_program_files: Optional[str] = kwargs.get(
+            "mymodel_program_files", None
+        )
+
     @staticmethod
     def _validate_type(
         value, expected_type: Union[List[type], Tuple[type], type], error_message: str
@@ -103,3 +112,11 @@ class Case:
 
     def __str__(self):
         return str(self.__dict__)
+
+    def to_dict(self) -> dict:
+        def serialize(obj):
+            if isinstance(obj, datetime):
+                return obj.isoformat()
+            return obj
+
+        return {k: serialize(v) for k, v in self.__dict__.items() if v is not None}

psr/cloud/version.py

@@ -2,4 +2,4 @@
 # Unauthorized copying of this file, via any medium is strictly prohibited
 # Proprietary and confidential
 
-__version__ = "0.3.3"
+__version__ = "0.3.7"

psr/cloud/xml.py

@@ -9,6 +9,8 @@ from xml.etree import ElementTree as ET
 def create_case_xml(parameters: Dict[str, Any]) -> str:
     root = ET.Element("ColecaoParametro")
     for name, value in parameters.items():
+        if value is None:
+            continue
         value = _handle_invalid_xml_chars(value)
         parameter = ET.SubElement(root, "Parametro", nome=name, tipo="System.String")
         parameter.text = value

psr/factory/__init__.py

@@ -2,6 +2,6 @@
 # Unauthorized copying of this file, via any medium is strictly prohibited
 # Proprietary and confidential
 
-__version__ = "5.0.0b1"
+__version__ = "5.0.0b2"
 
 from .api import *