Remove plugins from endpoint containers

model-engine/model_engine_server/core/docker/remote_build.py

@@ -70,7 +70,8 @@ def zip_context(
 
     assert len(folders_to_include) > 0
     assert s3_file_name.endswith(".gz")
-    print(f"Uploading to s3 at: {s3_file_name}")
+    s3_uri = f"s3://{S3_BUCKET}/{s3_file_name}"
+    print(f"Uploading to s3 at: {s3_uri}")
     try:
         # Need to gimme_okta_aws_creds (you can export AWS_PROFILE='ml-admin' right after)
         tar_command = _build_tar_cmd(context, ignore_file, folders_to_include)
@@ -83,7 +84,7 @@ def zip_context(
         ) as proc:
             assert proc.stdout is not None
             with storage_client.open(
-                f"s3://{S3_BUCKET}/{s3_file_name}",
+                s3_uri,
                 "wb",
             ) as out_file:
                 shutil.copyfileobj(proc.stdout, out_file)
@@ -429,6 +430,7 @@ def build_remote_block(
     :param ignore_file: File (e.g. .dockerignore) containing things to ignore when preparing docker context. Relative to context
     :return: BuildResult representing if docker image has successfully built/pushed
     """
+    logger.info(f"build_remote_block args {locals()}")
     job_name = build_remote(
         context,
         dockerfile,
@@ -439,6 +441,7 @@ def build_remote_block(
         build_args,
         custom_tags,
     )
+    logger.info(f"Waiting for job {job_name} to finish")
     result = get_pod_status_and_log(job_name)
     return result
 

model-engine/model_engine_server/core/loggers.py

@@ -10,7 +10,7 @@
 import ddtrace
 import json_log_formatter
 import tqdm
-from ddtrace.helpers import get_correlation_ids
+from ddtrace import tracer
 
 # DO NOT CHANGE LOGGING FORMAT
 LOG_FORMAT: str = "%(asctime)s %(levelname)s [%(name)s] [%(filename)s:%(lineno)d] - %(message)s"
@@ -82,11 +82,12 @@ def json_record(self, message: str, extra: dict, record: logging.LogRecord) -> d
         if request_id:
             extra["request_id"] = request_id
 
-        trace_id, span_id = get_correlation_ids()
+        context = tracer.current_trace_context()
+        trace_id, span_id = (context.trace_id, context.span_id) if context else (0, 0)
 
         # add ids to event dictionary
-        extra["dd.trace_id"] = trace_id or 0
-        extra["dd.span_id"] = span_id or 0
+        extra["dd.trace_id"] = trace_id
+        extra["dd.span_id"] = span_id
 
         # add the env, service, and version configured for the tracer.
         # If tracing is not set up, then this should pull values from DD_ENV, DD_SERVICE, and DD_VERSION.

model-engine/model_engine_server/inference/post_inference_hooks.py

@@ -1,32 +1,19 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional
-from uuid import uuid4
 
 import requests
-from model_engine_server.common.constants import (
-    BILLING_POST_INFERENCE_HOOK,
-    CALLBACK_POST_INFERENCE_HOOK,
-)
+from model_engine_server.common.constants import CALLBACK_POST_INFERENCE_HOOK
 from model_engine_server.common.dtos.tasks import EndpointPredictV1Request
 from model_engine_server.core.loggers import filename_wo_ext, make_logger
 from model_engine_server.domain.entities import CallbackAuth, CallbackBasicAuth
-from model_engine_server.inference.common import _write_to_s3
 from model_engine_server.inference.domain.gateways.inference_monitoring_metrics_gateway import (
     InferenceMonitoringMetricsGateway,
 )
-from model_engine_server.inference.domain.gateways.usage_metrics_gateway import UsageMetricsGateway
-from model_engine_server.inference.infra.gateways.fake_usage_metrics_gateway import (
-    FakeUsageMetricsGateway,
-)
 from tenacity import Retrying, stop_after_attempt, wait_exponential
 
 logger = make_logger(filename_wo_ext(__file__))
 
 
-def _upload_data(data: Any):
-    return _write_to_s3(data).get("result_url")
-
-
 class PostInferenceHook(ABC):
     def __init__(
         self,
@@ -48,41 +35,6 @@ def handle(
         pass
 
 
-class BillingHook(PostInferenceHook):
-    def __init__(
-        self,
-        endpoint_name: str,
-        bundle_name: str,
-        user_id: str,
-        billing_queue: Optional[str],
-        billing_tags: Optional[Dict[str, Any]],
-    ):
-        super().__init__(endpoint_name, bundle_name, user_id)
-        self._billing_queue = billing_queue
-        self._billing_tags = billing_tags or {}
-
-    def handle(
-        self,
-        request_payload: EndpointPredictV1Request,
-        response: Dict[str, Any],
-        task_id: Optional[str],
-    ):
-        if not self._user_id or not self._billing_queue:
-            logger.error("Usage inputs could not be found for billing hook, aborting")
-            return
-        if not task_id:
-            task_id = str(uuid4())
-
-        events_queue: UsageMetricsGateway
-        try:
-            from plugins.eventbridge_usage_metrics_gateway import EventbridgeUsageMetricsGateway
-
-            events_queue = EventbridgeUsageMetricsGateway(self._billing_queue)
-        except ModuleNotFoundError:
-            events_queue = FakeUsageMetricsGateway()
-        events_queue.emit_task_call_metric(idempotency_token=task_id, tags=self._billing_tags)
-
-
 class CallbackHook(PostInferenceHook):
     def __init__(
         self,
@@ -142,15 +94,7 @@ def __init__(
                 # TODO: Ensure that this process gracefully handles errors in
                 #   initializing each post-inference hook.
                 hook_lower = hook.lower()
-                if hook_lower == BILLING_POST_INFERENCE_HOOK:
-                    self._hooks[hook_lower] = BillingHook(
-                        endpoint_name,
-                        bundle_name,
-                        user_id,
-                        billing_queue,
-                        billing_tags,
-                    )
-                elif hook_lower == CALLBACK_POST_INFERENCE_HOOK:
+                if hook_lower == CALLBACK_POST_INFERENCE_HOOK:
                     self._hooks[hook_lower] = CallbackHook(
                         endpoint_name,
                         bundle_name,

model-engine/model_engine_server/inference/pytorch_or_tf.base.Dockerfile

@@ -57,7 +57,6 @@ COPY --chown=modelengine \
 RUN pip install -r /app/model-engine/model_engine_server/inference/requirements_base.txt
 
 COPY --chown=modelengine model-engine/setup.py /app/model-engine/setup.py
-COPY --chown=modelengine model-engine/model_engine_server.egg-info /app/model-engine/model_engine_server.egg-info
 COPY --chown=modelengine model-engine/model_engine_server/__init__.py /app/model-engine/model_engine_server/__init__.py
 COPY --chown=modelengine model-engine/model_engine_server/common /app/model-engine/model_engine_server/common
 COPY --chown=modelengine model-engine/model_engine_server/core /app/model-engine/model_engine_server/core

model-engine/model_engine_server/inference/requirements_base.txt

@@ -1,12 +1,22 @@
-aioredis==2.0.1
+aioredis~=2.0
+boto3>=1.28.38
 celery[redis,sqs,tblib]==5.3.1
+datadog-api-client==2.11.0
+datadog~=0.46.0
 fastapi==0.78.0
-gunicorn==20.1.0
 # Incompatibility between celery 5 and python 3.7 because of importlib-metadata 5, so we pin it
 importlib-metadata<5.0;python_version<"3.8"
-json-log-formatter==0.5.2
+scale-launch>=0.1.0
 smart_open==5.1.0
-tqdm==4.65.0
-# Pin typing-extensions so aioitertools doesn't break
 typing-extensions>=4.1.1
 uvicorn==0.17.6
+waitress==2.0.0
+
+# HACK: at time of adding, these deps are imported by model-engine/model_engine_server files
+# add here to to prevent `ModuleNotFoundError` error on container startup, these should be in sync with server reqs
+# long term: consider having slimmer deps and seperating inference container deps from server container deps
+ddtrace==1.8.3  # required for ddtrace-run entrypoint command as well
+json-log-formatter~=0.3  # model_engine_server/core/loggers.py
+tenacity>=6.0.0,<=6.2.0  # model_engine_server/core/loggers.py
+tqdm~=4.64  # model_engine_server/common/service_requests.py
+gunicorn~=20.0

model-engine/model_engine_server/infra/repositories/ecr_docker_repository.py

@@ -4,8 +4,11 @@
 from model_engine_server.core.config import infra_config
 from model_engine_server.core.docker.ecr import image_exists as ecr_image_exists
 from model_engine_server.core.docker.remote_build import build_remote_block
+from model_engine_server.core.loggers import logger_name, make_logger
 from model_engine_server.domain.repositories import DockerRepository
 
+logger = make_logger(logger_name())
+
 
 class ECRDockerRepository(DockerRepository):
     def image_exists(
@@ -21,6 +24,7 @@ def get_image_url(self, image_tag: str, repository_name: str) -> str:
         return f"{infra_config().docker_repo_prefix}/{repository_name}:{image_tag}"
 
     def build_image(self, image_params: BuildImageRequest) -> BuildImageResponse:
+        logger.info(f"build_image args {locals()}")
         folders_to_include = ["model-engine"]
         if image_params.requirements_folder:
             folders_to_include.append(image_params.requirements_folder)

model-engine/model_engine_server/infra/services/live_endpoint_builder_service.py

@@ -172,6 +172,7 @@ async def build_endpoint(
                         base_image_params = self.get_base_image_params(
                             build_endpoint_request, logger_adapter
                         )
+                        logger.info(f"base_image_params: {base_image_params}")
                         base_image = await self._build_image(
                             base_image_params,
                             build_endpoint_request,
@@ -490,6 +491,8 @@ def get_base_image_params(
         inference_folder = "model-engine/model_engine_server/inference"
         base_path: str = os.getenv("WORKSPACE")  # type: ignore
 
+        logger.info(f"inference_folder: {inference_folder}")
+        logger.info(f"dockerfile: {inference_folder}/{dockerfile}")
         return BuildImageRequest(
             repo="launch/inference",
             image_tag=resulting_image_tag[:MAX_IMAGE_TAG_LEN],

model-engine/model_engine_server/service_builder/tasks_v1.py

@@ -90,16 +90,7 @@ async def _build_endpoint(
     session = SessionAsyncNullPool
     pool = aioredis.BlockingConnectionPool.from_url(hmi_config.cache_redis_url)
     redis = aioredis.Redis(connection_pool=pool)
-
-    service: LiveEndpointBuilderService
-    try:
-        from plugins.dependencies import (
-            get_live_endpoint_builder_service as get_custom_live_endpoint_builder_service,
-        )
-
-        service = get_custom_live_endpoint_builder_service(session, redis)
-    except ModuleNotFoundError:
-        service = get_live_endpoint_builder_service(session, redis)
+    service: LiveEndpointBuilderService = get_live_endpoint_builder_service(session, redis)
 
     response = await service.build_endpoint(build_endpoint_request)
     await redis.close()

model-engine/requirements.in

@@ -15,7 +15,7 @@ croniter==1.4.1
 dataclasses-json>=0.5.7
 datadog-api-client==2.11.0
 datadog~=0.46.0
-ddtrace~=0.49.2
+ddtrace==1.8.3
 deprecation~=2.1
 docker~=5.0
 fastapi==0.78.0