New Model Training Refactor

[mborodii-prog]

This pull request fixes a bug (Bug #972) where models created using the name parameter for various training types (classify, extract, lookup,
and standardize) could be created in a broken state, causing errors on inference immediately afterward. The root cause was a malformed model-creation
request (e.g. missing variant for non-default extract variants) rather than backend eventual-consistency, so the fix corrects the create payload and
consolidates duplicated logic — it does not add client-side polling/retries on creation.

Bug Fixes

• Introduced _create_model_with_content, a single shared helper used by classify, extract, lookup, and standardize for model creation via
  name. It issues one POST to /model/content with the full training payload and correct params (including variant for extract), and parses/logs the
  new model_id from the response.
• This removes the previous per-type duplication of the POST + model-id-parsing logic, so all four training types stay in sync going forward.

Test Enhancements

• Added integration tests in tests/connectors/test_train.py verifying that models created via name work correctly, including an end-to-end
  regression test (test_extract_name_creates_working_model_after_short_wait) that trains an extract model by name and confirms it produces correct
  inference results shortly after creation — validating the fix without relying on a retry/polling mechanism.
• Added a _delete_model test helper to clean up models created during tests.

Logging

• Model-id logging on creation is now centralized in _create_model_with_content, with a fallback log message if the id can't be parsed from the
  response.

[Copilot]

Pull request overview

Refactors train.classify, train.extract, train.lookup, and train.standardize so that creation via name first POSTs to create the model, parses the returned model_id, then performs a follow-up PUT with retries to push the training data — addressing Bug #972 where freshly-created models were not immediately usable. New integration tests assert that name-trained models can be queried right after training.

Changes:

• Two-step (POST create + PUT update-with-retries) flow for new model training across all four model types, with model-id logging moved up-front.
• New integration tests covering classify/extract/lookup/standardize "trained-by-name" readiness, plus a _wait_for_model retry helper and a best-effort _delete_model cleanup utility.
• Removed redundant post-training logging blocks now handled earlier in the flow.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 8 comments.

File	Description
wrangles/train.py	Refactors the four train-by-name flows into create-then-update-with-retries, moves model-id logging earlier. Drops variant from the extract create call (regression).
tests/connectors/test_train.py	Adds four integration tests, retry helper, and model-cleanup helper; includes a stray tabnanny import and non-unique model names.

Comments suppressed due to low confidence (2)

tests/connectors/test_train.py:34

• time is imported at module level (line 11) and again inside _wait_for_model (line 34). The inner import is redundant and should be removed.

    import time

tests/connectors/test_train.py:229

• _delete_model is only invoked at the end of each test, which means if the test fails any assertion above it (e.g. assert new_model_id is not None, the inference assertions, etc.), the cleanup is skipped and the test model leaks. Use a try/finally, a pytest fixture with teardown, or addfinalizer so cleanup runs unconditionally. This affects all four new test_*_name_creates_working_model tests.

    _delete_model(new_model_id)

[mborodii-prog]

@thomasstvr pls review one more time

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.

Comments suppressed due to low confidence (6)

wrangles/train.py:71

• The two-step flow now sends the full training_data twice: once in the initial POST that creates the model, and again in the follow-up PUT to model_id. For large training sets this doubles the bandwidth and server-side processing on every new-model creation. If the backend supports a "create empty model" endpoint (or a POST that only accepts metadata), that should be used instead; otherwise consider documenting this intentional double-send. This pattern is repeated in extract (lines 127-166), lookup (lines 266-301), and standardize (lines 348-381).

            create_response = _requests.post(
                        f'{_config.api_host}/model/content',
                        params={'type':'classify', 'name': name},
                        headers={'Authorization': f'Bearer {access_token}'},
                        json=training_data,
                    )
            if not create_response.ok:
                return create_response
            try:
                _json = create_response.json()
                new_model_id = (
                    _json.get('model_id') or _json.get('id') or _json.get('modelId') or _json.get('model')
                )
                if new_model_id:
                    _logging.info(f": New classify model created :: {new_model_id}")
                else:
                    _logging.info(f": Classify model created. Response: {create_response.text}")
            except Exception:
                _logging.info(": Classify model created (could not parse model id from response)")
                new_model_id = None
            if new_model_id:
                response = _utils.request_retries(
                            request_type='PUT',
                            url=f'{_config.api_host}/model/content',
                            **{
                                'params': {'type':'classify', 'model_id': new_model_id},
                                'headers': {'Authorization': f'Bearer {access_token}'},
                                'json': training_data
                            }
                        )

wrangles/train.py:49

• Error handling is inconsistent across the four refactored functions. lookup raises RuntimeError on a non-OK create response (line 275), while classify (line 49), extract (line 140), and standardize (line 357) silently return create_response to the caller. Since the stated goal of the refactor is to guarantee that a returned model is immediately usable, a failed create should fail loudly in all four — otherwise callers (and the new integration tests that scan the log for "New … model created") will continue past a failed create as if everything is fine. Recommend unifying to either raise or to surface a clear error in all four paths.

            if not create_response.ok:
                return create_response

wrangles/train.py:71

• access_token is captured once at the start of the function and reused for the subsequent retrying PUT. request_retries performs up to 3 retries with backoff, and on a slow/initializing model the total elapsed time can be significant. If the access token expires between the POST and the final PUT retry, every retry will get a 401 (which is not in status_forcelist, so it is not retried by the session adapter either) and the new model will be left without training data. Consider re-fetching the token before the PUT, or letting request_retries lazily fetch it per attempt. Same concern in extract (line 128), lookup (line 267), and standardize (line 349).

            access_token = _auth.get_access_token()
            create_response = _requests.post(
                        f'{_config.api_host}/model/content',
                        params={'type':'classify', 'name': name},
                        headers={'Authorization': f'Bearer {access_token}'},
                        json=training_data,
                    )
            if not create_response.ok:
                return create_response
            try:
                _json = create_response.json()
                new_model_id = (
                    _json.get('model_id') or _json.get('id') or _json.get('modelId') or _json.get('model')
                )
                if new_model_id:
                    _logging.info(f": New classify model created :: {new_model_id}")
                else:
                    _logging.info(f": Classify model created. Response: {create_response.text}")
            except Exception:
                _logging.info(": Classify model created (could not parse model id from response)")
                new_model_id = None
            if new_model_id:
                response = _utils.request_retries(
                            request_type='PUT',
                            url=f'{_config.api_host}/model/content',
                            **{
                                'params': {'type':'classify', 'model_id': new_model_id},
                                'headers': {'Authorization': f'Bearer {access_token}'},
                                'json': training_data
                            }
                        )

wrangles/train.py:289

• The except RuntimeError: raise here is dead code: the only statements in the try block are create_response.json(), .get(...) calls, and _logging.info(...) — none of which raise RuntimeError. The earlier raise RuntimeError(...) at line 275 fires before this try is entered. The explicit re-raise just adds noise; either remove it or restructure so the intent (don't swallow a RuntimeError from elsewhere) is clearer.

            except RuntimeError:
                raise
            except Exception:
                _logging.info(": Lookup model created (could not parse model id from response)")
                new_model_id = None

tests/connectors/test_train.py:17

• import time is already done at module top (line 10). The redundant import time inside _wait_for_model should be removed.

    import time

wrangles/train.py:135

• The **({'variant': variant} if variant else {}) pattern excludes variant='' (empty string) as well as None. Compared to the previous code which always sent 'variant': variant, this is a subtle behavior change — if a caller passes variant='' (or any falsy non-None value), it will no longer be forwarded. If the intent is "omit when not provided", prefer if variant is not None to be explicit and to match the previous semantics for non-None falsy values. Same pattern in the PUT params at line 161.

                        params={
                            'type': 'extract',
                            'name': name,
                            **({'variant': variant} if variant else {}),
                        },

[thomasstvr]

@mborodii-prog sorry, I meant for us to plan to create a wrangle from that function at a later time, and not in this branch. Please revert the changes.

[mborodii-prog]

@thomasstvr changes for delete are reverted. also I added enhancement for deduplication code

[thomasstvr]

@mborodii-prog something happened with this branch, it looks like it is very far behind main. Please update with main and I will have another look

[mborodii-prog]

@thomasstvr branch is fixed

[mborodii-prog]

@ebhills @thomasstvr you can use recipe '972 fix demo' for testing in QA.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (2)

tests/connectors/test_train.py:33

• _delete_model silently swallows all exceptions and ignores non-2xx responses. Since these tests create real remote models, this can leave test artifacts behind without any signal when cleanup fails (e.g., auth issues, API errors). Consider logging a warning when deletion fails so failures are visible without making the test suite brittle.

def _delete_model(model_id, model_type=None):
    """Delete a model by id. Best-effort - silently ignores failures."""
    from wrangles import config as _config, auth as _auth
    try:
        params = {'model_id': model_id}
        if model_type:
            params['type'] = model_type
        _requests.delete(
            f'{_config.api_host}/model/content',
            params=params,
            headers={'Authorization': f'Bearer {_auth.get_access_token()}'},
        )
    except Exception:
        pass

tests/connectors/test_train.py:696

• test_extract_name_creates_working_model uses _wait_for_model(..., max_wait=300) which will pass even if a newly created model takes minutes to become usable. That doesn’t validate the stated requirement that models created via name are usable immediately after training; it can also make the test suite hang for a long time during service degradation. Consider reducing the timeout to a small window (or removing the retry once the training API itself blocks until ready).

        result = _wait_for_model(
            f"""
            wrangles:
                - extract.custom:
                    input: description
                    output: characters
                    model_id: {new_model_id}
            """,
            dataframe=pd.DataFrame({'description': [
                'Rachel is a replicant from Blade Runner',
                'Dolores woke up in Westworld',
                'No character mentioned here',
            ]}),
            max_wait=300,
        )

[ebhills]

Queue triage (2026-07-27)

• Disposition: Draft — delivery-owner action
• Delivery owner: @mborodii-prog
• Next action: Address or explicitly answer the outstanding requested changes, then re-request one primary reviewer. This is a near-term v1.20 candidate once the review decision is current.

GitHub is the status record; update this PR rather than the external spreadsheet.