Exemplo n.º 1
0
    def test_request_dict_without_user_args(self, MockDiscovery):
        self.setup(MockDiscovery)
        job_name = deploy.deploy_job(
            self.region,
            self.docker_img,
            self.chief_config,
            self.worker_count,
            self.worker_config,
            None,
            self.stream_logs,
        )
        build_ret_val = MockDiscovery.build.return_value
        proj_ret_val = build_ret_val.projects.return_value
        jobs_ret_val = proj_ret_val.jobs.return_value

        del self.expected_request_dict["trainingInput"]["args"]

        # Verify job creation args
        _, kwargs = jobs_ret_val.create.call_args
        self.assertDictEqual(
            kwargs,
            {
                "parent": "projects/" + self.mock_project_name,
                "body": self.expected_request_dict,
            },
        )
Exemplo n.º 2
0
    def test_deploy_job(self, MockDiscovery, MockStdOut):
        self.setup(MockDiscovery)

        job_name = deploy.deploy_job(
            self.region,
            self.docker_img,
            self.chief_config,
            self.worker_count,
            self.worker_config,
            self.entry_point_args,
            self.stream_logs,
        )

        self.assertEqual(job_name, self.mock_job_id)

        # Verify discovery API is invoked as expected.
        self.assertEqual(MockDiscovery.build.call_count, 1)
        args, _ = MockDiscovery.build.call_args
        self.assertListEqual(list(args), ["ml", "v1"])

        # Verify job is created as expected
        build_ret_val = MockDiscovery.build.return_value
        self.assertEqual(build_ret_val.projects.call_count, 1)
        proj_ret_val = build_ret_val.projects.return_value
        self.assertEqual(proj_ret_val.jobs.call_count, 1)
        jobs_ret_val = proj_ret_val.jobs.return_value
        self.assertEqual(jobs_ret_val.create.call_count, 1)

        # Verify job creation args
        _, kwargs = jobs_ret_val.create.call_args
        self.assertDictEqual(
            kwargs,
            {
                "parent": "projects/" + self.mock_project_name,
                "body": self.expected_request_dict,
            },
        )

        # Verify print statement
        self.assertEqual(
            MockStdOut.getvalue(),
            "Job submitted successfully.\nYour job ID is:  {}\nPlease access "
            "your job logs at the following URL:\nhttps://"
            "console.cloud.google.com/mlengine/jobs/{}?project={}\n".format(
                self.mock_job_id, self.mock_job_id, self.mock_project_name),
        )
Exemplo n.º 3
0
    def test_request_dict_with_TPU_worker(self, MockDiscovery):
        self.setup(MockDiscovery)
        chief_config = machine_config.COMMON_MACHINE_CONFIGS["CPU"]
        worker_config = machine_config.COMMON_MACHINE_CONFIGS["TPU"]
        worker_count = 1

        job_name = deploy.deploy_job(
            self.region,
            self.docker_img,
            chief_config,
            worker_count,
            worker_config,
            self.entry_point_args,
            self.stream_logs,
        )
        build_ret_val = MockDiscovery.build.return_value
        proj_ret_val = build_ret_val.projects.return_value
        jobs_ret_val = proj_ret_val.jobs.return_value

        self.expected_request_dict["trainingInput"]["workerCount"] = "1"
        self.expected_request_dict["trainingInput"]["workerType"] = "cloud_tpu"
        self.expected_request_dict["trainingInput"][
            "masterType"] = "n1-standard-4"
        self.expected_request_dict["trainingInput"]["workerConfig"][
            "acceleratorConfig"]["type"] = "TPU_V3"
        self.expected_request_dict["trainingInput"]["workerConfig"][
            "acceleratorConfig"]["count"] = "8"
        self.expected_request_dict["trainingInput"]["workerConfig"][
            "tpuTfVersion"] = "2.1"
        self.expected_request_dict["trainingInput"]["masterConfig"][
            "acceleratorConfig"]["type"] = "ACCELERATOR_TYPE_UNSPECIFIED"
        self.expected_request_dict["trainingInput"]["masterConfig"][
            "acceleratorConfig"]["count"] = "0"

        # Verify job creation args
        _, kwargs = jobs_ret_val.create.call_args
        self.assertDictEqual(
            kwargs,
            {
                "parent": "projects/" + self.mock_project_name,
                "body": self.expected_request_dict,
            },
        )
Exemplo n.º 4
0
    def test_request_dict_without_workers(self, MockDiscovery):
        self.setup(MockDiscovery)
        worker_count = 0

        job_name = deploy.deploy_job(self.region, self.docker_img,
                                     self.chief_config, worker_count, None,
                                     self.entry_point_args, self.stream_logs)
        build_ret_val = MockDiscovery.build.return_value
        proj_ret_val = build_ret_val.projects.return_value
        jobs_ret_val = proj_ret_val.jobs.return_value

        self.expected_request_dict['trainingInput']['workerCount'] = str(
            worker_count)
        del self.expected_request_dict['trainingInput']['workerType']
        del self.expected_request_dict['trainingInput']['workerConfig']

        # Verify job creation args
        _, kwargs = jobs_ret_val.create.call_args
        self.assertDictEqual(
            kwargs, {
                'parent': 'projects/' + self.mock_project_name,
                'body': self.expected_request_dict
            })