Async register API support

ads/aqua/common/task_status.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2025 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
+
+from ads.common.extended_enum import ExtendedEnum
+from ads.common.serializer import DataClassSerializable
+
+
+class TaskStatusEnum(ExtendedEnum):
+ MODEL_VALIDATION_SUCCESSFUL = "MODEL_VALIDATION_SUCCESSFUL"
+ MODEL_DOWNLOAD_STARTED = "MODEL_DOWNLOAD_STARTED"
+ MODEL_DOWNLOAD_SUCCESSFUL = "MODEL_DOWNLOAD_SUCCESSFUL"
+ MODEL_UPLOAD_STARTED = "MODEL_UPLOAD_STARTED"
+ MODEL_UPLOAD_SUCCESSFUL = "MODEL_UPLOAD_SUCCESSFUL"
+ DATASCIENCE_MODEL_CREATED = "DATASCIENCE_MODEL_CREATED"
+ MODEL_REGISTRATION_SUCCESSFUL = "MODEL_REGISTRATION_SUCCESSFUL"
+ REGISTRATION_FAILED = "REGISTRATION_FAILED"
+ MODEL_DOWNLOAD_INPROGRESS = "MODEL_DOWNLOAD_INPROGRESS"
+
+
+@dataclass
+class TaskStatus(DataClassSerializable):
+ state: TaskStatusEnum
+ message: str

ads/aqua/extension/aqua_ws_msg_handler.py

@@ -10,6 +10,7 @@
 from typing import List
 
 from tornado.web import HTTPError
+from tornado.websocket import WebSocketHandler
 
 from ads.aqua import logger
 from ads.aqua.common.decorator import handle_exceptions
@@ -53,6 +54,9 @@ def process(self) -> BaseResponse:
  """
  pass
 
+ def set_ws_connection(self, con: WebSocketHandler):
+ self.ws_connection = con
+
  def write_error(self, status_code, **kwargs):
  """AquaWSMSGhandler errors are JSON, not human pages."""
  reason = kwargs.get("reason")

ads/aqua/extension/model_handler.py

@@ -2,8 +2,11 @@
 # Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import threading
+from logging import getLogger
 from typing import Optional
 from urllib.parse import urlparse
+from uuid import uuid4
 
 from tornado.web import HTTPError
 
@@ -12,17 +15,29 @@
  CustomInferenceContainerTypeFamily,
 )
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
+from ads.aqua.common.task_status import TaskStatus, TaskStatusEnum
 from ads.aqua.common.utils import (
  get_hf_model_info,
  is_valid_ocid,
  list_hf_models,
 )
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
+from ads.aqua.extension.status_manager import (
+ RegistrationStatus,
+ StatusTracker,
+ TaskNameEnum,
+)
 from ads.aqua.model import AquaModelApp
-from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
+from ads.aqua.model.entities import (
+ AquaModel,
+ AquaModelSummary,
+ HFModelSummary,
+)
 from ads.aqua.ui import ModelFormat
 
+logger = getLogger(__name__)
+
 
 class AquaModelHandler(AquaAPIhandler):
  """Handler for Aqua Model REST APIs."""
@@ -108,6 +123,7 @@ def post(self, *args, **kwargs): # noqa: ARG002
  HTTPError
  Raises HTTPError if inputs are missing or are invalid
  """
+ task_id = str(uuid4())
  try:
  input_data = self.get_json_body()
  except Exception as ex:
@@ -145,27 +161,64 @@ def post(self, *args, **kwargs): # noqa: ARG002
  str(input_data.get("ignore_model_artifact_check", "false")).lower()
  == "true"
  )
+ async_mode = input_data.get("async_mode", False)
 
- return self.finish(
- AquaModelApp().register(
- model=model,
- os_path=os_path,
- download_from_hf=download_from_hf,
- local_dir=local_dir,
- cleanup_model_cache=cleanup_model_cache,
- inference_container=inference_container,
- finetuning_container=finetuning_container,
- compartment_id=compartment_id,
- project_id=project_id,
- model_file=model_file,
- inference_container_uri=inference_container_uri,
- allow_patterns=allow_patterns,
- ignore_patterns=ignore_patterns,
- freeform_tags=freeform_tags,
- defined_tags=defined_tags,
- ignore_model_artifact_check=ignore_model_artifact_check,
+ def register_model(callback=None) -> AquaModel:
+ """Wrapper method to help initialize callback in case of async mode"""
+ try:
+ registered_model = AquaModelApp().register(
+ model=model,
+ os_path=os_path,
+ download_from_hf=download_from_hf,
+ local_dir=local_dir,
+ cleanup_model_cache=cleanup_model_cache,
+ inference_container=inference_container,
+ finetuning_container=finetuning_container,
+ compartment_id=compartment_id,
+ project_id=project_id,
+ model_file=model_file,
+ inference_container_uri=inference_container_uri,
+ allow_patterns=allow_patterns,
+ ignore_patterns=ignore_patterns,
+ freeform_tags=freeform_tags,
+ defined_tags=defined_tags,
+ ignore_model_artifact_check=ignore_model_artifact_check,
+ callback=callback,
+ )
+ except Exception as e:
+ if async_mode:
+ StatusTracker.add_status(
+ RegistrationStatus(
+ task_id=task_id,
+ task_status=TaskStatus(
+ state=TaskStatusEnum.REGISTRATION_FAILED, message=str(e)
+ ),
+ )
+ )
+ raise
+ else:
+ raise
+ return registered_model
+
+ if async_mode:
+ t = threading.Thread(
+ target=register_model,
+ args=(
+ StatusTracker.prepare_status_callback(
+ TaskNameEnum.REGISTRATION_STATUS, task_id=task_id
+ ),
+ ),
+ daemon=True,
  )
- )
+ t.start()
+ output = {
+ "state": "ACCEPTED",
+ "task_id": task_id,
+ "progress_url": f"ws://host:port/aqua/ws/{task_id}",
+ }
+ else:
+ output = register_model()
+ return self.finish(output)
 
  @handle_exceptions
  def put(self, id):

ads/aqua/extension/models/ws_models.py

@@ -23,6 +23,7 @@ class RequestResponseType(ExtendedEnum):
  AdsVersion = "AdsVersion"
  CompatibilityCheck = "CompatibilityCheck"
  Error = "Error"
+ RegisterModelStatus = "RegisterModelStatus"
 
 
 @dataclass
@@ -141,3 +142,16 @@ class AquaWsError(DataClassSerializable):
 class ErrorResponse(BaseResponse):
  data: AquaWsError
  kind = RequestResponseType.Error
+
+
+@dataclass
+class RequestStatus(DataClassSerializable):
+ status: str
+ message: str
+
+
+@dataclass
+class ModelRegisterRequest(DataClassSerializable):
+ status: str
+ task_id: str
+ message: str = ""

ads/aqua/extension/models_ws_msg_handler.py

@@ -1,28 +1,46 @@
 #!/usr/bin/env python
 
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import json
+from logging import getLogger
 from typing import List, Union
 
 from ads.aqua.common.decorator import handle_exceptions
 from ads.aqua.extension.aqua_ws_msg_handler import AquaWSMsgHandler
 from ads.aqua.extension.models.ws_models import (
  ListModelsResponse,
  ModelDetailsResponse,
+ ModelRegisterRequest,
  RequestResponseType,
 )
+from ads.aqua.extension.status_manager import (
+ RegistrationSubscriber,
+ StatusTracker,
+ TaskNameEnum,
+)
 from ads.aqua.model import AquaModelApp
 
+logger = getLogger(__name__)
+
+REGISTRATION_STATUS = "registration_status"
+
 
 class AquaModelWSMsgHandler(AquaWSMsgHandler):
+ status_subscriber = {}
+ register_status = {} # Not threadsafe
+
  def __init__(self, message: Union[str, bytes]):
  super().__init__(message)
 
  @staticmethod
  def get_message_types() -> List[RequestResponseType]:
- return [RequestResponseType.ListModels, RequestResponseType.ModelDetails]
+ return [
+ RequestResponseType.ListModels,
+ RequestResponseType.ModelDetails,
+ RequestResponseType.RegisterModelStatus,
+ ]
 
  @handle_exceptions
  def process(self) -> Union[ListModelsResponse, ModelDetailsResponse]:
@@ -47,3 +65,26 @@ def process(self) -> Union[ListModelsResponse, ModelDetailsResponse]:
  kind=RequestResponseType.ModelDetails,
  data=response,
  )
+ elif request.get("kind") == "RegisterModelStatus":
+ task_id = request.get("task_id")
+ StatusTracker.add_subscriber(
+ subscriber=RegistrationSubscriber(
+ task_id=task_id, subscriber=self.ws_connection
+ ),
+ notify_latest_status=False,
+ )
+
+ latest_status = StatusTracker.get_latest_status(
+ TaskNameEnum.REGISTRATION_STATUS, task_id=task_id
+ )
+ logger.info(latest_status)
+ if latest_status:
+ return ModelRegisterRequest(
+ status=latest_status.state,
+ message=latest_status.message,
+ task_id=task_id,
+ )
+ else:
+ return ModelRegisterRequest(
+ status="SUBSCRIBED", task_id=task_id, message=""
+ )