Metadata-Version: 2.4
Name: altair-aitools-pyextension-wrapper
Version: 1.1b1
Summary: altair-aitools wrapper for internal AI Studio use
Author-email: Patrik Czako <pczako@altair.com>, Soma G Kontar <skontar@altair.com>
Requires-Python: >=3.9.23
Requires-Dist: altair-aitools-runtime>=1.0rc1
Requires-Dist: pandas>=2.0.3
Requires-Dist: pyarrow>=19.0.1
Requires-Dist: pydantic>=2.0.3
Requires-Dist: rapidminer>=2025.0.0
Requires-Dist: websockets>=15.0.1
Description-Content-Type: text/markdown

# Overview
The Python Extension Wrapper is a script which parses the arguments supplied by RapidMiner, delegates to the requested function calls, forwards intermediate results to subsequent function calls, and finally returns results to RapidMiner.

## Status
IMPLEMENTED: *PED - V2*

## Usage

There are two ways of executing python function calls. One can run the python script `run.py` to execute a given chain of functions. This option requires created the python environment at every execution.
Secondly, there is also a Websocket server mode that avoids the overhead of environment creation and supports long-running executions with a shutdown option.

### Script mode

Executes a given chain of functions defined by the [extension arguments file](#textual-description-of-the-argument-format).

Usage: run.py [-h] [-f FILE] [--temp-dir TEMP_DIR] [-v]

| Option                | Type     | Default Value | Description                           |
|-----------------------|----------|---------------|---------------------------------------|
| -f, --file            | Path     | -             | Specify the extension arguments file  |
| --temp-dir            | Path     | -             | Specify temporary directory           |
| -v, --verbose         | Boolean  | False         | Make the operation more descriptive   |

### Websocket Server mode

Starts a websocket server on the specified port.

Usage: start_ws_server.py [-h] [-p PORT] [-c CLIENT] [--pw PW] [--idle-timeout IDLE_TIMEOUT] [--temp-dir TEMP_DIR] [-v]

| Option                | Type     | Default Value | Description                                                   |
|-----------------------|----------|---------------|---------------------------------------------------------------|
| -p, --port            | int      | 8000          | Specify port number where the HTTP server will listen         |
| -c, --client          | str      | PEL           | Specify the username to access the API with Basic Auth        |
| --pw                  | str      | -             | Specify the password to access the API                        |
| --temp-dir            | Path     | -             | Specify temporary directory                                   |
| --idle-timeout        | int      | 360           | Specify the timeout in seconds for idle server (<0 to disable)|
| -v, --verbose         | Boolean  | False         | Make the operation more descriptive                           |

#### Authentication

The server supports HTTP Basic Authentication. Requests must have the `Authorization` header cointaining the username and the password in the following format: Basic [base64 encoded username:password]

#### Health check

This mode provides a HTTP endpoint to verify that the server is running.

| Endpoint | HTTP method | Expected response status |
|----------|-------------|--------------------------|
| /healthz | GET         | 200                      |

#### Communication protocol

Client-to-Server message structure:

```json
{
    "type": "START|STOP|EVICT_IO_CACHE|DATA_REQUEST|CONTROL",
    "execution_id": "uuid",
    "data": {...}
}
```

Server-to-Client message structure:

```json
{
    "type": "STARTED|STOPPED|OPERATOR_COMPLETED|FINISHED|LOG|ERROR",
    "execution_id": "uuid",
    "message": "error message or null",
    "data": {...}
}
```

**Running an operator chain:**

![Websocket Protocol](./assets/ws_sequence_normal.png)

Once the websocket connection is established, the client can request the execution of an operator chain sending a `START` message. This message must include the [extension arguments JSON object](#textual-description-of-the-argument-format) in its `data` attribute. The wrapper starts the execution on a worker thread, and notifies the client when it is finished (or failed).

**Inter-operator caching mechanism:**

Operator outputs can be stored in a cache to be processed by subsequent operators. The cache allows multiple operators (potentially from different chains) to consume the same cache items by passing deep-copied instances. This mechanism enables both intra- and inter-chain communication while keeping the objects in memory. Cached items can be evicted in the following ways:

- **Automatic eviction**: If the start message sets `max_hits` on the output, it will be evicted from the cache after the specified number of accesses.
- **Manual eviction**: A WebSocket message of type `EVICT_IO_CACHE` can be used to evict either a specified list of items or all of them. The `data` field of the message MUST carry the `clear_all` boolean. When this value is `false`, the system reads the list called `evict_keys` (otherwise optional) and evicts only the items it refers to.

**Requesting data from the cache:**

Clients can request objects stored in the cache by sending a message of type `DATA_REQUEST`. In this case, the `data` field must contain the following attributes:

- id (str | list[str]): Cache key of the requested object(s).
- data_type (str): Type of the requested object(s) as defined [below](#textual-description-of-the-argument-format).
- output_target (str): Identifier of the requester ("rapidminer" or "python").
- collection (bool): Indicates if the requested object is a collection.

The object(s) will then be sent over the websocket connection through which the request was made (see details [below](#textual-description-of-the-argument-format)).

**Operator logs:**

Logs created with the `logging` module inside operators are sent to PEL through the established WebSocket connection. In this case, `type` is `"LOG"` and the `data` attribute has the following structure:

- level (str): Log level (can be CRITICAL, ERROR, WARNING, INFO or DEBUG).
- message (str): User-defined log message.
- module (str): Name of Python module.
- function (str): Name of function.
- pathname (str): Path of file.
- filename (str): Name of file.
- lineno (int): Line number.

**Operator errors:**

Exceptions raised inside operators are also sent to PEL through the established WebSocket connection. In this case, `type` is `"ERROR"` and the `data` attribute has the following structure:

- classname (str): Name of exception class.
- stack_trace (list): Frames starting with the operator where the error occured, each specified by:
  - module (str): Name of Python module.
  - function (str): Name of function.
  - pathname (str): Path of file.
  - filename (str): Name of file.
  - lineno (int): Line number.

**Stopping the execution:**

![Websocket Protocol Stopped](./assets/ws_sequence_stopped.png)

In contrast to the HTTP interface, the Websocket API also makes it possible to stop a running execution. It can be requested by sending a `STOP` message where the `data` attribute carries the parameter:
- `patience` (float): The amount of time in seconds that the client is willing to wait for the execution to stop gracefully. In this case, the server waits for the current operator to complete and stops afterwards. If the executor thread is still alive after the specified duration, it is killed by raising a `SystemExit` exception. If `patience` <= 0, the execution is terminated immediately.


## Textual description of the argument format

- The `lib_zips` list is the list of *Path*s to each required library zip file for the functions, provided by the Python extension. 
- The `data_exchange` object specifies how the data exchange must be handled.
  - The `mode` attribute sets the type of the exchange ([see options below](#data-exchange-modes)).
  - Further configuration is done using the `parameters` attribute which must contain string &rarr; string mapping.
- The `environment` dict contains parameters related to the environment, passed from Studio, currently only containing the initial random seed, passing a new random seed to each function via the `**kwargs` mechanism, while ensuring total reproducibility as long as the chain of operators is unchanged ("doing the same thing twice results in the same output twice").
- The `additional_func_args` dict contains additional parameters passed from Studio, directly passing the parameters to the functions with the `**kwargs` mechanism.
- The `functions` list is the ordered list of all chained functions calls.

For each function call (aka Python operator in RapidMiner):
- `module`: The module to call
- `function`: The function to call
- List of inputs identified via parameter name, the type and their value:
  - The parameter `name` is the name of the argument in the function to call
  - The `source` show where the input comes from
    - *default*: The input value can be obtained using the default value from the function signature
    - *provided*: The value is provided int the `value` attribute (only applicable for primitives or enums)
    - *python*: Using the output value of a previous Python operator without any conversion. However, if the output of a Python operator has multiple consumers, deepcopied instances are passed.
    - *rapidminer*: The input value can be obtained by the configured data exchange method
    - *remote*: The input value is stored by another executor and can be obtained by network-based data exchange.
  - The `data_type` is the type of the argument, can be *int|float|str|bool|enum|file|dataframe|tensor|data_object|object*
    - The `object` type is only supported when the `source` is `python`
  - The `value` is being used in relation to the type
    - *int|float|str|bool* are obvious
    - *enum*: the string representation of an enum value
    - *file*, *dataframe*, *tensor*, *data_object*: depends on the chosen data exchane mode ([see details below](#data-exchange-modes))
    - *object*: the name of the output
  - `collection` indicating if the input is a collection
- List of outputs each with a name, a type and a target:
  - The `name` can be used to reference a result in a subsequent function
  - The `data_type` can be:
    - *file*
    - *dataframe*: pandas dataframes
    - *tensor*: numpy arrays
    - *data_object*: custom data objects
    - *object*: can be from any type, but cannot be passed back to RapidMiner
  - The `target` can be:
    - *discard*: not needed, will be thrown away, just here to correctly parse returned tuple
    - *python*: used as-is in a subsequent function as input
    - *rapidminer*: passed back to RapidMiner (by data exchange)
  - `collection` indicating if the output is a collection

## Data Exchange modes

This section describes the different options to exchange data objects between RapidMiner and the Wrapper.

### File system mode ('file_system')

Data exchange through serialization. This mode requires only the `working_dir` parameter to be set which determines the data exchange directory. Operator inputs are loaded from the `{working_dir}/inputs` directory, while the outputs are serialized into the `{working_dir}/outputs` directory.

In case of File system mode, the input.value attribute must carry a single string referring to the path of the serialized file. If this path is relative (filename), the wrapper attempts to load the object from `{working_dir}/inputs/{path}`. Otherwise, the object is loaded from a the specified (absolute) path.

The serialization of pandas.DataFrame objects must be done using the HDF5 data model.

### Network mode ('network')

Data exchange leveraging WebSockets, potentially using the already established connection for communication. The only required parameter is `remote_sources`, which holds a mapping between remote IDs and the necessary connection information. There are two additional optional parameters:

- `timeout`: Number of seconds the receiver is willing to wait for a chunk. Defaults to 30.
- `chunk_size`: Chunk size used during data transfer in bytes, should not exceed 1 MB. Defaults to 64 kB.

#### Example config and usage

```json
{
  "type": "START",
  "execution_id": "XXXXX",
  "data": {
    "lib_zips": [
      ...
    ],
    "data_exchange": {
      "mode": "network",
      "parameters": {
        "remote_sources": {
          "exec_a": {
            "uri": "ws://localhost:8765",
            "auth_token": "Basic XXXXXX"
          }
        }
      }
    },
    "additional_func_args": {
      ...
    },
    "environment": {
      ...
    },
    "functions": [
      {
        "function": "create_custom_array",
        "module": "samples",
        "inputs": [
          {
            "name": "data",
            "source": "remote",
            "data_type": "tensor",
            "value": {
              "object_id": "array0",
              "remote_id": "exec_a"
            },
            "collection": false
          }
        ],
        "outputs": [
          ...
        ]
      }
    ]
  }
}
```

In this example, one remote source is defined with ID `exec_a` that can be accessed at `localhost:8765` using HTTP Basic authentication. The `auth_token` field is optional, and it will be passed in the `Authorization` HTTP header when establishing the connection.

The specified resources can then be referenced in the `value` field of inputs, where the `remote_id` specifies the key of the remote executor in the `remote_sources` mapping, and the `object_id` identifies the requested object by its key used to store in the cache on the remote.

#### Protocol

Data transfer can be initiated in the following ways:
  - RAPIDMINER &rarr; PYTHON: If `rapidminer` is set as the input source, the executor will use the existing WS connection to request the data object by sending a message of type `DATA_REQUEST` (see details [above](#communication-protocol)).
  - RAPIDMINER &larr; PYTHON: The output target is set to `python` to store the object in the cache. Then a message of the same kind (`DATA_REQUEST`) can be sent to the executor requesting the object(s).
  - PYTHON &rarr; PYTHON: As seen in the example above, the executor attempts to fetch the object from the specified remote using the same requesting protocol.

From that point on, the server will send the object(s) in binary chunks, which the requester is supposed to consume. Apart from these binary chunks, there are three types of text-based control messages:
  - `BOF`: Sent before the first chunk, representing the beginning of the transfer.
  - `NEXT` (only for collections): As the items of a collection get transferred one-by-one, this message represents the end of an item, indicating that the next message will be the first chunk of the next item.
  - `EOF`: Sent after the last chunk, representing the end of the transfer.
  - `ERROR`: Transfer failed, no further chunks will be sent.

These control messages are wrapped up in the JSON structure as shown [above](#communication-protocol) where the type is `CONTROL`, and the control message itself travels in the data section under the `control_message` key.

The described method produces binaries (or a list of binaries in the case of collections) that then need to be deserialized on the receiver side. We list the serialization logic of the different supported types here:
  - Files: Binary content of the file with `utf-8` encoding.
  - Data objects: JSON-based representation constructed by the SDK, encoded with `utf-8`.
  - Dataframes: Apache Arrow table format.
  - Tensors: Apache Arrow tensor format (only supported in Python).

In the case of dataframes, column metadata is included in the Arrow table's metadata attached to the corresponding `field`. Currently, only `rm_role` is serialized, representing the column role.
