mpc_instance entity/repository migrate from fbpcp to fbpcs (#1910)

Summary:
Pull Request resolved: #1910

## Why
In next couple weeks, I'm going to migrate mpc service from fbpcp to fbpcs.
This is the first guard to avoid any further changes in fbpcp's mpc realted files
detailed plan: https://docs.google.com/document/d/1qR1XVVCA2By95tldl2Ey9m__GKX3raodGAZ5yK9SCEw/edit?usp=sharing

## What
- hg copy `mpc_instance.py` from fbpcp to fbpcs
- hg copy `repository/mpc_instance` as as well
- change import to fbpcs

Reviewed By: YigeZhu

Differential Revision: D41204111

fbshipit-source-id: 899c98e3f4fbb5321848fe936c3248931fe8ed50

fbpcs/common/entity/pcs_mpc_instance.py

@@ -10,9 +10,13 @@
 from typing import Any, Dict, List, Optional, Union
 
 from fbpcp.entity.container_instance import ContainerInstance
-from fbpcp.entity.mpc_instance import MPCInstance, MPCInstanceStatus, MPCParty
 from fbpcs.common.entity.instance_base import InstanceBase
 from fbpcs.common.entity.pcs_container_instance import PCSContainerInstance
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import (
+    MPCInstance,
+    MPCInstanceStatus,
+    MPCParty,
+)
 
 DEFAULT_SERVER_URIS = [
     "node0.study123.pci.facebook.com"

fbpcs/common/repository/mpc_instance_local.py

@@ -7,10 +7,12 @@
 # pyre-strict
 
 
-from fbpcp.entity.mpc_instance import MPCInstance
-from fbpcp.repository.mpc_instance import MPCInstanceRepository
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.common.repository.instance_local import LocalInstanceRepository
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCInstance
+from fbpcs.private_computation.service.mpc.repository.mpc_instance import (
+    MPCInstanceRepository,
+)
 
 
 class LocalMPCInstanceRepository(MPCInstanceRepository):

fbpcs/common/tests/repository/test_instance_local.py

@@ -9,10 +9,14 @@
 from pathlib import Path
 from unittest.mock import MagicMock, mock_open, patch
 
-from fbpcp.entity.mpc_instance import MPCInstanceStatus, MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.common.repository.instance_local import LocalInstanceRepository
 
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import (
+    MPCInstanceStatus,
+    MPCParty,
+)
+
 TEST_BASE_DIR = Path("./")
 TEST_INSTANCE_ID = "test-instance-id"
 TEST_GAME_NAME = "lift"

fbpcs/common/tests/repository/test_mpc_instance_local.py

@@ -9,10 +9,14 @@
 from pathlib import Path
 from unittest.mock import MagicMock, mock_open, patch
 
-from fbpcp.entity.mpc_instance import MPCInstanceStatus, MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.common.repository.mpc_instance_local import LocalMPCInstanceRepository
 
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import (
+    MPCInstanceStatus,
+    MPCParty,
+)
+
 TEST_BASE_DIR = Path("./")
 TEST_INSTANCE_ID = "test-instance-id"
 TEST_GAME_NAME = "lift"

fbpcs/private_computation/entity/private_computation_instance.py

@@ -26,7 +26,6 @@
 from pathlib import Path
 
 from fbpcp.entity.container_instance import ContainerInstance
-from fbpcp.entity.mpc_instance import MPCInstanceStatus
 from fbpcs.common.entity.instance_base import InstanceBase
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.common.entity.stage_state_instance import (
@@ -51,6 +50,7 @@
     PrivateIdDfcaConfig,
     ProductConfig,
 )
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCInstanceStatus
 
 UnionedPCInstanceStatus = Union[
     MPCInstanceStatus,

fbpcs/private_computation/service/mpc/entity/mpc_instance.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from fbpcp.entity.container_instance import ContainerInstance
+
+
+class MPCParty(Enum):
+    SERVER = "SERVER"
+    CLIENT = "CLIENT"
+
+
+class MPCInstanceStatus(Enum):
+    UNKNOWN = "UNKNOWN"
+    CREATED = "CREATED"
+    STARTED = "STARTED"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    CANCELED = "CANCELED"
+
+
+@dataclass
+class MPCInstance:
+    instance_id: str
+    game_name: str
+    mpc_party: MPCParty
+    num_workers: int
+    server_ips: Optional[List[str]]
+    containers: List[ContainerInstance]
+    status: MPCInstanceStatus
+    game_args: Optional[List[Dict[str, Any]]]
+    server_uris: Optional[List[str]] = None
+
+    def get_instance_id(self) -> str:
+        return self.instance_id

fbpcs/private_computation/service/mpc/mpc.py

@@ -13,13 +13,19 @@
 from fbpcp.entity.certificate_request import CertificateRequest
 
 from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus
-from fbpcp.entity.mpc_instance import MPCInstance, MPCInstanceStatus, MPCParty
 from fbpcp.error.pcp import PcpError
-from fbpcp.repository.mpc_instance import MPCInstanceRepository
 from fbpcp.service.container import ContainerService
 from fbpcp.service.onedocker import OneDockerService
 from fbpcp.util.typing import checked_cast
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import (
+    MPCInstance,
+    MPCInstanceStatus,
+    MPCParty,
+)
 from fbpcs.private_computation.service.mpc.mpc_game import MPCGameService
+from fbpcs.private_computation.service.mpc.repository.mpc_instance import (
+    MPCInstanceRepository,
+)
 
 DEFAULT_BINARY_VERSION = "latest"
 

fbpcs/private_computation/service/mpc/mpc_game.py

@@ -9,10 +9,11 @@
 import logging
 from typing import Any, Dict, Optional, Tuple
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcp.util.arg_builder import build_cmd_args
 
 from fbpcs.private_computation.service.mpc.entity.mpc_game_config import MPCGameConfig
+
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
 from fbpcs.private_computation.service.mpc.repository.mpc_game_repository import (
     MPCGameRepository,
 )

fbpcs/private_computation/service/mpc/repository/mpc_instance.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import abc
+
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCInstance
+
+
+class MPCInstanceRepository(abc.ABC):
+    @abc.abstractmethod
+    def create(self, instance: MPCInstance) -> None:
+        pass
+
+    @abc.abstractmethod
+    def read(self, instance_id: str) -> MPCInstance:
+        pass
+
+    @abc.abstractmethod
+    def update(self, instance: MPCInstance) -> None:
+        pass
+
+    @abc.abstractmethod
+    def delete(self, instance_id: str) -> None:
+        pass

fbpcs/private_computation/service/private_computation.py

@@ -13,7 +13,6 @@
 from typing import Any, DefaultDict, Dict, List, Optional, Set, Type, TypeVar, Union
 
 from fbpcp.entity.container_instance import ContainerInstanceStatus
-from fbpcp.entity.mpc_instance import MPCInstance
 from fbpcp.error.pcp import ThrottlingError
 from fbpcp.service.onedocker import OneDockerService
 from fbpcp.service.storage import StorageService
@@ -87,6 +86,7 @@
     PrivateComputationServiceInvalidStageError,
     PrivateComputationServiceValidationError,
 )
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCInstance
 from fbpcs.private_computation.service.mpc.mpc import MPCService
 from fbpcs.private_computation.service.pid_utils import (
     get_max_id_column_cnt,

fbpcs/private_computation/service/utils.py

@@ -15,7 +15,6 @@
 from typing import Any, Dict, List, Optional
 
 from fbpcp.entity.certificate_request import CertificateRequest
-from fbpcp.entity.mpc_instance import MPCInstance, MPCInstanceStatus, MPCParty
 from fbpcp.service.onedocker import OneDockerService
 from fbpcp.service.storage import StorageService
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
@@ -34,6 +33,11 @@
     SERVER_CERTIFICATE_ENV_VAR,
     SERVER_CERTIFICATE_PATH_ENV_VAR,
 )
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import (
+    MPCInstance,
+    MPCInstanceStatus,
+    MPCParty,
+)
 from fbpcs.private_computation.service.mpc.mpc import MPCService
 from fbpcs.private_computation.service.pid_utils import get_sharded_filepath
 

fbpcs/private_computation/test/entity/generate_instance_json.py

@@ -10,7 +10,6 @@
 import time
 
 from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus
-from fbpcp.entity.mpc_instance import MPCInstanceStatus, MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.common.entity.stage_state_instance import (
     StageStateInstance,
@@ -41,6 +40,10 @@
     LiftConfig,
     ProductConfig,
 )
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import (
+    MPCInstanceStatus,
+    MPCParty,
+)
 
 LIFT_PC_PATH: str = os.path.join(
     os.path.dirname(__file__),

fbpcs/private_computation/test/repository/test_private_computation_instance_local.py

@@ -10,7 +10,6 @@
 import string
 import unittest
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.private_computation.entity.infra_config import (
     InfraConfig,
@@ -30,6 +29,8 @@
     LocalPrivateComputationInstanceRepository,
 )
 
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
+
 
 class TestLocalPrivateComputationInstanceRepository(unittest.TestCase):
     def setUp(self) -> None:

fbpcs/private_computation/test/service/mpc/test_mpc.py

@@ -8,7 +8,11 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus
-from fbpcp.entity.mpc_instance import MPCInstance, MPCInstanceStatus, MPCParty
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import (
+    MPCInstance,
+    MPCInstanceStatus,
+    MPCParty,
+)
 from fbpcs.private_computation.service.mpc.mpc import MPCService
 
 
@@ -39,7 +43,9 @@
 class TestMPCService(IsolatedAsyncioTestCase):
     def setUp(self):
         cspatcher = patch("fbpcp.service.container.ContainerService")
-        irpatcher = patch("fbpcp.repository.mpc_instance.MPCInstanceRepository")
+        irpatcher = patch(
+            "fbpcs.private_computation.service.mpc.repository.mpc_instance.MPCInstanceRepository"
+        )
         gspatcher = patch(
             "fbpcs.private_computation.service.mpc.mpc_game.MPCGameService"
         )

fbpcs/private_computation/test/service/mpc/test_mpc_game.py

@@ -8,12 +8,12 @@
 from typing import List
 from unittest.mock import MagicMock, Mock
 
-from fbpcp.entity.mpc_instance import MPCParty
-
 from fbpcs.private_computation.service.mpc.entity.mpc_game_config import (
     MPCGameArgument,
     MPCGameConfig,
 )
+
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
 from fbpcs.private_computation.service.mpc.mpc_game import MPCGameService
 
 

fbpcs/private_computation/test/service/test_aggregate_shards_stage_service.py

@@ -8,7 +8,6 @@
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.infra.certificate.null_certificate_provider import NullCertificateProvider
 from fbpcs.onedocker_binary_config import OneDockerBinaryConfig
@@ -34,6 +33,8 @@
 )
 from fbpcs.private_computation.service.constants import NUM_NEW_SHARDS_PER_FILE
 
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
+
 
 class TestAggregateShardsStageService(IsolatedAsyncioTestCase):
     @patch("fbpcs.private_computation.service.mpc.mpc.MPCService")

fbpcs/private_computation/test/service/test_compute_metrics_stage_service.py

@@ -11,7 +11,6 @@
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.infra.certificate.null_certificate_provider import NullCertificateProvider
 from fbpcs.onedocker_binary_config import OneDockerBinaryConfig
@@ -35,6 +34,8 @@
     ComputeMetricsStageService,
 )
 from fbpcs.private_computation.service.constants import NUM_NEW_SHARDS_PER_FILE
+
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
 from fbpcs.private_computation.service.mpc.mpc import MPCService
 
 

fbpcs/private_computation/test/service/test_decoupled_aggregation_stage_service.py

@@ -8,7 +8,6 @@
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.infra.certificate.null_certificate_provider import NullCertificateProvider
 from fbpcs.onedocker_binary_config import OneDockerBinaryConfig
@@ -32,6 +31,8 @@
 from fbpcs.private_computation.service.decoupled_aggregation_stage_service import (
     AggregationStageService,
 )
+
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
 from fbpcs.private_computation.stage_flows.private_computation_decoupled_stage_flow import (
     PrivateComputationInstanceStatus,
 )

fbpcs/private_computation/test/service/test_decoupled_attribution_stage_service.py

@@ -8,7 +8,6 @@
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.infra.certificate.null_certificate_provider import NullCertificateProvider
 
@@ -35,6 +34,8 @@
     AttributionStageService,
 )
 
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
+
 
 class TestAttributionStageService(IsolatedAsyncioTestCase):
     @patch("fbpcs.private_computation.service.mpc.mpc.MPCService")

fbpcs/private_computation/test/service/test_pcf2_aggregation_stage_service.py

@@ -8,7 +8,6 @@
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.infra.certificate.null_certificate_provider import NullCertificateProvider
 from fbpcs.onedocker_binary_config import OneDockerBinaryConfig
@@ -29,6 +28,8 @@
 )
 from fbpcs.private_computation.repository.private_computation_game import GameNames
 from fbpcs.private_computation.service.constants import NUM_NEW_SHARDS_PER_FILE
+
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
 from fbpcs.private_computation.service.pcf2_aggregation_stage_service import (
     PCF2AggregationStageService,
 )

fbpcs/private_computation/test/service/test_pcf2_attribution_stage_service.py

@@ -8,7 +8,6 @@
 from unittest import IsolatedAsyncioTestCase
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from fbpcp.entity.mpc_instance import MPCParty
 from fbpcs.common.entity.pcs_mpc_instance import PCSMPCInstance
 from fbpcs.infra.certificate.null_certificate_provider import NullCertificateProvider
 from fbpcs.onedocker_binary_config import OneDockerBinaryConfig
@@ -30,6 +29,8 @@
 )
 from fbpcs.private_computation.repository.private_computation_game import GameNames
 from fbpcs.private_computation.service.constants import NUM_NEW_SHARDS_PER_FILE
+
+from fbpcs.private_computation.service.mpc.entity.mpc_instance import MPCParty
 from fbpcs.private_computation.service.pcf2_attribution_stage_service import (
     PCF2AttributionStageService,
 )