oracle-ads 2.11.9__py3-none-any.whl → 2.11.10__py3-none-any.whl

ads/aqua/extension/ui_websocket_handler.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+import concurrent.futures
+from asyncio.futures import Future
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, List, Type, Union
+
+import tornado
+from tornado import httputil
+from tornado.ioloop import IOLoop
+from tornado.websocket import WebSocketHandler
+
+from ads.aqua import logger
+from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
+from ads.aqua.extension.evaluation_ws_msg_handler import AquaEvaluationWSMsgHandler
+from ads.aqua.extension.models.ws_models import (
+    AquaWsError,
+    BaseRequest,
+    BaseResponse,
+    ErrorResponse,
+    RequestResponseType,
+)
+
+MAX_WORKERS = 20
+
+
+def get_aqua_internal_error_response(message_id: str) -> ErrorResponse:
+    error = AquaWsError(
+        status="500",
+        message="Internal Server Error",
+        service_payload={},
+        reason="",
+    )
+    return ErrorResponse(
+        message_id=message_id,
+        kind=RequestResponseType.Error,
+        data=error,
+    )
+
+
+class AquaUIWebSocketHandler(WebSocketHandler):
+    """Handler for Aqua Websocket."""
+
+    _handlers_: List[Type[AquaWSMsgHandler]] = [AquaEvaluationWSMsgHandler]
+
+    thread_pool: ThreadPoolExecutor
+
+    future_message_map: Dict[Future, BaseRequest]
+    message_type_handler_map: Dict[RequestResponseType, Type[AquaWSMsgHandler]]
+
+    def __init__(
+        self,
+        application: tornado.web.Application,
+        request: httputil.HTTPServerRequest,
+        **kwargs,
+    ):
+        # Create a mapping of message type to handler and check for duplicates
+        self.future_message_map = {}
+        self.message_type_handler_map = {}
+        for handler in self._handlers_:
+            for message_type in handler.get_message_types():
+                if message_type in self.message_type_handler_map:
+                    raise ValueError(
+                        f"Duplicate message type {message_type} in AQUA websocket handlers."
+                    )
+                else:
+                    self.message_type_handler_map[message_type] = handler
+
+        super().__init__(application, request, **kwargs)
+
+    def open(self, *args, **kwargs):
+        self.thread_pool = ThreadPoolExecutor(max_workers=MAX_WORKERS)
+        logger.info("AQUA WebSocket opened")
+
+    def on_message(self, message: Union[str, bytes]):
+        try:
+            request = BaseRequest.from_json(message, ignore_unknown=True)
+        except Exception as e:
+            logger.error(
+                f"Unable to parse WebSocket message {message}\nWith exception: {str(e)}"
+            )
+            raise e
+        # Find the handler for the message type.
+        # Each handler is responsible for some specific message types
+        handler = self.message_type_handler_map.get(request.kind, None)
+        if handler is None:
+            self.write_message(
+                get_aqua_internal_error_response(request.message_id).to_json()
+            )
+            raise ValueError(f"No handler found for message type {request.kind}")
+        else:
+            message_handler = handler(message)
+            future: Future = self.thread_pool.submit(message_handler.process)
+            self.future_message_map[future] = request
+            future.add_done_callback(self.on_message_processed)
+
+    def on_message_processed(self, future: concurrent.futures.Future):
+        """Callback function to handle the response from the various AquaWSMsgHandlers."""
+        try:
+            response: BaseResponse = future.result()
+
+        # Any exception coming here is an unhandled exception in the handler. We should log it and return an internal server error.
+        # In non WebSocket scenarios this would be handled by the tornado webserver
+        except Exception as e:
+            logger.error(
+                f"Unable to handle WebSocket message {self.future_message_map[future]}\nWith exception: {str(e)}"
+            )
+            response: BaseResponse = get_aqua_internal_error_response(
+                self.future_message_map[future].message_id
+            )
+            raise e
+        finally:
+            self.future_message_map.pop(future)
+            # Send the response back to the client on the event thread
+            IOLoop.current().run_sync(lambda: self.write_message(response.to_json()))
+
+    def on_close(self) -> None:
+        self.thread_pool.shutdown()
+        logger.info("AQUA WebSocket closed")
+
+
+__handlers__ = [("ws?([^/]*)", AquaUIWebSocketHandler)]

ads/aqua/extension/utils.py

@@ -7,7 +7,7 @@ from typing import Dict, Optional
 
 from tornado.web import HTTPError
 
-from ads.aqua.extension.base_handler import Errors
+from ads.aqua.extension.errors import Errors
 
 
 def validate_function_parameters(data_class, input_data: Dict):

ads/aqua/finetuning/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+from ads.aqua.finetuning.finetuning import AquaFineTuningApp
+
+__all__ = ["AquaFineTuningApp"]

ads/aqua/finetuning/constants.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+from ads.common.extended_enum import ExtendedEnumMeta
+
+
+class FineTuneCustomMetadata(str, metaclass=ExtendedEnumMeta):
+    FINE_TUNE_SOURCE = "fine_tune_source"
+    FINE_TUNE_SOURCE_NAME = "fine_tune_source_name"
+    FINE_TUNE_OUTPUT_PATH = "fine_tune_output_path"
+    FINE_TUNE_JOB_ID = "fine_tune_job_id"
+    FINE_TUNE_JOB_RUN_ID = "fine_tune_job_run_id"
+    SERVICE_MODEL_ARTIFACT_LOCATION = "artifact_location"
+    SERVICE_MODEL_DEPLOYMENT_CONTAINER = "deployment-container"
+    SERVICE_MODEL_FINE_TUNE_CONTAINER = "finetune-container"

ads/aqua/finetuning/entities.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from ads.aqua.data import AquaJobSummary
+from ads.common.serializer import DataClassSerializable
+
+
+@dataclass(repr=False)
+class AquaFineTuningParams(DataClassSerializable):
+    epochs: int
+    learning_rate: Optional[float] = None
+    sample_packing: Optional[bool] = "auto"
+    batch_size: Optional[
+        int
+    ] = None  # make it batch_size for user, but internally this is micro_batch_size
+    sequence_len: Optional[int] = None
+    pad_to_sequence_len: Optional[bool] = None
+    lora_r: Optional[int] = None
+    lora_alpha: Optional[int] = None
+    lora_dropout: Optional[float] = None
+    lora_target_linear: Optional[bool] = None
+    lora_target_modules: Optional[List] = None
+
+
+@dataclass(repr=False)
+class AquaFineTuningSummary(AquaJobSummary, DataClassSerializable):
+    parameters: AquaFineTuningParams = field(default_factory=AquaFineTuningParams)
+
+
+@dataclass(repr=False)
+class CreateFineTuningDetails(DataClassSerializable):
+    """Dataclass to create aqua model fine tuning.
+
+    Fields
+    ------
+    ft_source_id: str
+        The fine tuning source id. Must be model ocid.
+    ft_name: str
+        The name for fine tuning.
+    dataset_path: str
+        The dataset path for fine tuning. Could be either a local path from notebook session
+        or an object storage path.
+    report_path: str
+        The report path for fine tuning. Must be an object storage path.
+    ft_parameters: dict
+        The parameters for fine tuning.
+    shape_name: str
+        The shape name for fine tuning job infrastructure.
+    replica: int
+        The replica for fine tuning job runtime.
+    validation_set_size: float
+        The validation set size for fine tuning job. Must be a float in between [0,1).
+    ft_description: (str, optional). Defaults to `None`.
+        The description for fine tuning.
+    compartment_id: (str, optional). Defaults to `None`.
+        The compartment id for fine tuning.
+    project_id: (str, optional). Defaults to `None`.
+        The project id for fine tuning.
+    experiment_id: (str, optional). Defaults to `None`.
+        The fine tuning model version set id. If provided,
+        fine tuning model will be associated with it.
+    experiment_name: (str, optional). Defaults to `None`.
+        The fine tuning model version set name. If provided,
+        the fine tuning version set with the same name will be used if exists,
+        otherwise a new model version set will be created with the name.
+    experiment_description: (str, optional). Defaults to `None`.
+        The description for fine tuning model version set.
+    block_storage_size: (int, optional). Defaults to 256.
+        The storage for fine tuning job infrastructure.
+    subnet_id: (str, optional). Defaults to `None`.
+        The custom egress for fine tuning job.
+    log_group_id: (str, optional). Defaults to `None`.
+        The log group id for fine tuning job infrastructure.
+    log_id: (str, optional). Defaults to `None`.
+        The log id for fine tuning job infrastructure.
+    force_overwrite: (bool, optional). Defaults to `False`.
+        Whether to force overwrite the existing file in object storage.
+    """
+
+    ft_source_id: str
+    ft_name: str
+    dataset_path: str
+    report_path: str
+    ft_parameters: dict
+    shape_name: str
+    replica: int
+    validation_set_size: float
+    ft_description: Optional[str] = None
+    compartment_id: Optional[str] = None
+    project_id: Optional[str] = None
+    experiment_id: Optional[str] = None
+    experiment_name: Optional[str] = None
+    experiment_description: Optional[str] = None
+    block_storage_size: Optional[int] = None
+    subnet_id: Optional[str] = None
+    log_id: Optional[str] = None
+    log_group_id: Optional[str] = None
+    force_overwrite: Optional[bool] = False